Support direct (non-conservative) forces ensemble and uncertainty

Author: ppegolo

src/upet/calculator.py

@@ -200,6 +200,10 @@ def __init__(
         self._energy_ensemble_key: Optional[str] = (
             energy_ensemble_key if energy_ensemble_key in model_outputs else None
         )
+        self._direct_ensemble_key_BASE = "mtt::aux::non_conservative_forces_ensemble"
+        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.
@@ -427,36 +431,180 @@ def _run_forces_stress_uq(
             stress_ensemble if compute_stress else None,
         )
 
-    def get_forces_ensemble(self, atoms: Optional[Atoms] = None) -> np.ndarray:
+    @property
+    def _direct_ensemble_key(self) -> str:
+        return self._direct_ensemble_key_BASE + self._variant_postfix
+
+    @property
+    def _direct_uncertainty_key(self) -> str:
+        return self._direct_uncertainty_key_BASE + self._variant_postfix
+
+    def _run_direct_forces_uq(self, atoms: Optional[Atoms] = None) -> np.ndarray:
         """
-        Get the ensemble of forces for a given :py:class:`ase.Atoms` object.
+        Get the direct forces ensemble from the model.
+
+        :return: Forces ensemble as numpy.ndarray with shape [n_atoms, 3, n_ensemble].
+        """
+        if atoms is None:
+            if self.atoms is None:
+                raise ValueError(
+                    "No `atoms` provided and no previously calculated atoms found."
+                )
+            else:
+                atoms = self.atoms
+
+        calc = self.calculator
+        if isinstance(calc, SymmetrizedCalculator):
+            calc = calc._calculator
+
+        model_outputs = calc._model.capabilities().outputs
+        if self._direct_ensemble_key not in model_outputs:
+            raise NotImplementedError(
+                f"Direct forces ensemble ({self._direct_ensemble_key}) is not "
+                "available for the selected model."
+            )
+
+        outputs = calc.run_model(
+            atoms,
+            outputs={
+                self._direct_ensemble_key: ModelOutput(
+                    quantity="force", unit="eV/Angstrom", per_atom=True
+                )
+            },
+        )
+
+        # shape: [n_atoms, 3 * n_ensemble] -> [n_atoms, 3, n_ensemble]
+        values = (
+            outputs[self._direct_ensemble_key].block().values.detach().cpu().numpy()
+        )
+        return values.reshape(len(atoms), 3, -1)
+
+    def _run_direct_forces_uncertainty(
+        self, atoms: Optional[Atoms] = None
+    ) -> np.ndarray:
+        """
+        Get the built-in direct forces uncertainty from the model.
+
+        :return: Forces uncertainty as numpy.ndarray with shape [n_atoms, 3].
+        """
+        if atoms is None:
+            if self.atoms is None:
+                raise ValueError(
+                    "No `atoms` provided and no previously calculated atoms found."
+                )
+            else:
+                atoms = self.atoms
 
-        Forces are computed as the (negative) derivative of the energy ensemble
-        with respect to positions (conservative forces).
+        calc = self.calculator
+        if isinstance(calc, SymmetrizedCalculator):
+            calc = calc._calculator
+
+        model_outputs = calc._model.capabilities().outputs
+        if self._direct_uncertainty_key not in model_outputs:
+            raise NotImplementedError(
+                f"Direct forces uncertainty ({self._direct_uncertainty_key}) is not "
+                "available for the selected model."
+            )
+
+        outputs = calc.run_model(
+            atoms,
+            outputs={
+                self._direct_uncertainty_key: ModelOutput(
+                    quantity="force", unit="eV/Angstrom", per_atom=True
+                )
+            },
+        )
+
+        # shape: [n_atoms, 3, 1] -> [n_atoms, 3]
+        values = (
+            outputs[self._direct_uncertainty_key].block().values.detach().cpu().numpy()
+        )
+        return values.squeeze(-1)
+
+    def _resolve_forces_method(self, method: Optional[str]) -> str:
+        """Resolve the default forces method based on the calculator mode."""
+        if method is None:
+            return "direct" if self._non_conservative else "conservative"
+        return method
+
+    def get_forces_ensemble(
+        self, atoms: Optional[Atoms] = None, method: Optional[str] = None
+    ) -> np.ndarray:
+        """
+        Get the ensemble of forces for a given :py:class:`ase.Atoms` object.
 
         :param atoms: ASE atoms object. If ``None``, the last calculated atoms will be
             used.
+        :param method: Method to compute the forces ensemble. One of:
+
+            - ``"conservative"`` (default when ``non_conservative=False``): forces
+              derived from the energy ensemble via automatic differentiation.
+            - ``"direct"`` (default when ``non_conservative=True``): forces ensemble
+              predicted directly by the model. When the calculator
+              is in conservative mode (``non_conservative=False``), the ensemble is
+              shifted so its mean matches the conservative forces.
+
         :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
+        method = self._resolve_forces_method(method)
+
+        if method == "conservative":
+            if self._non_conservative:
+                raise ValueError(
+                    "method='conservative' is not available when the calculator was "
+                    "initialized with non_conservative=True."
+                )
+            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:
+        elif method == "direct":
+            direct_ensemble = self._run_direct_forces_uq(atoms=atoms)
+
+            if not self._non_conservative:
+                # Shift ensemble so its mean matches conservative forces
+                if atoms is None:
+                    atoms = self.atoms
+                atoms.calc = self
+                conservative_forces = atoms.get_forces()  # [n_atoms, 3]
+
+                direct_mean = np.mean(direct_ensemble, axis=2)  # [n_atoms, 3]
+                shift = conservative_forces - direct_mean  # [n_atoms, 3]
+                direct_ensemble = direct_ensemble + shift[:, :, np.newaxis]
+
+            return direct_ensemble
+
+        else:
+            raise ValueError(
+                f"Unknown method '{method}'. Must be one of: 'conservative', 'direct'."
+            )
+
+    def get_forces_uncertainty(
+        self, atoms: Optional[Atoms] = None, method: Optional[str] = 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.
+        :param method: Method to compute the forces uncertainty. One of:
+
+            - ``"conservative"``: standard deviation of the conservative forces
+              ensemble (derived from energy ensemble via autograd).
+            - ``"direct"``: built-in uncertainty from the model
+              (``mtt::aux::non_conservative_forces_uncertainty``).
+
         :return: Forces uncertainty as numpy.ndarray with shape [n_atoms, 3],
             in eV/Angstrom.
         """
-        forces_ensemble = self.get_forces_ensemble(atoms=atoms)
+        method = self._resolve_forces_method(method)
+
+        if method == "direct":
+            return self._run_direct_forces_uncertainty(atoms=atoms)
+
+        forces_ensemble = self.get_forces_ensemble(atoms=atoms, method=method)
         return np.std(forces_ensemble, axis=2)
 
     def get_stress_ensemble(

tests/upet/test_uncertainty_quantification.py

@@ -188,6 +188,34 @@ def test_forces_ensemble_cache():
     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")
+
+    # Conservative calculator defaults to method="conservative"
+    calc_cons = UPETCalculator(model="pet-mad-s", version="1.5.0")
+    forces_ensemble = calc_cons.get_forces_ensemble(atoms)
+    assert forces_ensemble.ndim == 3
+    assert forces_ensemble.shape[:2] == (len(atoms), 3)
+
+    # method="conservative" works explicitly
+    forces_ensemble_2 = calc_cons.get_forces_ensemble(atoms, method="conservative")
+    assert np.allclose(forces_ensemble, forces_ensemble_2, atol=1e-6)
+
+    # Non-conservative calculator: method="conservative" should raise
+    calc_nc = UPETCalculator(model="pet-mad-s", version="1.5.0", non_conservative=True)
+    with pytest.raises(ValueError, match="method='conservative' is not available"):
+        calc_nc.get_forces_ensemble(atoms, method="conservative")
+
+    # method="direct" requires non_conservative_forces_ensemble in model
+    with pytest.raises(NotImplementedError, match="non_conservative_forces_ensemble"):
+        calc_cons.get_forces_ensemble(atoms, method="direct")
+
+    # Invalid method
+    with pytest.raises(ValueError, match="Unknown method"):
+        calc_cons.get_forces_ensemble(atoms, method="invalid")
+
+
 def test_error_model_not_evaluated():
     atoms = bulk("Si", cubic=True, a=5.43, crystalstructure="diamond")
     calc = UPETCalculator(