Refactor ForceField class to facilitate future extensions

Author: tovrstra

CHANGELOG.md

@@ -9,7 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 
-- Refactored neighborlist API, to prepare for more efficient implementations.
+- Refactor `ForceField` class to facilitate future extensions.
+- Refactor neighborlist API, to prepare for more efficient implementations.
 
 
 ## [0.2.2] - 2024-10-09

README.md

@@ -58,15 +58,15 @@ This is the part that you are expected to write.
 The evaluation of the force field energy and its derivatives requires the following:
 
 ```python
-from tinyff.forcefield import CutoffWrapper, LennardJones, PairwiseForceField
+from tinyff.forcefield import CutoffWrapper, LennardJones, ForceField
 from tinyff.neighborlist import NBuildSimple
 
 # Define a pairwise potential, with energy and force shift
 rcut = 5.0
 lj = CutOffWrapper(LennardJones(2.5, 2.0), rcut)
 
 # Define a force field
-pwff = PairwiseForceField(lj, NBuildSimple(rcut))
+ff = ForceField([lj], NBuildSimple(rcut))
 
 # You need atomic positions and the length of a periodic cell edge.
 # The following line defines just two atomic positions.
@@ -80,11 +80,11 @@ cell_length = 20.0
 # - An array with Cartesian forces, same shape as `atpos`.
 # - The force contribution the pressure
 #   (often the written as the second term in the virial pressure).
-potential_energy, forces, frc_pressure = pwff(atpos, cell_length)
+potential_energy, forces, frc_pressure = ff(atpos, cell_length)
 ```
 
 This basic recipe can be extended by passing additional options
