Fix SMD hessian (#534)

* Fix SMD hessian

* lint

* import error

* Forcing an empty commit.

* Update SMD implementation to make it consistent with PySCF SMD module

* Mute some tests

* Update tests

* Fix to_gpu to_cpu issues

* lint

Author: sunqm

gpu4pyscf/df/hessian/rhf.py

@@ -643,7 +643,7 @@ def _get_jk_mo(hessobj, mol, dms, mo_coeff, mocc,
 class Hessian(rhf_hess.Hessian):
     '''Non-relativistic restricted Hartree-Fock hessian'''
 
-    from gpu4pyscf.lib.utils import to_gpu, device
+    _keys = {'auxbasis_response',}
 
     auxbasis_response = 1
     partial_hess_elec = partial_hess_elec

gpu4pyscf/df/hessian/rks.py

@@ -104,7 +104,8 @@ def make_h1(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
 
 class Hessian(rks_hess.Hessian):
     '''Non-relativistic RKS hessian'''
-    from gpu4pyscf.lib.utils import to_gpu, device
+
+    _keys = {'auxbasis_response',}
 
     auxbasis_response = 1
     partial_hess_elec = partial_hess_elec

gpu4pyscf/df/hessian/uhf.py

@@ -708,7 +708,7 @@ def _ao2mo(mat, mocc, mo):
 class Hessian(uhf_hess.Hessian):
     '''Non-relativistic restricted Hartree-Fock hessian'''
 
-    from gpu4pyscf.lib.utils import to_gpu, device
+    _keys = {'auxbasis_response',}
 
     auxbasis_response = 1
     partial_hess_elec = partial_hess_elec

gpu4pyscf/df/hessian/uks.py

@@ -132,7 +132,8 @@ def make_h1(hessobj, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
 
 class Hessian(uks_hess.Hessian):
     '''Non-relativistic RKS hessian'''
-    from gpu4pyscf.lib.utils import to_gpu, device
+
+    _keys = {'auxbasis_response',}
 
     auxbasis_response = 1
     partial_hess_elec = partial_hess_elec

gpu4pyscf/hessian/rhf.py

@@ -36,6 +36,7 @@
 from gpu4pyscf.__config__ import _streams, num_devices
 from gpu4pyscf.lib import logger
 from gpu4pyscf.lib import multi_gpu
+from gpu4pyscf.lib import utils
 from gpu4pyscf.scf.jk import (
     LMAX, QUEUE_DEPTH, SHM_SIZE, THREADS, GROUP_SIZE, libvhf_rys, _VHFOpt,
     init_constant, _make_tril_tile_mappings, _make_tril_pair_mappings,
@@ -911,6 +912,11 @@ def _get_veff_resp_mo(hessobj, mol, dms, mo_coeff, mo_occ, hermi=1, omega=None):
     return vj - 0.5 * vk
 
 class HessianBase(lib.StreamObject):
+
+    to_cpu = utils.to_cpu
+    to_gpu = utils.to_gpu
+    device = utils.device
+
     # attributes
     max_cycle   = rhf_hess_cpu.HessianBase.max_cycle
     level_shift = rhf_hess_cpu.HessianBase.level_shift
@@ -947,19 +953,9 @@ def dump_flags(self, verbose=None):
                  self.max_memory, lib.current_memory()[0])
         return self
 
-    def to_cpu(self):
-        mf = self.base.to_cpu()
-        from importlib import import_module
-        mod = import_module(self.__module__.replace('gpu4pyscf', 'pyscf'))
-        cls = getattr(mod, self.__class__.__name__)
-        obj = cls(mf)
-        return obj
-
 class Hessian(HessianBase):
     '''Non-relativistic restricted Hartree-Fock hessian'''
 
-    from gpu4pyscf.lib.utils import to_gpu, device
-
     def __init__(self, scf_method):
         self.verbose = scf_method.verbose
         self.stdout = scf_method.stdout

gpu4pyscf/hessian/rks.py

@@ -4250,8 +4250,6 @@ def get_veff_resp_mo(hessobj, mol, dms, mo_coeff, mo_occ, hermi=1, omega=None):
 class Hessian(rhf_hess.HessianBase):
     '''Non-relativistic RKS hessian'''
 
-    from gpu4pyscf.lib.utils import to_gpu, device
-
     _keys = {'grids', 'grid_response'}
 
     def __init__(self, mf):

gpu4pyscf/hessian/uhf.py

@@ -459,8 +459,6 @@ def _get_veff_resp_mo(hessobj, mol, dms, mo_coeff, mo_occ, hermi=1):
 class Hessian(rhf_hess_gpu.HessianBase):
     '''Non-relativistic unrestricted Hartree-Fock hessian'''
 
-    from gpu4pyscf.lib.utils import to_gpu, device
-
     __init__ = rhf_hess_gpu.Hessian.__init__
     partial_hess_elec = partial_hess_elec
     hess_elec = hess_elec

gpu4pyscf/hessian/uks.py

@@ -895,7 +895,6 @@ def get_veff_resp_mo(hessobj, mol, dms, mo_coeff, mo_occ, hermi=1):
 
 class Hessian(rhf_hess.HessianBase):
     '''Non-relativistic UKS hessian'''
-    from gpu4pyscf.lib.utils import to_gpu, device
 
     def __init__(self, mf):
         rhf_hess.Hessian.__init__(self, mf)

gpu4pyscf/lib/utils.py

@@ -87,7 +87,10 @@ def to_cpu(method, out=None):
             val = val.get()
         setattr(out, key, val)
     if hasattr(out, 'reset'):
-        out.reset()
+        try:
+            out.reset()
+        except NotImplementedError:
+            pass
     return out
 
 def to_gpu(method, device=None):

gpu4pyscf/scf/hf.py

@@ -763,6 +763,7 @@ def reset(self, mol=None):
             self.mol = mol
         self._opt_gpu = {None: None}
         self._opt_jengine = {None: None}
+        self._eri = None
         self.scf_summary = {}
         self.overlap_canonical_decomposed_x = None
         return self