"""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
# Enable float64 support
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."""
# Get mean-field data
_, cls.mf = get_be_ccpvdz()
# Create simple Jastrow factors for both implementations
cls.params_jax = {'alpha': jnp.array([1.0])}
cls.params_numpy = np.array([1.0])
cls.jastrow_jax = REXP_jax() # Remove params from constructor
cls.jastrow_numpy = REXP_np(cls.params_numpy) # Keep numpy version unchanged
# Initialize XTC calculators with low grid level for testing
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."""
# First test delta_U matrices
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) # Add params
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
)
# Test symmetry property
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 gradient calculations with defined points
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) # Add params
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
)
# Test hermiticity
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)
# Get standard h1e
h1e_std = tc_helper.get_hcore(self.mf)
# Combine
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)
# Get standard ERI
eri_std = tc_helper.get_eri(self.mf)
# Combine
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) # Pass mf and params
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.")
# get the hf energy using fock and eris
no = myrcc.nocc
# Reconstruct full h1e for checking
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()