fixed a bug in smd gradient

gpu4pyscf/solvent/grad/pcm.py

@@ -323,9 +323,24 @@ def make_grad_object(grad_method):
     name = (grad_method.base.with_solvent.__class__.__name__
             + grad_method.__class__.__name__)
     return lib.set_class(WithSolventGrad(grad_method),
-                         (grad_method.__class__, WithSolventGrad), name)
+                         (WithSolventGrad, grad_method.__class__), name)
+
+class WithSolventGrad:
+    _keys = {'de_solvent', 'de_solute'}
+
+    def __init__(self, grad_method):
+        self.__dict__.update(grad_method.__dict__)
+        self.de_solvent = None
+        self.de_solute = None
+
+    def undo_solvent(self):
+        cls = self.__class__
+        name_mixin = self.base.with_solvent.__class__.__name__
+        obj = lib.view(self, lib.drop_class(cls, WithSolventGrad, name_mixin))
+        del obj.de_solvent
+        del obj.de_solute
+        return obj
 
-class WithSolventGrad(ddcosmo_grad.WithSolventGrad):
     def kernel(self, *args, dm=None, atmlst=None, **kwargs):
         dm = kwargs.pop('dm', None)
         if dm is None:
@@ -335,7 +350,7 @@ def kernel(self, *args, dm=None, atmlst=None, **kwargs):
         self.de_solvent+= grad_solver(self.base.with_solvent, dm)
         self.de_solvent+= grad_nuc(self.base.with_solvent, dm)
 
-        self.de_solute = super().kernel(self, *args, **kwargs)
+        self.de_solute = super().kernel(*args, **kwargs)
         self.de = self.de_solute + self.de_solvent
 
         if self.verbose >= logger.NOTE:
@@ -345,3 +360,8 @@ def kernel(self, *args, dm=None, atmlst=None, **kwargs):
             rhf_grad._write(self, self.mol, self.de, self.atmlst)
             logger.note(self, '----------------------------------------------')
         return self.de
+
+    def _finalize(self):
+        # disable _finalize. It is called in grad_method.kernel method
+        # where self.de was not yet initialized.
+        pass

gpu4pyscf/solvent/grad/smd.py

@@ -215,8 +215,9 @@ def get_cds(smdobj):
     SASA *= radii.BOHR**2
     grad_SASA *= radii.BOHR**2
     mol_cds = mol_tension * np.sum(grad_SASA, axis=0) / 1000
-    atm_cds = SASA.dot(grad_tension) / 1000
-    atm_cds = 0.0*np.einsum('i,ijx->jx', atm_tension, grad_SASA) / 1000
+    grad_tension *= radii.BOHR
+    atm_cds = np.einsum('i,ijx->jx', SASA, grad_tension) / 1000
+    atm_cds+= np.einsum('i,ijx->jx', atm_tension, grad_SASA) / 1000
     return (mol_cds + atm_cds)/hartree2kcal # hartree
 
 def make_grad_object(grad_method):
@@ -229,7 +230,22 @@ def make_grad_object(grad_method):
     return lib.set_class(WithSolventGrad(grad_method),
                          (grad_method.__class__, WithSolventGrad), name)
 
-class WithSolventGrad(ddcosmo_grad.WithSolventGrad):
+class WithSolventGrad:
+    _keys = {'de_solvent', 'de_solute'}
+
+    def __init__(self, grad_method):
+        self.__dict__.update(grad_method.__dict__)
+        self.de_solvent = None
+        self.de_solute = None
+
+    def undo_solvent(self):
+        cls = self.__class__
+        name_mixin = self.base.with_solvent.__class__.__name__
+        obj = lib.view(self, lib.drop_class(cls, WithSolventGrad, name_mixin))
+        del obj.de_solvent
+        del obj.de_solute
+        return obj
+
     def kernel(self, *args, dm=None, atmlst=None, **kwargs):
         dm = kwargs.pop('dm', None)
         if dm is None:
@@ -239,10 +255,9 @@ def kernel(self, *args, dm=None, atmlst=None, **kwargs):
         self.de_solvent+= pcm_grad.grad_solver(self.base.with_solvent, dm)
         self.de_solvent+= pcm_grad.grad_nuc(self.base.with_solvent, dm)
 
-        self.de_solute = super().kernel(self, *args, **kwargs)
+        self.de_solute = super().kernel(*args, **kwargs)
         self.de = self.de_solute + self.de_solvent
         self.de += get_cds(self.base.with_solvent)
