Skip to content

Commit

Permalink
add the get_ab triplet for rks
Browse files Browse the repository at this point in the history
  • Loading branch information
@puzhichen
puzhichen committed Feb 28, 2025
1 parent 7a18372 commit e927cdf
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 14 deletions.
44 changes: 30 additions & 14 deletions gpu4pyscf/tdscf/rhf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,8 +45,8 @@ def get_ab(mf, mo_energy=None, mo_coeff=None, mo_occ=None, singlet=True):
'''
if hasattr(mf, 'with_df'):
raise NotImplementedError('DF-TDDFT is not implemented')
if not singlet:
raise NotImplementedError('Only singlet is implemented')
# if not singlet:
# raise NotImplementedError('Only singlet is implemented')
if mo_energy is None:
mo_energy = mf.mo_energy
if mo_coeff is None:
Expand Down Expand Up @@ -81,11 +81,15 @@ def add_hf_(a, b, hyb=1):
eri_mo = cp.einsum('ijpl,pk->ijkl', eri_mo, mo.conj())
eri_mo = cp.einsum('ijkp,pl->ijkl', eri_mo, mo)
eri_mo = eri_mo.reshape(nocc,nmo,nmo,nmo)
a += cp.einsum('iabj->iajb', eri_mo[:nocc,nocc:,nocc:,:nocc]) * 2
a -= cp.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb
if singlet:
a += cp.einsum('iabj->iajb', eri_mo[:nocc,nocc:,nocc:,:nocc]) * 2
a -= cp.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb

b += cp.einsum('iajb->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * 2
b -= cp.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb
b += cp.einsum('iajb->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * 2
b -= cp.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb
else:
a -= cp.einsum('ijba->iajb', eri_mo[:nocc,:nocc,nocc:,nocc:]) * hyb
b -= cp.einsum('jaib->iajb', eri_mo[:nocc,nocc:,:nocc,nocc:]) * hyb

if isinstance(mf, scf.hf.KohnShamDFT):
grids = mf.grids
Expand Down Expand Up @@ -128,8 +132,12 @@ def add_hf_(a, b, hyb=1):
mo_coeff_mask = mo_coeff[mask]
rho = ni.eval_rho2(_sorted_mol, ao, mo_coeff_mask,
mo_occ, mask, xctype, with_lapl=False)
fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
wfxc = fxc[0,0] * weight
if singlet or singlet is None:
fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
wfxc = fxc[0,0] * weight
else:
fxc = ni.eval_xc_eff(mf.xc, cp.stack((rho, rho)) * 0.5, deriv=2, xctype=xctype)[2]
wfxc = (fxc[0, 0, 0, 0] - fxc[1, 0, 0, 0]) * 0.5 * weight
orbo_mask = orbo[mask]
orbv_mask = orbv[mask]
rho_o = cp.einsum('pr,pi->ri', ao, orbo_mask)
Expand All @@ -147,8 +155,12 @@ def add_hf_(a, b, hyb=1):
mo_coeff_mask = mo_coeff[mask]
rho = ni.eval_rho2(_sorted_mol, ao, mo_coeff_mask,
mo_occ, mask, xctype, with_lapl=False)
fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
wfxc = fxc * weight
if singlet or singlet is None:
fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
wfxc = fxc * weight
else:
fxc = ni.eval_xc_eff(mf.xc, cp.stack((rho, rho)) * 0.5, deriv=2, xctype=xctype)[2]
wfxc = (fxc[0, :, 0, :] - fxc[1, :, 0, :]) * 0.5 * weight
orbo_mask = orbo[mask]
orbv_mask = orbv[mask]
rho_o = cp.einsum('xpr,pi->xri', ao, orbo_mask)
Expand All @@ -173,8 +185,12 @@ def add_hf_(a, b, hyb=1):
mo_coeff_mask = mo_coeff[mask]
rho = ni.eval_rho2(_sorted_mol, ao, mo_coeff_mask,
mo_occ, mask, xctype, with_lapl=False)
fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
wfxc = fxc * weight
if singlet or singlet is None:
fxc = ni.eval_xc_eff(mf.xc, rho, deriv=2, xctype=xctype)[2]
wfxc = fxc * weight
else:
fxc = ni.eval_xc_eff(mf.xc, cp.stack((rho, rho))*0.5, deriv=2, xctype=xctype)[2]
wfxc = (fxc[0, :, 0, :] - fxc[1, :, 0, :]) * 0.5 * weight
orbo_mask = orbo[mask]
orbv_mask = orbv[mask]
rho_o = cp.einsum('xpr,pi->xri', ao, orbo_mask)
Expand Down Expand Up @@ -254,10 +270,10 @@ class TDBase(lib.StreamObject):
_finalize = tdhf_cpu.TDBase._finalize

gen_vind = NotImplemented
def get_ab(self, mf=None, singlet=True):
def get_ab(self, mf=None):
if mf is None:
mf = self._scf
return get_ab(mf, singlet=singlet)
return get_ab(mf, singlet=self.singlet)

def get_precond(self, hdiag):
threshold_t=1.0e-4
Expand Down
41 changes: 41 additions & 0 deletions gpu4pyscf/tdscf/tests/test_tdrks.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -33,6 +33,16 @@ def diagonalize(a, b, nroots=4):
return lowest_e


def diagonalize_tda(a, nroots=5):
nocc, nvir = a.shape[:2]
nov = nocc * nvir
a = a.reshape(nov, nov)
e = np.linalg.eig(np.asarray(a))[0]
lowest_e = np.sort(e[e.real > 0].real)[:nroots]
lowest_e = lowest_e[lowest_e > 1e-3]
return lowest_e


class KnownValues(unittest.TestCase):
@classmethod
def setUpClass(cls):
Expand Down Expand Up @@ -214,6 +224,17 @@ def test_tda_b3lyp_triplet(self):
self.assertAlmostEqual(lib.fp(es), -1.4707787881198082, 6)
td.analyze()

mf_b3lyp_nodf = self.mf_b3lyp_nodf
td = mf_b3lyp_nodf.TDA()
assert td.device == 'gpu'
td.singlet = False
td.lindep=1.0E-6
es = td.kernel(nstates=5)[0]
a, b = td.get_ab()
ref = diagonalize_tda(a, nroots=5)
self.assertAlmostEqual(abs(es - ref).max(), 0, 8)


def test_tda_lda_triplet(self):
mf_lda = self.mf_lda
td = mf_lda.TDA()
Expand All @@ -225,6 +246,16 @@ def test_tda_lda_triplet(self):
self.assertAlmostEqual(abs(es - ref).max(), 0, 8)
self.assertAlmostEqual(lib.fp(es[[0,1,2,4,5]]), -1.4695846533898422, 6)

mf_lda_nodf = self.mf_lda_nodf
td = mf_lda_nodf.TDA()
assert td.device == 'gpu'
td.singlet = False
td.lindep=1.0E-6
es = td.kernel(nstates=5)[0]
a, b = td.get_ab()
ref = diagonalize_tda(a, nroots=5)
self.assertAlmostEqual(abs(es - ref).max(), 0, 8)

def test_tddft_b88p86_triplet(self):
mf_bp86 = self.mf_bp86
td = mf_bp86.TDDFT()
Expand All @@ -236,6 +267,16 @@ def test_tddft_b88p86_triplet(self):
self.assertAlmostEqual(abs(es - ref).max(), 0, 8)
self.assertAlmostEqual(lib.fp(es), -1.4412243124430528, 6)

mf_bp86_nodf = self.mf_bp86_nodf
td = mf_bp86_nodf.TDA()
assert td.device == 'gpu'
td.singlet = False
td.lindep=1.0E-6
es = td.kernel(nstates=5)[0]
a, b = td.get_ab()
ref = diagonalize_tda(a, nroots=5)
self.assertAlmostEqual(abs(es - ref).max(), 0, 8)

def test_tda_rsh(self):
mol = gto.M(atom='H 0 0 0.6; H 0 0 0', basis = "6-31g")
mf = mol.RKS()
Expand Down

0 comments on commit e927cdf

Please sign in to comment.