Cache last forces/stress ensemble to avoid redundant recomputation

Author: ppegolo

src/upet/calculator.py

@@ -203,6 +203,12 @@ def __init__(
         self._direct_uncertainty_key_BASE = (
             "mtt::aux::non_conservative_forces_uncertainty"
         )
+        # cache of the last conservative forces/stress ensemble computation, as
+        # (atoms, forces_ensemble, stress_ensemble); avoids recomputing the
+        # expensive Jacobian pass when called again for the same atoms.
+        self._uq_cache: Optional[
+            Tuple[Atoms, Optional[np.ndarray], Optional[np.ndarray]]
+        ] = None
 
         self.calculator = MetatomicCalculator(
             loaded_model,
@@ -340,6 +346,17 @@ def _run_forces_stress_uq(
             else:
                 atoms = self.atoms
 
+        cached = self._uq_cache
+        if cached is not None and cached[0] == atoms:
+            cached_forces, cached_stress = cached[1], cached[2]
+            if (not compute_forces or cached_forces is not None) and (
+                not compute_stress or cached_stress is not None
+            ):
+                return (
+                    cached_forces if compute_forces else None,
+                    cached_stress if compute_stress else None,
+                )
+
         calc = self.calculator
         # Unwrap SymmetrizedCalculator if present
         if isinstance(calc, SymmetrizedCalculator):
@@ -445,7 +462,16 @@ def _compute_ensemble(*positions_and_strain):
             # [3, 3, n_ensemble]
             stress_ensemble = stress_jac.transpose(1, 2, 0) / volume
 
-        return forces_ensemble, stress_ensemble
+        # keep whichever of forces/stress was already cached for these atoms but
+        # not recomputed here, instead of discarding it
+        if cached is not None and cached[0] == atoms:
+            forces_ensemble = forces_ensemble if compute_forces else cached[1]
+            stress_ensemble = stress_ensemble if compute_stress else cached[2]
+        self._uq_cache = (atoms.copy(), forces_ensemble, stress_ensemble)
+        return (
+            forces_ensemble if compute_forces else None,
+            stress_ensemble if compute_stress else None,
+        )
 
     @property
     def _direct_ensemble_key(self) -> str:

tests/upet/test_uncertainty_quantification.py

@@ -168,6 +168,26 @@ def test_forces_and_stress_ensemble():
         nc_calc.get_forces_and_stress_ensemble(atoms)
 
 
+def test_forces_ensemble_cache():
+    """Test that repeated calls with unchanged atoms reuse the cached ensemble."""
+    atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
+    calc = UPETCalculator(model="pet-mad-s", version="1.5.0")
+
+    forces_ensemble = calc.get_forces_ensemble(atoms)
+    # same atoms, unchanged: served from cache, not recomputed
+    assert calc.get_forces_ensemble(atoms) is forces_ensemble
+
+    # requesting stress next reuses the cached forces instead of discarding them
+    calc.get_stress_ensemble(atoms)
+    assert calc.get_forces_ensemble(atoms) is forces_ensemble
+
+    # moving the atoms invalidates the cache
+    atoms.positions[0, 0] += 0.1
+    forces_ensemble_moved = calc.get_forces_ensemble(atoms)
+    assert forces_ensemble_moved is not forces_ensemble
+    assert not np.allclose(forces_ensemble_moved, forces_ensemble)
+
+
 def test_forces_method_defaults():
     """Test that the default method depends on the non_conservative flag."""
     atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")