Simplify derivatives API

CHANGELOG.md

@@ -10,9 +10,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 - The `ForceField` object has a new method `compute` that can selectively compute
-  the energy, the forces and/or the force-contribution to the pressure.
-  This improves the efficiency in applications where not all quantities are of interest,
-  such as in Monte Carlo simulations.
+  the energy only or the energy, forces and the force-contribution to the pressure.
+  This improves the efficiency in applications where derivatives are of interest,
+  e.g. in Monte Carlo simulations.
 - The `ForceField` object is extended with `try_move` and `accept_move`
   to support (relatively) efficient Monte Carlo algorithms with TinyFF.
 - Basic analysis routines for radial distribution functions and autocorrelation functions.
@@ -24,12 +24,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - The `ForceField` class requires the `NBuild` instance to be provided as a keyword argument.
   For example: `ForceField([LennardJones()], nbuild=NBuildSimple)`.
 - The `ForceField.__call__` method is replaced by `ForceField.compute`,
-  which has a different API with new `do_*` arguments.
+  which has a different API with a new `nderiv=0` arguments.
   By default, only the energy is computed.
-  You must request forces and pressure explicitly.
-  The function always returns a list of results, even if there is only one.
+  You must request forces and pressure explicitly by setting `nderiv=1`.
+  The function always returns a list of results, even if `nderiv=0`.
 - The `PairwiseTerm.compute` (previously `PairwisePotential.compute`) has a new API:
-  it takes `do_*` arguments to decide what is computed (energy and or derivative).
+  it takes an `nderiv` arguments to decide what is computed
+  (energy or energy and derivative).
   It returns a list of requested results.
   By default, only the energy is computed.
 - The `ForceTerm.__call__` method has been replaced by `PairwiseTerm.compute_nlist`.

README.md

@@ -79,12 +79,11 @@ atpos = np.array([[0.0, 0.0, 1.0], [1.0, 2.0, 0.0]])
 cell_length = 20.0
 
 # Compute a selection of results with ff.compute.
-#   The ff.compute method has `do_*` arguments to compute only some results:
-#   - `do_energy` (default True)
-#   - `do_forces` (default False)
-#   - `do_press` (default False).
+#   The ff.compute method has an `nderiv` arguments to compute only some results:
+#   - `nderiv=0` (default): compute only the energy
+#   - `nderiv=1`: compute the energy, forces and force contribution to the pressure.
 #   Requested results are put in a list, even in only one result is requested.
