fixed a bug in smd gradient

benchmarks/smd/benchmark_smd.py

@@ -0,0 +1,64 @@
+# Copyright 2024 The GPU4PySCF Authors. All Rights Reserved.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <http://www.gnu.org/licenses/>.
+
+import os
+import numpy
+from pyscf import gto
+from gpu4pyscf.solvent import smd
+from gpu4pyscf.solvent.grad import smd as smd_grad
+
+def _check_grad(mol, solvent='water'):
+    natm = mol.natm
+    fd_cds = numpy.zeros([natm,3])
+    eps = 1e-4
+    for ia in range(mol.natm):
+        for j in range(3):
+            coords = mol.atom_coords(unit='B')
+            coords[ia,j] += eps
+            mol.set_geom_(coords, unit='B')
+            mol.build()
+
+            smdobj = smd.SMD(mol)
+            smdobj.solvent = solvent
+            e0_cds = smdobj.get_cds()
+
+            coords[ia,j] -= 2.0*eps
+            mol.set_geom_(coords, unit='B')
+            mol.build()
+
+            smdobj = smd.SMD(mol)
+            smdobj.solvent = solvent
+            e1_cds = smdobj.get_cds()
+
+            coords[ia,j] += eps
+            mol.set_geom_(coords, unit='B')
+            fd_cds[ia,j] = (e0_cds - e1_cds) / (2.0 * eps)
+
+    smdobj = smd.SMD(mol)
+    smdobj.solvent = solvent
+    grad_cds = smd_grad.get_cds(smdobj)
+    print(numpy.linalg.norm(fd_cds - grad_cds.get()))
+    assert numpy.linalg.norm(fd_cds - grad_cds.get()) < 1e-8
+
+if __name__ == "__main__":
+    path = '../molecules/organic/'
+    for filename in os.listdir(path):
+        f = os.path.join(path, filename)
+        mol = gto.Mole(atom=f)
+        mol.build()
+        print(f'benchmarking {f} in water')
+        _check_grad(mol, solvent='water')
+        print(f'benchmarking {f} in ethanol')
+        _check_grad(mol, solvent='ethanol')

examples/14-pcm_solvent.py

@@ -33,7 +33,7 @@
 mf.grids.atom_grid = (99,590)
 mf.small_rho_cutoff = 1e-10
 mf.with_solvent.lebedev_order = 29 # 302 Lebedev grids
-mf.with_solvent.method = 'C-PCM'
+mf.with_solvent.method = 'IEF-PCM'
 mf.with_solvent.eps = 78.3553
 mf.kernel()
 

gpu4pyscf/lib/cupy_helper.py

@@ -181,19 +181,23 @@ def add_sparse(a, b, indices):
         raise RuntimeError('failed in sparse_add2d')
     return a
 
-def dist_matrix(coords, out=None):
-    assert coords.flags.c_contiguous
-    n = coords.shape[0]
+def dist_matrix(x, y, out=None):
+    assert x.flags.c_contiguous
+    assert y.flags.c_contiguous
+
+    m = x.shape[0]
+    n = y.shape[0]
     if out is None:
-        out = cupy.empty([n,n])
+        out = cupy.empty([m,n])
 
     stream = cupy.cuda.get_current_stream()
     err = libcupy_helper.dist_matrix(
         ctypes.cast(stream.ptr, ctypes.c_void_p),
         ctypes.cast(out.data.ptr, ctypes.c_void_p),
-        ctypes.cast(coords.data.ptr, ctypes.c_void_p),
-        ctypes.cast(coords.data.ptr, ctypes.c_void_p),
-        ctypes.c_int(n),
+        ctypes.cast(x.data.ptr, ctypes.c_void_p),
+        ctypes.cast(y.data.ptr, ctypes.c_void_p),
+        ctypes.c_int(m),
+        ctypes.c_int(n)
     )
     if err != 0:
         raise RuntimeError('failed in calculating distance matrix')

gpu4pyscf/lib/cupy_helper/dist_matrix.cu

@@ -19,11 +19,11 @@
 #define THREADS        32
 
 __global__
