to_cpu && to_gpu for solvent

gpu4pyscf/lib/tests/test_to_gpu.py

@@ -128,9 +128,6 @@ def test_df_RKS(self):
         h = hobj.kernel()
         assert numpy.abs(lib.fp(h) - 2.187025544697092) < 1e-4
 
-
-
-
 if __name__ == "__main__":
     print("Full tests for to_gpu module")
     unittest.main()

gpu4pyscf/solvent/_attach_solvent.py

@@ -59,8 +59,9 @@ def undo_solvent(self):
         return obj
 
     def to_cpu(self):
+        from pyscf.solvent import _attach_solvent
         solvent_obj = self.with_solvent.to_cpu()
-        obj = _for_scf(self.undo_solvent().to_cpu(), solvent_obj)
+        obj = _attach_solvent._for_scf(self.undo_solvent().to_cpu(), solvent_obj)
         return obj
 
     def dump_flags(self, verbose=None):
@@ -83,7 +84,7 @@ def get_veff(self, mol=None, dm=None, *args, **kwargs):
 
     def get_fock(self, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1,
                  diis=None, diis_start_cycle=None,
-                 level_shift_factor=None, damp_factor=None):
+                 level_shift_factor=None, damp_factor=None, fock_last=None):
         # DIIS was called inside oldMF.get_fock. v_solvent, as a function of
         # dm, should be extrapolated as well. To enable it, v_solvent has to be
         # added to the fock matrix before DIIS was called.

gpu4pyscf/solvent/grad/pcm.py

@@ -143,6 +143,8 @@ def grad_nuc(pcmobj, dm):
     mol = pcmobj.mol
     log = logger.new_logger(mol, mol.verbose)
     t1 = log.init_timer()
+    if not pcmobj._intermediates:
+        pcmobj.build()
     dm_cache = pcmobj._intermediates.get('dm', None)
     if dm_cache is not None and cupy.linalg.norm(dm_cache - dm) < 1e-10:
         pass
@@ -180,6 +182,8 @@ def grad_qv(pcmobj, dm):
     '''
     contributions due to integrals
     '''
+    if not pcmobj._intermediates:
+        pcmobj.build()
     dm_cache = pcmobj._intermediates.get('dm', None)
     if dm_cache is not None and cupy.linalg.norm(dm_cache - dm) < 1e-10:
         pass
@@ -221,6 +225,8 @@ def grad_solver(pcmobj, dm):
     mol = pcmobj.mol
     log = logger.new_logger(mol, mol.verbose)
     t1 = log.init_timer()
+    if not pcmobj._intermediates:
+        pcmobj.build()
     dm_cache = pcmobj._intermediates.get('dm', None)
     if dm_cache is not None and cupy.linalg.norm(dm_cache - dm) < 1e-10:
         pass

gpu4pyscf/solvent/grad/smd.py

@@ -14,7 +14,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 '''
-Gradient of PCM family solvent model
+Gradient of SMD solvent model
 '''
 # pylint: disable=C0103
 

gpu4pyscf/solvent/hessian/pcm.py

@@ -14,7 +14,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 '''
-Gradient of PCM family solvent model
+Hessian of PCM family solvent model
 '''
 # pylint: disable=C0103
 

gpu4pyscf/solvent/hessian/smd.py

@@ -14,7 +14,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 '''
-Gradient of PCM family solvent model
+Hessian SMD solvent model
 '''
 # pylint: disable=C0103
 

gpu4pyscf/solvent/pcm.py

@@ -226,28 +226,42 @@ class PCM(lib.StreamObject):
     kernel = ddcosmo.DDCOSMO.kernel
 
     def __init__(self, mol):
-        ddcosmo.DDCOSMO.__init__(self, mol)
+        self.mol = mol
+        self.stdout = mol.stdout
+        self.verbose = mol.verbose
+        self.max_memory = mol.max_memory
         self.method = 'C-PCM'
+
         self.vdw_scale = 1.2 # default value in qchem
         self.surface = {}
         self.r_probe = 0.0
         self.radii_table = None
+        self.atom_radii = None
         self.lebedev_order = 29
         self._intermediates = {}
+        self.eps = 78.3553
+
+        self.max_cycle = 20
+        self.conv_tol = 1e-7
+        self.state_id = 0
+
+        self.frozen = False
+        self.equilibrium_solvation = False
+
+        self.e = None
+        self.v = None
+        self._dm = None
 
     def dump_flags(self, verbose=None):
         logger.info(self, '******** %s ********', self.__class__)
         logger.info(self, 'lebedev_order = %s (%d grids per sphere)',
                     self.lebedev_order, gen_grid.LEBEDEV_ORDER[self.lebedev_order])
-        logger.info(self, 'lmax = %s'         , self.lmax)
-        logger.info(self, 'eta = %s'          , self.eta)
         logger.info(self, 'eps = %s'          , self.eps)
         logger.info(self, 'frozen = %s'       , self.frozen)
         logger.info(self, 'equilibrium_solvation = %s', self.equilibrium_solvation)
         logger.debug2(self, 'radii_table %s', self.radii_table)
         if self.atom_radii:
             logger.info(self, 'User specified atomic radii %s', str(self.atom_radii))
