Skip to content

Commit

Permalink
Refactor _DFHF class. Add tests for to_cpu (pyscf#46)
Browse files Browse the repository at this point in the history 
* Refactor _DFHF class. Add tests for to_cpu

* Undefined variables

* Update df_jk.py

---------

Co-authored-by: Xiaojie Wu <[email protected]>
  • Loading branch information
@sunqm @wxj6000
sunqm and wxj6000 authored Oct 18, 2023
1 parent 3b7b091 commit 7c343e4
Show file tree
Hide file tree
Showing 8 changed files with 234 additions and 166 deletions.
1 change: 1 addition & 0 deletions gpu4pyscf/df/df.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,6 +47,7 @@ def __init__(self, mol, auxbasis=None):
def to_cpu(self):
from gpu4pyscf.lib.utils import to_cpu
obj = to_cpu(self)
del obj.intopt, obj.cd_low, obj.nao, obj.naux
return obj.reset()

def build(self, direct_scf_tol=1e-14, omega=None):
Expand Down
278 changes: 145 additions & 133 deletions gpu4pyscf/df/df_jk.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -96,8 +96,6 @@ def _density_fit(mf, auxbasis=None, with_df=None, only_dfj=False):
with_df.verbose = mf.verbose
with_df.auxbasis = auxbasis

mf_class = mf.__class__

if isinstance(mf, df_jk._DFHF):
if mf.with_df is None:
mf.with_df = with_df
Expand All @@ -108,139 +106,153 @@ def _density_fit(mf, auxbasis=None, with_df=None, only_dfj=False):
mf.only_dfj = only_dfj
return mf

class DensityFitting(df_jk._DFHF, mf_class):
__doc__ = '''
Density fitting SCF class
Attributes for density-fitting SCF:
auxbasis : str or basis dict
Same format to the input attribute mol.basis.
The default basis 'weigend+etb' means weigend-coulomb-fit basis
for light elements and even-tempered basis for heavy elements.
with_df : DF object
Set mf.with_df = None to switch off density fitting mode.
See also the documents of class %s for other SCF attributes.
''' % mf_class

from gpu4pyscf.lib.utils import to_cpu, to_gpu, device

def __init__(self, mf, dfobj, only_dfj):
self.__dict__.update(mf.__dict__)
self._eri = None
self.rhoj = None
self.rhok = None
self.direct_scf = False
self.with_df = dfobj
self.only_dfj = only_dfj
self._keys = self._keys.union(['with_df', 'only_dfj'])

init_workflow = init_workflow

def reset(self, mol=None):
self.with_df.reset(mol)
return mf_class.reset(self, mol)

def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
omega=None):
if dm is None: dm = self.make_rdm1()
if self.with_df and self.only_dfj:
vj = vk = None
if with_j:
vj, vk = self.with_df.get_jk(dm, hermi, True, False,
self.direct_scf_tol, omega)
if with_k:
vk = mf_class.get_jk(self, mol, dm, hermi, False, True, omega)[1]
elif self.with_df:
vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_k,
self.direct_scf_tol, omega)
else:
vj, vk = mf_class.get_jk(self, mol, dm, hermi, with_j, with_k, omega)
return vj, vk

def get_veff(self, mol=None, dm=None, dm_last=None, vhf_last=0, hermi=1):
'''
effective potential
'''
if mol is None: mol = self.mol
if dm is None: dm = self.make_rdm1()

# for DFT
if mf_class == rks.RKS:
return rks.get_veff(self, dm=dm)

if self.direct_scf:
ddm = cupy.asarray(dm) - dm_last
vj, vk = self.get_jk(mol, ddm, hermi=hermi)
return vhf_last + vj - vk * .5
else:
vj, vk = self.get_jk(mol, dm, hermi=hermi)
return vj - vk * .5

def energy_elec(self, dm=None, h1e=None, vhf=None):
'''
electronic energy
'''
if dm is None: dm = self.make_rdm1()
if h1e is None: h1e = self.get_hcore()
if vhf is None: vhf = self.get_veff(self.mol, dm)
# for DFT
if mf_class == rks.RKS:
e1 = cupy.sum(h1e*dm)
ecoul = self.ecoul
exc = self.exc
e2 = ecoul + exc
#logger.debug(self, f'E1 = {e1}, Ecoul = {ecoul}, Exc = {exc}')
return e1+e2, e2

e1 = cupy.einsum('ij,ji->', h1e, dm).real
e_coul = cupy.einsum('ij,ji->', vhf, dm).real * .5
self.scf_summary['e1'] = e1
self.scf_summary['e2'] = e_coul
#logger.debug(self, 'E1 = %s E_coul = %s', e1, e_coul)
return e1+e_coul, e_coul

