Update to_gpu

pyscf/cc/ccsd.py

@@ -1365,6 +1365,8 @@ def get_d2_diagnostic(self, t2=None):
         if t2 is None: t2 = self.t2
         return get_d2_diagnostic(t2)
 
+    to_gpu = lib.to_gpu
+
 CC = RCCSD = CCSD
 
 

pyscf/cc/dfccsd.py

@@ -48,6 +48,8 @@ def _add_vvvv(self, t1, t2, eris, out=None, with_ovvv=False, t2sym=None):
         assert (not self.direct)
         return ccsd.CCSD._add_vvvv(self, t1, t2, eris, out, with_ovvv, t2sym)
 
+    to_gpu = lib.to_gpu
+
 
 def _contract_vvvv_t2(mycc, mol, vvL, t2, out=None, verbose=None):
     '''Ht2 = numpy.einsum('ijcd,acdb->ijab', t2, vvvv)

pyscf/cc/gccsd.py

@@ -289,6 +289,8 @@ def spin2spatial(self, tx, orbspin=None):
                 orbspin = orbspin[self.get_frozen_mask()]
         return spin2spatial(tx, orbspin)
 
+    to_gpu = lib.to_gpu
+
 CCSD = GCCSD
 
 

pyscf/cc/uccsd.py

@@ -758,6 +758,8 @@ def vector_size(self, nmo=None, nocc=None):
     def amplitudes_from_rccsd(self, t1, t2):
         return amplitudes_from_rccsd(t1, t2)
 
+    to_gpu = lib.to_gpu
+
 CCSD = UCCSD
 
 

pyscf/ci/cisd.py

@@ -1131,6 +1131,8 @@ def nuc_grad_method(self):
         from pyscf.grad import cisd
         return cisd.Gradients(self)
 
+    to_gpu = lib.to_gpu
+
 class RCISD(CISD):
     pass
 

pyscf/df/addons.py

@@ -152,7 +152,7 @@ def make_auxbasis(mol, mp2fit=False):
         _basis.update(mol.basis)
         del (_basis['default'])
     else:
-        _basis = mol._basis
+        _basis = mol._basis or {}
 
     auxbasis = {}
     for k in _basis:

pyscf/df/df.py

@@ -308,9 +308,7 @@ def range_coulomb(self, omega):
             if auxmol_omega is not None:
                 auxmol.omega = auxmol_omega
 
-    def to_gpu(self):
-        from gpu4pyscf.df.df import DF as DF
-        return lib.to_gpu(self.__class__.reset(self.view(DF)))
+    to_gpu = lib.to_gpu
 
 GDF = DF
 

pyscf/df/df_jk.py

@@ -228,8 +228,7 @@ def CASSCF(self, ncas, nelecas, auxbasis=None, ncore=None, frozen=None):
 
     def to_gpu(self):
         obj = self.undo_df().to_gpu().density_fit()
-        obj.__dict__.update(self.__dict__)
-        return lib.to_gpu(obj)
+        return lib.to_gpu(self, obj)
 
 
 def get_jk(dfobj, dm, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-13):

pyscf/df/grad/casscf.py

@@ -224,6 +224,8 @@ def _finalize(self):
 
     as_scanner = as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 #from pyscf import mcscf

pyscf/df/grad/rhf.py

@@ -482,12 +482,14 @@ class Gradients(rhf_grad.Gradients):
     _keys = {'with_df', 'auxbasis_response'}
 
     def __init__(self, mf):
-        assert isinstance(mf, df.df_jk._DFHF)
         # Whether to include the response of DF auxiliary basis when computing
         # nuclear gradients of J/K matrices
         self.auxbasis_response = True
         rhf_grad.Gradients.__init__(self, mf)
 
+    def check_sanity(self):
+        assert isinstance(self.base, df.df_jk._DFHF)
+
     def get_jk(self, mol=None, dm=None, hermi=0, with_j=True, with_k=True,
                omega=None):
         if omega is None:
@@ -521,8 +523,6 @@ def extra_force(self, atom_id, envs):
         else:
             return 0
 
-    def to_gpu(self):
-        from gpu4pyscf.df.grad.rhf import Gradients
-        return lib.to_gpu(self.view(Gradients))
+    to_gpu = lib.to_gpu
 
 Grad = Gradients

pyscf/df/grad/rks.py

@@ -123,8 +123,6 @@ def extra_force(self, atom_id, envs):
             e1 += envs['vhf'].aux[atom_id]
         return e1
 
-    def to_gpu(self):
-        from gpu4pyscf.df.grad.rks import Gradients
-        return lib.to_gpu(self.view(Gradients))
+    to_gpu = lib.to_gpu
 
 Grad = Gradients

pyscf/df/grad/sacasscf.py

@@ -370,3 +370,5 @@ def kernel (self, **kwargs):
     def get_LdotJnuc (self, Lvec, **kwargs):
         with lib.temporary_env (sacasscf_grad, Lci_dot_dgci_dx=Lci_dot_dgci_dx, Lorb_dot_dgorb_dx=Lorb_dot_dgorb_dx):
             return sacasscf_grad.Gradients.get_LdotJnuc (self, Lvec, **kwargs)
+
+    to_gpu = lib.to_gpu

pyscf/df/grad/uhf.py

@@ -60,4 +60,6 @@ def extra_force(self, atom_id, envs):
         else:
             return 0
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients

pyscf/df/grad/uks.py

@@ -124,4 +124,6 @@ def extra_force(self, atom_id, envs):
             e1 += envs['vhf'].aux[atom_id]
         return e1
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients

pyscf/df/hessian/rhf.py

@@ -480,10 +480,7 @@ def __init__(self, mf):
 
     partial_hess_elec = partial_hess_elec
     make_h1 = make_h1
-
-    def to_gpu(self):
-        from gpu4pyscf.df.hessian.rhf import Hessian
-        return lib.to_gpu(self.view(Hessian))
+    to_gpu = lib.to_gpu
 
 #TODO: Insert into DF class
 

pyscf/df/hessian/rks.py

@@ -126,10 +126,7 @@ def __init__(self, mf):
 
     partial_hess_elec = partial_hess_elec
     make_h1 = make_h1
-
-    def to_gpu(self):
-        from gpu4pyscf.df.hessian.rks import Hessian
-        return lib.to_gpu(self.view(Hessian))
+    to_gpu = lib.to_gpu
 
 
 if __name__ == '__main__':

pyscf/df/hessian/uhf.py

@@ -531,6 +531,7 @@ def __init__(self, mf):
 
     partial_hess_elec = partial_hess_elec
     make_h1 = make_h1
+    to_gpu = lib.to_gpu
 
 #TODO: Insert into DF class
 

pyscf/df/hessian/uks.py

@@ -139,6 +139,7 @@ def __init__(self, mf):
 
     partial_hess_elec = partial_hess_elec
     make_h1 = make_h1
+    to_gpu = lib.to_gpu
 
 
 if __name__ == '__main__':

pyscf/dft/dks.py

@@ -143,6 +143,8 @@ def x2c1e(self):
         return x2chf
     x2c = x2c1e
 
+    to_gpu = lib.to_gpu
+
 UKS = UDKS = DKS
 
 class RDKS(DKS, dhf.RDHF):

pyscf/dft/gen_grid.py

@@ -587,9 +587,7 @@ def prune_by_density_(self, rho, threshold=0):
             self.screen_index = self.non0tab
         return self
 
-    def to_gpu(self):
-        from gpu4pyscf.dft.gen_grid import Grids
-        return lib.to_gpu(self.view(Grids))
+    to_gpu = lib.to_gpu
 
 
 def _default_rad(nuc, level=3):

pyscf/dft/gks.py

@@ -177,8 +177,7 @@ def to_hf(self):
         '''Convert to GHF object.'''
         return self._transfer_attrs_(self.mol.GHF())
 
-    def to_gpu(self):
-        raise NotImplementedError
+    to_gpu = lib.to_gpu
 
 
 if __name__ == '__main__':

pyscf/dft/gks_symm.py

@@ -57,6 +57,8 @@ def collinear(self, val):
     def nuc_grad_method(self):
         raise NotImplementedError
 
+    to_gpu = lib.to_gpu
+
 
 if __name__ == '__main__':
     import numpy

pyscf/dft/numint.py

@@ -2865,10 +2865,7 @@ def make_rho(idm, ao, sindex, xctype):
                                      with_lapl)
         return make_rho, ndms, nao
 
-    def to_gpu(self):
-        from gpu4pyscf.dft.numint import NumInt
-        # Note: gpu4pyscf NumInt initializes additional things in __init__.py
-        return NumInt()
+    to_gpu = lib.to_gpu
 
 _NumInt = NumInt
 

pyscf/dft/rks.py

@@ -531,9 +531,4 @@ def to_hf(self):
         '''Convert to RHF object.'''
         return self._transfer_attrs_(self.mol.RHF())
 
-    def to_gpu(self):
-        from gpu4pyscf.dft.rks import RKS
-        obj = lib.to_gpu(hf.SCF.reset(self.view(RKS)))
-        # Attributes only defined in gpu4pyscf.RKS
-        obj.screen_tol = 1e-14
-        return obj
+    to_gpu = lib.to_gpu

pyscf/dft/rks_symm.py

@@ -46,12 +46,14 @@ def nuc_grad_method(self):
         from pyscf.grad import rks
         return rks.Gradients(self)
 
+    to_gpu = lib.to_gpu
+
 RKS = SymAdaptedRKS
 
 
 class SymAdaptedROKS(rks.KohnShamDFT, hf_symm.SymAdaptedROHF):
     ''' Restricted Kohn-Sham '''
-    def __init__(self, mol, xc='LDA,VWN'):
+    def __init__(self, mol=None, xc='LDA,VWN'):
         hf_symm.ROHF.__init__(self, mol)
         rks.KohnShamDFT.__init__(self, xc)
 
@@ -70,6 +72,8 @@ def nuc_grad_method(self):
         from pyscf.grad import roks
         return roks.Gradients(self)
 
+    to_gpu = lib.to_gpu
+
 ROKS = SymAdaptedROKS
 
 

pyscf/dft/roks.py

@@ -65,14 +65,7 @@ def to_hf(self):
         '''Convert to ROHF object.'''
         return self._transfer_attrs_(self.mol.ROHF())
 
-    def to_gpu(self):
-        from pyscf.scf.hf import SCF
-        from gpu4pyscf.dft.roks import ROKS
-        obj = lib.to_gpu(SCF.reset(self.view(ROKS)))
-        # Attributes only defined in gpu4pyscf.RKS
-        obj.screen_tol = 1e-14
-        obj.disp = None
-        return obj
+    to_gpu = lib.to_gpu
 
 
 if __name__ == '__main__':

pyscf/dft/uks.py

@@ -197,10 +197,4 @@ def to_hf(self):
         '''Convert to UHF object.'''
         return self._transfer_attrs_(self.mol.UHF())
 
-    def to_gpu(self):
-        from pyscf.scf.hf import SCF
-        from gpu4pyscf.dft.uks import UKS
-        obj = lib.to_gpu(SCF.reset(self.view(UKS)))
-        # Attributes only defined in gpu4pyscf.RKS
-        obj.screen_tol = 1e-14
-        return obj
+    to_gpu = lib.to_gpu

pyscf/dft/uks_symm.py

@@ -47,6 +47,8 @@ def nuc_grad_method(self):
         from pyscf.grad import uks
         return uks.Gradients(self)
 
+    to_gpu = lib.to_gpu
+
 UKS = SymAdaptedUKS
 
 

pyscf/fci/direct_spin1.py

@@ -727,13 +727,7 @@ class FCIBase(lib.StreamObject):
 
     def __init__(self, mol=None):
         if mol is None:
-            self.stdout = sys.stdout
-            self.verbose = logger.NOTE
-            self.max_memory = lib.param.MAX_MEMORY
-        else:
-            self.stdout = mol.stdout
-            self.verbose = mol.verbose
-            self.max_memory = mol.max_memory
+            mol = lib.omniobj
         self.mol = mol
         self.nroots = 1
         self.spin = None
@@ -943,6 +937,8 @@ def transform_ci_for_orbital_rotation(self, fcivec, norb, nelec, u):
         nelec = _unpack_nelec(nelec, self.spin)
         return addons.transform_ci_for_orbital_rotation(fcivec, norb, nelec, u)
 
+    to_gpu = lib.to_gpu
+
 FCI = FCISolver
 
 class FCIvector(numpy.ndarray):

pyscf/grad/casci.py

@@ -342,6 +342,8 @@ def _finalize(self):
 
     as_scanner = as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 from pyscf import mcscf

pyscf/grad/casscf.py

@@ -220,6 +220,8 @@ def _finalize(self):
 
     as_scanner = as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 from pyscf import mcscf

pyscf/grad/ccsd.py

@@ -456,6 +456,8 @@ def grad_nuc(self, mol=None, atmlst=None):
 
     as_scanner = as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 ccsd.CCSD.Gradients = lib.class_as_method(Gradients)

pyscf/grad/cisd.py

@@ -203,6 +203,8 @@ def _finalize(self):
 
     as_scanner = as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 cisd.CISD.Gradients = lib.class_as_method(Gradients)

pyscf/grad/dhf.py

@@ -217,6 +217,8 @@ def kernel(self, mo_energy=None, mo_coeff=None, mo_occ=None, atmlst=None):
 
     as_scanner = rhf_grad.as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 from pyscf import scf

pyscf/grad/mp2.py

@@ -309,6 +309,8 @@ def grad_nuc(self, mol=None, atmlst=None):
 
     as_scanner = as_scanner
 
+    to_gpu = lib.to_gpu
+
 Grad = Gradients
 
 # Inject to RMP2 class

pyscf/grad/rhf.py

@@ -463,9 +463,7 @@ def make_rdm1e(self, mo_energy=None, mo_coeff=None, mo_occ=None):
 
     grad_elec = grad_elec
 
-    def to_gpu(self):
-        from gpu4pyscf.grad.rhf import Gradients
-        return lib.to_gpu(self.view(Gradients))
+    to_gpu = lib.to_gpu
 
 Grad = Gradients