-        self.grids.dump_flags(verbose)
         return self
 
     def to_cpu(self):
@@ -391,7 +405,6 @@ def Hessian(self, hess_method):
     def reset(self, mol=None):
         if mol is not None:
             self.mol = mol
-            self.grids.reset(mol)
         self._intermediates = None
         self.surface = None
         self.intopt = None

gpu4pyscf/solvent/smd.py

@@ -559,7 +559,6 @@ def dump_flags(self, verbose=None):
         logger.info(self, 'radii_table %s', self.radii_table*radii.BOHR)
         if self.atom_radii:
             logger.info(self, 'User specified atomic radii %s', str(self.atom_radii))
-        self.grids.dump_flags(verbose)
         return self
 
     def get_cds(self):

gpu4pyscf/solvent/tests/test_pcm.py

@@ -15,9 +15,14 @@
 
 import unittest
 import numpy
+import pyscf
+import pytest
 from pyscf import gto
 from gpu4pyscf import scf, dft
 from gpu4pyscf.solvent import pcm
+from packaging import version
+
+pyscf_25 = version.parse(pyscf.__version__) <= version.parse('2.5.0')
 
 def setUpModule():
     global mol, epsilon, lebedev_order
@@ -86,16 +91,55 @@ def test_uks(self):
 
     def test_dfrks(self):
         e_tot = _energy_with_solvent(dft.RKS(mol, xc='b3lyp').density_fit(), 'IEF-PCM')
-        print(e_tot)
         print(f"Energy error in DFRKS with IEF-PCM: {numpy.abs(e_tot - -75.31863727142068)}")
         assert numpy.abs(e_tot -  -75.31863727142068) < 1e-9
 
     def test_dfuks(self):
         e_tot = _energy_with_solvent(dft.UKS(mol, xc='b3lyp').density_fit(), 'IEF-PCM')
-        print(e_tot)
         print(f"Energy error in DFUKS with IEF-PCM: {numpy.abs(e_tot - -75.31863727142068)}")
         assert numpy.abs(e_tot - -75.31863727142068) < 1e-9
 
+    def test_to_cpu(self):
+        mf = dft.RKS(mol, xc='b3lyp')
+        e_gpu = mf.kernel()
+        mf = mf.to_cpu()
+        e_cpu = mf.kernel()
+        assert abs(e_cpu - e_gpu) < 1e-8
+
+        mf = dft.RKS(mol, xc='b3lyp').density_fit()
+        e_gpu = mf.kernel()
+        mf = mf.to_cpu()
+        e_cpu = mf.kernel()
+        assert abs(e_cpu - e_gpu) < 1e-8
+
+    @pytest.mark.skipif(pyscf_25, reason='requires pyscf 2.6 or higher')
+    def test_to_gpu(self):
+        import pyscf
+        mf = pyscf.dft.RKS(mol, xc='b3lyp').PCM()
+        e_cpu = mf.kernel()
+        mf = mf.to_gpu()
+        e_gpu = mf.kernel()
+        assert abs(e_cpu - e_gpu) < 1e-8
+
+        mf = pyscf.dft.RKS(mol, xc='b3lyp').density_fit().PCM()
+        e_cpu = mf.kernel()
+        mf = mf.to_gpu()
+        e_gpu = mf.kernel()
+        assert abs(e_cpu - e_gpu) < 1e-8
+
+    @pytest.mark.skipif(pyscf_25, reason='requires pyscf 2.6 or higher')
+    def test_to_cpu(self):
+        mf = dft.RKS(mol, xc='b3lyp').PCM()
+        e_gpu = mf.kernel()
+        mf = mf.to_cpu()
+        e_cpu = mf.kernel()
+        assert abs(e_cpu - e_gpu) < 1e-8
+
+        mf = dft.RKS(mol, xc='b3lyp').density_fit().PCM()
+        e_gpu = mf.kernel()
+        mf = mf.to_cpu()
+        e_cpu = mf.kernel()
+        assert abs(e_cpu - e_gpu) < 1e-8
 if __name__ == "__main__":
     print("Full Tests for PCMs")
     unittest.main()

gpu4pyscf/solvent/tests/test_pcm_grad.py

@@ -15,10 +15,15 @@
 
 import unittest
 import numpy
+import pyscf
+import pytest
 from pyscf import gto
 from gpu4pyscf import scf
 from gpu4pyscf.solvent import pcm
 from gpu4pyscf.solvent.grad import pcm as pcm_grad
+from packaging import version
+
+pyscf_25 = version.parse(pyscf.__version__) <= version.parse('2.5.0')
 
 def setUpModule():
     global mol, epsilon, lebedev_order
@@ -58,6 +63,7 @@ def _grad_with_solvent(method, unrestricted=False):
     return grad
 
 class KnownValues(unittest.TestCase):
+
     def test_dA_dF(self):
         cm = pcm.PCM(mol)
         cm.lebedev_order = 3
@@ -165,6 +171,56 @@ def test_uhf_grad_IEFPCM(self):
         print(f"Gradient error in UHF with IEFPCM: {numpy.linalg.norm(g0 - grad)}")
         assert numpy.linalg.norm(g0 - grad) < 1e-6
 
