Support conservative forces and stress ensembles

Force and stress ensembles are obtained as explicit positions/strain
gradients of the model's "energy_ensemble" output, requested through
MetatomicCalculator.run_model. The Jacobian machinery lives inside the
model (metatrain's LLPRUncertaintyModel, so no engine-side autograd
setup and no private metatomic attributes are needed; the ensemble
output key is resolved from the model's own capabilities, following the
same variant convention as the non-conservative output keys.

- get_forces_ensemble / get_stress_ensemble, and the corresponding
  uncertainties as the standard deviation over ensemble members
- get_forces_and_stress_ensemble: both from a single backward pass
- forces are recentered per ensemble member (translational invariance)
- the last computed ensembles are cached and reused for unchanged atoms

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
EOF
)

Author: ppegolo

src/upet/calculator.py

@@ -6,7 +6,11 @@
 import torch
 from ase import Atoms
 from metatomic.torch import ModelOutput
-from metatomic_ase import MetatomicCalculator, SymmetrizedCalculator
+from metatomic_ase import (
+    MetatomicCalculator,
+    SymmetrizedCalculator,
+)
+from metatomic_ase._calculator import _full_3x3_to_voigt_6_stress
 from packaging.version import Version
 
 from ._models import (
@@ -40,6 +44,17 @@
 }
 
 
+def _stress_ensemble_to_voigt(stress_ensemble: np.ndarray) -> np.ndarray:
+    """Convert a [3, 3, n_ensemble] stress ensemble to Voigt [6, n_ensemble]."""
+    n_ensemble = stress_ensemble.shape[2]
+    return np.array(
+        [
+            _full_3x3_to_voigt_6_stress(stress_ensemble[:, :, i])
+            for i in range(n_ensemble)
+        ]
+    ).T
+
+
 class UPETCalculator(ase.calculators.calculator.Calculator):
     """
     ASE Calculator for universal MLIPs based on the PET architecture.
@@ -174,6 +189,24 @@ def __init__(
             loaded_model._capabilities.dtype = DTYPE_TO_STR[dtype]
             loaded_model = loaded_model.to(dtype=dtype, device=device)
 
+        self._non_conservative = non_conservative
+        selected_variant = None if variants is None else variants.get("energy")
+        variant_postfix = f"/{selected_variant}" if selected_variant else ""
+        self._variant_postfix = variant_postfix
+        # resolved here from the model's own outputs (same variant convention as the
+        # non-conservative keys above) so that no private MetatomicCalculator
+        # attribute is needed to request the ensemble
+        energy_ensemble_key = "energy_ensemble" + variant_postfix
+        self._energy_ensemble_key: Optional[str] = (
+            energy_ensemble_key if energy_ensemble_key in model_outputs else None
+        )
+        # 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,
             extensions_directory=None,
@@ -274,9 +307,232 @@ def get_energy_ensemble(
         :param per_atom: Whether to return the energies per atom.
         :return: Energy uncertainty in numpy.ndarray format.
         """
-        key = self.calculator._energy_uq_key.replace("_uncertainty", "_ensemble")
+        key = self._energy_ensemble_key
+        if key is None:
+            raise NotImplementedError(
+                "Energy ensemble is not available for the selected model. For "
+                "uncertainty estimates, please use one of the following models: "
+                f"{UPET_UQ_SUPPORTED_MODELS}"
+            )
         return self._run_uq(atoms=atoms, per_atom=per_atom, key=key)
 
