From 6ed9ebdfad4c7a7ec0b6872e03d10a3af9f7e053 Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Fri, 24 Apr 2026 08:38:15 +0200 Subject: [PATCH 1/2] Change `non_conservative` to 4-way string selector --- python/metatomic_ase/CHANGELOG.md | 6 + .../src/metatomic_ase/_calculator.py | 109 +++++++++++------- python/metatomic_ase/tests/calculator.py | 50 ++++++++ python/metatomic_torchsim/CHANGELOG.md | 6 + .../metatomic_torchsim/_model.py | 101 ++++++++++------ python/metatomic_torchsim/tests/torchsim.py | 51 ++++++++ 6 files changed, 246 insertions(+), 77 deletions(-) diff --git a/python/metatomic_ase/CHANGELOG.md b/python/metatomic_ase/CHANGELOG.md index 5ac049d8..54af70a1 100644 --- a/python/metatomic_ase/CHANGELOG.md +++ b/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 diff --git a/python/metatomic_ase/src/metatomic_ase/_calculator.py b/python/metatomic_ase/src/metatomic_ase/_calculator.py index 0973c712..ac7e7991 100644 --- a/python/metatomic_ase/src/metatomic_ase/_calculator.py +++ b/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" ) diff --git a/python/metatomic_ase/tests/calculator.py b/python/metatomic_ase/tests/calculator.py index 1dd62f4f..d0d54ea3 100644 --- a/python/metatomic_ase/tests/calculator.py +++ b/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 diff --git a/python/metatomic_torchsim/CHANGELOG.md b/python/metatomic_torchsim/CHANGELOG.md index 561dbc5a..40caeeb3 100644 --- a/python/metatomic_torchsim/CHANGELOG.md +++ b/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 diff --git a/python/metatomic_torchsim/metatomic_torchsim/_model.py b/python/metatomic_torchsim/metatomic_torchsim/_model.py index 500aa37e..ff5a0101 100644 --- a/python/metatomic_torchsim/metatomic_torchsim/_model.py +++ b/python/metatomic_torchsim/metatomic_torchsim/_model.py @@ -13,7 +13,7 @@ import os import pathlib import warnings -from typing import Dict, List, Optional, Union +from typing import Dict, List, Literal, Optional, Union import torch from metatensor.torch import TensorMap @@ -73,7 +73,7 @@ def __init__( compute_forces: bool = True, compute_stress: bool = True, variants: Optional[Dict[str, Optional[str]]] = None, - non_conservative: bool = False, + non_conservative: Union[bool, Literal["forces", "stress"]] = False, uncertainty_threshold: Optional[float] = 0.1, additional_outputs: Optional[Dict[str, ModelOutput]] = None, ) -> None: @@ -95,10 +95,23 @@ def __init__( non-conservative outputs unless overridden (e.g. ``{"energy": "pbe", "energy_uncertainty": "r2scan"}`` would select ``energy/pbe`` and ``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 uncertainty_threshold: Threshold for per-atom energy uncertainty in eV. When the model supports ``energy_uncertainty`` with ``sample_kind="atom"``, atoms exceeding this threshold trigger a warning. @@ -189,8 +202,17 @@ def __init__( self._energy_uq_key = None # Non-conservative outputs - self._non_conservative = non_conservative - if non_conservative: + _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._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 @@ -205,18 +227,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 # Additional outputs @@ -266,15 +292,14 @@ def __init__( run_outputs[self._energy_uq_key] = ModelOutput( unit="eV", sample_kind="atom" ) - if self._non_conservative: - if self._compute_forces: - run_outputs[self._nc_forces_key] = ModelOutput( - unit="eV/Angstrom", sample_kind="atom" - ) - if self._compute_stress: - run_outputs[self._nc_stress_key] = ModelOutput( - unit="eV/Angstrom^3", sample_kind="system" - ) + if self._nc_forces and self._compute_forces: + run_outputs[self._nc_forces_key] = ModelOutput( + unit="eV/Angstrom", sample_kind="atom" + ) + if self._nc_stress and self._compute_stress: + run_outputs[self._nc_stress_key] = ModelOutput( + unit="eV/Angstrom^3", sample_kind="system" + ) run_outputs.update(self._additional_output_requests) self._evaluation_options = ModelEvaluationOptions( @@ -314,8 +339,8 @@ def forward(self, state: "ts.SimState") -> Dict[str, torch.Tensor]: ) # Determine whether autograd is needed - do_autograd_forces = self._compute_forces and not self._non_conservative - do_autograd_stress = self._compute_stress and not self._non_conservative + do_autograd_forces = self._compute_forces and not self._nc_forces + do_autograd_stress = self._compute_stress and not self._nc_stress # Build per-system System objects. Metatomic expects a list of System # rather than a single batched graph. @@ -398,24 +423,24 @@ def forward(self, state: "ts.SimState") -> Dict[str, torch.Tensor]: stacklevel=2, ) - # Forces and stresses - if self._non_conservative: - if self._compute_forces: - nc_forces = model_outputs[self._nc_forces_key].block().values.detach() - nc_forces = nc_forces.reshape(-1, 3) - # Remove spurious net force per system - for sys_idx in range(n_systems): - mask = state.system_idx == sys_idx - sys_forces = nc_forces[mask] - nc_forces[mask] = sys_forces - sys_forces.mean(dim=0, keepdim=True) - results["forces"] = nc_forces - - if self._compute_stress: - nc_stress = model_outputs[self._nc_stress_key].block().values.detach() - nc_stress = nc_stress.reshape(n_systems, 3, 3) - results["stress"] = nc_stress - - elif do_autograd_forces or do_autograd_stress: + # Forces and stresses; NC outputs are read directly from model + if self._compute_forces and self._nc_forces: + nc_forces = model_outputs[self._nc_forces_key].block().values.detach() + nc_forces = nc_forces.reshape(-1, 3) + # Remove spurious net force per system + for sys_idx in range(n_systems): + mask = state.system_idx == sys_idx + sys_forces = nc_forces[mask] + nc_forces[mask] = sys_forces - sys_forces.mean(dim=0, keepdim=True) + results["forces"] = nc_forces + + if self._compute_stress and self._nc_stress: + nc_stress = model_outputs[self._nc_stress_key].block().values.detach() + nc_stress = nc_stress.reshape(n_systems, 3, 3) + results["stress"] = nc_stress + + # Forces and stresses; autograd outputs + if do_autograd_forces or do_autograd_stress: grad_inputs: List[torch.Tensor] = [] if do_autograd_forces: for system in systems: diff --git a/python/metatomic_torchsim/tests/torchsim.py b/python/metatomic_torchsim/tests/torchsim.py index 518da07f..df398950 100644 --- a/python/metatomic_torchsim/tests/torchsim.py +++ b/python/metatomic_torchsim/tests/torchsim.py @@ -557,3 +557,54 @@ def test_additional_outputs_invalid_raises(lj_model): device=DEVICE, additional_outputs={"bad": "not a ModelOutput"}, ) + + +@pytest.mark.parametrize("non_conservative", ["on", "off", "invalid"]) +def test_non_conservative_invalid_raises(lj_model, non_conservative): + """Passing an invalid non_conservative value raises ValueError.""" + with pytest.raises(ValueError, match="non_conservative must be one of"): + MetatomicModel(model=lj_model, device=DEVICE, non_conservative=non_conservative) + + +@pytest.mark.parametrize("non_conservative", ["forces", "stress"]) +def test_non_conservative_mixed_modes(lj_model, ni_atoms, non_conservative): + """'forces' mode uses NC forces + autograd stress; 'stress' mode is the reverse. + + The autograd side must match the conservative reference; the NC side has zero + net force (forces mode) or is simply present with the correct shape (stress mode). + """ + model_ref = MetatomicModel( + model=lj_model, device=DEVICE, uncertainty_threshold=None + ) + model_nc = MetatomicModel( + model=lj_model, + device=DEVICE, + non_conservative=non_conservative, + uncertainty_threshold=None, + ) + + sim_state = ts.io.atoms_to_state([ni_atoms], DEVICE, DTYPE) + out_ref = model_ref(sim_state) + out_nc = model_nc(sim_state) + + assert "energy" in out_nc + assert "forces" in out_nc + assert "stress" in out_nc + assert out_nc["forces"].shape == (len(ni_atoms), 3) + assert out_nc["stress"].shape == (1, 3, 3) + + if non_conservative == "forces": + # Stress comes from autograd — must match conservative reference + torch.testing.assert_close( + out_nc["stress"], out_ref["stress"], atol=1e-8, rtol=0 + ) + # Forces come from NC output — zero net force + net_force = out_nc["forces"].sum(dim=0) + torch.testing.assert_close( + net_force, torch.zeros(3, dtype=DTYPE), atol=1e-6, rtol=0 + ) + elif non_conservative == "stress": + # Forces come from autograd — must match conservative reference + torch.testing.assert_close( + out_nc["forces"], out_ref["forces"], atol=1e-8, rtol=0 + ) From 8642b5b1cc2ec1dee3253e5276179daee9414bb7 Mon Sep 17 00:00:00 2001 From: Philip Loche Date: Fri, 24 Apr 2026 09:34:47 +0200 Subject: [PATCH 2/2] Update lammps docs --- docs/src/engines/lammps.rst | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/docs/src/engines/lammps.rst b/docs/src/engines/lammps.rst index 62cb7adf..9fe0d4c1 100644 --- a/docs/src/engines/lammps.rst +++ b/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