From de99b9d0064dbe5aa14748c74e1e641d6b6e74f4 Mon Sep 17 00:00:00 2001 From: "xiaojie.wu" Date: Sun, 21 Jan 2024 07:50:33 +0800 Subject: [PATCH] fixed a bug in smd gradient --- gpu4pyscf/solvent/grad/pcm.py | 26 +++++++- gpu4pyscf/solvent/grad/smd.py | 30 +++++++-- gpu4pyscf/solvent/hessian/pcm.py | 83 +++++++++++++----------- gpu4pyscf/solvent/tests/test_smd_grad.py | 2 +- 4 files changed, 93 insertions(+), 48 deletions(-) diff --git a/gpu4pyscf/solvent/grad/pcm.py b/gpu4pyscf/solvent/grad/pcm.py index fb90d56f..33c5d8ce 100644 --- a/gpu4pyscf/solvent/grad/pcm.py +++ b/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 \ No newline at end of file diff --git a/gpu4pyscf/solvent/grad/smd.py b/gpu4pyscf/solvent/grad/smd.py index 68f89626..8f55e9f2 100644 --- a/gpu4pyscf/solvent/grad/smd.py +++ b/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 + diff --git a/gpu4pyscf/solvent/hessian/pcm.py b/gpu4pyscf/solvent/hessian/pcm.py index ae28942c..2b849003 100644 --- a/gpu4pyscf/solvent/hessian/pcm.py +++ b/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 + diff --git a/gpu4pyscf/solvent/tests/test_smd_grad.py b/gpu4pyscf/solvent/tests/test_smd_grad.py index 6b24d6bd..45706f43 100644 --- a/gpu4pyscf/solvent/tests/test_smd_grad.py +++ b/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')