Automatically apply SCF initial guess from existing wavefunction (#269)

* Automatically apply SCF initial guess from existing wavefunction

* indention issue

* Undefined variable

* Bug fixes and update tests

* Bug fixes

---------

Co-authored-by: Qiming Sun <qiming.sun@bytedance.com>

Author: sunqm

gpu4pyscf/df/df.py

@@ -123,7 +123,7 @@ def get_jk(self, dm, hermi=1, with_j=True, with_k=True,
         if key in self._rsh_df:
             rsh_df = self._rsh_df[key]
         else:
-            rsh_df = self._rsh_df[key] = copy.copy(self).reset()
+            rsh_df = self._rsh_df[key] = self.copy().reset()
             logger.info(self, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
 
         return df_jk.get_jk(rsh_df, dm, hermi, with_j, with_k, direct_scf_tol, omega=omega)

gpu4pyscf/df/df_jk.py

@@ -47,7 +47,7 @@ def build_df():
             if key in mf.with_df._rsh_df:
                 rsh_df = mf.with_df._rsh_df[key]
             else:
-                rsh_df = mf.with_df._rsh_df[key] = copy.copy(mf.with_df).reset()
+                rsh_df = mf.with_df._rsh_df[key] = mf.with_df.copy().reset()
             rsh_df.build(omega=omega)
         return
 
@@ -101,7 +101,7 @@ def _density_fit(mf, auxbasis=None, with_df=None, only_dfj=False):
             mf.with_df = with_df
         elif getattr(mf.with_df, 'auxbasis', None) != auxbasis:
             #logger.warn(mf, 'DF might have been initialized twice.')
-            mf = copy.copy(mf)
+            mf = mf.copy()
             mf.with_df = with_df
             mf.only_dfj = only_dfj
         return mf
@@ -528,7 +528,7 @@ def _get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True,
     if key in dfobj._rsh_df:
         rsh_df = dfobj._rsh_df[key]
     else:
-        rsh_df = dfobj._rsh_df[key] = copy.copy(dfobj).reset()
+        rsh_df = dfobj._rsh_df[key] = dfobj.copy().reset()
         logger.info(dfobj, 'Create RSH-DF object %s for omega=%s', rsh_df, omega)
 
     with rsh_df.mol.with_range_coulomb(omega):

gpu4pyscf/df/grad/rhf.py

@@ -71,7 +71,7 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
             with_df = mf_grad.base.with_df._rsh_df[key]
         else:
             dfobj = mf_grad.base.with_df
-            with_df = dfobj._rsh_df[key] = copy.copy(dfobj).reset()
+            with_df = dfobj._rsh_df[key] = dfobj.copy().reset()
 
     auxmol = with_df.auxmol
     if not hasattr(with_df, 'intopt') or with_df._cderi is None:
@@ -282,4 +282,4 @@ def extra_force(self, atom_id, envs):
         else:
             return 0
 
-Grad = Gradients
+Grad = Gradients

gpu4pyscf/df/grad/uhf.py

@@ -51,7 +51,7 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True,
             with_df = mf_grad.base.with_df._rsh_df[key]
         else:
             dfobj = mf_grad.base.with_df
-            with_df = dfobj._rsh_df[key] = copy.copy(dfobj).reset()
+            with_df = dfobj._rsh_df[key] = dfobj.copy().reset()
 
     auxmol = with_df.auxmol
     if not hasattr(with_df, 'intopt') or with_df._cderi is None:

gpu4pyscf/df/tests/test_df_geomopt.py

@@ -82,7 +82,8 @@ def test_df_rhf_geomopt(self):
         assert np.linalg.norm(coords - coords_qchem) < 1e-4
 
     def test_df_uks_geomopt(self):
-        mf = uks.UKS(mol, xc=xc, disp=disp).density_fit()
+        mf = uks.UKS(mol, xc=xc).density_fit()
+        mf.disp = disp
         mf.grids.level = grids_level
         mf.kernel()
         mol_eq = optimize(mf, maxsteps=20)

gpu4pyscf/dft/roks.py

@@ -16,38 +16,42 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import numpy as np
-from pyscf.dft import roks
-from gpu4pyscf.dft import numint
+import cupy as cp
+from pyscf.dft import roks as roks_cpu
 from gpu4pyscf.scf.rohf import ROHF
-from gpu4pyscf.dft import uks, gen_grid
+from gpu4pyscf.dft import rks, uks
 from gpu4pyscf.lib.cupy_helper import tag_array
+from gpu4pyscf.lib import utils
 
-class ROKS(roks.ROKS, ROHF):
-    from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
+class ROKS(rks.KohnShamDFT, ROHF):
 
     def __init__(self, mol, xc='LDA,VWN'):
-        super().__init__(mol, xc)
-        self._numint = numint.NumInt()
-        self.disp = None
-        self.screen_tol = 1e-14
-
-        grids_level = self.grids.level
-        self.grids = gen_grid.Grids(mol)
-        self.grids.level = grids_level
-
-        nlcgrids_level = self.nlcgrids.level
-        self.nlcgrids = gen_grid.Grids(mol)
-        self.nlcgrids.level = nlcgrids_level
+        ROHF.__init__(self, mol)
+        rks.KohnShamDFT.__init__(self, xc)
 
     def get_veff(self, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
-        if getattr(dm, 'mo_coeff', None) is not None:
+        if dm is None:
+            dm = self.make_rdm1()
+        elif getattr(dm, 'mo_coeff', None) is not None:
             mo_coeff = dm.mo_coeff
             mo_occ_a = (dm.mo_occ > 0).astype(np.double)
             mo_occ_b = (dm.mo_occ ==2).astype(np.double)
-            dm = tag_array(dm, mo_coeff=(mo_coeff,mo_coeff),
-                           mo_occ=(mo_occ_a,mo_occ_b))
+            if dm.ndim == 2:
+                dm = cp.repeat(dm[None]*.5, 2, axis=0)
+            dm = tag_array(dm, mo_coeff=cp.asarray((mo_coeff,mo_coeff)),
+                           mo_occ=cp.asarray((mo_occ_a,mo_occ_b)))
+        elif dm.ndim == 2:
+            dm = cp.repeat(dm[None]*.5, 2, axis=0)
         return uks.get_veff(self, mol, dm, dm_last, vhf_last, hermi)
 
     energy_elec = uks.UKS.energy_elec
     nuc_grad_method = NotImplemented
     to_hf = NotImplemented
+
+    to_gpu = utils.to_gpu
+    device = utils.device
+
+    def to_cpu(self):
+        mf = roks_cpu.ROKS(self.mol, xc=self.xc)
+        utils.to_cpu(self, mf)
+        return mf

gpu4pyscf/dft/uks.py

@@ -16,12 +16,13 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import cupy
-from pyscf.dft import uks
+from pyscf.dft import uks as uks_cpu
 from pyscf import lib
 from gpu4pyscf.lib import logger
 from gpu4pyscf.dft import numint, gen_grid, rks
 from gpu4pyscf.scf import hf, uhf
 from gpu4pyscf.lib.cupy_helper import tag_array
+from gpu4pyscf.lib import utils
 
 
 def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
@@ -30,8 +31,9 @@ def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
     '''
     if mol is None: mol = ks.mol
     if dm is None: dm = ks.make_rdm1()
+    assert dm.ndim == 3
     t0 = logger.init_timer(ks)
-    rks.initialize_grids(ks, mol, dm)
+    rks.initialize_grids(ks, mol, cupy.asarray(dm[0]+dm[1]))
 
     if hasattr(ks, 'screen_tol') and ks.screen_tol is not None:
         ks.direct_scf_tol = ks.screen_tol
@@ -42,7 +44,7 @@ def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
         n, exc, vxc = (0,0), 0, 0
     else:
         max_memory = ks.max_memory - lib.current_memory()[0]
-        n, exc, vxc = ni.nr_uks(mol, ks.grids, ks.xc, dm, max_memory=max_memory)
+        n, exc, vxc = ni.nr_uks(mol, ks.grids, ks.xc, dm.view(cupy.ndarray), max_memory=max_memory)
         logger.debug(ks, 'nelec by numeric integration = %s', n)
         if ks.do_nlc():
             if ni.libxc.is_nlc(ks.xc):
@@ -61,7 +63,10 @@ def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
         vk = None
         if (ks._eri is None and ks.direct_scf and
             getattr(vhf_last, 'vj', None) is not None):
-            ddm = cupy.asarray(dm) - cupy.asarray(dm_last)
+            dm_last = cupy.asarray(dm_last)
+            dm = cupy.asarray(dm)
+            assert dm_last.ndim == 0 or dm_last.ndim == dm.ndim
+            ddm = dm - dm_last
             vj = ks.get_j(mol, ddm[0]+ddm[1], hermi)
             vj += vhf_last.vj
         else:
@@ -71,7 +76,10 @@ def get_veff(ks, mol=None, dm=None, dm_last=0, vhf_last=0, hermi=1):
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(ks.xc, spin=mol.spin)
         if (ks._eri is None and ks.direct_scf and
             getattr(vhf_last, 'vk', None) is not None):
-            ddm = cupy.asarray(dm) - cupy.asarray(dm_last)
+            dm_last = cupy.asarray(dm_last)
+            dm = cupy.asarray(dm)
+            assert dm_last.ndim == 0 or dm_last.ndim == dm.ndim
+            ddm = dm - dm_last
             vj, vk = ks.get_jk(mol, ddm, hermi)
             vk *= hyb
             if abs(omega) > 1e-10:  # For range separated Coulomb operator
@@ -113,19 +121,15 @@ def energy_elec(ks, dm=None, h1e=None, vhf=None):
 
 
 class UKS(rks.KohnShamDFT, uhf.UHF):
-    from gpu4pyscf.lib.utils import to_gpu, device
-    _keys = {'disp', 'screen_tol'}
-
-    def __init__(self, mol, xc='LDA,VWN', disp=None):
+    def __init__(self, mol, xc='LDA,VWN'):
         uhf.UHF.__init__(self, mol)
         rks.KohnShamDFT.__init__(self, xc)
-        self.disp = disp
 
     get_veff = get_veff
-    get_vasp = uks.get_vsap
+    get_vasp = uks_cpu.get_vsap
     energy_elec = energy_elec
     energy_tot = hf.RHF.energy_tot
-    init_guess_by_vsap = uks.UKS.init_guess_by_vsap
+    init_guess_by_vsap = uks_cpu.UKS.init_guess_by_vsap
 
     to_hf = NotImplemented
 
@@ -141,9 +145,10 @@ def nuc_grad_method(self):
         from gpu4pyscf.grad import uks as uks_grad
         return uks_grad.Gradients(self)
 
+    to_gpu = utils.to_gpu
+    device = utils.device
+
     def to_cpu(self):
-        from gpu4pyscf.lib import utils
-        mf = uks.UKS(self.mol, xc=self.xc)
-        mf.disp = self.disp
+        mf = uks_cpu.UKS(self.mol, xc=self.xc)
         utils.to_cpu(self, mf)
         return mf

gpu4pyscf/grad/tests/test_geomopt.py

@@ -81,7 +81,8 @@ def test_rhf_geomopt(self):
         assert np.linalg.norm(coords - coords_qchem) < 1e-4
 
     def test_uks_geomopt(self):
-        mf = uks.UKS(mol, xc=xc, disp=disp)
+        mf = uks.UKS(mol, xc=xc)
+        mf.disp = disp
         mf.grids.level = grids_level
         mf.kernel()
         mol_eq = optimize(mf, maxsteps=20)

gpu4pyscf/gto/mole.py

@@ -82,7 +82,7 @@ def basis_seg_contraction(mol, allow_replica=False):
             bas_templates[key] = bas_of_ia
         _bas.append(bas_of_ia)
 
-    pmol = copy.copy(mol)
+    pmol = mol.copy()
     pmol.output = mol.output
     pmol.verbose = mol.verbose
     pmol.stdout = mol.stdout

gpu4pyscf/scf/_response_functions.py

@@ -16,7 +16,7 @@
 import cupy
 from pyscf import lib
 from gpu4pyscf.lib import logger
-from gpu4pyscf.scf import hf, uhf
+from gpu4pyscf.scf import hf, uhf, rohf
 
 def _gen_rhf_response(mf, mo_coeff=None, mo_occ=None,
                       singlet=None, hermi=0, grids=None, max_memory=None):
@@ -133,7 +133,7 @@ def _gen_uhf_response(mf, mo_coeff=None, mo_occ=None,
     '''Generate a function to compute the product of UHF response function and
     UHF density matrices.
     '''
-    assert isinstance(mf, uhf.UHF)
+    assert isinstance(mf, (uhf.UHF, rohf.ROHF))
     if mo_coeff is None: mo_coeff = mf.mo_coeff
     if mo_occ is None: mo_occ = mf.mo_occ
     mol = mf.mol
@@ -194,3 +194,4 @@ def vind(dm1):
 
 hf.RHF.gen_response = _gen_rhf_response
 uhf.UHF.gen_response = _gen_uhf_response
+rohf.ROHF.gen_response = _gen_uhf_response