Atomic Simulation Environment (ASE) integration

The code in metatensor.torch.atomistic.ase_calculator defines a class that allow using MetatensorAtomisticModel which predict the energy of a system as an ASE calculator; enabling the use of machine learning interatomic potentials to drive simulations inside ASE.

Additionally, it allow using arbitrary models with prediction targets which are not just the energy, through the ase_calculator.MetatensorCalculator.run_model() function.

class metatensor.torch.atomistic.ase_calculator.MetatensorCalculator(model: str | bytes | PurePath | MetatensorAtomisticModel, *, extensions_directory=None, check_consistency=False, device=None)[source]

Bases: Calculator

The MetatensorCalculator class implements ASE’s ase.calculators.calculator.Calculator API using metatensor atomistic models to compute energy, forces and any other supported property.

This class can be initialized with any MetatensorAtomisticModel, and used to run simulations using ASE’s MD facilities.

Parameters:

  • model (str | bytes | PurePath | MetatensorAtomisticModel) – model to use for the calculation. This can be a file path, a Python instance of MetatensorAtomisticModel, or the output of torch.jit.script() on MetatensorAtomisticModel.

  • extensions_directory – if the model uses extensions, we will try to load them from this directory

  • check_consistency – should we check the model for consistency when running, defaults to False.

  • device – torch device to use for the calculation. If None, we will try the options in the model’s supported_device in order.

metadata() → ModelMetadata[source]

Get the metadata of the underlying model

Return type:

ModelMetadata

run_model(atoms: Atoms, outputs: Dict[str, ModelOutput], selected_atoms: Labels | None = None) → Dict[str, TensorBlock][source]

Run the model on the given atoms, computing properties according to the outputs and selected_atoms options.

The output of the model is returned directly, and as such the blocks’ values will be torch.Tensor.

This is intended as an easy way to run metatensor models on ase.Atoms when the model can predict properties not supported by the usual ASE’s calculator interface.

All the parameters have the same meaning as the corresponding ones in metatensor.torch.atomistic.ModelInterface.forward().

Parameters:

  • atoms (Atoms) – system on which to run the model

  • outputs (Dict[str, ModelOutput]) – outputs of the model that should be predicted

  • selected_atoms (Labels | None) – subset of atoms on which to run the calculation

Return type:

Dict[str, TensorBlock]

calculate(atoms: Atoms, properties: List[str], system_changes: List[str]) → Dict[str, ndarray][source]

Compute some properties with this calculator, and return them in the format expected by ASE.

This is not intended to be called directly by users, but to be an implementation detail of atoms.get_energy() and related functions. See ase.calculators.calculator.Calculator.calculate() for more information.

Parameters:
Return type:

Dict[str, ndarray]