import unittest
import jax
import jax.numpy as jnp
from jax import random
import numpy as np
from pyscf import gto, scf
import time
import psutil
import os
from pytc.ansatz.sj import SlaterJastrow
from pytc.ansatz.det import SlaterDet
from pytc.jastrow import REXP, CompositeJastrow, NuclearCusp, BoysHandy, Poly, NeuralEEN
from pytc.vmc.walker import initialize_walker_state, initialize_walkers
from pytc.vmc.hamiltonian import (
compute_jastrow_terms,
compute_potential_matrix,
compute_single_walker_energy,
eval_local_energy
)
from pytc.vmc.loss import make_energy_loss, make_variance_loss
[docs]
class TestHamiltonian(unittest.TestCase):
[docs]
def setUp(self):
self.mol = gto.Mole()
self.mol.atom = 'H 0 0 0; H 0 0 1.4'
self.mol.unit = 'Bohr'
self.mol.basis = 'sto-3g'
self.mol.build()
mf = scf.RHF(self.mol)
mf.kernel()
self.hf_energy = mf.e_tot
det = SlaterDet.create(self.mol, mf.mo_coeff)
jastrow = Poly()
self.ansatz = SlaterJastrow.create(self.mol, jastrow, [det])
self.jastrow_params = jnp.zeros(1)
self.linear_coeffs = jnp.array([1.0])
self.params = (self.jastrow_params, self.linear_coeffs)
key = random.PRNGKey(42)
self.n_electrons = self.mol.nelectron
self.positions = random.normal(key, (1, self.n_electrons, 3))
self.walker = initialize_walker_state(self.ansatz, self.positions)
# Update walker with determinant values (required for energy calc)
# self.walker is batched (1, n, 3), so we need vmapped ansatz
batch_ansatz = jax.vmap(self.ansatz, in_axes=(0, None))
_, self.walker = batch_ansatz(self.walker, self.params)
# We need a single walker for the functions in hamiltonian.py,
# initialize_walker_state returns batched walker (n_walkers=1 here)
# Extract single walker data for functions that expect single walker input
# Walker dataclass fields are batched, so we index [0]
from pytc.vmc.walker import Walker
self.single_walker = Walker(
positions=self.walker.positions[0],
slater_up=self.walker.slater_up[0],
slater_down=self.walker.slater_down[0],
inv_up=self.walker.inv_up[0],
inv_down=self.walker.inv_down[0],
grad_up=self.walker.grad_up[0],
grad_down=self.walker.grad_down[0],
lap_up=self.walker.lap_up[0],
lap_down=self.walker.lap_down[0],
det_up=self.walker.det_up[0],
det_down=self.walker.det_down[0],
move_mask=self.walker.move_mask[0],
log_psi=self.walker.log_psi[0],
psi_sign=self.walker.psi_sign[0],
log_jastrow=self.walker.log_jastrow[0],
)
[docs]
def test_compute_jastrow_terms_shape(self):
"""Test output shapes of compute_jastrow_terms."""
grad_J_over_J, lap_J_over_J = compute_jastrow_terms(
self.ansatz, self.single_walker.positions, self.jastrow_params
)
self.assertEqual(grad_J_over_J.shape, (self.n_electrons, 3))
self.assertEqual(lap_J_over_J.shape, (self.n_electrons,))
[docs]
def test_compute_jastrow_terms_zero_params(self):
"""Test that zero Jastrow params result in zero gradients/laplacians."""
# Poly jastrow with zero params should be 1 (log is 0)
grad_J_over_J, lap_J_over_J = compute_jastrow_terms(
self.ansatz, self.single_walker.positions, jnp.zeros(1)
)
np.testing.assert_allclose(grad_J_over_J, 0.0, atol=1e-10)
np.testing.assert_allclose(lap_J_over_J, 0.0, atol=1e-10)
[docs]
def test_compute_potential_matrix_shape(self):
"""Test output shapes of compute_potential_matrix."""
B_alpha, B_beta = compute_potential_matrix(
self.ansatz,
self.single_walker.positions,
self.single_walker.slater_up,
self.single_walker.slater_down
)
n_alpha = self.ansatz.n_alpha
n_beta = self.ansatz.n_beta
self.assertEqual(B_alpha.shape, (n_alpha, n_alpha))
self.assertEqual(B_beta.shape, (n_beta, n_beta))
[docs]
def test_compute_single_walker_energy(self):
"""Test compute_single_walker_energy returns a scalar."""
energy = compute_single_walker_energy(
self.ansatz, self.single_walker, self.jastrow_params
)
self.assertEqual(energy.shape, ())
self.assertTrue(jnp.isfinite(energy))
[docs]
def test_eval_local_energy(self):
"""Test eval_local_energy wrapper."""
energy, walker = eval_local_energy(
self.ansatz, self.single_walker, self.params
)
self.assertEqual(energy.shape, ())
np.testing.assert_array_equal(walker.positions, self.single_walker.positions)
[docs]
def test_potential_matrix_values(self):
"""Test potential matrix values for a simple H2 case."""
# H2 at 0 and 1.4
# Place one electron at 0.5 (near first H) and one at 0.9 (near second H)
# 1D along z-axis for simplicity in manual check, but coords are 3D
pos = jnp.array([[0.0, 0.0, 0.5], [0.0, 0.0, 0.9]])
# Manually compute potentials
# Nuclei at (0,0,0) and (0,0,1.4) with charge 1
r1 = pos[0] # (0,0,0.5)
r2 = pos[1] # (0,0,0.9)
d1_n1 = 0.5
d1_n2 = 1.4 - 0.5 # 0.9
v_en_1 = -1/d1_n1 - 1/d1_n2 # -2 - 1.111... = -3.111...
d2_n1 = 0.9
d2_n2 = 1.4 - 0.9 # 0.5
v_en_2 = -1/d2_n1 - 1/d2_n2 # -1.111... - 2 = -3.111...
r12 = jnp.linalg.norm(r1 - r2) # 0.4
v_ee = 1.0 / r12 # 2.5
# Total potential for electron 1 (including half of ee)
pot1 = v_en_1 + 0.5 * v_ee
# Total potential for electron 2 (including half of ee)
pot2 = v_en_2 + 0.5 * v_ee
# In the code, B matrix includes potential * slater_matrix
# But we can check if B / slater matches potential if we use identity slater or just check scaling
# Let's mock slater matrices as identity to easily extract potential
slater_up = jnp.eye(1) # 1 alpha electron
slater_down = jnp.eye(1) # 1 beta electron
# H2 has 2 electrons. In this setup n_alpha=1, n_beta=1.
# compute_potential_matrix expects full coords
B_alpha, B_beta = compute_potential_matrix(
self.ansatz, pos, slater_up, slater_down
)
self.assertTrue(jnp.allclose(B_alpha[0,0], pot1, rtol=1e-4))
self.assertTrue(jnp.allclose(B_beta[0,0], pot2, rtol=1e-4))
[docs]
def test_be_hf_energy(self):
"""Test that Be atom with zero Jastrow gives HF energy."""
mol = gto.Mole()
mol.atom = 'Be 0 0 0'
mol.unit = 'Bohr' # Use Bohr to avoid unit issues
mol.basis = 'sto-3g'
mol.build()
mf = scf.RHF(mol)
mf.kernel()
hf_energy = mf.e_tot
print(f"Be HF Energy: {hf_energy}")
det = SlaterDet.create(mol, mf.mo_coeff)
jastrow = Poly() # Zero params = identity
ansatz = SlaterJastrow.create(mol, jastrow, [det])
key = random.PRNGKey(123)
walker_batch = initialize_walkers(ansatz, 1, key=key)
jastrow_params = jnp.zeros(1)
params = (jastrow_params, jnp.array([1.0]))
# walker_batch is batched, need vmap
batch_ansatz = jax.vmap(ansatz, in_axes=(0, None))
_, walker_batch = batch_ansatz(walker_batch, params)
from pytc.vmc.walker import Walker
single_walker = Walker(
positions=walker_batch.positions[0],
slater_up=walker_batch.slater_up[0],
slater_down=walker_batch.slater_down[0],
inv_up=walker_batch.inv_up[0],
inv_down=walker_batch.inv_down[0],
grad_up=walker_batch.grad_up[0],
grad_down=walker_batch.grad_down[0],
lap_up=walker_batch.lap_up[0],
lap_down=walker_batch.lap_down[0],
det_up=walker_batch.det_up[0],
det_down=walker_batch.det_down[0],
move_mask=walker_batch.move_mask[0],
log_psi=walker_batch.log_psi[0],
psi_sign=walker_batch.psi_sign[0],
log_jastrow=walker_batch.log_jastrow[0],
)
energy = compute_single_walker_energy(ansatz, single_walker, jastrow_params)
print(f"Be Local Energy at random config: {energy}")
self.assertTrue(jnp.isfinite(energy))
self.assertTrue(energy > -50.0 and energy < -5.0)
[docs]
class TestMemoryUsage(unittest.TestCase):
[docs]
def setUp(self):
self.mol = gto.Mole()
# Use Be atom instead of Benzene for speed/memory test
self.mol.atom = 'Be 0 0 0'
self.mol.basis = 'sto-3g'
self.mol.build()
mf = scf.RHF(self.mol)
mf.kernel()
det = SlaterDet.create(self.mol, mf.mo_coeff)
jastrow = Poly()
self.ansatz = SlaterJastrow.create(self.mol, jastrow, [det])
self.jastrow_params = jnp.zeros(1)
self.linear_coeffs = jnp.array([1.0])
self.params = (self.jastrow_params, self.linear_coeffs)
self.n_electrons = self.mol.nelectron
[docs]
def test_memory_scaling(self):
"""Test memory usage with increasing number of walkers."""
process = psutil.Process(os.getpid())
initial_memory = process.memory_info().rss / 1024 / 1024 # MB
dummy_pos = jnp.zeros((self.n_electrons, 3))
dummy_walker = initialize_walker_state(self.ansatz, dummy_pos[None, ...])
from pytc.vmc.hamiltonian import compute_single_walker_energy
vmapped_energy = jax.vmap(
lambda w: compute_single_walker_energy(self.ansatz, w, self.jastrow_params)
)
vmapped_energy(dummy_walker).block_until_ready()
post_compile_memory = process.memory_info().rss / 1024 / 1024
print(f"Memory after compilation: {post_compile_memory:.2f} MB")
n_walkers = 100
key = random.PRNGKey(123)
positions = random.normal(key, (n_walkers, self.n_electrons, 3))
walkers = initialize_walker_state(self.ansatz, positions)
start_mem = process.memory_info().rss / 1024 / 1024
energies = vmapped_energy(walkers)
energies.block_until_ready()
end_mem = process.memory_info().rss / 1024 / 1024
peak_mem_increase = end_mem - start_mem
print(f"Memory increase for {n_walkers} walkers (Benzene): {peak_mem_increase:.2f} MB")
# Assert memory increase is less than 1.5 GB for this workload
self.assertLess(peak_mem_increase, 1500, "Memory usage seems excessive (>1.5GB)")
[docs]
class TestHamiltonianGrad(unittest.TestCase):
[docs]
def setUp(self):
self.mol = gto.Mole()
self.mol.atom = 'Be 0 0 0'
self.mol.basis = 'sto-3g'
self.mol.build()
mf = scf.RHF(self.mol)
mf.kernel()
det = SlaterDet.create(self.mol, mf.mo_coeff)
ncusp = NuclearCusp.create(self.mol)
bh_new = BoysHandy.create(self.mol)
jastrow = CompositeJastrow.create([ncusp, bh_new])
self.ansatz = SlaterJastrow.create(self.mol, jastrow, [det])
self.jastrow_params = jastrow.init_params()
self.n_electrons = self.mol.nelectron
[docs]
class TestJastrowTermsPolymorphicDispatch(unittest.TestCase):
"""BoysHandyAnalytical's get_pair_grid_grad_lap override (the fast
whole-electron-set path) must give bit-identical full-E_L results to
the Jastrow base class's default per-pair implementation (inherited
unchanged by generic BoysHandy), on multi-type systems with the same
CompositeJastrow (NuclearCusp + BH) structure production uses. There
is no flag -- class choice (BoysHandyAnalytical vs BoysHandy) is the
only dispatch. See pytc/jastrow/test/test_bha.py for
the lower-level, same-instance base-vs-override check.
"""
def _build(self, mol, jastrow_cls):
mf = scf.RHF(mol).density_fit()
mf.kernel()
det = SlaterDet.create(mol, mf.mo_coeff)
bh = jastrow_cls.create(mol)
ncusp = NuclearCusp.create(mol, name="ncusp")
jastrow = CompositeJastrow.create([ncusp, bh])
params = jastrow.init_params()
ansatz = SlaterJastrow.create(mol, jastrow, [det])
return ansatz, params
def _check(self, mol, n_walkers=4, seed=11):
from pytc.jastrow.bha import BoysHandyAnalytical
ansatz_base, params_base = self._build(mol, BoysHandy)
ansatz_override, params_override = self._build(mol, BoysHandyAnalytical)
key = random.PRNGKey(seed)
walkers = initialize_walkers(ansatz_base, n_walkers, key=key)
for w in range(n_walkers):
walker_w = jax.tree_util.tree_map(lambda x: x[w], walkers)
e_base = compute_single_walker_energy(ansatz_base, walker_w, params_base)
e_override = compute_single_walker_energy(
ansatz_override, walker_w, params_override)
# Tight but not bit-exact: the two paths use different op
# orderings, which coincide only for special parameter values.
np.testing.assert_allclose(
float(e_override), float(e_base), rtol=1e-12, atol=1e-12,
err_msg=f"walker {w}: BoysHandyAnalytical E_L != generic BoysHandy E_L")
[docs]
def test_lih(self):
mol = gto.M(atom="Li 0 0 0; H 0 0 1.6", basis="sto-3g", unit="Bohr", verbose=0)
self._check(mol)
[docs]
def test_h2o(self):
mol = gto.M(atom="O 0 0 0; H 0 -1.4 1.1; H 0 1.4 1.1",
basis="sto-3g", unit="Bohr", verbose=0)
self._check(mol)
[docs]
def test_water_dimer(self):
mol = gto.M(atom="""
O 0.000 0.000 0.000
H 0.757 0.586 0.000
H -0.757 0.586 0.000
O 0.000 0.000 2.900
H 0.757 -0.586 2.900
H -0.757 -0.586 2.900
""", basis="cc-pVDZ", unit="Angstrom", verbose=0)
self._check(mol)
[docs]
class TestJastrowTermsCompositeLengthMismatch(unittest.TestCase):
"""A CompositeJastrow/params length mismatch must raise, not silently
truncate via zip() and return a wrong (partial) energy."""
if __name__ == "__main__":
unittest.main()