From fbe0cddcd049ac516959fa0afdc6ce3a6210d26b Mon Sep 17 00:00:00 2001 From: "puzhichen.996" Date: Wed, 15 Jan 2025 09:43:56 +0800 Subject: [PATCH] Finish debugging and unit tests. --- gpu4pyscf/tdscf/_lr_eig.py | 7 +++++-- gpu4pyscf/tdscf/tests/test_tdrhf.py | 8 ++++---- gpu4pyscf/tdscf/tests/test_tdrks.py | 6 +++--- gpu4pyscf/tdscf/tests/test_tduhf.py | 4 ++-- gpu4pyscf/tdscf/tests/test_tduks.py | 6 +++--- gpu4pyscf/tdscf/uks.py | 2 +- 6 files changed, 18 insertions(+), 15 deletions(-) diff --git a/gpu4pyscf/tdscf/_lr_eig.py b/gpu4pyscf/tdscf/_lr_eig.py index 1616cd71..59f206ff 100644 --- a/gpu4pyscf/tdscf/_lr_eig.py +++ b/gpu4pyscf/tdscf/_lr_eig.py @@ -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): @@ -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) @@ -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] diff --git a/gpu4pyscf/tdscf/tests/test_tdrhf.py b/gpu4pyscf/tdscf/tests/test_tdrhf.py index 4e6761ca..f1a0774a 100644 --- a/gpu4pyscf/tdscf/tests/test_tdrhf.py +++ b/gpu4pyscf/tdscf/tests/test_tdrhf.py @@ -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): @@ -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__": diff --git a/gpu4pyscf/tdscf/tests/test_tdrks.py b/gpu4pyscf/tdscf/tests/test_tdrks.py index cebae87c..233b6463 100644 --- a/gpu4pyscf/tdscf/tests/test_tdrks.py +++ b/gpu4pyscf/tdscf/tests/test_tdrks.py @@ -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): @@ -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): @@ -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__": diff --git a/gpu4pyscf/tdscf/tests/test_tduhf.py b/gpu4pyscf/tdscf/tests/test_tduhf.py index 884249e4..0acd8578 100644 --- a/gpu4pyscf/tdscf/tests/test_tduhf.py +++ b/gpu4pyscf/tdscf/tests/test_tduhf.py @@ -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): @@ -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__": diff --git a/gpu4pyscf/tdscf/tests/test_tduks.py b/gpu4pyscf/tdscf/tests/test_tduks.py index ad1e3f45..34c214b8 100644 --- a/gpu4pyscf/tdscf/tests/test_tduks.py +++ b/gpu4pyscf/tdscf/tests/test_tduks.py @@ -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): @@ -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): @@ -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__": diff --git a/gpu4pyscf/tdscf/uks.py b/gpu4pyscf/tdscf/uks.py index 8ea2f2cf..c2c48d22 100644 --- a/gpu4pyscf/tdscf/uks.py +++ b/gpu4pyscf/tdscf/uks.py @@ -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