def energy_tot(self, dm, h1e, vhf=None):
'''
compute tot energy
'''
nuc = self.energy_nuc()
e_tot = self.energy_elec(dm, h1e, vhf)[0] + nuc
self.scf_summary['nuc'] = nuc.real
return e_tot

def nuc_grad_method(self):
if mf_class == rks.RKS:
from gpu4pyscf.df.grad import rks as rks_grad
return rks_grad.Gradients(self)
if mf_class == hf.RHF:
from gpu4pyscf.df.grad import rhf as rhf_grad
return rhf_grad.Gradients(self)
raise NotImplementedError()


def Hessian(self):
from gpu4pyscf.df.hessian import rhf, rks
if isinstance(self, scf.rhf.RHF):
if isinstance(self, scf.hf.KohnShamDFT):
return rks.Hessian(self)
else:
return rhf.Hessian(self)
else:
raise NotImplementedError
dfmf = _DFHF(mf, with_df, only_dfj)
return lib.set_class(dfmf, (_DFHF, mf.__class__))

# for pyscf 1.0, 1.1 compatibility
@property
def _cderi(self):
naux = self.with_df.get_naoaux()
return next(self.with_df.loop(blksize=naux))
@_cderi.setter
def _cderi(self, x):
self.with_df._cderi = x
class _DFHF(df_jk._DFHF):
'''
Density fitting SCF class
Attributes for density-fitting SCF:
auxbasis : str or basis dict
Same format to the input attribute mol.basis.
The default basis 'weigend+etb' means weigend-coulomb-fit basis
for light elements and even-tempered basis for heavy elements.
with_df : DF object
Set mf.with_df = None to switch off density fitting mode.
'''

@property
def auxbasis(self):
return getattr(self.with_df, 'auxbasis', None)
from gpu4pyscf.lib.utils import to_gpu, device

def __init__(self, mf, dfobj, only_dfj):
self.__dict__.update(mf.__dict__)
self._eri = None
self.rhoj = None
self.rhok = None
self.direct_scf = False
self.with_df = dfobj
self.only_dfj = only_dfj
self._keys = self._keys.union(['with_df', 'only_dfj'])

def undo_df(self):
'''Remove the DFHF Mixin'''
obj = lib.view(self, lib.drop_class(self.__class__, _DFHF))
del obj.rhoj, obj.rhok, obj.with_df, obj.only_dfj
return obj

def reset(self, mol=None):
self.with_df.reset(mol)
return super().reset(mol)

init_workflow = init_workflow

def get_jk(self, mol=None, dm=None, hermi=1, with_j=True, with_k=True,
omega=None):
if dm is None: dm = self.make_rdm1()
if self.with_df and self.only_dfj:
vj = vk = None
if with_j:
vj, vk = self.with_df.get_jk(dm, hermi, True, False,
self.direct_scf_tol, omega)
if with_k:
vk = super().get_jk(mol, dm, hermi, False, True, omega)[1]
elif self.with_df:
vj, vk = self.with_df.get_jk(dm, hermi, with_j, with_k,
self.direct_scf_tol, omega)
else:
vj, vk = super().get_jk(mol, dm, hermi, with_j, with_k, omega)
return vj, vk

return DensityFitting(mf, with_df, only_dfj)
def nuc_grad_method(self):
if isinstance(self, rks.RKS):
from gpu4pyscf.df.grad import rks as rks_grad
return rks_grad.Gradients(self)
if isinstance(self, hf.RHF):
from gpu4pyscf.df.grad import rhf as rhf_grad
return rhf_grad.Gradients(self)
raise NotImplementedError()

def Hessian(self):
from pyscf.dft.rks import KohnShamDFT
from gpu4pyscf.df.hessian import rhf, rks
if isinstance(self, scf.rhf.RHF):
if isinstance(self, KohnShamDFT):
return rks.Hessian(self)
else:
return rhf.Hessian(self)
else:
raise NotImplementedError

@property
def auxbasis(self):
return getattr(self.with_df, 'auxbasis', None)

def get_veff(self, mol=None, dm=None, dm_last=None, vhf_last=0, hermi=1):
'''
effective potential
'''
if mol is None: mol = self.mol
if dm is None: dm = self.make_rdm1()

# for DFT
if super() == rks.RKS:
return rks.get_veff(self, dm=dm)

if self.direct_scf:
ddm = cupy.asarray(dm) - dm_last
vj, vk = self.get_jk(mol, ddm, hermi=hermi)
return vhf_last + vj - vk * .5
else:
vj, vk = self.get_jk(mol, dm, hermi=hermi)
return vj - vk * .5

