Atomic Simulation Environment (ASE) integration

The code in metatomic.torch.ase_calculator defines a class that allow using AtomisticModel 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.MetatomicCalculator.run_model() function.

class metatomic.torch.ase_calculator.MetatomicCalculator(model: str | bytes | PurePath | AtomisticModel, *, additional_outputs: Dict[str, ModelOutput] | None = None, extensions_directory=None, check_consistency=False, device=None, non_conservative=False)

Bases: Calculator

The MetatomicCalculator 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 AtomisticModel, and used to run simulations using ASE’s MD facilities.

Neighbor lists are computed using the fast vesin neighbor list library, unless the system has mixed periodic and non-periodic boundary conditions (which are not yet supported by vesin), in which case the slower ASE neighbor list is used.

Parameters:

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

• additional_outputs (Dict[str, <torch.ScriptClass object at 0x7fd40c1f1330>]) – Dictionary of additional outputs to be computed by the model. These outputs will always be computed whenever the calculate() function is called (e.g. by ase.Atoms.get_potential_energy(), ase.optimize.optimize.Dynamics.run(), etc.) and stored in the additional_outputs attribute. If you want more control over when and how to compute specific outputs, you should use run_model() instead.

• 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.

• 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).

additional_outputs: ScriptClass object at 0x7fd40c1f1330>]

Additional outputs computed by calculate() are stored in this dictionary.

The keys will match the keys of the additional_outputs parameters to the constructor; and the values will be the corresponding raw metatensor.torch.TensorMap produced by the model.

todict()

Obtain a dictionary of parameter information

metadata() → ModelMetadata

Get the metadata of the underlying model

Return type:

ModelMetadata

run_model(atoms: ~ase.atoms.Atoms | ~typing.List[~ase.atoms.Atoms], outputs: ~typing.Dict[str, ~metatomic.torch.documentation.ModelOutput], selected_atoms: <torch.ScriptClass object at 0x7fd40c149e30> | None = None) → ScriptClass object at 0x7fd40c1f1330>]

Run the model on the given atoms, computing the requested outputs and only these.

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 compute outputs not supported by the standard ASE’s calculator interface.

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

Parameters:

• atoms (Atoms | List[Atoms]) – ase.Atoms, or list of ase.Atoms, on which to run the model

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

• selected_atoms (<torch.ScriptClass object at 0x7fd40c149e30> | None) – subset of atoms on which to run the calculation

Return type:

Dict[str, <torch.ScriptClass object at 0x7fd40c1f1330>]

calculate(atoms: Atoms, properties: List[str], system_changes: List[str]) → None

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:

• atoms (Atoms)

• properties (List[str])

• system_changes (List[str])

Return type:

None

compute_energy(atoms: Atoms | List[Atoms], compute_forces_and_stresses: bool = False) → Dict[str, float | ndarray | List[float | ndarray]]

Compute the energy of the given atoms.

Energies are computed in eV, forces in eV/Å, and stresses in 3x3 tensor format and in units of eV/Å^3.

Parameters:

• atoms (Atoms | List[Atoms]) – ase.Atoms, or list of ase.Atoms, on which to run the model

• compute_forces_and_stresses (bool) – if True, the model will also compute forces and stresses. IMPORTANT: stresses will only be computed if all provided systems have periodic boundary conditions in all directions.

Returns:

A dictionary with the computed properties. The dictionary will contain the energy as a float, and, if requested, the forces and stress as numpy arrays. In case of a list of ase.Atoms, the dictionary values will instead be lists of the corresponding properties, in the same format.

Return type:

Dict[str, float | ndarray | List[float | ndarray]]