TorchSim

Official website	How is metatomic supported?
https://torchsim.github.io/torch-sim/	Via the metatomic-torchsim package

How to install the code

Install the integration package from PyPI:

pip install metatomic-torchsim

For the full TorchSim documentation, see https://torchsim.github.io/torch-sim/.

Supported model outputs

The energy output is the primary output. Forces and stresses are derived via autograd by default. The wrapper also supports:

• Non-conservative forces/stress: use direct prediction of gradients instead of autograd (non_conservative=True)

• Energy uncertainty: per-atom uncertainty warnings when the model provides an energy_uncertainty output

• Additional outputs: request arbitrary extra model outputs via additional_outputs; results are stored as metatensor.torch.TensorMap in the additional_outputs attribute

See the MetatomicModel API documentation below for details on all parameters, and the tutorials for worked examples:

• Getting started with TorchSim – loading a model and running NVE dynamics

• Batched simulations with TorchSim – evaluating multiple systems in a single call

How to use the code

import ase.build
import torch_sim as ts
from metatomic_torchsim import MetatomicModel

model = MetatomicModel("model.pt", device="cpu")

atoms = ase.build.bulk("Si", "diamond", a=5.43, cubic=True)
sim_state = ts.initialize_state(atoms, device=model.device, dtype=model.dtype)

results = model(sim_state)
print(results["energy"])   # shape [1]
print(results["forces"])   # shape [n_atoms, 3]
print(results["stress"])   # shape [1, 3, 3]

API documentation

class metatomic_torchsim.MetatomicModel(model: str | bytes | PurePath | AtomisticModel | RecursiveScriptModule, *, extensions_directory: str | bytes | PurePath | None = None, device: device | str | None = None, check_consistency: bool = False, compute_forces: bool = True, compute_stress: bool = True, variants: Dict[str, str | None] | None = None, non_conservative: bool = False, uncertainty_threshold: float | None = 0.1, additional_outputs: Dict[str, ModelOutput] | None = None)

Bases: ModelInterface

TorchSim wrapper for metatomic atomistic models.

Wraps a metatomic model to compute energies, forces, and stresses within the TorchSim framework. Handles the translation between TorchSim’s batched SimState and metatomic’s list-of-System convention, and uses autograd for force/stress derivatives.

Neighbor lists are computed with vesin, or with nvalchemiops on CUDA when available and the model requests full neighbor lists.

Parameters:

• model (str | bytes | PurePath | AtomisticModel | ForwardRef('torch.jit.RecursiveScriptModule')) – Model to use. Accepts a file path to a .pt saved model, a Python AtomisticModel instance, or a TorchScript torch.jit.RecursiveScriptModule.

• extensions_directory (str | bytes | PurePath | None) – Directory containing compiled TorchScript extensions required by the model, if any.

• device (device | str | None) – Torch device for evaluation. When None, the best device is selected from the model’s supported_devices.

• check_consistency (bool) – Run consistency checks during model evaluation. Useful for debugging but hurts performance.

• compute_forces (bool) – Compute atomic forces via autograd.

• compute_stress (bool) – Compute stress tensors via the strain trick.

• variants (Dict[str, str | None] | None) – Dictionary mapping output names to a variant that should be used. Setting {"energy": "pbe"} selects the "energy/pbe" output. The energy variant propagates to uncertainty and non-conservative outputs unless overridden (e.g. {"energy": "pbe", "energy_uncertainty": "r2scan"} would select energy/pbe and energy_uncertainty/r2scan).

• non_conservative (bool) – 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).

• uncertainty_threshold (float | None) – Threshold for per-atom energy uncertainty in eV. When the model supports energy_uncertainty with per_atom=True, atoms exceeding this threshold trigger a warning. Set to None to disable.

• additional_outputs (Dict[str, ModelOutput] | None) – Dictionary of extra ModelOutput to request from the model. Results are stored in additional_outputs after each forward call.

additional_outputs: Dict[str, TensorMap]

Additional outputs computed by forward() are stored here. Keys match the additional_outputs parameter to the constructor; values are raw metatensor.torch.TensorMap from the model.

forward(state: SimState) → Dict[str, Tensor]

Compute energies, forces, and stresses for the given simulation state.

Parameters:

state (SimState) – TorchSim simulation state

Returns:

Dictionary with "energy" (shape [n_systems]), "forces" (shape [n_atoms, 3], if compute_forces), and "stress" (shape [n_systems, 3, 3], if compute_stress).

Return type:

Dict[str, Tensor]