Extend non_conservative to a 4-way string mode selector in ASE and TorchSim

docs/src/engines/lammps.rst

@@ -290,13 +290,23 @@ documentation.
         directory = path to a directory containing TorchScript extensions as shared
         libraries. If the model uses extensions, we will try to load them from this
         directory first
-      **non_conservative** values = on or off
-        set this to on to use non-conservative forces and stresses in your simulation,
-        typically affording a speedup factor between 2 and 3. We recommend using this in
-        combination with RESPA to obtain physically correct observables (see
-        https://arxiv.org/abs/2412.11569 for more information, and
-        https://atomistic-cookbook.org/examples/pet-mad-nc/pet-mad-nc.html for an
-        example of how to set up the RESPA run). Default to off.
+      **non_conservative** values = on or off or forces or stress
+        controls which outputs are read directly from the model rather than computed
+        via autograd on the energy:
+
+        - ``off`` (default): conservative mode; forces and stress are both derived
+          from the gradient of the energy.
+        - ``on``: both forces and stress are read directly from the model's
+          non-conservative outputs, typically affording a speedup factor between 2
+          and 3. We recommend using this in combination with RESPA to obtain
+          physically correct observables (see https://arxiv.org/abs/2412.11569 for
+          more information, and
+          https://atomistic-cookbook.org/examples/pet-mad-nc/pet-mad-nc.html for an
+          example of how to set up the RESPA run).
+        - ``forces``: forces are read directly from the model's
+          ``non_conservative_forces`` output; stress is still obtained via autograd.
+        - ``stress``: stress is read directly from the model's
+          ``non_conservative_stress`` output; forces are still obtained via autograd.
       **scale** values = float
         multiplies the contribution of the potential by a scaling factor. Defaults to 1.
       **check_consistency** values = on or off

python/metatomic_ase/CHANGELOG.md