-static void _calc_distances(double *dist, const double *x, const double *y, int n)
+static void _calc_distances(double *dist, const double *x, const double *y, int m, int n)
 {
     int i = blockIdx.x * blockDim.x + threadIdx.x;
     int j = blockIdx.y * blockDim.y + threadIdx.y;
-    if (i >= n || j >= n){
+    if (i >= m || j >= n){
         return;
     }
 
@@ -34,12 +34,13 @@ static void _calc_distances(double *dist, const double *x, const double *y, int
 }
 
 extern "C" {
-int dist_matrix(cudaStream_t stream, double *dist, const double *x, const double *y, int n)
+int dist_matrix(cudaStream_t stream, double *dist, const double *x, const double *y, int m, int n)
 {
-    int ntile = (n + THREADS - 1) / THREADS;
+    int ntilex = (m + THREADS - 1) / THREADS;
+    int ntiley = (n + THREADS - 1) / THREADS;
     dim3 threads(THREADS, THREADS);
-    dim3 blocks(ntile, ntile);
-    _calc_distances<<<blocks, threads, 0, stream>>>(dist, x, y, n);
+    dim3 blocks(ntilex, ntiley);
+    _calc_distances<<<blocks, threads, 0, stream>>>(dist, x, y, m, n);
     cudaError_t err = cudaGetLastError();
     if (err != cudaSuccess) {
         return 1;

gpu4pyscf/lib/tests/test_cupy_helper.py

@@ -72,8 +72,7 @@ def test_sparse(self):
     def test_dist_matrix(self):
         a = cupy.random.rand(4, 3)
         rij = cupy.sum((a[:,None,:] - a[None,:,:])**2, axis=2)**0.5
-
-        rij0 = dist_matrix(a)
+        rij0 = dist_matrix(a, a)
         assert cupy.linalg.norm(rij - rij0) < 1e-10
 
     def test_takebak(self):

gpu4pyscf/lib/tests/test_cutensor.py

@@ -16,7 +16,7 @@
 import unittest
 import numpy
 import cupy
-from gpu4pyscf.lib.cupy_helper import *
+from gpu4pyscf.lib.cupy_helper import contract
 
 class KnownValues(unittest.TestCase):
     def test_contract(self):
@@ -49,7 +49,7 @@ def test_cache(self):
         c = contract('ijkl,jl->ik', a, b)
         c_einsum = cupy.einsum('ijkl,jl->ik', a, b)
         assert cupy.linalg.norm(c - c_einsum) < 1e-10
-    
+
 if __name__ == "__main__":
     print("Full tests for cutensor module")
     unittest.main()

gpu4pyscf/solvent/grad/smd.py

@@ -87,12 +87,12 @@ def get_atom_tension(sym_i):
             dt_HO = np.zeros([natm,3])
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C':
-                    r, dr = r_zz['H','C']
+                    r, dr = r_zz.get(('H','C'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_HC[i] += dt_drij
                     dt_HC[j] -= dt_drij
                 if sym_j == 'O':
-                    r, dr = r_zz['H','O']
+                    r, dr = r_zz.get(('H','O'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_HO[i] += dt_drij
                     dt_HO[j] -= dt_drij
@@ -107,12 +107,12 @@ def get_atom_tension(sym_i):
             t_CN = 0.0
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C' and i != j:
-                    r, dr = r_zz['C', 'C']
+                    r, dr = r_zz.get(('C', 'C'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_CC[i] += dt_drij
                     dt_CC[j] -= dt_drij
                 if sym_j == 'N':
-                    r, dr = r_zz['C', 'N']
+                    r, dr = r_zz.get(('C', 'N'), (0.0, 0.0))
                     t_CN += swtich_function(rij[i,j], r, dr)
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_CN[i] += dt_drij
@@ -128,12 +128,12 @@ def get_atom_tension(sym_i):
             dt_NC3 = np.zeros([natm,3])
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C':
-                    r, dr = r_zz['N','C']
+                    r, dr = r_zz.get(('N','C'), (0.0, 0.0))
                     tk = 0.0
-                    dtk = cupy.zeros([natm,3])
+                    dtk = np.zeros([natm,3])
                     for k, sym_k in enumerate(symbols):
                         if k != i and k != j:
-                            rjk, drjk = r_zz['C', sym_k]
+                            rjk, drjk = r_zz.get(('C', sym_k), (0.0, 0.0))
                             tk += swtich_function(rij[j,k], rjk, drjk)
                             dtk_rjk = grad_swtich_function(rij[j,k], rjk, drjk) * drij[j,k]
                             dtk[j] += dtk_rjk
@@ -144,10 +144,10 @@ def get_atom_tension(sym_i):
                     dt_NC[j] -= dt_drij
 
                     t = swtich_function(rij[i,j], r, dr)
-                    dt_NC += t * (2 * dtk)
+                    dt_NC += t * (2 * tk * dtk)
                     t_NC += t * tk**2
 
-                    r, dr = r_zz['N','C3']
+                    r, dr = r_zz.get(('N','C3'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_NC3[i] += dt_drij
                     dt_NC3[j] -= dt_drij
@@ -163,22 +163,22 @@ def get_atom_tension(sym_i):
             dt_OP = np.zeros([natm,3])
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C':
-                    r, dr = r_zz['O','C']
+                    r, dr = r_zz.get(('O','C'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_OC[i] += dt_drij
                     dt_OC[j] -= dt_drij
                 if sym_j == 'N':
-                    r, dr = r_zz['O','N']
+                    r, dr = r_zz.get(('O','N'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_ON[i] += dt_drij
                     dt_ON[j] -= dt_drij
                 if sym_j == 'O' and j != i:
-                    r, dr = r_zz['O','O']
+                    r, dr = r_zz.get(('O','O'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_OO[i] += dt_drij
                     dt_OO[j] -= dt_drij
                 if sym_j == 'P':
-                    r, dr = r_zz['O','P']
+                    r, dr = r_zz.get(('O','P'), (0.0, 0.0))
                     dt_drij = grad_swtich_function(rij[i,j], r, dr) * drij[i,j]
                     dt_OP[i] += dt_drij
                     dt_OP[j] -= dt_drij

gpu4pyscf/solvent/pcm.py

@@ -122,7 +122,8 @@ def gen_surface(mol, ng=302, rad=modified_Bondi, vdw_scale=1.2, r_probe=0.0):
 
         atom_grid = r_vdw * unit_sphere[:,:3] + atom_coords[ia,:]
         #riJ = scipy.spatial.distance.cdist(atom_grid[:,:3], atom_coords)
-        riJ = cupy.sum((atom_grid[:,None,:] - atom_coords[None,:,:])**2, axis=2)**0.5
+        #riJ = cupy.sum((atom_grid[:,None,:] - atom_coords[None,:,:])**2, axis=2)**0.5
+        riJ = dist_matrix(atom_grid, atom_coords)
         diJ = (riJ - R_in_J) / R_sw_J
         diJ[:,ia] = 1.0
         diJ[diJ<1e-8] = 0.0
@@ -193,16 +194,21 @@ def get_D_S(surface, with_S=True, with_D=False):
     xi_ij = xi_i * xi_j / (xi_i**2 + xi_j**2)**0.5
     #rij = scipy.spatial.distance.cdist(grid_coords, grid_coords)
     #rij = cupy.sum((grid_coords[:,None,:] - grid_coords[None,:,:])**2, axis=2)**0.5
-    rij = dist_matrix(grid_coords)
+    rij = dist_matrix(grid_coords, grid_coords)
     xi_r_ij = xi_ij * rij
     cupy.fill_diagonal(rij, 1)
     S = scipy.special.erf(xi_r_ij) / rij
     cupy.fill_diagonal(S, charge_exp * (2.0 / PI)**0.5 / switch_fun)
 
     D = None
     if with_D:
-        drij = cupy.expand_dims(grid_coords, axis=1) - grid_coords
-        nrij = cupy.sum(drij * norm_vec, axis=-1)
+        #drij = cupy.expand_dims(grid_coords, axis=1) - grid_coords
+        #nri = cupy.sum(grid_coords * norm_vec, axis=-1)
+        #nrij = cupy.expand_dims(nri, axis=1) - nri
+        #cupy.expand_dims(grid_coords, axis=1) * norm_vec
+        nrij = grid_coords.dot(norm_vec.T) - cupy.sum(grid_coords * norm_vec, axis=-1)
+
+        #nrij = cupy.sum(drij * norm_vec, axis=-1)
 
         D = S*nrij/rij**2 -2.0*xi_r_ij/PI**0.5*cupy.exp(-xi_r_ij**2)*nrij/rij**3
         cupy.fill_diagonal(D, -charge_exp * (2.0 / PI)**0.5 / (2.0 * R_vdw))

gpu4pyscf/solvent/smd.py

@@ -374,10 +374,10 @@ def get_atom_tension(sym_i):
             t_HO = 0.0
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C':
-                    r, dr = r_zz['H','C']
+                    r, dr = r_zz.get(('H','C'), (0.0, 0.0))
                     t_HC += swtich_function(rij[i,j], r, dr)
                 if sym_j == 'O':
-                    r, dr = r_zz['H','O']
+                    r, dr = r_zz.get(('H','O'), (0.0, 0.0))
                     t_HO += swtich_function(rij[i,j], r, dr)
             sig_HC = get_bond_tension(('H','C'))
             sig_HO = get_bond_tension(('H','O'))
@@ -389,10 +389,10 @@ def get_atom_tension(sym_i):
             t_CN = 0.0
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C' and i != j:
-                    r, dr = r_zz['C', 'C']
+                    r, dr = r_zz.get(('C', 'C'), (0.0, 0.0))
                     t_CC += swtich_function(rij[i,j], r, dr)
                 if sym_j == 'N':
-                    r, dr = r_zz['C', 'N']
+                    r, dr = r_zz.get(('C', 'N'), (0.0, 0.0))
                     t_CN += swtich_function(rij[i,j], r, dr)
             sig_CC = get_bond_tension(('C','C'))
             sig_CN = get_bond_tension(('C','N'))
@@ -404,15 +404,15 @@ def get_atom_tension(sym_i):
             t_NC3 = 0.0
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C':
-                    r, dr = r_zz['N','C']
+                    r, dr = r_zz.get(('N','C'), (0.0,0.0))
                     tk = 0.0
                     for k, sym_k in enumerate(symbols):
                         if k != i and k != j:
-                            rjk, drjk = r_zz['C', sym_k]
+                            rjk, drjk = r_zz.get(('C', sym_k), (0.0,0.0))
                             tk += swtich_function(rij[j,k], rjk, drjk)
                     t_NC += swtich_function(rij[i,j], r, dr) * tk**2
 
-                    r, dr = r_zz['N','C3']
+                    r, dr = r_zz.get(('N','C3'), (0.0, 0.0))
                     t_NC3 += swtich_function(rij[i,j], r, dr)
             sig_NC = get_bond_tension(('N','C'))
             sig_NC3= get_bond_tension(('N','C3'))
@@ -426,16 +426,16 @@ def get_atom_tension(sym_i):
             t_OP = 0.0
             for j, sym_j in enumerate(symbols):
                 if sym_j == 'C':
-                    r, dr = r_zz['O','C']
+                    r, dr = r_zz.get(('O','C'), (0.0, 0.0))
                     t_OC += swtich_function(rij[i,j], r, dr)
                 if sym_j == 'N':
-                    r, dr = r_zz['O','N']
+                    r, dr = r_zz.get(('O','N'), (0.0, 0.0))
                     t_ON += swtich_function(rij[i,j], r, dr)
                 if sym_j == 'O' and j != i:
-                    r, dr = r_zz['O','O']
+                    r, dr = r_zz.get(('O','O'), (0.0, 0.0))
                     t_OO += swtich_function(rij[i,j], r, dr)
                 if sym_j == 'P':
-                    r, dr = r_zz['O','P']
+                    r, dr = r_zz.get(('O','P'), (0.0, 0.0))
                     t_OP += swtich_function(rij[i,j], r, dr)
             sig_OC = get_bond_tension(('O','C'))
             sig_ON = get_bond_tension(('O','N'))