"""Tests for JAX implementation of Extended Transcorrelated methods."""
import unittest
import numpy as np
from functools import partial
import jax
import jax.numpy as jnp
from pyscf import gto, scf
from pytc.legacy.xtc import XTC as XTC_numpy
from pytc.legacy.jastrow import REXP as REXP_np
from pytc.xtc import XTC as XTC_jax
from pytc.jastrow import REXP as REXP_jax
from pytc import tc_helper
jax.config.update("jax_enable_x64", True)
[docs]
def get_be_ccpvdz():
"""Return a Be atom with cc-pVDZ basis for testing."""
mol = gto.M(atom='Be 0 0 0', basis='ccpvdz', unit='Bohr')
mf = scf.RHF(mol)
mf.kernel()
return mol, mf
[docs]
class TestXTC(unittest.TestCase):
"""Test JAX implementation of Extended Transcorrelated methods."""
[docs]
@classmethod
def setUpClass(cls):
"""Set up test case using Be atom."""
_, cls.mf = get_be_ccpvdz()
cls.params_jax = {'alpha': jnp.array([1.0])}
cls.params_numpy = np.array([1.0])
cls.jastrow_jax = REXP_jax()
cls.jastrow_numpy = REXP_np(cls.params_numpy)
cls.xtc_jax = XTC_jax.from_pyscf(cls.mf, cls.jastrow_jax, grid_lvl=2)
cls.xtc_numpy = XTC_numpy(cls.mf, cls.jastrow_numpy, grid_lvl=2)
[docs]
def test_grid_initialization(self):
"""Test grid initialization and conversion to JAX arrays."""
np.testing.assert_allclose(
np.asarray(self.xtc_jax.grid_points),
self.xtc_numpy.grid_points,
rtol=1e-7, atol=1e-7
)
np.testing.assert_allclose(
np.asarray(self.xtc_jax.weights),
self.xtc_numpy.weights,
rtol=1e-7, atol=1e-7
)
[docs]
def test_basis_evaluation(self):
"""Test basis function evaluation on grid."""
phi_jax = self.xtc_jax.phi
rho_numpy, _ = self.xtc_numpy._eval_basis_on_grid()
np.testing.assert_allclose(
np.asarray(phi_jax),
rho_numpy,
rtol=1e-5, atol=1e-5
)
[docs]
def test_delta_U(self):
"""Test delta_U calculation."""
dm1_jax = self.xtc_jax._get_mf_dm()
# dm1_numpy = self.xtc_numpy._get_mf_dm() # Not needed if we trust JAX dm1
delta_U_jax = self.xtc_jax.get_delta_U(self.params_jax, dm1_jax)
delta_U_numpy = self.xtc_numpy.get_delta_U()
np.testing.assert_allclose(
np.asarray(delta_U_jax),
delta_U_numpy,
rtol=1e-6, atol=1e-6
)
np.testing.assert_allclose(
np.asarray(delta_U_jax),
np.asarray(delta_U_jax.transpose(2,3,0,1)),
rtol=1e-7, atol=1e-7
)
test_r1 = np.array([[0.0, 0.0, 0.0]])
test_r2 = np.array([[0.0, 0.0, 1.0]])
grad_jax = self.xtc_jax.jastrow_factor.grad_r(test_r1, test_r2, self.params_jax)
grad_numpy = self.jastrow_numpy.grad(test_r1, test_r2)
np.testing.assert_allclose(
np.asarray(grad_jax[:, None, :]), # Add middle dimension to match numpy shape
grad_numpy,
rtol=1e-5, atol=1e-5,
err_msg="Gradients don't match"
)
[docs]
def test_delta_h(self):
"""Test delta_h calculation."""
delta_h_jax = self.xtc_jax.get_delta_h(self.params_jax)
delta_h_numpy = self.xtc_numpy._calc_delta_h()
np.testing.assert_allclose(
np.asarray(delta_h_jax),
delta_h_numpy,
rtol=1e-5, atol=1e-5
)
np.testing.assert_allclose(
np.asarray(delta_h_jax),
np.asarray(delta_h_jax.T.conj()),
rtol=1e-7, atol=1e-7
)
[docs]
def test_one_body(self):
"""Test one-body operator calculation."""
# JAX returns only delta_h
delta_h_jax = self.xtc_jax.get_1b(self.params_jax)
h1e_std = tc_helper.get_hcore(self.mf)
h1e_jax = h1e_std + delta_h_jax
# Numpy returns full h1e
h1e_numpy = self.xtc_numpy.get_1b()
np.testing.assert_allclose(
np.asarray(h1e_jax),
h1e_numpy,
rtol=1e-5, atol=1e-5
)
[docs]
def test_two_body(self):
"""Test two-body term calculation."""
# JAX returns correction + delta_U
correction_jax = self.xtc_jax.get_2b(self.params_jax)
eri_std = tc_helper.get_eri(self.mf)
v2e_jax = eri_std + correction_jax
# Numpy returns full effective ERI
v2e_numpy = self.xtc_numpy.get_2b()
np.testing.assert_allclose(
np.asarray(v2e_jax),
v2e_numpy,
rtol=1e-5, atol=1e-5
)
[docs]
def test_nhccsd(self):
from pyscf.cc import rccsd, CCSD
import numpy as np
from functools import reduce
mycc = CCSD(self.mf)
e_corr, t1, t2 = mycc.kernel()
mycc.verbose = 5
print("E_CCSD = ", e_corr)
#self.assertAlmostEqual(e_corr, -0.04503138331130402, places=6)
t = mycc.amplitudes_to_vector(t1, t2)
print("|t2| = ", np.linalg.norm(t2))
print("|t1+t2| = ", np.linalg.norm(t))
eri1 = tc_helper.get_eri(self.mf, self.xtc_jax.mo_coeff)
h1e = mycc._scf.get_hcore()
h1e = reduce(np.dot, (self.xtc_jax.mo_coeff.T, h1e, self.xtc_jax.mo_coeff))
e_hf_0 = 2. * np.einsum('ii->', h1e[:mycc.nocc, :mycc.nocc])
e_dir = 2. * np.einsum('jjii->', eri1[:mycc.nocc, :mycc.nocc, :mycc.nocc, :mycc.nocc])
e_ex = -1. * np.einsum('ijji->', eri1[:mycc.nocc, :mycc.nocc, :mycc.nocc, :mycc.nocc])
e_hf_0 += (e_dir + e_ex) + mycc._scf.energy_nuc()
print("Check e_hf = ", e_hf_0)
self.assertAlmostEqual(e_hf_0, -14.57233763095337, places=6)
myrcc = rccsd.RCCSD(self.mf)
#myrcc.verbose = 5
eris = self.xtc_jax.make_eris(self.mf, self.params_jax)
tc_e_corr, t1, t2 = myrcc.kernel(eris=eris)
t = myrcc.amplitudes_to_vector(t1, t2)
print("|t2| = ", np.linalg.norm(t2))
print("|t1+t2| = ", np.linalg.norm(t))
print("corr E_XTC_CCSD = ", myrcc.e_corr)
# Check against expected value (either custom PySCF or standard PySCF)
# Custom PySCF (tc-ccsd branch) gives -0.0327...
# Standard PySCF gives -0.0537...
expected_custom = -0.03271941719618445
expected_standard = -0.053729178994280744
if np.isclose(tc_e_corr, expected_custom, atol=1e-5):
self.assertAlmostEqual(tc_e_corr, expected_custom, places=5)
elif np.isclose(tc_e_corr, expected_standard, atol=1e-5):
print("Warning: Using standard PySCF result. For correct tc-ccsd results, install https://github.com/nickirk/pyscf/tree/tc-ccsd")
self.assertAlmostEqual(tc_e_corr, expected_standard, places=5)
else:
self.fail(f"Correlation energy {tc_e_corr} does not match expected custom ({expected_custom}) or standard ({expected_standard}) values.")
no = myrcc.nocc
tc_h1e_corr = self.xtc_jax.get_1b(self.params_jax)
h1e_std = tc_helper.get_hcore(self.mf, self.xtc_jax.mo_coeff)
tc_h1e = h1e_std + tc_h1e_corr
tc_e_hf = 2. * np.einsum('ii->', tc_h1e[:no, :no])
tc_e_dir = 2. * np.einsum('jjii->', eris.oooo)
tc_e_ex = -1. * np.einsum('ijji->', eris.oooo)
tc_e_hf += (tc_e_dir + tc_e_ex) + eris.e_core
print("E_XTC_CCSD = ", myrcc.e_corr + tc_e_hf)
expected_total_custom = -14.65640358838446
expected_total_standard = -14.677381617968336
if np.isclose(tc_e_hf + tc_e_corr, expected_total_custom, atol=1e-5):
self.assertAlmostEqual(tc_e_hf + tc_e_corr, expected_total_custom, places=5)
elif np.isclose(tc_e_hf + tc_e_corr, expected_total_standard, atol=1e-5):
self.assertAlmostEqual(tc_e_hf + tc_e_corr, expected_total_standard, places=5)
else:
self.fail(f"Total energy {tc_e_hf + tc_e_corr} does not match expected custom ({expected_total_custom}) or standard ({expected_total_standard}) values.")
if __name__ == '__main__':
unittest.main()