@@ -16,6 +16,12 @@ follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased](https://github.com/metatensor/metatomic/)
 
+### Added
+
+- `non_conservative` in `MetatomicCalculator` now also accepts `"forces"` and
+  `"stress"`. `"forces"` reads forces directly from the model while still computing
+  stress via autograd; `"stress"` does the reverse.
+
 ## [Version 0.1.0](https://github.com/metatensor/metatomic/releases/tag/metatomic-ase-v0.1.0) - 2026-03-25
 
 - `metatomic-ase` is now a standalone package, containing the ASE integration

python/metatomic_ase/src/metatomic_ase/_calculator.py

@@ -2,7 +2,7 @@
 import os
 import pathlib
 import warnings
-from typing import Dict, List, Optional, Union
+from typing import Dict, List, Literal, Optional, Union
 
 import metatensor.torch as mts
 import numpy as np
@@ -125,7 +125,7 @@ def __init__(
         check_consistency=False,
         device=None,
         variants: Optional[Dict[str, Optional[str]]] = None,
-        non_conservative=False,
+        non_conservative: Union[bool, Literal["forces", "stress"]] = False,
         do_gradients_with_energy=True,
         uncertainty_threshold=0.1,
     ):
@@ -154,10 +154,23 @@ def __init__(
             taken from this variant. This behaviour can be overriden by setting the
             corresponding keys explicitly to ``None`` or to another value (e.g.
             ``{"energy_uncertainty": "r2scan"}``).
-        :param non_conservative: if ``True``, the model will be asked to compute
-            non-conservative forces and stresses. This can afford a speed-up,
-            potentially at the expense of physical correctness (especially in molecular
-            dynamics simulations).
+        :param non_conservative: controls which outputs are obtained directly from the
+            model rather than via autograd on the energy. Accepted values are:
+
+            - ``False`` (default): conservative mode; forces and stress are both
+              derived from the gradient of the energy.
+            - ``True``: both forces and stress are read directly from the model's
+              ``non_conservative_forces`` and ``non_conservative_stress`` outputs.
+            - ``"forces"``: forces come from the model's
+              ``non_conservative_forces`` output; stress is still obtained via
+              autograd.
+            - ``"stress"``: stress comes from the model's
+              ``non_conservative_stress`` output; forces are still obtained via
+              autograd.
+
+            Using any value other than ``False`` can afford a speed-up, potentially
+            at the expense of physical correctness (especially in molecular dynamics
+            simulations).
         :param do_gradients_with_energy: if ``True``, this calculator will always
             compute the energy gradients (forces and stress) when the energy is
             requested (e.g. through ``atoms.get_potential_energy()``). Because the
@@ -174,11 +187,18 @@ def __init__(
         """
         super().__init__()
 
+        _valid_nc = (True, False, "forces", "stress")
+        if non_conservative not in _valid_nc:
+            raise ValueError(
+                f"non_conservative must be one of {list(_valid_nc)}, "
+                f"got {non_conservative!r}"
+            )
+
         self.parameters = {
             "extensions_directory": extensions_directory,
             "check_consistency": bool(check_consistency),
             "variants": variants,
-            "non_conservative": bool(non_conservative),
+            "non_conservative": non_conservative,
             "do_gradients_with_energy": bool(do_gradients_with_energy),
             "additional_outputs": additional_outputs,
             "uncertainty_threshold": uncertainty_threshold,
@@ -258,7 +278,10 @@ def __init__(
         else:
             self._energy_uq_key = None
 
-        if non_conservative:
+        self._nc_forces = non_conservative in (True, "forces")
+        self._nc_stress = non_conservative in (True, "stress")
+
+        if self._nc_forces and self._nc_stress:
             if (
                 "non_conservative_stress" in variants
                 and "non_conservative_forces" in variants
@@ -273,18 +296,22 @@ def __init__(
                     "must either be both `None` or both not `None`."
                 )
 
+        if self._nc_forces:
             self._nc_forces_key = pick_output(
                 "non_conservative_forces",
                 outputs,
                 resolved_variants["non_conservative_forces"],
             )
+        else:
+            self._nc_forces_key = None
+
+        if self._nc_stress:
             self._nc_stress_key = pick_output(
                 "non_conservative_stress",
                 outputs,
                 resolved_variants["non_conservative_stress"],
             )
         else:
-            self._nc_forces_key = None
             self._nc_stress_key = None
 
         if additional_outputs is None:
@@ -512,20 +539,21 @@ def calculate(
                 atoms=atoms, dtype=self._dtype, device=self._device
             )
 
-            do_backward = False
-            if calculate_forces and not self.parameters["non_conservative"]:
-                do_backward = True
-                positions.requires_grad_(True)
+            do_backward_for_forces = calculate_forces and not self._nc_forces
+            do_backward_for_stress = calculate_stress and not self._nc_stress
+            do_backward = do_backward_for_forces or do_backward_for_stress
 
-            if calculate_stress and not self.parameters["non_conservative"]:
-                do_backward = True
+            if do_backward_for_forces:
+                positions.requires_grad_(True)
 
+            if do_backward_for_stress:
                 strain = torch.eye(
                     3, requires_grad=True, device=self._device, dtype=self._dtype
                 )
 
                 positions = positions @ strain
-                positions.retain_grad()
+                if do_backward_for_forces:
+                    positions.retain_grad()
 
                 cell = cell @ strain
 
@@ -616,7 +644,7 @@ def calculate(
                 self.results["energy"] = energy_values.numpy()[0, 0]
 
             if calculate_forces:
-                if self.parameters["non_conservative"]:
+                if self._nc_forces:
                     forces_values = outputs[self._nc_forces_key].block().values.detach()
                     # remove any spurious net force
                     forces_values = forces_values - forces_values.mean(
@@ -629,7 +657,7 @@ def calculate(
                 self.results["forces"] = forces_values.numpy()
 
             if calculate_stress:
-                if self.parameters["non_conservative"]:
+                if self._nc_stress:
                     stress_values = outputs[self._nc_stress_key].block().values.detach()
                 else:
                     stress_values = strain.grad / atoms.cell.volume
@@ -704,13 +732,16 @@ def compute_energy(
             types, positions, cell, pbc = _ase_to_torch_data(
                 atoms=atoms, dtype=self._dtype, device=self._device
             )
-            if compute_forces_and_stresses and not self.parameters["non_conservative"]:
+            if compute_forces_and_stresses and not self._nc_forces:
                 positions.requires_grad_(True)
+
+            if compute_forces_and_stresses and not self._nc_stress:
                 strain = torch.eye(
                     3, requires_grad=True, device=self._device, dtype=self._dtype
                 )
                 positions = positions @ strain
-                positions.retain_grad()
+                if not self._nc_forces:
+                    positions.retain_grad()
                 cell = cell @ strain
                 strains.append(strain)
             system = System(types, positions, cell, pbc)
@@ -770,8 +801,13 @@ def compute_energy(
             }
 
         if compute_forces_and_stresses:
-            if self.parameters["non_conservative"]:
-                results_as_numpy_arrays["forces"] = (
+            # Run backward if any output still requires autograd
+            if (not self._nc_forces) or (not self._nc_stress):
+                energy_tensor = energies.block().values
+                energy_tensor.backward(torch.ones_like(energy_tensor))
+
+            if self._nc_forces:
+                forces_values = (
                     predictions[self._nc_forces_key]
                     .block()
                     .values.squeeze(-1)
@@ -783,17 +819,18 @@ def compute_energy(
                 # split them into the original systems
                 split_sizes = [len(system) for system in systems]
                 split_indices = np.cumsum(split_sizes[:-1])
-                results_as_numpy_arrays["forces"] = np.split(
-                    results_as_numpy_arrays["forces"], split_indices, axis=0
-                )
-
+                forces_values = np.split(forces_values, split_indices, axis=0)
                 # remove net forces
                 results_as_numpy_arrays["forces"] = [
-                    f - f.mean(axis=0, keepdims=True)
-                    for f in results_as_numpy_arrays["forces"]
+                    f - f.mean(axis=0, keepdims=True) for f in forces_values
+                ]
+            else:
+                results_as_numpy_arrays["forces"] = [
+                    -system.positions.grad.cpu().numpy() for system in systems
                 ]
 
-                if all(atoms.pbc.all() for atoms in atoms_list):
+            if all(atoms.pbc.all() for atoms in atoms_list):
+                if self._nc_stress:
                     results_as_numpy_arrays["stress"] = [
                         s
                         for s in predictions[self._nc_stress_key]
@@ -803,13 +840,7 @@ def compute_energy(
                         .cpu()
                         .numpy()
                     ]
-            else:
-                energy_tensor = energies.block().values
-                energy_tensor.backward(torch.ones_like(energy_tensor))
-                results_as_numpy_arrays["forces"] = [
-                    -system.positions.grad.cpu().numpy() for system in systems
-                ]
-                if all(atoms.pbc.all() for atoms in atoms_list):
+                else:
                     results_as_numpy_arrays["stress"] = [
                         strain.grad.cpu().numpy() / atoms.cell.volume
                         for strain, atoms in zip(strains, atoms_list, strict=True)
@@ -852,19 +883,19 @@ def _ase_properties_to_metatensor_outputs(
                 output.sample_kind = "system"
 
             metatensor_outputs[self._energy_key] = output
-        if calculate_forces and self.parameters["non_conservative"]:
+        if calculate_forces and self._nc_forces:
             metatensor_outputs[self._nc_forces_key] = ModelOutput(
                 unit="eV/Angstrom",
                 sample_kind="atom",
             )
 
-        if calculate_stress and self.parameters["non_conservative"]:
+        if calculate_stress and self._nc_stress:
             metatensor_outputs[self._nc_stress_key] = ModelOutput(
                 unit="eV/Angstrom^3",
                 sample_kind="system",
             )
 
-        if calculate_stresses and self.parameters["non_conservative"]:
+        if calculate_stresses and (self._nc_forces or self._nc_stress):
             raise NotImplementedError(
                 "non conservative, per-atom stress is not yet implemented"
             )

python/metatomic_ase/tests/calculator.py

@@ -695,6 +695,56 @@ def test_variant_non_conservative_error(atoms, model, force_is_None):
         )
 
 
+@pytest.mark.parametrize("non_conservative", ["on", "off", "invalid"])
+def test_non_conservative_invalid_raises(model, non_conservative):
+    """Passing an invalid non_conservative value raises ValueError."""
+    with pytest.raises(ValueError, match="non_conservative must be one of"):
+        MetatomicCalculator(model, non_conservative=non_conservative)
+
+
+@pytest.mark.parametrize("non_conservative", ["forces", "stress"])
+def test_non_conservative_mixed_modes(tmpdir, model, atoms, non_conservative):
+    """'forces' mode uses NC forces + autograd stress; 'stress' mode is the reverse"""
+    ref = atoms.copy()
+    ref.calc = ase.calculators.lj.LennardJones(
+        sigma=SIGMA, epsilon=EPSILON, rc=CUTOFF, ro=CUTOFF, smooth=False
+    )
+
+    path = os.path.join(tmpdir, "exported-model.pt")
+    model.save(path)
+
+    # Conservative reference via calculate()
+    calc_ref = MetatomicCalculator(
+        path, check_consistency=True, non_conservative=False, uncertainty_threshold=None
+    )
+    atoms.calc = calc_ref
+    ref_forces = atoms.get_forces()
+    ref_stress = atoms.get_stress()
+
+    # Mixed-mode calculator
+    calc = MetatomicCalculator(
+        path,
+        check_consistency=True,
+        non_conservative=non_conservative,
+        uncertainty_threshold=None,
+    )
+    atoms.calc = calc
+    energy = atoms.get_potential_energy()
+    forces = atoms.get_forces()
+    stress = atoms.get_stress()
+
+    np.testing.assert_allclose(ref.get_potential_energy(), energy, rtol=1e-5)
+    assert forces.shape == ref_forces.shape
+    assert stress.shape == ref_stress.shape
+
+    if non_conservative == "stress":
+        # Forces come from autograd — must match conservative reference
+        np.testing.assert_allclose(ref_forces, forces, rtol=1e-5, atol=1e-8)
+    elif non_conservative == "forces":
+        # Stress comes from autograd — must match conservative reference
+        np.testing.assert_allclose(ref_stress, stress, rtol=1e-5, atol=1e-8)
+
+
 def test_model_without_energy(atoms):
     """
     Test that a MetatomicCalculator can be created with a model without energy

python/metatomic_torchsim/CHANGELOG.md

@@ -16,6 +16,12 @@ follows [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
 ## [Unreleased](https://github.com/metatensor/metatomic/)
 
+### Added
+
+- `non_conservative` in `MetatomicModel` now also accepts `"forces"` and `"stress"`.
+  `"forces"` reads forces directly from the model while still computing stress via
+  autograd; `"stress"` does the reverse.
+
 ## [Version 0.1.2](https://github.com/metatensor/metatomic/releases/tag/metatomic-torchsim-v0.1.2) - 2026-04-22
 
 ### Changed