Skip to content

Commit

Permalink
Finish debugging and unit tests.
Browse files Browse the repository at this point in the history
  • Loading branch information
puzhichen committed Jan 15, 2025
1 parent 8d234d8 commit fbe0cdd
Show file tree
Hide file tree
Showing 6 changed files with 18 additions and 15 deletions.
7 changes: 5 additions & 2 deletions gpu4pyscf/tdscf/_lr_eig.py
Original file line number Diff line number Diff line change
Expand Up @@ -244,7 +244,7 @@ def eigh(aop, x0, precond, tol_residual=1e-5, lindep=1e-12, nroots=1,
if len(x0) < min(x0_size, nroots):
log.warn(f'Not enough eigenvectors (len(x0)={len(x0)}, nroots={nroots})')

return conv, e, x0
return conv, e.get(), x0.get()

def eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=None,
max_cycle=50, max_memory=MAX_MEMORY, lindep=1e-12, verbose=logger.WARN):
Expand Down Expand Up @@ -759,7 +759,7 @@ def real_eig(aop, x0, precond, tol_residual=1e-5, nroots=1, x0sym=None, pick=Non
if len(x0[0]) < min(A_size, nroots):
log.warn(f'Not enough eigenvectors (len(x0)={len(x0[0])}, nroots={nroots})')

return conv[:nroots], e[:nroots], cp.hstack(x0)
return conv[:nroots], e[:nroots].get(), cp.hstack(x0).get()

def _gen_x0(v, xs):
out = _outprod_to_subspace(v[::2], xs)
Expand Down Expand Up @@ -1140,6 +1140,9 @@ def VW_Gram_Schmidt_fill_holder_gpu(V_holder, W_holder, X_new, Y_new, lindep=1e-
csc = c_orth.T.dot(s21).dot(c_orth)
w, u = cp.linalg.eigh(csc)
c_orth = c_orth.dot(u)
mask = 1 - abs(w) > lindep
w = w[mask]
c_orth = c_orth[:,mask]

# Symmetric diagonalize
# [1 w] => c = [a b]
Expand Down
8 changes: 4 additions & 4 deletions gpu4pyscf/tdscf/tests/test_tdrhf.py
Original file line number Diff line number Diff line change
Expand Up @@ -125,12 +125,12 @@ def test_tda_vind(self):
nvir = nmo - nocc
zs = np.random.rand(3,nocc,nvir)
ref = mf.to_cpu().TDA().set(singlet=False).gen_vind()[0](zs)
dat = mf.TDA().set(singlet=False).gen_vind()[0](cp.asarray(zs))
dat = mf.TDA().set(singlet=False).gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

df_mf = self.df_mf
ref = df_mf.to_cpu().TDA().set(singlet=True).gen_vind()[0](zs)
dat = df_mf.TDA().set(singlet=True).gen_vind()[0](cp.asarray(zs))
dat = df_mf.TDA().set(singlet=True).gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

def test_tdhf_vind(self):
Expand All @@ -140,12 +140,12 @@ def test_tdhf_vind(self):
nvir = nmo - nocc
zs = np.random.rand(3,2,nocc,nvir)
ref = mf.to_cpu().TDHF().set(singlet=True).gen_vind()[0](zs)
dat = mf.TDHF().set(singlet=True).gen_vind()[0](zs)
dat = mf.TDHF().set(singlet=True).gen_vind()[0](zs).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

df_mf = self.df_mf
ref = df_mf.to_cpu().TDHF().set(singlet=False).gen_vind()[0](zs)
dat = df_mf.TDHF().set(singlet=False).gen_vind()[0](zs)
dat = df_mf.TDHF().set(singlet=False).gen_vind()[0](zs).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

if __name__ == "__main__":
Expand Down
6 changes: 3 additions & 3 deletions gpu4pyscf/tdscf/tests/test_tdrks.py
Original file line number Diff line number Diff line change
Expand Up @@ -251,7 +251,7 @@ def test_tda_vind(self):
nvir = nmo - nocc
zs = np.random.rand(3,nocc,nvir)
ref = mf.to_cpu().TDA().set(singlet=False).gen_vind()[0](zs)
dat = mf.TDA().set(singlet=False).gen_vind()[0](cp.asarray(zs))
dat = mf.TDA().set(singlet=False).gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

def test_tddft_vind(self):
Expand All @@ -261,7 +261,7 @@ def test_tddft_vind(self):
nvir = nmo - nocc
zs = np.random.rand(3,2,nocc,nvir)
ref = mf.to_cpu().TDDFT().set(singlet=True).gen_vind()[0](zs)
dat = mf.TDDFT().set(singlet=True).gen_vind()[0](zs)
dat = mf.TDDFT().set(singlet=True).gen_vind()[0](zs).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

def test_casida_tddft_vind(self):
Expand All @@ -271,7 +271,7 @@ def test_casida_tddft_vind(self):
nvir = nmo - nocc
zs = np.random.rand(3,nocc,nvir)
ref = mf.to_cpu().CasidaTDDFT().gen_vind()[0](zs)
dat = mf.CasidaTDDFT().gen_vind()[0](cp.asarray(zs))
dat = mf.CasidaTDDFT().gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

if __name__ == "__main__":
Expand Down
4 changes: 2 additions & 2 deletions gpu4pyscf/tdscf/tests/test_tduhf.py
Original file line number Diff line number Diff line change
Expand Up @@ -100,7 +100,7 @@ def test_tda_vind(self):
nvirb = nmo - noccb
zs = np.random.rand(3,nocca*nvira+noccb*nvirb)
ref = mf.to_cpu().TDA().set().gen_vind()[0](zs)
dat = mf.TDA().set().gen_vind()[0](cp.asarray(zs))
dat = mf.TDA().set().gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

def test_tdhf_vind(self):
Expand All @@ -111,7 +111,7 @@ def test_tdhf_vind(self):
nvirb = nmo - noccb
zs = np.random.rand(3,2,nocca*nvira+noccb*nvirb)
ref = mf.to_cpu().TDHF().set().gen_vind()[0](zs)
dat = mf.TDHF().set().gen_vind()[0](zs)
dat = mf.TDHF().set().gen_vind()[0](zs).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

if __name__ == "__main__":
Expand Down
6 changes: 3 additions & 3 deletions gpu4pyscf/tdscf/tests/test_tduks.py
Original file line number Diff line number Diff line change
Expand Up @@ -187,7 +187,7 @@ def test_tda_vind(self):
nvirb = nmo - noccb
zs = np.random.rand(3,nocca*nvira+noccb*nvirb)
ref = mf.to_cpu().TDA().gen_vind()[0](zs)
dat = mf.TDA().gen_vind()[0](cp.asarray(zs))
dat = mf.TDA().gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

def test_tddft_vind(self):
Expand All @@ -198,7 +198,7 @@ def test_tddft_vind(self):
nvirb = nmo - noccb
zs = np.random.rand(3,2,nocca*nvira+noccb*nvirb)
ref = mf.to_cpu().TDDFT().gen_vind()[0](zs)
dat = mf.TDDFT().gen_vind()[0](cp.asarray(zs))
dat = mf.TDDFT().gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

def test_casida_tddft_vind(self):
Expand All @@ -209,7 +209,7 @@ def test_casida_tddft_vind(self):
nvirb = nmo - noccb
zs = np.random.rand(3,nocca*nvira+noccb*nvirb)
ref = mf.to_cpu().CasidaTDDFT().gen_vind()[0](zs)
dat = mf.CasidaTDDFT().gen_vind()[0](cp.asarray(zs))
dat = mf.CasidaTDDFT().gen_vind()[0](cp.asarray(zs)).get()
self.assertAlmostEqual(abs(ref - dat).max(), 0, 9)

if __name__ == "__main__":
Expand Down
2 changes: 1 addition & 1 deletion gpu4pyscf/tdscf/uks.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
import cupy as cp
from pyscf import symm
from pyscf import lib
from pyscf.tdscf._lr_eig import eigh as lr_eigh
from gpu4pyscf.tdscf._lr_eig import eigh as lr_eigh
from gpu4pyscf.dft.rks import KohnShamDFT
from gpu4pyscf.lib.cupy_helper import contract, tag_array, transpose_sum
from gpu4pyscf.lib import logger
Expand Down

0 comments on commit fbe0cdd

Please sign in to comment.