-into the `PairwiseForceField` constructor:
+into the `ForceField` constructor:
 
 - Linear-scaling neighbor lists with the
   [cell lists](https://en.wikipedia.org/wiki/Cell_lists) method:
@@ -93,7 +93,7 @@ into the `PairwiseForceField` constructor:
     from tinyff.neighborlist import NBuildCellLists
 
     # Construct your force field object as follows:
-    pwff = PairwiseForceField(lj, NBuildCellLists(rcut))
+    ff = ForceField([lj], NBuildCellLists(rcut))
     ```
 
     Note that the current cell lists implementation is not very efficient (yet),
@@ -104,7 +104,7 @@ into the `PairwiseForceField` constructor:
 
     ```python
     rmax = 6.0  # > rcut, so buffer of 1.0
-    pwff = PairwiseForceField(lj, NBuildSimple(rmax, nlist_reuse=16))
+    ff = ForceField([lj], NBuildSimple(rmax, nlist_reuse=16))
     ```
 
 

src/tinyff/atomsmithy.py

@@ -23,7 +23,7 @@
 from numpy.typing import ArrayLike, NDArray
 from scipy.optimize import minimize
 
-from .forcefield import PairPotential, PairwiseForceField
+from .forcefield import ForceField, PairPotential
 from .neighborlist import NBuild, NBuildSimple
 
 __all__ = (
@@ -86,7 +86,7 @@ class PushPotential(PairPotential):
 
     rcut: float = attrs.field(converter=float, validator=attrs.validators.gt(0))
 
-    def __call__(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
+    def compute(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
         """Compute pair potential energy and its derivative towards distance."""
         dist = np.asarray(dist, dtype=float)
         x = dist / self.rcut
@@ -113,11 +113,11 @@ def build_random_cell(
     # Define cost function to push the atoms appart.
     if nbuild is None:
         nbuild = NBuildSimple(rcut)
-    pwff = PairwiseForceField(PushPotential(rcut), nbuild)
+    ff = ForceField([PushPotential(rcut)], nbuild)
 
     def costgrad(atpos_raveled):
         atpos = atpos_raveled.reshape(-1, 3)
-        energy, force, _ = pwff(atpos, cell_length)
+        energy, force, _ = ff(atpos, cell_length)
         return energy, -force.ravel()
 
     # Optimize and return structure

src/tinyff/forcefield.py

@@ -22,14 +22,25 @@
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 
-from .neighborlist import NBuild
+from .neighborlist import NLIST_DTYPE, NBuild
 
-__all__ = ("PairwiseForceField",)
+__all__ = ("ForceTerm", "PairPotential", "LennardJones", "CutOffWrapper", "ForceField")
+
+
+@attrs.define
+class ForceTerm:
+    def __call__(self, nlist: NDArray[NLIST_DTYPE]):
+        raise NotImplementedError  # pragma: nocover
 
 
 @attrs.define
 class PairPotential:
-    def __call__(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
+    def __call__(self, nlist: NDArray[NLIST_DTYPE]):
+        energy, gdist = self.compute(nlist["dist"])
+        nlist["energy"] += energy
+        nlist["gdist"] += gdist
+
+    def compute(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
         """Compute pair potential energy and its derivative towards distance."""
         raise NotImplementedError  # pragma: nocover
 
@@ -39,7 +50,7 @@ class LennardJones(PairPotential):
     epsilon: float = attrs.field(converter=float)
     sigma: float = attrs.field(converter=float)
 
-    def __call__(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
+    def compute(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
         """Compute pair potential energy and its derivative towards distance."""
         dist = np.asarray(dist, dtype=float)
         x = self.sigma / dist
@@ -57,23 +68,23 @@ class CutOffWrapper(PairPotential):
 
     def __attrs_post_init__(self):
         """Post initialization changes."""
-        self.ecut, self.gcut = self.original(self.rcut)
+        self.ecut, self.gcut = self.original.compute(self.rcut)
 
-    def __call__(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
+    def compute(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
         """Compute pair potential energy and its derivative towards distance."""
         dist = np.asarray(dist, dtype=float)
         mask = dist < self.rcut
         if mask.ndim == 0:
             # Deal with non-array case
             if mask:
-                energy, gdist = self.original(dist)
+                energy, gdist = self.original.compute(dist)
                 energy -= self.ecut + self.gcut * (dist - self.rcut)
                 gdist -= self.gcut
             else:
                 energy = 0.0
                 gdist = 0.0
         else:
-            energy, gdist = self.original(dist)
+            energy, gdist = self.original.compute(dist)
             energy[mask] -= self.ecut + self.gcut * (dist[mask] - self.rcut)
             energy[~mask] = 0.0
             gdist[mask] -= self.gcut
@@ -82,11 +93,9 @@ def __call__(self, dist: ArrayLike) -> tuple[NDArray, NDArray]:
 
 
 @attrs.define
-class PairwiseForceField:
-    pair_potential: PairPotential = attrs.field(
-        validator=attrs.validators.instance_of(PairPotential)
-    )
-    """A definition of the pair potential."""
+class ForceField:
+    force_terms: list[ForceTerm] = attrs.field()
+    """A list of contributions to the potential energy."""
 
     nbuild: NBuild = attrs.field(validator=attrs.validators.instance_of(NBuild))
     """Algorithm to build the neigborlist."""
@@ -116,10 +125,11 @@ def __call__(self, atpos: NDArray, cell_length: float):
         self.nbuild.update(atpos, cell_length)
         nlist = self.nbuild.nlist
         # Compute all pairwise quantities
-        nlist["energy"], nlist["gdist"] = self.pair_potential(nlist["dist"])
-        nlist["gdelta"] = (nlist["gdist"] / nlist["dist"]).reshape(-1, 1) * nlist["delta"]
+        for force_term in self.force_terms:
+            force_term(nlist)
         # Compute the totals
         energy = nlist["energy"].sum()
+        nlist["gdelta"] = (nlist["gdist"] / nlist["dist"]).reshape(-1, 1) * nlist["delta"]
         forces = np.zeros(atpos.shape, dtype=float)
         np.add.at(forces, nlist["iatom0"], nlist["gdelta"])
         np.add.at(forces, nlist["iatom1"], -nlist["gdelta"])

tests/test_atomsmithy.py

@@ -112,18 +112,18 @@ def test_fcc_lattice():
 def test_push_derivative():
     pp = PushPotential(2.5)
     dist = np.linspace(0.4, 3.0, 50)
-    gdist1 = pp(dist)[1]
-    gdist2 = nd.Derivative(lambda dist: pp(dist)[0])(dist)
+    gdist1 = pp.compute(dist)[1]
+    gdist2 = nd.Derivative(lambda dist: pp.compute(dist)[0])(dist)
     assert gdist1 == pytest.approx(gdist2)
 
 
 def test_push_cutoff():
     pp = PushPotential(2.5)
     eps = 1e-13
-    assert abs(pp(2.5 - 0.1)[0]) > eps
-    assert abs(pp(2.5 - 0.1)[1]) > eps
-    assert abs(pp(2.5 - eps)[0]) < eps
-    assert abs(pp(2.5 - eps)[1]) < eps
+    assert abs(pp.compute(2.5 - 0.1)[0]) > eps
+    assert abs(pp.compute(2.5 - 0.1)[1]) > eps
+    assert abs(pp.compute(2.5 - eps)[0]) < eps
+    assert abs(pp.compute(2.5 - eps)[1]) < eps
 
 
 def test_random_box():

tests/test_forcefield.py

@@ -22,36 +22,37 @@
 import numpy as np
 import pytest
 
-from tinyff.forcefield import CutOffWrapper, LennardJones, PairwiseForceField
+from tinyff.atomsmithy import PushPotential
+from tinyff.forcefield import CutOffWrapper, ForceField, LennardJones
 from tinyff.neighborlist import NBuildCellLists, NBuildSimple
 
 
 def test_lennard_jones_derivative():
     lj = LennardJones(2.5, 0.5)
     dist = np.linspace(0.4, 3.0, 50)
-    gdist1 = lj(dist)[1]
-    gdist2 = nd.Derivative(lambda dist: lj(dist)[0])(dist)
+    gdist1 = lj.compute(dist)[1]
+    gdist2 = nd.Derivative(lambda dist: lj.compute(dist)[0])(dist)
     assert gdist1 == pytest.approx(gdist2)
 
 
 def test_lennard_jones_cut_derivative():
     lj = CutOffWrapper(LennardJones(2.5, 0.5), 3.5)
     dist = np.linspace(0.4, 5.0, 50)
-    gdist1 = lj(dist)[1]
-    gdist2 = nd.Derivative(lambda x: lj(x)[0])(dist)
+    gdist1 = lj.compute(dist)[1]
+    gdist2 = nd.Derivative(lambda x: lj.compute(x)[0])(dist)
     assert gdist1 == pytest.approx(gdist2)
 
 
 def test_lennard_jones_cut_zero_array():
     lj = CutOffWrapper(LennardJones(2.5, 0.5), 3.5)
-    e, g = lj([5.0, 3.6])
+    e, g = lj.compute([5.0, 3.6])
     assert (e == 0.0).all()
     assert (g == 0.0).all()
 
 
 def test_lennard_jones_cut_zero_scalar():
     lj = CutOffWrapper(LennardJones(2.5, 0.5), 3.5)
-    e, g = lj(5.0)
+    e, g = lj.compute(5.0)
     assert e == 0.0
     assert g == 0.0
 
@@ -65,13 +66,12 @@ def test_pairwise_force_field_two(nbuild_class):
     # Define the force field.
     rcut = 8.0
     lj = CutOffWrapper(LennardJones(2.5, 1.3), rcut)
-    nbuild = nbuild_class(rcut)
-    pwff = PairwiseForceField(lj, nbuild)
+    ff = ForceField([lj], nbuild_class(rcut))
 
     # Compute and check against manual result
-    energy, forces, frc_press = pwff(atpos, cell_length)
+    energy, forces, frc_press = ff(atpos, cell_length)
     d = np.linalg.norm(atpos[0] - atpos[1])
-    e, g = lj(d)
+    e, g = lj.compute(d)
     assert energy == pytest.approx(e)
     assert forces == pytest.approx(np.array([[g, 0.0, 0.0], [-g, 0.0, 0.0]]))
     assert frc_press == pytest.approx(-g * d / (3 * cell_length**3))
@@ -86,31 +86,30 @@ def test_pairwise_force_field_three(nbuild_class):
     # Define the force field.
     rcut = 8.0
     lj = CutOffWrapper(LennardJones(2.5, 1.3), rcut)
-    nbuild = nbuild_class(rcut)
-    pwff = PairwiseForceField(lj, nbuild)
+    ff = ForceField([lj], nbuild_class(rcut))
 
     # Compute the energy, the forces and the force contribution pressure.
-    energy1, forces1, frc_press1 = pwff(atpos, cell_length)
+    energy1, forces1, frc_press1 = ff(atpos, cell_length)
 
     # Compute the energy manually and compare.
     dists = [
         np.linalg.norm(atpos[1] - atpos[2]),
         np.linalg.norm(atpos[2] - atpos[0]),
         np.linalg.norm(atpos[0] - atpos[1]),
     ]
-    energy2 = lj(dists)[0].sum()
+    energy2 = lj.compute(dists)[0].sum()
     assert energy1 == pytest.approx(energy2)
 
     # Test forces with numdifftool
-    forces2 = -nd.Gradient(lambda x: pwff(x.reshape(-1, 3), cell_length)[0])(atpos)
+    forces2 = -nd.Gradient(lambda x: ff(x.reshape(-1, 3), cell_length)[0])(atpos)
     forces2.shape = (-1, 3)
     assert forces1 == pytest.approx(forces2.reshape(-1, 3))
 
     # Test pressure with numdifftool
     def energy_volume(volume):
         my_cell_length = volume ** (1.0 / 3.0)
         scale = my_cell_length / cell_length
-        return pwff(atpos * scale, my_cell_length)[0]
+        return ff(atpos * scale, my_cell_length)[0]
 
     frc_press2 = -nd.Derivative(energy_volume)(cell_length**3)
     assert frc_press1 == pytest.approx(frc_press2)
@@ -143,23 +142,43 @@ def test_pairwise_force_field_fifteen(nbuild_class):
     # Define the force field.
     rcut = 8.0
     lj = CutOffWrapper(LennardJones(2.5, 1.3), rcut)
-    nbuild = nbuild_class(rcut)
-    pwff = PairwiseForceField(lj, nbuild)
+    ff = ForceField([lj], nbuild_class(rcut))
 
     # Compute the energy, the forces and the force contribution to the pressure.
-    energy, forces1, frc_press1 = pwff(atpos, cell_length)
+    energy, forces1, frc_press1 = ff(atpos, cell_length)
     assert energy < 0
 
     # Test forces with numdifftool
-    forces2 = -nd.Gradient(lambda x: pwff(x.reshape(-1, 3), cell_length)[0])(atpos)
+    forces2 = -nd.Gradient(lambda x: ff(x.reshape(-1, 3), cell_length)[0])(atpos)
     forces2.shape = (-1, 3)
     assert forces1 == pytest.approx(forces2.reshape(-1, 3))
 
     # Test pressure with numdifftool
     def energy_volume(volume):
         my_cell_length = volume ** (1.0 / 3.0)
         scale = my_cell_length / cell_length
-        return pwff(atpos * scale, my_cell_length)[0]
+        return ff(atpos * scale, my_cell_length)[0]
 
     frc_press2 = -nd.Derivative(energy_volume)(cell_length**3)
     assert frc_press1 == pytest.approx(frc_press2)
+
+
+@pytest.mark.parametrize("nbuild_class", [NBuildSimple, NBuildCellLists])
+def test_superposition(nbuild_class):
+    atpos = np.array([[0.0, 0.0, 5.0], [0.0, 0.0, 0.0], [0.0, 3.0, 0.0]])
+    cell_length = 10.0
+
+    # Define the force field.
+    rcut = 4.9
+    lj = CutOffWrapper(LennardJones(2.5, 1.3), rcut)
+    pp = PushPotential(rcut)
+    ff_lj = ForceField([lj], nbuild_class(rcut))
+    ff_pp = ForceField([pp], nbuild_class(rcut))
+    ff_su = ForceField([lj, pp], nbuild_class(rcut))
+
+    energy_lj, forces_lj, frc_press_lj = ff_lj(atpos, cell_length)
+    energy_pp, forces_pp, frc_press_pp = ff_pp(atpos, cell_length)
+    energy_su, forces_su, frc_press_su = ff_su(atpos, cell_length)
+    assert energy_lj + energy_pp == pytest.approx(energy_su)
+    assert forces_lj + forces_pp == pytest.approx(forces_su)
+    assert frc_press_lj + frc_press_pp == pytest.approx(frc_press_su)