Change all K^-1 algorithm from LU solve to SVD pseudo inverse

gpu4pyscf/solvent/grad/pcm.py

@@ -276,10 +276,10 @@ def grad_solver(pcmobj, dm):
     A            = pcmobj._intermediates['A']
     D            = pcmobj._intermediates['D']
     S            = pcmobj._intermediates['S']
-    K            = pcmobj._intermediates['K']
+    inverse_K    = pcmobj._intermediates['inverse_K']
     q            = pcmobj._intermediates['q']
 
-    vK_1 = cupy.linalg.solve(K.T, v_grids)
+    vK_1 = cupy.dot(inverse_K.T, v_grids)
     epsilon = pcmobj.eps
 
     de = cupy.zeros([pcmobj.mol.natm,3])

gpu4pyscf/solvent/hessian/pcm.py

@@ -153,7 +153,7 @@ def pcm_grad_scanner(mol):
     pcmobj.reset(pmol)
     return de
 
-def get_dqsym_dx_fix_vgrids(pcmobj, atmlst, inverse_K):
+def get_dqsym_dx_fix_vgrids(pcmobj, atmlst):
     assert pcmobj._intermediates is not None
 
     gridslice    = pcmobj.surface['gslice_by_atom']
@@ -162,6 +162,7 @@ def get_dqsym_dx_fix_vgrids(pcmobj, atmlst, inverse_K):
     D            = pcmobj._intermediates['D']
     S            = pcmobj._intermediates['S']
     R            = pcmobj._intermediates['R']
+    inverse_K    = pcmobj._intermediates['inverse_K']
     q_sym        = pcmobj._intermediates['q_sym']
     f_epsilon    = pcmobj._intermediates['f_epsilon']
 
@@ -317,14 +318,15 @@ def dK_dot_q(q):
 
     return dqdx_fix_Vq
 
-def get_dqsym_dx_fix_K_R(pcmobj, dm, atmlst, inverse_K, intopt_derivative):
+def get_dqsym_dx_fix_K_R(pcmobj, dm, atmlst, intopt_derivative):
     assert pcmobj._intermediates is not None
 
     mol = pcmobj.mol
     gridslice    = pcmobj.surface['gslice_by_atom']
     charge_exp   = pcmobj.surface['charge_exp']
     grid_coords  = pcmobj.surface['grid_coords']
     R            = pcmobj._intermediates['R']
+    inverse_K    = pcmobj._intermediates['inverse_K']
 
     atom_coords = mol.atom_coords(unit='B')
     atom_charges = numpy.asarray(mol.atom_charges(), dtype=numpy.float64)
@@ -357,9 +359,7 @@ def get_dqsym_dx_fix_K_R(pcmobj, dm, atmlst, inverse_K, intopt_derivative):
     return dqdx_fix_K_R
 
 def get_dqsym_dx(pcmobj, dm, atmlst, intopt_derivative):
