"""Tests for DTN (Drummond-Towler-Needs) 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.dtn import DTN, DTNTermEE, DTNTermEN, DTNTermEEN
# 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]
class TestDTN(unittest.TestCase):
"""Test DTN Jastrow implementation."""
[docs]
def setUp(self):
self.mol = get_h2_molecule()
self.key = random.PRNGKey(0)
ee = [DTNTermEE(1, 0.5), DTNTermEE(2, -0.1)]
en = [[DTNTermEN(2, -0.1)]]
een = [[DTNTermEEN(2, 0, 2, 0.05)]]
self.jastrow = DTN.create(
self.mol, ee_terms=ee, en_terms=en, een_terms=een
)
self.params = self.jastrow.init_params(key=self.key)
[docs]
def test_init(self):
"""Test initialization with default terms."""
jastrow = DTN.create(self.mol)
params = jastrow.init_params()
# Check parameter structure
self.assertIn('rc_en_raw', params)
self.assertIn('rc_ee_raw', params)
self.assertIn('c_ee_raw', params)
self.assertIn('c_en_raw', params)
self.assertIn('c_een_raw', params)
self.assertNotIn('b_raw', params)
self.assertNotIn('d_raw', params)
# H2 has 1 atom type
self.assertEqual(params['rc_en_raw'].shape, (1,))
self.assertEqual(params['rc_ee_raw'].shape, (1,)) # Global
self.assertEqual(params['c_ee_raw'].shape, (9,)) # 9 default EE terms (global)
self.assertEqual(params['c_en_raw'].shape, (1, 8)) # 8 default EN terms
self.assertEqual(params['c_een_raw'].shape, (1, 13)) # 13 cusp-safe EEN terms
[docs]
def test_parameter_handling(self):
"""Test parameter flattening/unflattening roundtrip."""
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_decay_to_zero(self):
"""Test that DTN decays to zero beyond cutoff."""
r1 = jnp.array([100., 100., 100.])
r2 = jnp.array([-100., -100., -100.])
value = self.jastrow._compute(r1, r2, self.params)
np.testing.assert_allclose(
float(value), 0.0, atol=1e-10,
err_msg="DTN should decay to zero far beyond cutoff"
)
[docs]
def test_nonzero_inside_cutoff(self):
"""Test that DTN gives non-zero values inside the cutoff radius."""
r1 = jnp.array([0., 0., 0.1])
r2 = jnp.array([0., 0., -0.5])
value = self.jastrow._compute(r1, r2, self.params)
self.assertNotEqual(float(value), 0.0)
self.assertTrue(jnp.isfinite(value))
[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))
[docs]
def test_gradient_vs_finite_diff(self):
"""Test autodiff gradients against numerical finite differences."""
r1 = jnp.array([0.3, -0.2, 0.5])
r2 = jnp.array([-0.1, 0.4, -0.3])
grad_fn = jax.grad(lambda r: self.jastrow._compute(r, r2, self.params))
autodiff_grad = grad_fn(r1)
eps = 1e-5
fd_grad = np.zeros(3)
for i in range(3):
r1_plus = r1.at[i].add(eps)
r1_minus = r1.at[i].add(-eps)
f_plus = float(self.jastrow._compute(r1_plus, r2, self.params))
f_minus = float(self.jastrow._compute(r1_minus, r2, self.params))
fd_grad[i] = (f_plus - f_minus) / (2 * eps)
np.testing.assert_allclose(
np.array(autodiff_grad), fd_grad, rtol=1e-4, atol=1e-8
)
[docs]
def test_laplacian_vs_finite_diff(self):
"""Test autodiff Laplacian against numerical finite differences."""
r1 = jnp.array([0.3, -0.2, 0.5])
r2 = jnp.array([-0.1, 0.4, -0.3])
_, lap = self.jastrow.get_log_grads_r1(r1, r2, self.params)
eps = 1e-4
fd_lap = 0.0
f0 = float(self.jastrow._compute(r1, r2, self.params))
for i in range(3):
r1_plus = r1.at[i].add(eps)
r1_minus = r1.at[i].add(-eps)
f_plus = float(self.jastrow._compute(r1_plus, r2, self.params))
f_minus = float(self.jastrow._compute(r1_minus, r2, self.params))
fd_lap += (f_plus - 2 * f0 + f_minus) / (eps ** 2)
np.testing.assert_allclose(float(lap), fd_lap, rtol=1e-3, atol=1e-7)
[docs]
def test_parameter_gradient_nan_check(self):
"""Test that parameter gradients have no NaNs beyond cutoff."""
r1 = jnp.array([10.0, 10.0, 10.0])
r2 = jnp.array([-10.0, -10.0, -10.0])
def energy_fn(params):
g, l = self.jastrow.get_log_grads_r1(r1, r2, params)
return jnp.sum(g) + l
grads = jax.grad(energy_fn)(self.params)
for key in grads:
self.assertFalse(
jnp.any(jnp.isnan(grads[key])),
f"NaN found in parameter gradient for {key}"
)
[docs]
def test_multi_atom_type(self):
"""Test with a molecule having multiple atom types (LiH)."""
mol = gto.M(atom='Li 0 0 0; H 0 0 3.0', basis='sto-3g', unit='bohr')
ee = [DTNTermEE(1, 0.5)]
en = [[DTNTermEN(2, 0.1)], [DTNTermEN(2, 0.1)]]
een = [[DTNTermEEN(2, 0, 2, 0.1)], [DTNTermEEN(2, 0, 2, 0.1)]]
jastrow = DTN.create(mol, ee_terms=ee, en_terms=en, een_terms=een)
params = jastrow.init_params()
self.assertEqual(jastrow.n_types, 2)
self.assertEqual(params['rc_en_raw'].shape, (2,))
self.assertEqual(params['rc_ee_raw'].shape, (1,))
self.assertEqual(params['c_ee_raw'].shape, (1,))
r1 = jnp.array([0., 0., 1.0])
r2 = jnp.array([0., 0., 2.0])
value = jastrow._compute(r1, r2, params)
self.assertTrue(jnp.isfinite(value))
[docs]
class TestDTNCuspPreservation(unittest.TestCase):
"""Test that the e-e cusp is exactly 0.5 and preserved during optimization.
With separated EE/EN/EEN channels, the e-e cusp comes from the
pure EE term: c_1 * r12 * C_ee(r12). Since C_ee(0) = 1, the
cusp slope du/dr12|_{r12→0} = c_1 = 0.5.
For any molecule, this is exactly 0.5 because EE is a global term.
Note: _safe_norm uses epsilon=1e-8, so finite differences must use
separations >> sqrt(epsilon) ≈ 1e-4.
"""
def _cusp_slope(self, jastrow, params, r1, eps=1e-3):
"""Numerical du/dr12 at small but finite r12 separation."""
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))
return (ub - ua) / eps
[docs]
def test_cusp_exactly_half(self):
"""Verify du/dr12 ≈ 0.5 for a multi-atom system.
With EE evaluated globally, the slope should be exactly 0.5
for H2 as well (not 1.0).
"""
mol = gto.M(atom='H 0 0 -1; H 0 0 1', basis='sto-3g', unit='bohr')
# Only EE terms to isolate the cusp
ee = [DTNTermEE(1, 0.5)]
jastrow = DTN.create(mol, ee_terms=ee, en_terms=[[], []], een_terms=[[], []])
params = jastrow.init_params()
r1 = jnp.array([0.3, 0.0, 0.0])
du = self._cusp_slope(jastrow, params, r1)
np.testing.assert_allclose(
du, 0.5, rtol=0.01,
err_msg=f"EE cusp should be exactly 0.5 for H2, got {du:.6f}"
)
[docs]
def test_cusp_preserved_after_optimize_ref_var(self):
"""Run optimize_ref_var for a few steps, verify cusp invariance.
Evaluates the actual c_ee array used in _compute_forward to verify
the cusp mask correctly clamps the coefficient to 0.5 even if the
optimizer updates c_ee_raw.
"""
from pyscf import scf
from pytc.vmc import optimize_ref_var
from pytc.ansatz.sj import SlaterJastrow
from pytc.ansatz.det import SlaterDet
mol = gto.M(atom='He 0 0 0', basis='sto-3g', unit='bohr')
mf = scf.RHF(mol)
mf.kernel()
det = SlaterDet.create(mol, mf.mo_coeff)
# Mix of terms
ee = [DTNTermEE(1, 0.5), DTNTermEE(2, 0.1)]
en = [[DTNTermEN(2, 0.1)]]
een = [[DTNTermEEN(2, 0, 2, 0.1)]]
jastrow = DTN.create(mol, ee_terms=ee, en_terms=en, een_terms=een)
ansatz = SlaterJastrow.create(mol, jastrow, [det])
params_init = jastrow.init_params()
linear_coeffs = jnp.ones(1)
# Run a few optimization steps
key = random.PRNGKey(123)
opt_results = optimize_ref_var(
ansatz,
params=[params_init, linear_coeffs],
n_walkers=200,
n_steps=5,
step_size=0.1,
burn_in_steps=100,
n_opt_steps=3,
optimizer_type='adam',
learning_rate=0.01,
key=key,
)
optimized_params = opt_results['params'][-1][0]
# Verify raw parameters changed
for key_name in ['c_ee_raw', 'c_en_raw', 'c_een_raw']:
if optimized_params[key_name].size > 0:
changed = not np.allclose(
params_init[key_name], optimized_params[key_name], atol=1e-10
)
if changed:
break
self.assertTrue(changed, "At least some parameters should change")
# The cusp coefficient is the first EE term (index 0).
# It must be exactly 0.5 because of the Jastrow masking logic.
# Note: Because of jnp.where, the gradient w.r.t this raw parameter is 0,
# so it stays exactly at its initialized value.
actual_c_ee = jnp.where(jastrow._ee_cusp_mask, 0.5, optimized_params['c_ee_raw'])
# Check that the evaluated parameter is strictly clamped to 0.5
clamped_cusp_val = actual_c_ee[0]
np.testing.assert_allclose(
float(clamped_cusp_val), 0.5, atol=1e-12,
err_msg="The effective cusp coefficient must be strictly 0.5"
)
if __name__ == '__main__':
unittest.main()