"""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
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."""
mol = get_atom_molecule(atom_symbol)
bh_terms = sm7_coeffs_to_bh_terms(atom_symbol)
bh = BoysHandy.create(mol, terms_per_nucleus=[bh_terms])
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
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')
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
bh_value = bh._compute(r1, r2, bh_params)
# 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])
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')
key1, key2 = random.split(self.key)
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"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')
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')
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')
key1, key2 = random.split(self.key)
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 returns shape (1, 1, 3) for single points
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"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')
key1, key2 = random.split(self.key)
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"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')
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')
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)
# 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."""
jastrow = BoysHandy.create(self.mol, epsilon=1e-16)
params = jastrow.init_params()
self.assertIn('b_raw', params)
self.assertIn('d_raw', params)
self.assertIn('c_raw', params)
self.assertEqual(params['b_raw'].shape, (1,)) # 1 atom type
self.assertEqual(params['d_raw'].shape, (1,)) # 1 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)
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.
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)
self.assertEqual(grad1.shape, (3,))
self.assertEqual(grad2.shape, (3,))
self.assertEqual(lap1.shape, ())
self.assertEqual(lap2.shape, ())
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()