-potential_energy, forces = ff.compute(atpos, cell_length, do_forces=True)
+potential_energy, forces, press = ff.compute(atpos, cell_length, nderiv=1)
 ```
 
 This basic recipe can be extended by passing additional options

src/tinyff/atomsmithy.py

@@ -102,7 +102,7 @@ def build_random_cell(
 
     def costgrad(atpos_raveled):
         atpos = atpos_raveled.reshape(-1, 3)
-        energy, force = ff.compute(atpos, cell_length, do_forces=True)
+        energy, force, _ = ff.compute(atpos, cell_length, nderiv=1)
         return energy, -force.ravel()
 
     # Optimize and return structure

src/tinyff/forcefield.py

@@ -47,15 +47,7 @@ class ForceField:
     nbuild: NBuild = attrs.field(validator=attrs.validators.instance_of(NBuild), kw_only=True)
     """Algorithm to build the neigborlist."""
 
-    def compute(
-        self,
-        atpos: NDArray,
-        cell_lengths: ArrayLike | float,
-        *,
-        do_energy: bool = True,
-        do_forces: bool = False,
-        do_press: bool = False,
-    ):
+    def compute(self, atpos: NDArray, cell_lengths: ArrayLike | float, nderiv: int = 0):
         """Compute microscopic properties related to the potential energy.
 
         Parameters
@@ -66,12 +58,9 @@ def compute(
         cell_length
             The length of the edge of the cubic simulation cell,
             or an array of lengths of three cell vectors.
-        do_energy
-            if True, the energy is returned.
-        do_forces
-            if True, the atomic forces are returned.
-        do_press
-            if True, the force contribution to the pressure is returned.
+        nderiv
+            The order of derivatives to compute, either 0 (energy)
+            or 1 (energy, forces and pressure).
 
         Returns
         -------
@@ -84,23 +73,20 @@ def compute(
 
         # Compute all pairwise quantities, if needed with derivatives.
         for pairwise_term in self.pairwise_terms:
-            pairwise_term.compute_nlist(nlist, do_energy, do_forces | do_press)
+            pairwise_term.compute_nlist(nlist, nderiv)
 
         # Compute the totals
         results = []
-        if do_energy:
-            energy = nlist["energy"].sum()
-            results.append(energy)
-        if do_forces or do_press:
+        energy = nlist["energy"].sum()
+        results.append(energy)
+        if nderiv >= 1:
             nlist["gdelta"] = (nlist["gdist"] / nlist["dist"]).reshape(-1, 1) * nlist["delta"]
-            if do_forces:
-                atfrc = np.zeros(atpos.shape, dtype=float)
-                np.subtract.at(atfrc, nlist["iatom1"], nlist["gdelta"])
-                np.add.at(atfrc, nlist["iatom0"], nlist["gdelta"])
-                results.append(atfrc)
-            if do_press:
-                frc_press = -np.dot(nlist["gdist"], nlist["dist"]) / (3 * cell_lengths.prod())
-                results.append(frc_press)
+            atfrc = np.zeros(atpos.shape, dtype=float)
+            np.subtract.at(atfrc, nlist["iatom1"], nlist["gdelta"])
+            np.add.at(atfrc, nlist["iatom0"], nlist["gdelta"])
+            results.append(atfrc)
+            frc_press = -np.dot(nlist["gdist"], nlist["dist"]) / (3 * cell_lengths.prod())
+            results.append(frc_press)
 
         return results
 

src/tinyff/pairwise.py

@@ -29,23 +29,33 @@
 
 @attrs.define
 class PairwiseTerm:
-    def compute_nlist(
-        self, nlist: NDArray[NLIST_DTYPE], do_energy: bool = True, do_gdist: bool = False
-    ):
-        """Compute energies and derivatives and add them to the neighborlist."""
-        results = self.compute(nlist["dist"], do_energy, do_gdist)
-        if do_energy:
-            nlist["energy"] += results.pop(0)
-        if do_gdist:
+    def compute_nlist(self, nlist: NDArray[NLIST_DTYPE], nderiv: int = 0):
+        """Compute energies and derivatives and add them to the neighborlist.
+
+        Parameters
+        ----------
+        nlist
+            The neighborlist to which energies (and derivatives) must be added.
+        nderiv
+            The order of derivatives to compute, either 0 (energy)
+            or 1 (energy and its derivative).
+        """
+        results = self.compute(nlist["dist"], nderiv)
+        nlist["energy"] += results.pop(0)
+        if nderiv >= 1:
             nlist["gdist"] += results.pop(0)
 
-    def compute(
-        self,
-        dist: NDArray[float],
-        do_energy: bool = True,
-        do_gdist: bool = False,
-    ) -> list[NDArray]:
-        """Compute pair potential energy and its derivative towards distance."""
+    def compute(self, dist: NDArray[float], nderiv: int = 0) -> list[NDArray]:
+        """Compute pair potential energy and its derivative towards distance.
+
+        Parameters
+        ----------
+        dist
+            The interatomic distances.
+        nderiv
+            The order of derivatives to compute, either 0 (energy)
+            or 1 (energy and its derivative).
+        """
         raise NotImplementedError  # pragma: nocover
 
 
@@ -54,23 +64,17 @@ class LennardJones(PairwiseTerm):
     epsilon: float = attrs.field(default=1.0, converter=float)
     sigma: float = attrs.field(default=1.0, converter=float)
 
-    def compute(
-        self,
-        dist: NDArray[float],
-        do_energy: bool = True,
-        do_gdist: bool = False,
-    ) -> list[NDArray]:
+    def compute(self, dist: NDArray[float], nderiv: int = 0) -> list[NDArray]:
         """Compute pair potential energy and its derivative towards distance."""
         results = []
         dist = np.asarray(dist, dtype=float)
         x = self.sigma / dist
         x3 = x * x * x
         x6 = x3 * x3
-        if do_energy:
-            energy = (x6 - 1) * x6
-            energy *= 4 * self.epsilon
-            results.append(energy)
-        if do_gdist:
+        energy = (x6 - 1) * x6
+        energy *= 4 * self.epsilon
+        results.append(energy)
+        if nderiv >= 1:
             gdist = (x6 - 0.5) * x6 / dist
             gdist *= -48 * self.epsilon
             results.append(gdist)
@@ -86,43 +90,35 @@ class CutOffWrapper(PairwiseTerm):
 
     def __attrs_post_init__(self):
         """Post initialization changes."""
-        self.ecut, self.gcut = self.original.compute(self.rcut, do_gdist=True)
-
-    def compute(
-        self,
-        dist: NDArray[float],
-        do_energy: bool = True,
-        do_gdist: bool = False,
-    ) -> list[NDArray]:
+        self.ecut, self.gcut = self.original.compute(self.rcut, nderiv=1)
+
+    def compute(self, dist: NDArray[float], nderiv: int = 0) -> list[NDArray]:
         """Compute pair potential energy and its derivative towards distance."""
         dist = np.asarray(dist, dtype=float)
         mask = dist < self.rcut
         results = []
         if mask.ndim == 0:
             # Deal with non-array case
             if mask:
-                orig_results = self.original.compute(dist, do_energy, do_gdist)
-                if do_energy:
-                    energy = orig_results.pop(0)
-                    energy -= self.ecut + self.gcut * (dist - self.rcut)
-                    results.append(energy)
-                if do_gdist:
+                orig_results = self.original.compute(dist, nderiv)
+                energy = orig_results.pop(0)
+                energy -= self.ecut + self.gcut * (dist - self.rcut)
+                results.append(energy)
+                if nderiv >= 1:
                     gdist = orig_results.pop(0)
                     gdist -= self.gcut
                     results.append(gdist)
             else:
-                if do_energy:
-                    results.append(0.0)
-                if do_gdist:
+                results.append(0.0)
+                if nderiv >= 1:
                     results.append(0.0)
         else:
-            orig_results = self.original.compute(dist, do_energy, do_gdist)
-            if do_energy:
-                energy = orig_results.pop(0)
-                energy -= self.ecut + self.gcut * (dist - self.rcut)
-                energy *= mask
-                results.append(energy)
-            if do_gdist:
+            orig_results = self.original.compute(dist, nderiv)
+            energy = orig_results.pop(0)
+            energy -= self.ecut + self.gcut * (dist - self.rcut)
+            energy *= mask
+            results.append(energy)
+            if nderiv >= 1:
                 gdist = orig_results.pop(0)
                 gdist -= self.gcut
                 gdist *= mask
@@ -136,21 +132,15 @@ class CheapRepulsion(PairwiseTerm):
 
     rcut: float = attrs.field(converter=float, validator=attrs.validators.gt(0))
 
-    def compute(
-        self,
-        dist: NDArray[float],
-        do_energy: bool = True,
-        do_gdist: bool = False,
-    ) -> list[NDArray]:
+    def compute(self, dist: NDArray[float], nderiv: int = 0) -> list[NDArray]:
         """Compute pair potential energy and its derivative towards distance."""
         dist = np.asarray(dist, dtype=float)
         x = dist / self.rcut
         results = []
         common = (x - 1) * (x < 1)
-        if do_energy:
-            energy = common * common
-            results.append(energy)
-        if do_gdist:
+        energy = common * common
+        results.append(energy)
+        if nderiv >= 1:
             gdist = (2 / self.rcut) * common
             results.append(gdist)
         return results

tests/test_forcefield.py

@@ -39,9 +39,9 @@ def test_pairwise_force_field_two(nbuild):
     ff = ForceField([lj], nbuild=nbuild)
 
     # Compute and check against manual result
-    energy, forces, frc_press = ff.compute(atpos, cell_length, do_forces=True, do_press=True)
+    energy, forces, frc_press = ff.compute(atpos, cell_length, nderiv=1)
     d = np.linalg.norm(atpos[0] - atpos[1])
-    e, g = lj.compute(d, do_gdist=True)
+    e, g = lj.compute(d, nderiv=1)
     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))
@@ -59,7 +59,7 @@ def test_pairwise_force_field_three(nbuild):
     ff = ForceField([lj], nbuild=nbuild)
 
     # Compute the energy, the forces and the force contribution pressure.
-    energy1, forces1, frc_press1 = ff.compute(atpos, cell_length, do_forces=True, do_press=True)
+    energy1, forces1, frc_press1 = ff.compute(atpos, cell_length, nderiv=1)
 
     # Compute the energy manually and compare.
     dists = [
@@ -115,7 +115,7 @@ def test_pairwise_force_field_fifteen(nbuild):
     ff = ForceField([lj], nbuild=nbuild)
 
     # Compute the energy, the forces and the force contribution to the pressure.
-    energy, forces1, frc_press1 = ff.compute(atpos, cell_length, do_forces=True, do_press=True)
+    energy, forces1, frc_press1 = ff.compute(atpos, cell_length, nderiv=1)
     assert energy < 0
 
     # Test forces with numdifftool
@@ -148,9 +148,9 @@ def test_superposition(nbuild_class, kwargs):
     ff_pp = ForceField([cr], nbuild=nbuild_class(rcut, **kwargs))
     ff_su = ForceField([lj, cr], nbuild=nbuild_class(rcut, **kwargs))
 
-    ener_lj, forces_lj, press_lj = ff_lj.compute(atpos, cell_length, do_forces=True, do_press=True)
-    ener_pp, forces_pp, press_pp = ff_pp.compute(atpos, cell_length, do_forces=True, do_press=True)
-    ener_su, forces_su, press_su = ff_su.compute(atpos, cell_length, do_forces=True, do_press=True)
+    ener_lj, forces_lj, press_lj = ff_lj.compute(atpos, cell_length, nderiv=1)
+    ener_pp, forces_pp, press_pp = ff_pp.compute(atpos, cell_length, nderiv=1)
+    ener_su, forces_su, press_su = ff_su.compute(atpos, cell_length, nderiv=1)
     assert ener_lj + ener_pp == pytest.approx(ener_su)
     assert forces_lj + forces_pp == pytest.approx(forces_su)
     assert press_lj + press_pp == pytest.approx(press_su)

tests/test_pairwise.py

@@ -28,45 +28,51 @@
 def test_lennard_jones_derivative():
     lj = LennardJones(2.5, 0.5)
     dist = np.linspace(0.4, 3.0, 50)
-    gdist1 = lj.compute(dist, do_energy=False, do_gdist=True)[0]
+    gdist1 = lj.compute(dist, nderiv=1)[1]
     gdist2 = nd.Derivative(lambda x: lj.compute(x)[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.compute(dist, do_energy=False, do_gdist=True)[0]
+    gdist1 = lj.compute(dist, nderiv=1)[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.compute([5.0, 3.6], do_gdist=True)
+    e, g = lj.compute([5.0, 3.6], nderiv=1)
     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.compute(5.0, do_gdist=True)
+    e, g = lj.compute(5.0, nderiv=1)
     assert e == 0.0
     assert g == 0.0
 
 
 def test_cheap_repulsion_derivative():
     cr = CheapRepulsion(2.5)
     dist = np.linspace(0.4, 3.0, 50)
-    gdist1 = cr.compute(dist, do_energy=False, do_gdist=True)[0]
+    gdist1 = cr.compute(dist, nderiv=1)[1]
     gdist2 = nd.Derivative(lambda x: cr.compute(x)[0])(dist)
     assert gdist1 == pytest.approx(gdist2)
 
 
 def test_cheap_repulsion_cutoff():
-    cr = CheapRepulsion(2.5)
+    rcut = 2.5
+    cr = CheapRepulsion(rcut)
     eps = 1e-13
-    assert abs(cr.compute(2.5 - 0.1)[0]) > eps
-    assert abs(cr.compute(2.5 - 0.1, do_energy=False, do_gdist=True)[0]) > eps
-    assert abs(cr.compute(2.5 - eps)[0]) < eps
-    assert abs(cr.compute(2.5 - eps, do_energy=False, do_gdist=True)[0]) < eps
+    e, g = cr.compute(rcut - 0.1, nderiv=1)
+    assert abs(e) > eps
+    assert abs(g) > eps
+    e, g = cr.compute(rcut - eps, nderiv=1)
+    assert abs(e) < eps
+    assert abs(g) < eps
+    e, g = cr.compute(rcut + 0.2, nderiv=1)
+    assert e == 0.0
+    assert g == 0.0