"""Tests for Boys-Handy Jastrow implementation."""
import unittest
import numpy as np
import jax
import jax.numpy as jnp
from jax import random
from pyscf import gto
from pytc.jastrow.bh import BoysHandy, BHTerm
from pytc.legacy.jastrow.sm7 import SM7
# Enable float64 support
jax.config.update("jax_enable_x64", True)
[docs]
def get_h2_molecule(bond_length=1.4):
"""Create H2 molecule."""
mol = gto.M(
atom=f'H 0 0 0; H 0 0 {bond_length}',
basis='sto-3g',
unit='bohr'
)
return mol
[docs]
def get_atom_molecule(atom_symbol):
"""Create single atom molecule at origin."""
mol = gto.M(
atom=f'{atom_symbol} 0 0 0',
basis='sto-3g',
unit='bohr'
)
return mol
[docs]
def sm7_coeffs_to_bh_terms(atom_symbol):
"""Convert SM7 coefficients to BH terms.
SM7 uses coefficients indexed by (m, n, o) tuples.
BH uses BHTerm(m, n, o, c) where c is the coefficient.
Note: SM7 applies a factor of 0.5 when m == n, but BH handles
this internally via delta_factor, so we keep the full coefficient.
"""
sm7_coeffs = SM7._coeff_table[atom_symbol]
terms = []
for (m, n, o), coeff in sm7_coeffs.items():
terms.append(BHTerm(m, n, o, coeff))
return terms
[docs]
class TestBoysHandyVsSM7(unittest.TestCase):
"""Test Boys-Handy implementation against SM7 reference."""
[docs]
def setUp(self):
"""Set up test cases."""
self.key = random.PRNGKey(42)
def _setup_atom_comparison(self, atom_symbol):
"""Set up BH and SM7 for a given atom."""
# Create molecule with atom at origin
mol = get_atom_molecule(atom_symbol)
# Get SM7 coefficients and convert to BH terms
bh_terms = sm7_coeffs_to_bh_terms(atom_symbol)
# Create BH Jastrow with single nucleus, so terms_per_nucleus is a list with one element
# (It naturally has no cutoffs now as we removed them)
bh = BoysHandy.create(mol, terms_per_nucleus=[bh_terms])
# Initialize parameters
bh_params = bh.init_params(key=self.key)
# SM7 uses fixed scaling with b=1.0 and d=1.0
# We need to set BH parameters to match this for comparison
# softplus(x) = 1.0 => x = log(exp(1) - 1)
val_1 = 1.0
raw_val = float(np.log(np.exp(val_1) - 1.0))
bh_params['b_raw'] = jnp.ones_like(bh_params['b_raw']) * raw_val
bh_params['d_raw'] = jnp.ones_like(bh_params['d_raw']) * raw_val
# Create SM7 Jastrow
sm7 = SM7(atom=atom_symbol)
return bh, bh_params, sm7
[docs]
def test_he_function_evaluation(self):
"""Test He atom: compare BH and SM7 function values."""
bh, bh_params, sm7 = self._setup_atom_comparison('He')
# Generate random electron positions
key1, key2 = random.split(self.key)
r1 = random.normal(key1, (3,)) * 2.0 # Scale to ~[-2, 2] bohr
r2 = random.normal(key2, (3,)) * 2.0
# Compute BH value
bh_value = bh._compute(r1, r2, bh_params)
# Compute SM7 value
# SM7 expects numpy arrays; handle both scalar and array outputs
sm7_output = sm7(np.array(r1), np.array(r2))
sm7_value = float(np.atleast_1d(sm7_output).flat[0])
# Compare values
np.testing.assert_allclose(
float(bh_value), sm7_value, rtol=1e-5, atol=1e-8,
err_msg=f"BH vs SM7 mismatch for He at r1={r1}, r2={r2}"
)
[docs]
def test_be_function_evaluation(self):
"""Test Be atom: compare BH and SM7 function values."""
bh, bh_params, sm7 = self._setup_atom_comparison('Be')
# Generate random electron positions
key1, key2 = random.split(self.key)
r1 = random.normal(key1, (3,)) * 2.0
r2 = random.normal(key2, (3,)) * 2.0
# Compute BH value
bh_value = bh._compute(r1, r2, bh_params)
# Compute SM7 value
sm7_output = sm7(np.array(r1), np.array(r2))
sm7_value = float(np.atleast_1d(sm7_output).flat[0])
# Compare values
np.testing.assert_allclose(
float(bh_value), sm7_value, rtol=1e-5, atol=1e-8,
err_msg=f"BH vs SM7 mismatch for Be at r1={r1}, r2={r2}"
)
[docs]
def test_he_function_evaluation_multiple_points(self):
"""Test He atom with multiple random positions."""
bh, bh_params, sm7 = self._setup_atom_comparison('He')
# Test multiple random point pairs
n_tests = 10
for i in range(n_tests):
key_i = random.fold_in(self.key, i)
key1, key2 = random.split(key_i)
r1 = random.normal(key1, (3,)) * 2.0
r2 = random.normal(key2, (3,)) * 2.0
bh_value = bh._compute(r1, r2, bh_params)
sm7_output = sm7(np.array(r1), np.array(r2))
sm7_value = float(np.atleast_1d(sm7_output).flat[0])
np.testing.assert_allclose(
float(bh_value), sm7_value, rtol=1e-5, atol=1e-8,
err_msg=f"Test {i}: BH vs SM7 mismatch for He"
)
[docs]
def test_be_function_evaluation_multiple_points(self):
"""Test Be atom with multiple random positions."""
bh, bh_params, sm7 = self._setup_atom_comparison('Be')
# Test multiple random point pairs
n_tests = 10
for i in range(n_tests):
key_i = random.fold_in(self.key, i)
key1, key2 = random.split(key_i)
r1 = random.normal(key1, (3,)) * 2.0
r2 = random.normal(key2, (3,)) * 2.0
bh_value = bh._compute(r1, r2, bh_params)
sm7_output = sm7(np.array(r1), np.array(r2))
sm7_value = float(np.atleast_1d(sm7_output).flat[0])
np.testing.assert_allclose(
float(bh_value), sm7_value, rtol=1e-5, atol=1e-8,
err_msg=f"Test {i}: BH vs SM7 mismatch for Be"
)
[docs]
def test_he_gradient_r1(self):
"""Test He atom: compare gradients with respect to r1."""
bh, bh_params, sm7 = self._setup_atom_comparison('He')
# Generate random electron positions
key1, key2 = random.split(self.key)
r1 = random.normal(key1, (3,)) * 2.0
r2 = random.normal(key2, (3,)) * 2.0
# Compute BH gradient using JAX autodiff on _compute
bh_grad_fn = jax.grad(lambda r: bh._compute(r, r2, bh_params))
bh_grad = bh_grad_fn(r1)
# Compute SM7 gradient
# SM7.grad returns shape (1, 1, 3) for single points
sm7_grad = sm7.grad(np.array(r1), np.array(r2))[0, 0, :]
# Compare gradients
np.testing.assert_allclose(
np.array(bh_grad), sm7_grad, rtol=1e-4, atol=1e-7,
err_msg=f"BH vs SM7 gradient mismatch for He at r1={r1}, r2={r2}"
)
[docs]
def test_be_gradient_r1(self):
"""Test Be atom: compare gradients with respect to r1."""
bh, bh_params, sm7 = self._setup_atom_comparison('Be')
# Generate random electron positions
key1, key2 = random.split(self.key)
r1 = random.normal(key1, (3,)) * 2.0
r2 = random.normal(key2, (3,)) * 2.0
# Compute BH gradient using JAX autodiff
bh_grad_fn = jax.grad(lambda r: bh._compute(r, r2, bh_params))
bh_grad = bh_grad_fn(r1)
# Compute SM7 gradient
sm7_grad = sm7.grad(np.array(r1), np.array(r2))[0, 0, :]
# Compare gradients
np.testing.assert_allclose(
np.array(bh_grad), sm7_grad, rtol=1e-4, atol=1e-7,
err_msg=f"BH vs SM7 gradient mismatch for Be at r1={r1}, r2={r2}"
)
[docs]
def test_he_gradient_r1_multiple_points(self):
"""Test He atom gradients with multiple random positions."""
bh, bh_params, sm7 = self._setup_atom_comparison('He')
# Test multiple random point pairs
n_tests = 10
for i in range(n_tests):
key_i = random.fold_in(self.key, i)
key1, key2 = random.split(key_i)
r1 = random.normal(key1, (3,)) * 2.0
r2 = random.normal(key2, (3,)) * 2.0
bh_grad_fn = jax.grad(lambda r: bh._compute(r, r2, bh_params))
bh_grad = bh_grad_fn(r1)
sm7_grad = sm7.grad(np.array(r1), np.array(r2))[0, 0, :]
np.testing.assert_allclose(
np.array(bh_grad), sm7_grad, rtol=1e-4, atol=1e-7,
err_msg=f"Test {i}: BH vs SM7 gradient mismatch for He"
)
[docs]
def test_be_gradient_r1_multiple_points(self):
"""Test Be atom gradients with multiple random positions."""
bh, bh_params, sm7 = self._setup_atom_comparison('Be')
# Test multiple random point pairs
n_tests = 10
for i in range(n_tests):
key_i = random.fold_in(self.key, i)
key1, key2 = random.split(key_i)
r1 = random.normal(key1, (3,)) * 2.0
r2 = random.normal(key2, (3,)) * 2.0
bh_grad_fn = jax.grad(lambda r: bh._compute(r, r2, bh_params))
bh_grad = bh_grad_fn(r1)
sm7_grad = sm7.grad(np.array(r1), np.array(r2))[0, 0, :]
np.testing.assert_allclose(
np.array(bh_grad), sm7_grad, rtol=1e-4, atol=1e-7,
err_msg=f"Test {i}: BH vs SM7 gradient mismatch for Be"
)
[docs]
class TestBoysHandy(unittest.TestCase):
"""Test Boys-Handy Jastrow implementation."""
[docs]
def setUp(self):
self.mol = get_h2_molecule()
self.key = random.PRNGKey(0)
# Create Boys-Handy Jastrow with custom terms
# Use negative coefficients for e-n terms for expected decay behavior
# Now terms_per_nucleus should be per atom TYPE, not per atom
# H2 has only 1 atom type (H), so only 1 list of terms
terms = [
[BHTerm(0, 0, 1, 0.5), # e-e cusp term
BHTerm(1, 0, 0, -0.1), # e-n term (attractive)
BHTerm(2, 0, 0, -0.1)] # higher order term (attractive)
]
self.jastrow = BoysHandy.create(self.mol, terms_per_nucleus=terms, epsilon=1e-16)
self.params = self.jastrow.init_params(key=self.key)
[docs]
def test_init(self):
"""Test initialization."""
# Test default initialization
jastrow = BoysHandy.create(self.mol, epsilon=1e-16)
params = jastrow.init_params()
# Check parameter structure
self.assertIn('b_raw', params)
self.assertIn('d_raw', params)
self.assertIn('c_raw', params)
# Check shapes - H2 has 2 atoms but only 1 atom type (H)
self.assertEqual(params['b_raw'].shape, (1,)) # 1 atom type
self.assertEqual(params['d_raw'].shape, (1,)) # 1 atom type
# Default has 17 terms per atom type
self.assertEqual(params['c_raw'].shape[0], 1) # 1 atom type
self.assertEqual(params['c_raw'].shape[1], 17) # 17 default terms
[docs]
def test_parameter_handling(self):
"""Test parameter flattening/unflattening."""
flat_params = self.jastrow.flatten_params(self.params)
restored_params = self.jastrow.unflatten_params(flat_params)
# Check that parameters are correctly restored
for key in self.params:
np.testing.assert_allclose(self.params[key], restored_params[key])
[docs]
def test_electron_symmetry(self):
"""Test symmetry with respect to electron exchange."""
r1 = jnp.array([0., 0., 0.])
r2 = jnp.array([0., 0., 1.0])
value1 = self.jastrow._compute(r1, r2, self.params)
value2 = self.jastrow._compute(r2, r1, self.params)
np.testing.assert_allclose(value1, value2, rtol=1e-7)
[docs]
def test_electron_cusp(self):
"""Test electron-electron cusp condition."""
# Shift electrons far from nuclei (at origin) to minimize e-n slope interference
r1 = jnp.array([10.0, 10.0, 0.0])
eps = 1e-3
r2_a = r1 + jnp.array([eps, 0.0, 0.0])
r2_b = r1 + jnp.array([2.0*eps, 0.0, 0.0])
ua = float(self.jastrow._compute(r1, r2_a, self.params))
ub = float(self.jastrow._compute(r1, r2_b, self.params))
du = (ub - ua) / eps
np.testing.assert_allclose(
du, 0.5, rtol=0.03,
err_msg=f"EE cusp should be exactly 0.5, got {du:.6f}"
)
[docs]
def test_electron_cusp_h2o(self):
"""Test electron-electron cusp condition for a molecule with multiple atom types."""
from pyscf import gto
mol = gto.M(
atom='O 0 0 0; H 0 1 1; H 0 -1 1',
basis='sto-3g',
unit='bohr'
)
jastrow = BoysHandy.create(mol, epsilon=1e-16)
params = jastrow.init_params()
# Shift electrons far from nuclei to minimize e-n slope interference
r1 = jnp.array([10.0, 10.0, 0.0])
eps = 1e-6
r2_a = r1 + jnp.array([eps, 0.0, 0.0])
r2_b = r1 + jnp.array([2.0*eps, 0.0, 0.0])
ua = float(jastrow._compute(r1, r2_a, params))
ub = float(jastrow._compute(r1, r2_b, params))
du = (ub - ua) / eps
np.testing.assert_allclose(
du, 0.5, rtol=0.03,
err_msg=f"EE cusp should be exactly 0.5 for H2O, got {du:.6f}"
)
[docs]
def test_nuclear_decay(self):
"""Test decay of correlation with nuclear distance."""
r1_near = jnp.array([0., 0., 0.1])
r1_far = jnp.array([0., 0., 5.0])
r2 = jnp.array([0., 0., -0.5])
value_near = self.jastrow._compute(r1_near, r2, self.params)
value_far = self.jastrow._compute(r1_far, r2, self.params)
# With attractive e-n terms (negative coefficients),
# the Jastrow value can become more negative at larger distances
# depending on the interplay of e-n and e-e terms.
# The key physical requirement is that the function remains finite
# and well-behaved at all distances.
# Let's just check that both values are finite
self.assertTrue(jnp.isfinite(value_near), "Value at near distance should be finite")
self.assertTrue(jnp.isfinite(value_far), "Value at far distance should be finite")
[docs]
def test_gradient_computation(self):
"""Test gradient and laplacian computation."""
r1 = jnp.array([0., 0., 0.5])
r2 = jnp.array([0., 0., -0.5])
grad1, lap1 = self.jastrow.get_log_grads_r1(r1, r2, self.params)
grad2, lap2 = self.jastrow.get_log_grads_r2(r1, r2, self.params)
# Check shapes
self.assertEqual(grad1.shape, (3,))
self.assertEqual(grad2.shape, (3,))
self.assertEqual(lap1.shape, ())
self.assertEqual(lap2.shape, ())
# Check finite values
self.assertTrue(jnp.all(jnp.isfinite(grad1)))
self.assertTrue(jnp.all(jnp.isfinite(grad2)))
self.assertTrue(jnp.isfinite(lap1))
self.assertTrue(jnp.isfinite(lap2))
if __name__ == '__main__':
unittest.main()