-    K = pcmobj._intermediates['K']
-    inverse_K = cupy.linalg.inv(K)
-    return get_dqsym_dx_fix_vgrids(pcmobj, atmlst, inverse_K) + get_dqsym_dx_fix_K_R(pcmobj, dm, atmlst, inverse_K, intopt_derivative)
+    return get_dqsym_dx_fix_vgrids(pcmobj, atmlst) + get_dqsym_dx_fix_K_R(pcmobj, dm, atmlst, intopt_derivative)
 
 def analytic_grad_vmat(pcmobj, dm, mo_coeff, mo_occ, atmlst=None, verbose=None):
     '''

gpu4pyscf/solvent/pcm.py

@@ -331,11 +331,15 @@ def build(self, ng=None):
         else:
             raise RuntimeError(f"Unknown implicit solvent model: {self.method}")
 
+        # Notice the SVD in pseudo inverse scales badly with ngrids
+        inverse_K = cupy.linalg.pinv(K)
+
         intermediates = {
             'S': cupy.asarray(S),
             'D': cupy.asarray(D),
             'A': cupy.asarray(A),
             'K': cupy.asarray(K),
+            'inverse_K': cupy.asarray(inverse_K),
             'R': cupy.asarray(R),
             'f_epsilon': f_epsilon
         }
@@ -366,23 +370,21 @@ def _get_vind(self, dms):
             dms = (dms[0] + dms[1]).reshape(-1,nao,nao)
         if not isinstance(dms, cupy.ndarray):
             dms = cupy.asarray(dms)
-        K = self._intermediates['K']
+        inverse_K = self._intermediates['inverse_K']
         R = self._intermediates['R']
         v_grids_e = self._get_v(dms)
         v_grids = self.v_grids_n - v_grids_e
 
         b = cupy.dot(R, v_grids.T)
-        q = cupy.linalg.solve(K, b).T
+        q = cupy.dot(inverse_K, b).T
 
-        vK_1 = cupy.linalg.solve(K.T, v_grids.T)
+        vK_1 = cupy.dot(inverse_K.T, v_grids.T)
         qt = cupy.dot(R.T, vK_1).T
         q_sym = (q + qt)/2.0
 
         vmat = self._get_vmat(q_sym)
         epcm = 0.5 * cupy.dot(v_grids[0], q_sym[0])
 
-        self._intermediates['K'] = K
-        self._intermediates['R'] = R
         self._intermediates['q'] = q[0]
         self._intermediates['q_sym'] = q_sym[0]
         self._intermediates['v_grids'] = v_grids[0]
@@ -439,14 +441,14 @@ def _B_dot_x(self, dms):
         nao = dms.shape[-1]
         dms = dms.reshape(-1,nao,nao)
 
-        K = self._intermediates['K']
+        inverse_K = self._intermediates['inverse_K']
         R = self._intermediates['R']
         v_grids = -self._get_v(dms)
 
         b = cupy.dot(R, v_grids.T)
-        q = cupy.linalg.solve(K, b).T
+        q = cupy.dot(inverse_K, b).T
 
-        vK_1 = cupy.linalg.solve(K.T, v_grids.T)
+        vK_1 = cupy.dot(inverse_K.T, v_grids.T)
         qt = cupy.dot(R.T, vK_1).T
         q_sym = (q + qt)/2.0
 

gpu4pyscf/solvent/smd.py

@@ -381,11 +381,15 @@ def build(self, ng=None):
         K = S - f_epsilon/(2.0*np.pi) * DAS
         R = -f_epsilon * (cupy.eye(K.shape[0]) - 1.0/(2.0*np.pi)*DA)
 
+        # Notice the SVD in pseudo inverse scales badly with ngrids
+        inverse_K = cupy.linalg.pinv(K)
+
         intermediates = {
             'S': cupy.asarray(S),
             'D': cupy.asarray(D),
             'A': cupy.asarray(A),
             'K': cupy.asarray(K),
+            'inverse_K': cupy.asarray(inverse_K),
             'R': cupy.asarray(R),
             'f_epsilon': f_epsilon
         }

gpu4pyscf/solvent/tests/test_pcm_hessian.py

@@ -195,7 +195,7 @@ def test_grad_vmat_cpcm(self):
         test_grad_vmat = analytic_grad_vmat(hobj.base.with_solvent, dm, mo_coeff, mo_occ)
         ref_grad_vmat = _fd_grad_vmat(hobj.base.with_solvent, dm, mo_coeff, mo_occ)
 
-        cp.testing.assert_allclose(ref_grad_vmat, test_grad_vmat, atol = 1e-10)
+        cp.testing.assert_allclose(ref_grad_vmat, test_grad_vmat, atol = 3e-10)
 
     def test_grad_vmat_iefpcm(self):
         print("testing IEF-PCM dV_solv/dx")
@@ -209,7 +209,7 @@ def test_grad_vmat_iefpcm(self):
         test_grad_vmat = analytic_grad_vmat(hobj.base.with_solvent, dm, mo_coeff, mo_occ)
         ref_grad_vmat = _fd_grad_vmat(hobj.base.with_solvent, dm, mo_coeff, mo_occ)
 
-        cp.testing.assert_allclose(ref_grad_vmat, test_grad_vmat, atol = 1e-10)
+        cp.testing.assert_allclose(ref_grad_vmat, test_grad_vmat, atol = 3e-10)
 
     def test_grad_vmat_ssvpe(self):
         print("testing SS(V)PE dV_solv/dx")
@@ -223,7 +223,7 @@ def test_grad_vmat_ssvpe(self):
         test_grad_vmat = analytic_grad_vmat(hobj.base.with_solvent, dm, mo_coeff, mo_occ)
         ref_grad_vmat = _fd_grad_vmat(hobj.base.with_solvent, dm, mo_coeff, mo_occ)
 
-        cp.testing.assert_allclose(ref_grad_vmat, test_grad_vmat, atol = 1e-10)
+        cp.testing.assert_allclose(ref_grad_vmat, test_grad_vmat, atol = 3e-10)
 
     @pytest.mark.skipif(pyscf_25, reason='requires pyscf 2.6 or higher')
     def test_to_gpu(self):