+    def _run_forces_stress_uq(
+        self,
+        atoms: Optional[Atoms] = None,
+        compute_forces: bool = True,
+        compute_stress: bool = True,
+    ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
+        """
+        Compute force and/or stress ensembles via gradients of the energy ensemble.
+
+        Returns a tuple (forces_ensemble, stress_ensemble) where each is None if not
+        requested. Forces ensemble has shape [n_atoms, 3, n_ensemble], stress ensemble
+        has shape [3, 3, n_ensemble].
+        """
+        if not self.calculator._calculate_uncertainty:
+            raise NotImplementedError(
+                "Forces/stress uncertainty and ensemble are not available for the "
+                "selected model. For uncertainty estimates, please use one of the "
+                f"following models: {UPET_UQ_SUPPORTED_MODELS}"
+            )
+
+        if not compute_forces and not compute_stress:
+            raise ValueError(
+                "At least one of compute_forces or compute_stress must be True."
+            )
+
+        if atoms is None:
+            if self.atoms is None:
+                raise ValueError(
+                    "No `atoms` provided and no previously calculated atoms found."
+                )
+            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):
+            calc = calc._calculator
+
+        ensemble_key = self._energy_ensemble_key
+        if ensemble_key is None:
+            raise NotImplementedError(
+                "Energy ensemble is not available for the selected model. For "
+                "uncertainty estimates, please use one of the following models: "
+                f"{UPET_UQ_SUPPORTED_MODELS}"
+            )
+
+        explicit_gradients = []
+        if compute_forces:
+            explicit_gradients.append("positions")
+        if compute_stress:
+            explicit_gradients.append("strain")
+
+        outputs_request = {
+            ensemble_key: ModelOutput(
+                unit="eV", sample_kind="system", explicit_gradients=explicit_gradients
+            )
+        }
+        if calc._energy_key is not None:
+            # some models require the base "energy" output to be requested
+            # alongside the ensemble; this is cheap to compute (it is an
+            # intermediate step of the ensemble itself) and its value is unused here
+            outputs_request[calc._energy_key] = ModelOutput(
+                unit="eV", sample_kind="system"
+            )
+
+        block = calc.run_model(atoms, outputs_request)[ensemble_key].block()
+
+        forces_ensemble = None
+        stress_ensemble = None
+
+        if compute_forces:
+            # gradient shape: [n_atoms, 3, n_ensemble]
+            forces_ensemble = (
+                -block.gradient("positions").values.detach().cpu().double().numpy()
+            )
+            # remove the mean over atoms for each ensemble member to impose
+            # translational invariance
+            forces_ensemble = forces_ensemble - np.mean(
+                forces_ensemble, axis=0, keepdims=True
+            )
+
+        if compute_stress:
+            # gradient shape: [1, 3, 3, n_ensemble] (single system)
+            # -> [3, 3, n_ensemble]
+            stress_ensemble = (
+                block.gradient("strain").values.detach().cpu().double().numpy()[0]
+                / atoms.cell.volume
+            )
+
+        # 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,
+        )
+
+    def get_forces_ensemble(self, atoms: Optional[Atoms] = None) -> np.ndarray:
+        """
+        Get the ensemble of forces for a given :py:class:`ase.Atoms` object.
+
+        Forces are computed as the (negative) derivative of the energy ensemble
+        with respect to positions (conservative forces).
+
+        :param atoms: ASE atoms object. If ``None``, the last calculated atoms will be
+            used.
+        :return: Forces ensemble as numpy.ndarray with shape [n_atoms, 3, n_ensemble],
+            in eV/Angstrom.
+        """
+        forces_ensemble, _ = self._run_forces_stress_uq(
+            atoms=atoms, compute_forces=True, compute_stress=False
+        )
+        return forces_ensemble
+
+    def get_forces_uncertainty(self, atoms: Optional[Atoms] = None) -> np.ndarray:
+        """
+        Get the forces uncertainty for a given :py:class:`ase.Atoms` object.
+
+        Uncertainty is computed as the standard deviation of the conservative
+        forces ensemble (derived from the energy ensemble).
+
+        :param atoms: ASE atoms object. If ``None``, the last calculated atoms will be
+            used.
+        :return: Forces uncertainty as numpy.ndarray with shape [n_atoms, 3],
+            in eV/Angstrom.
+        """
+        forces_ensemble = self.get_forces_ensemble(atoms=atoms)
+        return np.std(forces_ensemble, axis=2)
+
+    def get_stress_ensemble(
+        self, atoms: Optional[Atoms] = None, voigt: bool = True
+    ) -> np.ndarray:
+        """
+        Get the ensemble of stresses for a given :py:class:`ase.Atoms` object.
+
+        Stresses are computed as the derivative of the energy ensemble with
+        respect to strain (conservative stresses).
+
+        :param atoms: ASE atoms object. If ``None``, the last calculated atoms will be
+            used.
+        :param voigt: If ``True`` (default), return stresses in Voigt notation with
+            shape [6, n_ensemble] (xx, yy, zz, yz, xz, xy). If ``False``, return
+            full 3x3 tensors with shape [3, 3, n_ensemble].
+        :return: Stress ensemble in eV/Angstrom^3.
+        """
+        _, stress_ensemble = self._run_forces_stress_uq(
+            atoms=atoms, compute_forces=False, compute_stress=True
+        )
+        if voigt:
+            assert stress_ensemble is not None  # for mypy
+            stress_ensemble = _stress_ensemble_to_voigt(stress_ensemble)
+        return stress_ensemble
+
+    def get_stress_uncertainty(
+        self, atoms: Optional[Atoms] = None, voigt: bool = True
+    ) -> np.ndarray:
+        """
+        Get the stress uncertainty for a given :py:class:`ase.Atoms` object.
+
+        Uncertainty is computed as the standard deviation of the stress ensemble
+        (conservative stresses derived from the energy ensemble).
+
+        :param atoms: ASE atoms object. If ``None``, the last calculated atoms will be
+            used.
+        :param voigt: If ``True`` (default), return uncertainty in Voigt notation with
+            shape [6] (xx, yy, zz, yz, xz, xy). If ``False``, return full 3x3 tensor
+            with shape [3, 3].
+        :return: Stress uncertainty in eV/Angstrom^3.
+        """
+        stress_ensemble = self.get_stress_ensemble(atoms=atoms, voigt=voigt)
+        return np.std(stress_ensemble, axis=-1)
+
+    def get_forces_and_stress_ensemble(
+        self, atoms: Optional[Atoms] = None, voigt: bool = True
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        Get force and stress ensembles together (conservative method only).
+
+        Equivalent to calling :py:meth:`get_forces_ensemble` and
+        :py:meth:`get_stress_ensemble` separately, but cheaper: both are
+        derived from the same backward pass through the energy ensemble.
+
+        :param atoms: ASE atoms object. If ``None``, the last calculated atoms will be
+            used.
+        :param voigt: If ``True`` (default), return stress in Voigt notation with
+            shape [6, n_ensemble]. If ``False``, return full 3x3 tensors with shape
+            [3, 3, n_ensemble].
+        :return: Tuple ``(forces_ensemble, stress_ensemble)``. Forces ensemble has
+            shape [n_atoms, 3, n_ensemble], in eV/Angstrom.
+        """
+        if self._non_conservative:
+            raise ValueError(
+                "get_forces_and_stress_ensemble is not available when the "
+                "calculator was initialized with non_conservative=True."
+            )
+        forces_ensemble, stress_ensemble = self._run_forces_stress_uq(
+            atoms=atoms, compute_forces=True, compute_stress=True
+        )
+        if voigt:
+            assert stress_ensemble is not None  # for mypy
+            stress_ensemble = _stress_ensemble_to_voigt(stress_ensemble)
+        return forces_ensemble, stress_ensemble
+
 
 # For PET-MAD-DOS predictions
 ENERGY_INTERVAL = 0.05  # Interval of the energy grid for DOS