simplify

Author: jeanwsr

pyscf/dh/grad/dfdh.py

@@ -19,23 +19,19 @@
 #
 
 from __future__ import annotations
-# dh import
 from pyscf.dh.resp import RDHRespMixin
 from pyscf.dh.resp import _r_get_eri_cpks, _u_get_eri_cpks
-from pyscf.dh.resp import _get_Y_mo, _prepare_pt2_r, _prepare_pt2_u
+from pyscf.dh.resp import _get_Y_mo
 from pyscf.dh.dh import DHBase
 from pyscf.dh.dhutil import calc_batch_size, gen_batch, gen_shl_batch, timing, as_scanner_grad, available_memory
-# pyscf import
 from pyscf import gto, lib, df
 from pyscf.df.grad.rhf import _int3c_wrapper as int3c_wrapper
-# other import
 import numpy as np
 
 einsum = lib.einsum
 
 
 def kernel(mf_dh: Gradients, **kwargs):
-    # unrestricted method requires dump t_ijab_αβ to disk; controling αβ and SS dumping is too hard for me
     dump_t_ijab = mf_dh.with_t_ijab
 
     mf_dh.base.build()
@@ -68,10 +64,7 @@ def kernel(mf_dh: Gradients, **kwargs):
     mf_dh.tensors.consume(Y_mo).consume(eri)
 
     mf_dh.prepare_xc_kernel()
-    if mf_dh.D.ndim == 2:
-        pt2_res, eng_bi = _prepare_pt2_r(mf_dh, dump_t_ijab=dump_t_ijab)
-    else:
-        pt2_res, eng_bi = _prepare_pt2_u(mf_dh, dump_t_ijab=dump_t_ijab)
+    pt2_res, eng_bi = mf_dh._prepare_pt2(dump_t_ijab=dump_t_ijab)
     mf_dh.tensors.consume(pt2_res)
     if mf_dh.base.eng_tot is None and eng_bi is not None:
         DHBase.kernel(mf_dh.base, eng_bi=(None,) + tuple(eng_bi))

pyscf/dh/polar/rdfdh.py

@@ -22,7 +22,7 @@
 # dh import
 from pyscf.dh.resp import RDHRespMixin
 from pyscf.dh.resp import _r_get_eri_cpks, _u_get_eri_cpks
-from pyscf.dh.resp import _get_Y_mo, _prepare_pt2_r, _prepare_pt2_u
+from pyscf.dh.resp import _get_Y_mo
 from pyscf.dh.dh import DHBase
 from pyscf.dh.dhutil import gen_batch, get_rho_from_dm_gga, restricted_biorthogonalize, hermi_sum_last2dim
 from pyscf.dh.xccode import xc_equal
@@ -55,10 +55,7 @@ def kernel(mf: Polar):
     mf.tensors.consume(Y_mo).consume(eri)
 
     mf.prepare_xc_kernel()
-    if mf.D.ndim == 2:
-        pt2_res, eng_bi = _prepare_pt2_r(mf, dump_t_ijab=True)
-    else:
-        pt2_res, eng_bi = _prepare_pt2_u(mf, dump_t_ijab=True)
+    pt2_res, eng_bi = mf._prepare_pt2(dump_t_ijab=True)
     mf.tensors.consume(pt2_res)
     if mf.base.eng_tot is None and eng_bi is not None:
         DHBase.kernel(mf.base, eng_bi=(None,) + tuple(eng_bi))

pyscf/dh/resp.py

@@ -370,10 +370,7 @@ def make_rdm1_relaxed(self, ao_repr=False):
                                        self.nocc, self.cx, self._incore_Y_mo, self.max_memory)
             self.tensors.consume(Y_mo).consume(eri)
             self.prepare_xc_kernel()
-            if self.D.ndim == 2:
-                pt2_res, _ = _prepare_pt2_r(self, dump_t_ijab=True)
-            else:
-                pt2_res, _ = _prepare_pt2_u(self, dump_t_ijab=True)
+            pt2_res, _ = self._prepare_pt2(dump_t_ijab=True)
             self.tensors.consume(pt2_res)
             self.prepare_lagrangian()
             self.prepare_D_r()
@@ -449,6 +446,9 @@ def build(sI, t_ijab, g_ijab):
 class RDHRespMixin(RespMixin):
     """Restricted response mixin: Ax0_*, prepare_lagrangian."""
 
+    def _prepare_pt2(self, dump_t_ijab=True):
+        return _prepare_pt2_r(self, dump_t_ijab=dump_t_ijab)
+
     def Ax0_Core_HF(self, si, sa, sj, sb, cx=None):
         Y_mo_jk = self.tensors["Y_mo_jk"]
         cx = cx if cx else self.cx
@@ -605,6 +605,9 @@ def _prepare_pt2_u(mf, dump_t_ijab=True):
 class UDHRespMixin(RespMixin):
     """Unrestricted response mixin: Ax0_*, prepare_lagrangian."""
 
+    def _prepare_pt2(self, dump_t_ijab=True):
+        return _prepare_pt2_u(self, dump_t_ijab=dump_t_ijab)
+
     def Ax0_Core_HF(self, si, sa, sj, sb, cx=None):
         Y_mo_jk = [self.tensors["Y_mo_jk" + str(σ)] for σ in (α, β)]
         cx = cx if cx else self.cx