fixed bugs

pyscf/solvent/grad/pcm.py

@@ -252,10 +252,10 @@ def grad_solver(pcmobj, dm):
 
     dF, dA = get_dF_dA(pcmobj.surface)
 
-    with_D = pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SS(V)PE', 'SMD']
+    with_D = pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SS(V)PE']
     dD, dS, dSii = get_dD_dS(pcmobj.surface, dF, with_D=with_D, with_S=True)
 
-    if pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SS(V)PE', 'SMD']:
+    if pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SS(V)PE']:
         DA = D*A
 
     epsilon = pcmobj.eps
@@ -273,7 +273,7 @@ def grad_solver(pcmobj, dm):
         de -= numpy.asarray([numpy.sum(de_dS[p0:p1], axis=0) for p0,p1, in gridslice])
         de -= 0.5*numpy.einsum('i,xij->jx', vK_1*q, dSii) # 0.5*cupy.einsum('i,xij,i->jx', vK_1, dSii, q)
 
-    elif pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SS(V)PE', 'SMD']:
+    elif pcmobj.method.upper() in ['IEF-PCM', 'IEFPCM', 'SS(V)PE']:
         dD = dD.transpose([2,0,1])
         dS = dS.transpose([2,0,1])
         dSii = dSii.transpose([2,0,1])
@@ -286,7 +286,7 @@ def grad_solver(pcmobj, dm):
         fac = f_epsilon/(2.0*PI)
 
         Av = A*v_grids
-        de_dR  = 0.5*fac * contract_ket(vK_1, dD, Av)
+        de_dR  = 0.5*fac * numpy.einsum('i,xij,j->ix', vK_1, dD, Av)
         de_dR -= 0.5*fac * numpy.einsum('i,xij,j->jx', vK_1, dD, Av)
         de_dR  = numpy.asarray([numpy.sum(de_dR[p0:p1], axis=0) for p0,p1 in gridslice])
 

pyscf/solvent/grad/smd.py

@@ -26,6 +26,92 @@
 from pyscf.solvent.grad import pcm as pcm_grad
 from pyscf.lib import logger
 
+def grad_solver(pcmobj, dm):
+    '''
+    dE = 0.5*v* d(K^-1 R) *v + q*dv
+    v^T* d(K^-1 R)v = v^T*K^-1(dR - dK K^-1R)v = v^T K^-1(dR - dK q)
+    '''
+    mol = pcmobj.mol
+    log = logger.new_logger(mol, mol.verbose)
+    t1 = (logger.process_clock(), logger.perf_counter())
+    if not pcmobj._intermediates:
+        pcmobj.build()
+    dm_cache = pcmobj._intermediates.get('dm', None)
+    if dm_cache is not None and np.linalg.norm(dm_cache - dm) < 1e-10:
+        pass
+    else:
+        pcmobj._get_vind(dm)
+
+    gridslice    = pcmobj.surface['gslice_by_atom']
+    v_grids      = pcmobj._intermediates['v_grids']
+    A            = pcmobj._intermediates['A']
+    D            = pcmobj._intermediates['D']
+    S            = pcmobj._intermediates['S']
+    K            = pcmobj._intermediates['K']
+    q            = pcmobj._intermediates['q']
+
+    vK_1 = np.linalg.solve(K.T, v_grids)
+
+    dF, dA = pcm_grad.get_dD_dSt_dF_dA(pcmobj.surface)
+    dD, dS, dSii = pcm_grad.get_dD_dS(pcmobj.surface, dF, with_D=True, with_S=True)
+    DA = D*A
+
+    epsilon = pcmobj.eps
+
+    #de_dF = v0 * -dSii_dF * q
+    #de += 0.5*numpy.einsum('i,inx->nx', de_dF, dF)
+    # dQ = v^T K^-1 (dR - dK K^-1 R) v
+    de = np.zeros([pcmobj.mol.natm,3])
+    dD = dD.transpose([2,0,1])
+    dS = dS.transpose([2,0,1])
+    dSii = dSii.transpose([2,0,1])
+    dA = dA.transpose([2,0,1])
+
+    # IEF-PCM and SS(V)PE formally are the same in gradient calculation
+    # dR = f_eps/(2*pi) * (dD*A + D*dA),
+    # dK = dS - f_eps/(2*pi) * (dD*A*S + D*dA*S + D*A*dS)
+    f_epsilon = (epsilon - 1.0)/(epsilon + 1.0)
+    fac = f_epsilon/(2.0*PI)
+
+    Av = A*v_grids
+    de_dR  = 0.5*fac * np.einsum('i,xij,j->ix', vK_1, dD, Av)
+    de_dR -= 0.5*fac * np.einsum('i,xij,j->jx', vK_1, dD, Av)
+    de_dR  = np.asarray([np.sum(de_dR[p0:p1], axis=0) for p0,p1 in gridslice])
+
+    vK_1_D = vK_1.dot(D)
+    vK_1_Dv = vK_1_D * v_grids
+    de_dR += 0.5*fac * np.einsum('j,xjn->nx', vK_1_Dv, dA)
+
+    de_dS0  = 0.5*np.einsum('i,xij,j->ix', vK_1, dS, q)
+    de_dS0 -= 0.5*np.einsum('i,xij,j->jx', vK_1, dS, q)
+    de_dS0  = np.asarray([np.sum(de_dS0[p0:p1], axis=0) for p0,p1 in gridslice])
+
+    vK_1_q = vK_1 * q
+    de_dS0 += 0.5*np.einsum('i,xin->nx', vK_1_q, dSii)
+
+    vK_1_DA = np.dot(vK_1, DA)
+    de_dS1  = 0.5*np.einsum('i,xij,j->ix', vK_1_DA, dS, q)
+    de_dS1 -= 0.5*np.einsum('i,xij,j->jx', vK_1_DA, dS, q)
+    de_dS1  = np.asarray([np.sum(de_dS1[p0:p1], axis=0) for p0,p1 in gridslice])
+
+    vK_1_DAq = vK_1_DA*q
+    de_dS1 += 0.5*np.einsum('j,xjn->nx', vK_1_DAq, dSii)
+
+    Sq = np.dot(S,q)
+    ASq = A*Sq
+    de_dD  = 0.5*np.einsum('i,xij,j->ix', vK_1, dD, ASq)
+    de_dD -= 0.5*np.einsum('i,xij,j->jx', vK_1, dD, ASq)
+    de_dD  = np.asarray([np.sum(de_dD[p0:p1], axis=0) for p0,p1 in gridslice])
+
+    vK_1_D = np.dot(vK_1, D)
+    de_dA = 0.5*np.einsum('j,xjn->nx', vK_1_D*Sq, dA)   # 0.5*cupy.einsum('j,xjn,j->nx', vK_1_D, dA, Sq)
+
+    de_dK = de_dS0 - fac * (de_dD + de_dA + de_dS1)
+    de += de_dR - de_dK
+
+    t1 = log.timer_debug1('grad solver', *t1)
+    return de
+
 def get_cds(smdobj):
     return smd.get_cds_legacy(smdobj)[1]
 
