refactor dh

Author: jeanwsr

pyscf/dh/dh.py

@@ -82,6 +82,23 @@ def __init__(self,
         self.e_tot = NotImplemented
         self.eng_tot = self.eng_nc = self.eng_pt2 = self.eng_nuc = self.eng_os = self.eng_ss = NotImplemented
 
+    def _init_common(self):
+        """Shared initialization after subclass creates mf_s / mf_n."""
+        mf_s = self.mf_s
+        mol = self.mol
+        if not hasattr(self, 'grids') or self.grids is None:
+            self.grids = getattr(mf_s, 'grids', dft.Grids(mol))
+        if not hasattr(self, 'grids_cpks') or self.grids_cpks is None:
+            self.grids_cpks = self.grids
+        self.mf_s.grids = self.grids
+        self.ni = getattr(mf_s, '_numint', dft.numint.NumInt())
+        self.cx = self.ni.hybrid_coeff(self.xc)
+        self.cx_n = self.ni.hybrid_coeff(self.xc_n)
+        self.df_jk = mf_s.with_df
+        self.df_ri = df.DF(mol, self.auxbasis_ri) if not self.same_aux else self.df_jk
+        if self._scf is not None and self._scf.e_tot != 0:
+            self.run_scf()
+
     @property
     def base(self):
         return self
@@ -163,47 +180,8 @@ def build(self):
         self.mf_s.grids = self.mf_n.grids = self.grids
         if self.df_jk.auxmol is None:
             self.df_jk.build()
-            self.aux_jk = self.df_jk.auxmol
         if self.df_ri.auxmol is None:
             self.df_ri.build()
-            self.aux_ri = self.df_ri.auxmol
-
-    def set_frozen(self, method='auto'):
-        if method != 'auto':
-            raise NotImplementedError("Only method='auto' is supported.")
-        return _ccsd_set_frozen(self)
-
-    def dump_intermediates(self, dir_path="scratch"):
-        os.makedirs(dir_path, exist_ok=True)
-        tensors = self.tensors
-        h5_path = dir_path + "/tensors.h5"
-        dat_path = dir_path + "/tensors.dat"
-        tensors.dump(h5_path, dat_path)
-        att_path = dir_path + "/attributes.dat"
-        dct = {
-            "mo_coeff": self.mo_coeff,
-            "mo_energy": self.mo_energy,
-            "D": self.D,
-            "mo_occ": self.mo_occ,
-            "mf_s_e_tot": self.mf_s.e_tot,
-        }
-        with open(att_path, "wb") as f:
-            pickle.dump(dct, f)
-
-    def load_intermediates(self, dir_path="scratch", rerun_scf=False):
-        h5_path = dir_path + "/tensors.h5"
-        dat_path = dir_path + "/tensors.dat"
-        self.tensors = HybridDict.pick(h5_path, dat_path)
-        att_path = dir_path + "/attributes.dat"
-        with open(att_path, "rb") as f:
-            dct = pickle.load(f)
-        self.mf_s.mo_coeff = dct["mo_coeff"]
-        self.mf_s.mo_energy = dct["mo_energy"]
-        self.mf_s.mo_occ = dct["mo_occ"]
-        self.mf_s.e_tot = dct["mf_s_e_tot"]
-        if rerun_scf:
-            self.mf_s.kernel(dm=self.mf_s.make_rdm1())
-        self.run_scf()
         return self
 
 

pyscf/dh/grad/dfdh.py

@@ -407,7 +407,7 @@ def prepare_gradient_gga(self):
     def prepare_gradient_pt2(self):
         tensors = self.tensors
         C, e = self.mo_coeff, self.mo_energy
-        mol, aux_ri = self.mol, self.aux_ri
+        mol, aux_ri = self.mol, self.df_ri.auxmol
         natm, nao, nmo, nocc, nvir, naux = mol.natm, self.nao, self.nmo, self.nocc, self.nvir, self.df_ri.get_naoaux()
         # this algorithm asserts naux = aux.nao, i.e. no linear dependency in auxiliary basis
         assert naux == aux_ri.nao

pyscf/dh/grad/udfdh.py

@@ -164,7 +164,7 @@ def prepare_gradient_gga(self):
     def prepare_gradient_pt2(self):
         tensors = self.tensors
         C, e = self.mo_coeff, self.mo_energy
-        mol, aux_ri = self.mol, self.aux_ri
+        mol, aux_ri = self.mol, self.df_ri.auxmol
         natm, nao, nocc, nvir, naux = mol.natm, self.nao, self.nocc, self.nvir, self.df_ri.get_naoaux()
         mocc = max(nocc)
         # this algorithm asserts naux = aux.nao, i.e. no linear dependency in auxiliary basis

pyscf/dh/rdfdh.py

@@ -68,25 +68,15 @@ def __init__(self,
         self.grids = grids if grids else (getattr(mf_s, 'grids', dft.Grids(mol)))
         self.grids_cpks = grids_cpks if grids_cpks else self.grids
         self.mf_s = mf_s
-        self.mf_s.grids = self.grids
         self.xc_n = None if xc_equal(self.xc_n, self.xc) else self.xc_n
         self.mf_n = self.mf_s
         if self.xc_n:
             self.mf_n = dft.KS(mol, xc=self.xc_n).density_fit(auxbasis=self.auxbasis_jk)
-            self.mf_n.grids = self.mf_s.grids
-            self.mf_n.grids = self.grids
-        self.ni = getattr(self.mf_s, '_numint', dft.numint.NumInt())
-        self.cx = self.ni.hybrid_coeff(self.xc)
-        self.cx_n = self.ni.hybrid_coeff(self.xc_n)
-        self.df_jk = mf_s.with_df
-        self.aux_jk = self.df_jk.auxmol
-        self.df_ri = df.DF(mol, self.auxbasis_ri) if not self.same_aux else self.df_jk
-        self.aux_ri = self.df_ri.auxmol
+            self.mf_n.grids = self.mf_s.grids = self.grids
+        self._init_common()
         self.nocc = mol.nelec[0]
         self.nmo = self.nao
         self.nvir = self.nmo - self.nocc
-        if self._scf is not None and self._scf.e_tot != 0:
-            self.run_scf()
         if mp2_backend == "dfmp2_native":
             self.energy_elec_mp2 = partial(energy_elec_mp2_dfmp2_native, self)
         elif mp2_backend == "dfmp2":

pyscf/dh/udfdh.py

@@ -71,27 +71,16 @@ def __init__(self,
         self.grids = grids if grids else (getattr(mf_s, 'grids', dft.Grids(mol)))
         self.grids_cpks = grids_cpks if grids_cpks else self.grids
         self.mf_s = mf_s
-        self.mf_s.grids = self.grids
         self.xc_n = None if xc_equal(self.xc_n, self.xc) else self.xc_n
         self.mf_n = self.mf_s
         if self.xc_n:
             self.mf_n = dft.UKS(mol, xc=self.xc_n).density_fit(auxbasis=self.auxbasis_jk)
-            self.mf_n.grids = self.mf_s.grids
-            self.mf_n.grids = self.grids
-        self.ni = getattr(self.mf_s, '_numint', dft.numint.NumInt())
-        self.cx = self.ni.hybrid_coeff(self.xc)
-        self.cx_n = self.ni.hybrid_coeff(self.xc_n)
-        self.df_jk = mf_s.with_df
-        self.aux_jk = self.df_jk.auxmol
-        self.df_ri = df.DF(mol, self.auxbasis_ri) if not self.same_aux else self.df_jk
-        self.aux_ri = self.df_ri.auxmol
+            self.mf_n.grids = self.mf_s.grids = self.grids
+        self._init_common()
         self.nocc = mol.nelec
-        self.mvir = NotImplemented
         self.mocc = max(max(self.nocc), 1)
         self.nmo = self.nao
         self.nvir = (self.nmo - self.nocc[α], self.nmo - self.nocc[β])
-        if self._scf is not None and self._scf.e_tot != 0:
-            self.run_scf()
         if mp2_backend == "dfmp2_native":
             self.energy_elec_mp2 = partial(energy_elec_mp2_dfump2_native, self)
         elif mp2_backend == "dfmp2":