gpu4pyscf 0.7.2__tar.gz → 0.7.3__tar.gz

{gpu4pyscf-0.7.2/gpu4pyscf.egg-info → gpu4pyscf-0.7.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpu4pyscf
-Version: 0.7.2
+Version: 0.7.3
 Summary: GPU extensions for PySCF
 Home-page: UNKNOWN
 Author: Qiming Sun

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/__init__.py

@@ -1,6 +1,6 @@
 from . import lib, grad, hessian, solvent, scf, dft
 
-__version__ = '0.7.2'
+__version__ = '0.7.3'
 
 # monkey patch libxc reference due to a bug in nvcc
 from pyscf.dft import libxc

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/df/grad/uhf.py

@@ -254,6 +254,9 @@ def get_jk(mf_grad, mol=None, dm0=None, hermi=0, with_j=True, with_k=True, omega
 
 class Gradients(uhf_grad.Gradients):
     '''Unrestricted density-fitting Hartree-Fock gradients'''
+
+    _keys = {'with_df', 'auxbasis_response'}
+
     def __init__(self, mf):
         # Whether to include the response of DF auxiliary basis when computing
         # nuclear gradients of J/K matrices
@@ -262,7 +265,7 @@ class Gradients(uhf_grad.Gradients):
         uhf_grad.Gradients.__init__(self, mf)
 
     get_jk = get_jk
-    
+
     # TODO: finish these two functions
     def get_j(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None):
         vj, _, vjaux, _ =  self.get_jk(mol, dm, with_k=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2)
@@ -271,7 +274,7 @@ class Gradients(uhf_grad.Gradients):
     def get_k(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None):
         _, vk, _, vkaux = self.get_jk(mol, dm, with_j=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2)
         return vk, vkaux
-        
+
 
     def get_veff(self, mol=None, dm=None):
         vj0, vk0, vjaux0, vkaux0 = self.get_jk(mol, dm[0], mo_coeff=self.base.mo_coeff[0], mo_occ=self.base.mo_occ[0])

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/df/grad/uks.py

@@ -100,7 +100,7 @@ def get_veff(ks_grad, mol=None, dm=None):
         vxc += vj0 + vj1 + vj0_m1 + vj1_m0
     else:
         omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=mol.spin)
-        
+
         vj0, vk0, vjaux0, vkaux0 = ks_grad.get_jk(mol, dm[0], mo_coeff=ks_grad.base.mo_coeff[0], mo_occ=ks_grad.base.mo_occ[0])
         vj1, vk1, vjaux1, vkaux1 = ks_grad.get_jk(mol, dm[1], mo_coeff=ks_grad.base.mo_coeff[1], mo_occ=ks_grad.base.mo_occ[1])
         vj0_m1, vjaux0_m1 = ks_grad.get_j(mol, dm[0], mo_coeff=ks_grad.base.mo_coeff[0], mo_occ=ks_grad.base.mo_occ[0], dm2=dm[1])
@@ -110,29 +110,32 @@ def get_veff(ks_grad, mol=None, dm=None):
         if ks_grad.auxbasis_response:
             vj_aux = vjaux0 + vjaux1 + vjaux0_m1 + vjaux1_m0
             vk_aux = (vkaux0+vkaux1) * hyb
-        
+
         if omega != 0:
             vk_lr0, vkaux_lr0 = ks_grad.get_k(mol, dm[0], mo_coeff=ks_grad.base.mo_coeff[0], mo_occ=ks_grad.base.mo_occ[0], omega=omega)
-            vk_lr1, vkaux_lr1 = ks_grad.get_k(mol, dm[1], mo_coeff=ks_grad.base.mo_coeff[1], mo_occ=ks_grad.base.mo_occ[1], omega=omega) 
+            vk_lr1, vkaux_lr1 = ks_grad.get_k(mol, dm[1], mo_coeff=ks_grad.base.mo_coeff[1], mo_occ=ks_grad.base.mo_occ[1], omega=omega)
             vk += (vk_lr0 + vk_lr1) * (alpha-hyb)
             if ks_grad.auxbasis_response:
                 vk_aux += (vkaux_lr0 + vkaux_lr1) * (alpha-hyb)
-                
+
         vxc += vj - vk
         if ks_grad.auxbasis_response:
             e1_aux = vj_aux - vk_aux
-            
+
     if ks_grad.auxbasis_response:
         logger.debug1(ks_grad, 'sum(auxbasis response) %s', e1_aux.sum(axis=0))
     else:
-        e1_aux = None   
-             
+        e1_aux = None
+
     vxc = tag_array(vxc, aux=e1_aux)
-    
+
     return vxc
 
 
 class Gradients(uks_grad.Gradients):
+
+    _keys = {'with_df', 'auxbasis_response'}
+
     def __init__(self, mf):
         # Whether to include the response of DF auxiliary basis when computing
         # nuclear gradients of J/K matrices
@@ -140,7 +143,7 @@ class Gradients(uks_grad.Gradients):
         uks_grad.Gradients.__init__(self, mf)
 
     get_jk = get_jk
-    
+
     def get_j(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None, omega=None):
         vj, _, vjaux, _ = self.get_jk(mol, dm, with_k=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2, omega=omega)
         return vj, vjaux
@@ -148,7 +151,7 @@ class Gradients(uks_grad.Gradients):
     def get_k(self, mol=None, dm=None, hermi=0, mo_coeff=None, mo_occ=None, dm2 = None, omega=None):
         _, vk, _, vkaux = self.get_jk(mol, dm, with_j=False, mo_coeff=mo_coeff, mo_occ=mo_occ, dm2=dm2, omega=omega)
         return vk, vkaux
-    
+
     get_veff = get_veff
 
     def extra_force(self, atom_id, envs):

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/scf/uhf.py

@@ -113,7 +113,7 @@ def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
 class UHF(uhf.UHF):
     from gpu4pyscf.lib.utils import to_cpu, to_gpu, device
 
-    _keys = {'e_disp', 'screen_tol', 'conv_tol_cpscf'}
+    _keys = {'e_disp', 'screen_tol', 'conv_tol_cpscf', 'h1e', 's1e'}
 
     screen_tol = 1e-14
     conv_tol_cpscf = 1e-3

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/solvent/grad/pcm.py

@@ -344,7 +344,8 @@ class WithSolventGrad:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
-
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
         self.de_solute = super().kernel(*args, **kwargs)
         self.de_solvent = grad_qv(self.base.with_solvent, dm)
         self.de_solvent+= grad_solver(self.base.with_solvent, dm)

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/solvent/grad/smd.py

@@ -255,14 +255,15 @@ class WithSolventGrad:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
-
-        self.de_solute = super().kernel(*args, **kwargs)
+        if dm.ndim == 3:
+            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+= pcm_grad.grad_nuc(self.base.with_solvent, dm)
+        self.de_cds     = get_cds(self.base.with_solvent)
+        self.de = self.de_solute + self.de_solvent + self.de_cds
 
-        self.de = self.de_solute + self.de_solvent
-        self.de += get_cds(self.base.with_solvent)
         if self.verbose >= logger.NOTE:
             logger.note(self, '--------------- %s (+%s) gradients ---------------',
                         self.base.__class__.__name__,

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/solvent/hessian/pcm.py

@@ -21,8 +21,7 @@ Gradient of PCM family solvent model
 import numpy
 import cupy
 from cupyx import scipy
-from pyscf import lib
-from pyscf import gto, df
+from pyscf import lib, scf, gto, df
 from pyscf.grad import rhf as rhf_grad
 from gpu4pyscf.solvent.pcm import PI, switch_h
 from gpu4pyscf.solvent.grad.pcm import grad_switch_h, get_dF_dA, get_dD_dS, grad_qv, grad_solver, grad_nuc
@@ -155,7 +154,7 @@ def hess_elec(pcmobj, dm, verbose=None):
     pcmobj.reset(pmol)
     return de
 
-def fd_grad_vmat(pcmobj, mo_coeff, mo_occ, atmlst=None, verbose=None):
+def fd_grad_vmat(pcmobj, dm, mo_coeff, mo_occ, atmlst=None, verbose=None):
     '''
     dv_solv / da
     slow version with finite difference
@@ -169,7 +168,6 @@ def fd_grad_vmat(pcmobj, mo_coeff, mo_occ, atmlst=None, verbose=None):
     nao, nmo = mo_coeff.shape
     mocc = mo_coeff[:,mo_occ>0]
     nocc = mocc.shape[1]
-    dm = cupy.dot(mocc, mocc.T) * 2
     coords = mol.atom_coords(unit='Bohr')
     def pcm_vmat_scanner(mol):
         pcmobj.reset(mol)
@@ -250,6 +248,8 @@ class WithSolventHess:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
+        if dm.ndim == 3:
+            dm = dm[0] + dm[1]
         is_equilibrium = self.base.with_solvent.equilibrium_solvation
         self.base.with_solvent.equilibrium_solvation = True
         self.de_solvent = hess_elec(self.base.with_solvent, dm, verbose=self.verbose)
@@ -263,11 +263,25 @@ class WithSolventHess:
         if atmlst is None:
             atmlst = range(self.mol.natm)
         h1ao = super().make_h1(mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        dv = fd_grad_vmat(self.base.with_solvent, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        for i0, ia in enumerate(atmlst):
-            h1ao[i0] += dv[i0]
-        return h1ao
-
+        if isinstance(self.base, scf.hf.RHF):
+            dm = self.base.make_rdm1(ao_repr=True)
+            dv = fd_grad_vmat(self.base.with_solvent, dm, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1ao[i0] += dv[i0]
+            return h1ao
+        elif isinstance(self.base, scf.uhf.UHF):
+            h1aoa, h1aob = h1ao
+            solvent = self.base.with_solvent
+            dm = self.base.make_rdm1(ao_repr=True)
+            dm = dm[0] + dm[1]
+            dva = fd_grad_vmat(solvent, dm, mo_coeff[0], mo_occ[0], atmlst=atmlst, verbose=verbose)
+            dvb = fd_grad_vmat(solvent, dm, mo_coeff[1], mo_occ[1], atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1aoa[i0] += dva[i0]
+                h1aob[i0] += dvb[i0]
+            return h1aoa, h1aob
+        else:
+            raise NotImplementedError('Base object is not supported')
     def _finalize(self):
         # disable _finalize. It is called in grad_method.kernel method
         # where self.de was not yet initialized.

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/solvent/hessian/smd.py

@@ -20,9 +20,7 @@ Gradient of PCM family solvent model
 
 import numpy as np
 import cupy
-#from cupyx import scipy, jit
-from pyscf import lib
-from pyscf import gto, df
+from pyscf import gto, df, scf, lib
 from pyscf.grad import rhf as rhf_grad
 from pyscf.data import radii
 from pyscf.solvent import ddcosmo_grad
@@ -278,13 +276,15 @@ class WithSolventHess:
         dm = kwargs.pop('dm', None)
         if dm is None:
             dm = self.base.make_rdm1(ao_repr=True)
+        if dm.ndim == 3:
+            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_solute = super().kernel(*args, **kwargs)
-        self.de = self.de_solute + self.de_solvent
-        self.de += get_cds(self.base.with_solvent)
+        self.de_cds = get_cds(self.base.with_solvent)
+        self.de = self.de_solute + self.de_solvent + self.de_cds
         self.base.with_solvent.equilibrium_solvation = is_equilibrium
         return self.de
 
@@ -292,11 +292,25 @@ class WithSolventHess:
         if atmlst is None:
             atmlst = range(self.mol.natm)
         h1ao = super().make_h1(mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        dv = pcm_hess.fd_grad_vmat(self.base.with_solvent, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
-        for i0, ia in enumerate(atmlst):
-            h1ao[i0] += dv[i0]
-        return h1ao
-
+        if isinstance(self.base, scf.hf.RHF):
+            dm = self.base.make_rdm1(ao_repr=True)
+            dv = pcm_hess.fd_grad_vmat(self.base.with_solvent, dm, mo_coeff, mo_occ, atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1ao[i0] += dv[i0]
+            return h1ao
+        elif isinstance(self.base, scf.uhf.UHF):
+            h1aoa, h1aob = h1ao
+            solvent = self.base.with_solvent
+            dm = self.base.make_rdm1(ao_repr=True)
+            dm = dm[0] + dm[1]
+            dva = pcm_hess.fd_grad_vmat(solvent, dm, mo_coeff[0], mo_occ[0], atmlst=atmlst, verbose=verbose)
+            dvb = pcm_hess.fd_grad_vmat(solvent, dm, mo_coeff[1], mo_occ[1], atmlst=atmlst, verbose=verbose)
+            for i0, ia in enumerate(atmlst):
+                h1aoa[i0] += dva[i0]
+                h1aob[i0] += dvb[i0]
+            return h1aoa, h1aob
+        else:
+            raise NotImplementedError('Base object is not supported')
     def _finalize(self):
         # disable _finalize. It is called in grad_method.kernel method
         # where self.de was not yet initialized.

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/solvent/pcm.py

@@ -310,7 +310,6 @@ class PCM(ddcosmo.DDCOSMO):
     def _get_vind(self, dms):
         if not self._intermediates:
             self.build()
-
         nao = dms.shape[-1]
         dms = dms.reshape(-1,nao,nao)
         if dms.shape[0] == 2:
@@ -335,7 +334,6 @@ class PCM(ddcosmo.DDCOSMO):
         self._intermediates['q'] = q[0]
         self._intermediates['q_sym'] = q_sym[0]
         self._intermediates['v_grids'] = v_grids[0]
-
         return epcm, vmat[0]
 
     def _get_v(self, dms):
@@ -363,7 +361,7 @@ class PCM(ddcosmo.DDCOSMO):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(grad_method.base, scf.hf.RHF):
+        if isinstance(grad_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return pcm_grad.make_grad_object(grad_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')
@@ -373,7 +371,7 @@ class PCM(ddcosmo.DDCOSMO):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(hess_method.base, scf.hf.RHF):
+        if isinstance(hess_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return pcm_hess.make_hess_object(hess_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')
@@ -387,7 +385,7 @@ class PCM(ddcosmo.DDCOSMO):
     def _B_dot_x(self, dms):
         if not self._intermediates:
             self.build()
-
+        out_shape = dms.shape
         nao = dms.shape[-1]
         dms = dms.reshape(-1,nao,nao)
 
@@ -403,6 +401,5 @@ class PCM(ddcosmo.DDCOSMO):
         q_sym = (q + qt)/2.0
 
         vmat = self._get_vmat(q_sym)
-
-        return vmat
+        return vmat.reshape(out_shape)
 

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3}/gpu4pyscf/solvent/smd.py

@@ -572,7 +572,7 @@ class SMD(pcm.PCM):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(grad_method.base, scf.hf.RHF):
+        if isinstance(grad_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return smd_grad.make_grad_object(grad_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')
@@ -582,7 +582,7 @@ class SMD(pcm.PCM):
         if self.frozen:
             raise RuntimeError('Frozen solvent model is not supported')
         from gpu4pyscf import scf
-        if isinstance(hess_method.base, scf.hf.RHF):
+        if isinstance(hess_method.base, (scf.hf.RHF, scf.uhf.UHF)):
             return smd_hess.make_hess_object(hess_method)
         else:
             raise RuntimeError('Only SCF gradient is supported')

{gpu4pyscf-0.7.2 → gpu4pyscf-0.7.3/gpu4pyscf.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: gpu4pyscf
-Version: 0.7.2
+Version: 0.7.3
 Summary: GPU extensions for PySCF
 Home-page: UNKNOWN
 Author: Qiming Sun