torch-sim

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

Only the energy output is supported. Forces and stresses are derived via autograd.

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)

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.

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]