+    @pytest.mark.skipif(pyscf_25, reason='requires pyscf 2.6 or higher')
+    def test_to_cpu(self):
+        mf = scf.RHF(mol).PCM()
+        mf.verbose = 0
+        mf.conv_tol = 1e-12
+        mf.kernel()
+
+        gradobj = mf.nuc_grad_method()
+        grad_gpu = gradobj.kernel()
+        gradobj = gradobj.to_cpu()
+        grad_cpu = gradobj.kernel()
+        assert numpy.linalg.norm(grad_gpu - grad_cpu) < 1e-8
+
+        mf = scf.RHF(mol).density_fit().PCM()
+        mf.verbose = 0
+        mf.conv_tol = 1e-12
+        mf.kernel()
+
+        gradobj = mf.nuc_grad_method()
+        grad_gpu = gradobj.kernel()
+        gradobj = gradobj.to_cpu()
+        grad_cpu = gradobj.kernel()
+        assert numpy.linalg.norm(grad_gpu - grad_cpu) < 1e-8
+
+    @pytest.mark.skipif(pyscf_25, reason='requires pyscf 2.6 or higher')
+    def test_to_gpu(self):
+        mf = pyscf.scf.RHF(mol).PCM()
+        mf.verbose = 0
+        mf.conv_tol = 1e-12
+        mf.kernel()
+
+        g = mf.nuc_grad_method()
+        grad_cpu = g.kernel()
+
+        g = g.to_gpu()
+        grad_gpu = g.kernel()
+        assert numpy.linalg.norm(grad_gpu - grad_cpu) < 1e-8
+
+        mf = pyscf.scf.RHF(mol).density_fit().PCM()
+        mf.verbose = 0
+        mf.conv_tol = 1e-12
+        mf.kernel()
+
+        g = mf.nuc_grad_method()
+        grad_cpu = g.kernel()
+
+        g = g.to_gpu()
+        grad_gpu = g.kernel()
+        assert numpy.linalg.norm(grad_gpu - grad_cpu) < 1e-8
+
 if __name__ == "__main__":
     print("Full Tests for Gradient of PCMs")
     unittest.main()

gpu4pyscf/solvent/tests/test_pcm_hessian.py

@@ -15,10 +15,13 @@
 
 import unittest
 import numpy as np
+import pyscf
+import pytest
 from pyscf import gto
 from gpu4pyscf import scf, dft
-from gpu4pyscf.solvent import pcm
-from gpu4pyscf.solvent.grad import pcm as pcm_grad
+from packaging import version
+
+pyscf_25 = version.parse(pyscf.__version__) <= version.parse('2.5.0')
 
 def setUpModule():
     global mol, epsilon, lebedev_order, eps, xc, tol
@@ -110,6 +113,48 @@ def test_uhf_hess_iefpcm(self):
         _check_hessian(mf, h, ix=0, iy=0)
         _check_hessian(mf, h, ix=0, iy=1)
 
+    @pytest.mark.skipif(pyscf_25, reason='requires pyscf 2.6 or higher')
+    def test_to_gpu(self):
+        import pyscf
+        # Not implemented yet
+        '''
+        mf = pyscf.dft.RKS(mol, xc='b3lyp').SMD()
+        mf.kernel()
+        hessobj = mf.Hessian()
+        hess_cpu = hessobj.kernel()
+        hessobj = hessobj.to_gpu()
+        hess_gpu = hessobj.kernel()
+        assert np.linalg.norm(hess_cpu - hess_gpu) < 1e-8
+        '''
+        mf = pyscf.dft.RKS(mol, xc='b3lyp').density_fit().PCM()
+        mf.conv_tol = 1e-12
+        mf.conv_tol_cpscf = 1e-7
+        mf.kernel()
+        hessobj = mf.Hessian()
+        hess_cpu = hessobj.kernel()
+        hessobj = hessobj.to_gpu()
+        hess_gpu = hessobj.kernel()
+        assert np.linalg.norm(hess_cpu - hess_gpu) < 1e-5
+
+    @pytest.mark.skipif(pyscf_25, reason='requires pyscf 2.6 or higher')
+    def test_to_cpu(self):
+        # Not implemented yet
+        '''
+        mf = dft.RKS(mol, xc='b3lyp').SMD()
+        e_gpu = mf.kernel()
+        mf = mf.to_cpu()
+        e_cpu = mf.kernel()
+        assert abs(e_cpu - e_gpu) < 1e-8
+        '''
+        mf = dft.RKS(mol, xc='b3lyp').density_fit().PCM()
+        mf.conv_tol = 1e-12
+        mf.conv_tol_cpscf = 1e-7
+        mf.kernel()
+        hessobj = mf.Hessian()
+        hess_gpu = hessobj.kernel()
+        hessobj = hessobj.to_cpu()
+        hess_cpu = hessobj.kernel()
+        assert np.linalg.norm(hess_cpu - hess_gpu) < 1e-5
 if __name__ == "__main__":
     print("Full Tests for Hessian of PCMs")
     unittest.main()