def energy_elec(self, dm=None, h1e=None, vhf=None):
'''
electronic energy
'''
if dm is None: dm = self.make_rdm1()
if h1e is None: h1e = self.get_hcore()
if vhf is None: vhf = self.get_veff(self.mol, dm)
# for DFT
if super() == rks.RKS:
e1 = cupy.sum(h1e*dm)
ecoul = self.ecoul
exc = self.exc
e2 = ecoul + exc
#logger.debug(self, f'E1 = {e1}, Ecoul = {ecoul}, Exc = {exc}')
return e1+e2, e2

e1 = cupy.einsum('ij,ji->', h1e, dm).real
e_coul = cupy.einsum('ij,ji->', vhf, dm).real * .5
self.scf_summary['e1'] = e1
self.scf_summary['e2'] = e_coul
#logger.debug(self, 'E1 = %s E_coul = %s', e1, e_coul)
return e1+e_coul, e_coul

def energy_tot(self, dm, h1e, vhf=None):
'''
compute tot energy
'''
nuc = self.energy_nuc()
e_tot = self.energy_elec(dm, h1e, vhf)[0] + nuc
self.scf_summary['nuc'] = nuc.real
return e_tot


def to_cpu(self):
obj = self.undo_df().to_cpu().density_fit()
keys = dir(obj)
obj.__dict__.update(self.__dict__)

for key in set(dir(self)).difference(keys):
delattr(obj, key)

for key in keys:
val = getattr(obj, key)
if isinstance(val, cupy.ndarray):
setattr(obj, key, cupy.asnumpy(val))
elif hasattr(val, 'to_cpu'):
setattr(obj, key, val.to_cpu())
return obj

def get_jk(dfobj, dms_tag, hermi=1, with_j=True, with_k=True, direct_scf_tol=1e-14, omega=None):
'''
Expand Down Expand Up @@ -387,4 +399,4 @@ def get_j(dfobj, dm, hermi=1, direct_scf_tol=1e-13):
vj = int3c2e.get_j_int3c2e_pass2(intopt, rhoj)
return vj

density_fit = _density_fit
density_fit = _density_fit
25 changes: 20 additions & 5 deletions gpu4pyscf/df/tests/test_df_scf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@
import numpy as np
import pyscf
from pyscf import lib
from pyscf.df import df_jk as cpu_df_jk
from gpu4pyscf import scf
from gpu4pyscf.df import df_jk
from gpu4pyscf.dft import rks

lib.num_threads(8)
Expand All @@ -37,7 +39,7 @@ def setUpModule():
mol.output = '/dev/null'
mol.build()
mol.verbose = 1

def tearDownModule():
global mol
mol.stdout.close()
Expand All @@ -57,7 +59,7 @@ def test_rhf(self):
mf = scf.RHF(mol).density_fit(auxbasis='def2-tzvpp-jkfit')
e_tot = mf.kernel()
assert np.allclose(e_tot, -76.0624582299)

def test_rks_lda(self):
print('------- LDA ----------------')
e_tot = run_dft("LDA_X,LDA_C_VWN")
Expand All @@ -67,17 +69,17 @@ def test_rks_pbe(self):
print('------- PBE ----------------')
e_tot = run_dft('PBE')
assert np.allclose(e_tot, -76.3800181250)

def test_rks_b3lyp(self):
print('-------- B3LYP -------------')
e_tot = run_dft('B3LYP')
assert np.allclose(e_tot, -76.4666493796)

def test_rks_m06(self):
print('--------- M06 --------------')
e_tot = run_dft("M06")
assert np.allclose(e_tot, -76.4265841359)

def test_rks_wb97(self):
print('-------- wB97 --------------')
e_tot = run_dft("HYB_GGA_XC_WB97")
Expand All @@ -88,6 +90,19 @@ def test_rks_wb97(self):
e_tot = run_dft("HYB_MGGA_XC_WB97M_V")
assert np.allclose(e_tot, -76.4334567297)

def test_to_cpu(self):
mf = scf.RHF(mol).density_fit().to_cpu()
assert isinstance(mf, cpu_df_jk._DFHF)
mf = mf.to_gpu()
assert isinstance(mf, df_jk._DFHF)

mf = rks.RKS(mol).density_fit().to_cpu()
assert isinstance(mf, cpu_df_jk._DFHF)
assert 'gpu' not in mf.grids.__module__
mf = mf.to_gpu()
assert isinstance(mf, df_jk._DFHF)
assert 'gpu' in mf.grids.__module__

if __name__ == "__main__":
print("Full Tests for SCF")
unittest.main()
3 changes: 1 addition & 2 deletions gpu4pyscf/dft/gks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -23,8 +23,7 @@ class GKS(gks.GKS):
from gpu4pyscf.lib.utils import to_cpu, to_gpu, device

def __init__(self, mol, xc='LDA,VWN'):
super().__init__(mol, xc)
self._numint = numint.NumInt()
raise NotImplementedError

get_jk = GHF.get_jk
_eigh = GHF._eigh
Loading

0 comments on commit 7c343e4

Please sign in to comment.