-
         if self.verbose >= logger.NOTE:
             logger.note(self, '--------------- %s (+%s) gradients ---------------',
                         self.base.__class__.__name__,
@@ -251,4 +266,9 @@ def kernel(self, *args, dm=None, atmlst=None, **kwargs):
             logger.note(self, '----------------------------------------------')
         return self.de
 
+    def _finalize(self):
+        # disable _finalize. It is called in grad_method.kernel method
+        # where self.de was not yet initialized.
+        pass
+
 

gpu4pyscf/solvent/hessian/pcm.py

@@ -224,44 +224,49 @@ def analytic_grad_vmat(pcmobj, mo_coeff, mo_occ, atmlst=None, verbose=None):
     pcmobj.reset(pmol)
     return vmat
 """
+
 def make_hess_object(hess_method):
-    '''
-    return solvent hessian object
-    '''
-    hess_method_class = hess_method.__class__
-    class WithSolventHess(hess_method_class):
-        def __init__(self, hess_method):
-            self.__dict__.update(hess_method.__dict__)
-            self.de_solvent = None
-            self.de_solute = None
-            self._keys = self._keys.union(['de_solvent', 'de_solute'])
-
-        def kernel(self, *args, dm=None, atmlst=None, **kwargs):
-            dm = kwargs.pop('dm', None)
-            if dm is None:
-                dm = self.base.make_rdm1(ao_repr=True)
-            is_equilibrium = self.base.with_solvent.equilibrium_solvation
-            self.base.with_solvent.equilibrium_solvation = True
-            self.de_solvent = hess_elec(self.base.with_solvent, dm, verbose=self.verbose)
-            #self.de_solvent+= hess_nuc(self.base.with_solvent)
-            self.de_solute = hess_method_class.kernel(self, *args, **kwargs)
-            self.de = self.de_solute + self.de_solvent
-            self.base.with_solvent.equilibrium_solvation = is_equilibrium
-            return self.de
-
-        def make_h1(self, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
-            if atmlst is None:
-                atmlst = range(self.mol.natm)
-            h1ao = hess_method_class.make_h1(self, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-            dv = fd_grad_vmat(self.base.with_solvent, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-            for i0, ia in enumerate(atmlst):
-                h1ao[i0] += dv[i0]
-            return h1ao
-
-        def _finalize(self):
-            # disable _finalize. It is called in grad_method.kernel method
-            # where self.de was not yet initialized.
-            pass
-
-    return WithSolventHess(hess_method)
+    if hess_method.base.with_solvent.frozen:
+        raise RuntimeError('Frozen solvent model is not avialbe for energy hessian')
+
+    name = (hess_method.base.with_solvent.__class__.__name__
+            + hess_method.__class__.__name__)
+    return lib.set_class(WithSolventHess(hess_method),
+                         (WithSolventHess, hess_method.__class__), name)
+
+class WithSolventHess:
+    _keys = {'de_solvent', 'de_solute'}
+
+    def __init__(self, hess_method):
+        self.__dict__.update(hess_method.__dict__)
+        self.de_solvent = None
+        self.de_solute = None
+
+    def kernel(self, *args, dm=None, atmlst=None, **kwargs):
+        dm = kwargs.pop('dm', None)
+        if dm is None:
+            dm = self.base.make_rdm1(ao_repr=True)
+        is_equilibrium = self.base.with_solvent.equilibrium_solvation
+        self.base.with_solvent.equilibrium_solvation = True
+        self.de_solvent = hess_elec(self.base.with_solvent, dm, verbose=self.verbose)
+        #self.de_solvent+= hess_nuc(self.base.with_solvent)
+        self.de_solute = super().kernel(*args, **kwargs)
+        self.de = self.de_solute + self.de_solvent
+        self.base.with_solvent.equilibrium_solvation = is_equilibrium
+        return self.de
+
+    def make_h1(self, mo_coeff, mo_occ, chkfile=None, atmlst=None, verbose=None):
+        if atmlst is None:
+            atmlst = range(self.mol.natm)
+        h1ao = super().make_h1(mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
+        dv = fd_grad_vmat(self.base.with_solvent, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
+        for i0, ia in enumerate(atmlst):
+            h1ao[i0] += dv[i0]
+        return h1ao
+
+    def _finalize(self):
+        # disable _finalize. It is called in grad_method.kernel method
+        # where self.de was not yet initialized.
+        pass
+
 

gpu4pyscf/solvent/tests/test_smd_grad.py

@@ -43,7 +43,7 @@ def tearDownModule():
 def _check_grad(mol, solvent='water'):
     natm = mol.natm
     fd_cds = numpy.zeros([natm,3])
-    eps = 1e-5
+    eps = 1e-4
     for ia in range(mol.natm):
         for j in range(3):
             coords = mol.atom_coords(unit='B')