Add function for computing the nuclear repulsion energy (#127)

* Add function for computing the nuclear repulsion energy

* explicit axis

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))

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)