From fdc30def900d9cb1fef04cc0c5e50048bc1b10dd Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Sun, 15 Oct 2023 19:03:10 +0000 Subject: [PATCH 1/2] Add function for computing the nuclear repulsion energy --- pyscf_ipu/experimental/structure.py | 21 +++++++++++++++++++++ test/test_interop.py | 10 +++++++++- 2 files changed, 30 insertions(+), 1 deletion(-) diff --git a/pyscf_ipu/experimental/structure.py b/pyscf_ipu/experimental/structure.py index 7977151..b800efe 100644 --- a/pyscf_ipu/experimental/structure.py +++ b/pyscf_ipu/experimental/structure.py @@ -73,3 +73,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) From 1f915056c8d4beb0f147d438cdb8d079e39fa2e0 Mon Sep 17 00:00:00 2001 From: Hatem Helal Date: Wed, 18 Oct 2023 11:55:56 +0000 Subject: [PATCH 2/2] explicit axis --- pyscf_ipu/experimental/structure.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyscf_ipu/experimental/structure.py b/pyscf_ipu/experimental/structure.py index b800efe..81411da 100644 --- a/pyscf_ipu/experimental/structure.py +++ b/pyscf_ipu/experimental/structure.py @@ -93,4 +93,4 @@ def nuclear_energy(structure: Structure) -> float: 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)) + return np.sum(u / np.linalg.norm(rij, axis=1))