"""Tests for JAX implementation of Transcorrelated method."""
import unittest
import numpy as np
import jax
import jax.numpy as jnp
from pyscf import gto, scf
from pytc.legacy.tc import TC as TC_numpy
from pytc.tc import TC as TC_jax
from pytc.jastrow import Poly
from pytc import tc_helper
# Enable float64 support
jax.config.update("jax_enable_x64", True)
[docs]
class TestTC(unittest.TestCase):
"""Test JAX implementation of TC method."""
[docs]
def setUp(self):
"""Set up test fixtures."""
# Create a simple molecule
self.mol = gto.M(atom='He 0 0 0', basis='sto-3g')
self.mf = scf.RHF(self.mol)
self.mf.kernel()
# Create simple Jastrow factors for both implementations
self.params = jnp.array([1.0])
self.jastrow_jax = Poly() # No params in constructor
# Create numpy version of same jastrow for comparison
class PolyNumpy:
"""Numpy version of Poly for comparison."""
def __init__(self, params):
self.params = params
def __call__(self, r1, r2):
r12 = np.sqrt(np.sum((r1 - r2)**2, axis=-1))
return self.params[0] * r12
def grad(self, r1, r2):
"""Combined gradient calculation for comparison."""
diff = r1[:, None, :] - r2[None, :, :]
r12_sq = np.sum(diff * diff, axis=-1)
r12 = np.sqrt(r12_sq + 1e-10) # Add epsilon for stability
cutoff = 1.0 / (1.0 + np.exp(-(r12 - 1e-5) * 1e6)) # Sigmoid cutoff
grad = np.where(r12_sq[..., None] > 1e-10,
diff * self.params[0] * cutoff[..., None] / r12[..., None],
np.zeros_like(diff))
return grad
self.jastrow_numpy = PolyNumpy(self.params)
# Create TC objects
self.tc_jax = TC_jax.from_pyscf(self.mf, self.jastrow_jax)
self.tc_numpy = TC_numpy(self.mf, self.jastrow_numpy)
[docs]
def test_grid_initialization(self):
"""Test grid initialization and conversion to JAX arrays."""
self.assertIsNotNone(self.tc_jax.grid_points)
self.assertIsNotNone(self.tc_jax.weights)
self.assertTrue(isinstance(self.tc_jax.grid_points, jnp.ndarray))
self.assertTrue(isinstance(self.tc_jax.weights, jnp.ndarray))
[docs]
def test_basis_evaluation(self):
"""Test basis function evaluation on grid."""
# In new design, phi is pre-computed and stored in struct
phi_jax = self.tc_jax.phi
grad_phi_jax = self.tc_jax.grad_phi
rho_numpy, nabla_rho_numpy = self.tc_numpy._eval_basis_on_grid()
np.testing.assert_allclose(
np.asarray(phi_jax), rho_numpy,
rtol=1e-5, atol=1e-5,
err_msg="JAX and numpy basis evaluations don't match"
)
np.testing.assert_allclose(
np.asarray(grad_phi_jax), nabla_rho_numpy,
rtol=1e-5, atol=1e-5,
err_msg="JAX and numpy basis gradients don't match"
)
[docs]
def test_get_2b_against_numpy(self):
"""Test two-body term calculation against numpy version."""
# Compute two-body correction with explicit parameter passing
correction_jax = self.tc_jax.get_2b(self.params)
# Get standard ERI
eri1 = tc_helper.get_eri(self.mf)
# Combine to get full effective ERI
result_jax = eri1 + correction_jax
# Numpy version returns full effective ERI
result_numpy = self.tc_numpy.get_2b()
# Convert JAX array to numpy for comparison
result_jax = np.asarray(result_jax)
np.testing.assert_allclose(
result_jax, result_numpy,
rtol=1e-5, atol=1e-5,
err_msg="JAX and numpy two-body terms don't match"
)
[docs]
def test_get_eri_without_incore_cache(self):
"""Test AO-to-MO transformation when PySCF does not cache AO ERIs."""
expected = tc_helper.get_eri(self.mf)
self.mf._eri = None
actual = tc_helper.get_eri(self.mf)
np.testing.assert_allclose(actual, expected, rtol=1e-12, atol=1e-12)
[docs]
def test_mo_coeff_handling(self):
"""Test handling of molecular orbital coefficients."""
# Test with explicit mo_coeff
new_mo = self.mf.mo_coeff + 0.1
tc_jax_new = TC_jax.from_pyscf(self.mf, self.jastrow_jax, mo_coeff=new_mo)
tc_numpy_new = TC_numpy(self.mf, self.jastrow_numpy, mo_coeff=new_mo)
correction_jax = tc_jax_new.get_2b(self.params)
eri1 = tc_helper.get_eri(self.mf, mo_coeff=new_mo)
result_jax = eri1 + correction_jax
result_numpy = tc_numpy_new.get_2b()
np.testing.assert_allclose(
np.asarray(result_jax), result_numpy,
rtol=1e-5, atol=1e-5,
err_msg="Results don't match with explicit mo_coeff"
)
[docs]
def test_two_body_terms(self):
"""Test two-body term calculation."""
r1 = np.array([[0.0, 0.0, 0.0]])
r2 = np.array([[0.0, 0.0, 1.0]])
grad_jax = self.jastrow_jax.grad_r(r1, r2, self.params)[:, None, :]
grad_numpy = self.jastrow_numpy.grad(r1, r2)
np.testing.assert_allclose(
grad_jax, grad_numpy,
rtol=1e-5, atol=1e-5,
err_msg="Gradients don't match"
)
if __name__ == '__main__':
unittest.main()