@@ -68,7 +154,7 @@ def kernel(self, *args, dm=None, atmlst=None, **kwargs):
             dm = dm[0] + dm[1]
         self.de_solute  = super().kernel(*args, **kwargs)
         self.de_solvent = pcm_grad.grad_qv(self.base.with_solvent, dm)
-        self.de_solvent+= pcm_grad.grad_solver(self.base.with_solvent, dm)
+        self.de_solvent+= grad_solver(self.base.with_solvent, dm)
         self.de_solvent+= pcm_grad.grad_nuc(self.base.with_solvent, dm)
         #self.de_cds     = get_cds(self.base.with_solvent)
         self.de_cds     = smd.get_cds_legacy(self.base.with_solvent)[1]

pyscf/solvent/hessian/smd.py

@@ -22,9 +22,47 @@
 from pyscf import scf, lib
 from pyscf.solvent import smd
 from pyscf.solvent.grad import smd as smd_grad
+from pyscf.solvent.grad import pcm as pcm_grad
 from pyscf.solvent.hessian import pcm as pcm_hess
 from pyscf.lib import logger
 
+def hess_elec(smdobj, dm, verbose=None):
+    '''
+    slow version with finite difference
+    TODO: use analytical hess_nuc
+    '''
+    log = logger.new_logger(smdobj, verbose)
+    t1 = (logger.process_clock(), logger.perf_counter())
+    pmol = smdobj.mol.copy()
+    mol = pmol.copy()
+    coords = mol.atom_coords(unit='Bohr')
+
+    def smd_grad_scanner(mol):
+        # TODO: use more analytical forms
+        smdobj.reset(mol)
+        smdobj._get_vind(dm)
+        grad = pcm_grad.grad_nuc(smdobj, dm)
+        grad+= smd_grad.grad_solver(smdobj, dm)
+        grad+= pcm_grad.grad_qv(smdobj, dm)
+        return grad
+
+    mol.verbose = 0
+    de = np.zeros([mol.natm, mol.natm, 3, 3])
+    eps = 1e-3
+    for ia in range(mol.natm):
+        for ix in range(3):
+            dv = np.zeros_like(coords)
+            dv[ia,ix] = eps
+            mol.set_geom_(coords + dv, unit='Bohr')
+            g0 = smd_grad_scanner(mol)
+
+            mol.set_geom_(coords - dv, unit='Bohr')
+            g1 = smd_grad_scanner(mol)
+            de[ia,:,ix] = (g0 - g1)/2.0/eps
+    t1 = log.timer_debug1('solvent energy', *t1)
+    smdobj.reset(pmol)
+    return de
+
 def get_cds(smdobj):
     mol = smdobj.mol
     solvent = smdobj.solvent
@@ -97,8 +135,7 @@ def kernel(self, *args, dm=None, atmlst=None, **kwargs):
             dm = dm[0] + dm[1]
         is_equilibrium = self.base.with_solvent.equilibrium_solvation
         self.base.with_solvent.equilibrium_solvation = True
-        self.de_solvent = pcm_hess.hess_elec(self.base.with_solvent, dm, verbose=self.verbose)
-        #self.de_solvent+= hess_nuc(self.base.with_solvent)
+        self.de_solvent = hess_elec(self.base.with_solvent, dm, verbose=self.verbose)
         self.de_solute = super().kernel(*args, **kwargs)
         self.de_cds = get_cds(self.base.with_solvent)
         self.de = self.de_solute + self.de_solvent + self.de_cds

pyscf/solvent/smd.py

@@ -18,8 +18,7 @@
 '''
 
 import numpy as np
-import scipy
-from pyscf import lib
+from pyscf import lib, gto
 from pyscf.data import radii
 from pyscf.dft import gen_grid
 from pyscf.solvent import pcm
@@ -349,7 +348,7 @@ def build(self, ng=None):
         fakemol = gto.fakemol_for_charges(grid_coords, expnt=charge_exp**2)
         fakemol_nuc = gto.fakemol_for_charges(atom_coords)
         v_ng = gto.mole.intor_cross(int2c2e, fakemol_nuc, fakemol)
-        self.v_grids_n = numpy.dot(atom_charges, v_ng)
+        self.v_grids_n = np.dot(atom_charges, v_ng)
 
     @property
     def solvent(self):