diff --git a/pyscf_ipu/experimental/structure.py b/pyscf_ipu/experimental/structure.py index 28c0316..4e894ff 100644 --- a/pyscf_ipu/experimental/structure.py +++ b/pyscf_ipu/experimental/structure.py @@ -76,3 +76,24 @@ def molecule(name: str): ) raise NotImplementedError(f"No structure registered for: {name}") + + +def nuclear_energy(structure: Structure) -> float: + """Nuclear electrostatic interaction energy + + Evaluated by taking sum over all unique pairs of atom centers: + + sum_{j > i} z_i z_j / |r_i - r_j| + + where z_i is the charge of the ith atom (the atomic number). + + Args: + structure (Structure): input structure + + Returns: + float: the total nuclear repulsion energy + """ + idx, jdx = np.triu_indices(structure.num_atoms, 1) + u = structure.atomic_number[idx] * structure.atomic_number[jdx] + rij = structure.position[idx, :] - structure.position[jdx, :] + return np.sum(u / np.linalg.norm(rij, axis=1)) diff --git a/test/test_interop.py b/test/test_interop.py index 164254a..8ed52cf 100644 --- a/test/test_interop.py +++ b/test/test_interop.py @@ -7,7 +7,7 @@ from pyscf_ipu.experimental.basis import basisset from pyscf_ipu.experimental.interop import to_pyscf from pyscf_ipu.experimental.mesh import electron_density, uniform_mesh -from pyscf_ipu.experimental.structure import molecule +from pyscf_ipu.experimental.structure import molecule, nuclear_energy @pytest.mark.parametrize("basis_name", ["sto-3g", "6-31g**"]) @@ -44,3 +44,11 @@ def test_gto(): actual = electron_density(basis, mesh, C) expect = eval_rho(mol, expect_ao, mf.make_rdm1(), "lda") assert_allclose(actual, expect, atol=1e-6) + + +@pytest.mark.parametrize("name", ["water", "h2"]) +def test_nuclear_energy(name): + mol = molecule(name) + actual = nuclear_energy(mol) + expect = to_pyscf(mol).energy_nuc() + assert_allclose(actual, expect)