Models

class metatensor.torch.atomistic.MetatensorAtomisticModel(module: Module, capabilities: ModelCapabilities)

MetatensorAtomisticModel is the main entry point for atomistic machine learning based on metatensor. It is the interface between custom, user-defined models and simulation engines. Users should wrap their models with this class, and use export() to save and export the model to a file. The exported models can then be loaded by a simulation engine to compute properties of atomistic systems.

When wrapping a module, you should declare what the model is capable of (using ModelCapabilities). This includes what units the model expects as input and what properties the model can compute (using ModelOutput). The simulation engine will then ask the model to compute some subset of these properties (through a metatensor.torch.atomistic.ModelEvaluationOptions), on all or a subset of atoms of an atomistic system.

Additionally, the wrapped module can request neighbors list to be computed by the simulation engine, and stored inside the input System. This is done by defining requested_neighbors_lists(self) -> List[NeighborsListOptions] on the wrapped model or any of it’s sub-module. MetatensorAtomisticModel will unify identical requests before storing them and exposing it’s own requested_neighbors_lists() that should be used by the engine to know what it needs to compute.

There are several requirements on the wrapped module must satisfy. The main one is concerns the forward() function, which must have the following signature:

>>> import torch
>>> from typing import List, Dict, Optional
>>> from metatensor.torch import Labels, TensorBlock
>>> from metatensor.torch.atomistic import ModelOutput, System
>>> class CustomModel(torch.nn.Module):
...     def forward(
...         self,
...         systems: List[System],
...         outputs: Dict[str, ModelOutput],
...         selected_atoms: Optional[Labels] = None,
...     ) -> Dict[str, TensorMap]: ...
...

The returned dictionary should have the same keys as outputs, and the values should contains the corresponding properties of the systems, as computed for the subset of atoms defined in selected_atoms. For some specific outputs, there are additional constrains on how the associated metadata should look like, documented in the Standard model outputs section.

Additionally, the wrapped module should not already be compiled by TorchScript, and should be in “eval” mode (i.e. module.training should be False).

For example, a custom module predicting the energy as a constant time the number of atoms could look like this

>>> class ConstantEnergy(torch.nn.Module):
...     def __init__(self, constant: float):
...         super().__init__()
...         self.constant = torch.tensor(constant).reshape(1, 1)
...
...     def forward(
...         self,
...         systems: List[System],
...         outputs: Dict[str, ModelOutput],
...         selected_atoms: Optional[Labels] = None,
...     ) -> Dict[str, TensorMap]:
...         results: Dict[str, TensorMap] = {}
...         if "energy" in outputs:
...             if outputs["energy"].per_atom:
...                 raise NotImplementedError("per atom energy is not implemented")
...
...             dtype = systems[0].positions.dtype
...             energies = torch.zeros(len(systems), 1, dtype=dtype)
...             for i, system in enumerate(systems):
...                 if selected_atoms is None:
...                     n_atoms = len(system)
...                 else:
...                     n_atoms = len(selected_atoms)
...
...                 energies[i] = self.constant * n_atoms
...
...             systems_idx = torch.tensor([[i] for i in range(len(systems))])
...             energy_block = TensorBlock(
...                 values=energies,
...                 samples=Labels(["system"], systems_idx.to(torch.int32)),
...                 components=[],
...                 properties=Labels(["energy"], torch.IntTensor([[0]])),
...             )
...
...             results["energy"] = TensorMap(
...                 keys=Labels(["_"], torch.IntTensor([[0]])),
...                 blocks=[energy_block],
...             )
...
...         return results
...

Wrapping and exporting this model would then look like this:

>>> import os
>>> import tempfile
>>> from metatensor.torch.atomistic import MetatensorAtomisticModel
>>> from metatensor.torch.atomistic import ModelCapabilities, ModelOutput
>>> model = ConstantEnergy(constant=3.141592)
>>> # put the model in inference mode
>>> model = model.eval()
>>> # Define the model capabilities
>>> capabilities = ModelCapabilities(
...     length_unit="angstrom",
...     species=[1, 2, 6, 8, 12],
...     outputs={
...         "energy": ModelOutput(
...             quantity="energy",
...             unit="eV",
...             per_atom=False,
...             explicit_gradients=[],
...         ),
...     },
... )
>>> # wrap the model
>>> wrapped = MetatensorAtomisticModel(model, capabilities)
>>> # export the model
>>> with tempfile.TemporaryDirectory() as directory:
...     wrapped.export(os.path.join(directory, "constant-energy-model.pt"))
...

Parameters:

• module (Module) – The torch module to wrap and export.

• capabilities (ModelCapabilities) – Description of the model capabilities.

wrapped_module() → Module

Get the module wrapped in this MetatensorAtomisticModel

Return type:

Module

capabilities() → ModelCapabilities

Get the capabilities of the wrapped model

Return type:

ModelCapabilities

requested_neighbors_lists(length_unit: str | None = None) → List[NeighborsListOptions]

Get the neighbors lists required by the wrapped model or any of the child module.

Parameters:

length_unit (str | None) – If not None, this should contain a known unit of length. The returned neighbors lists will use this to set the engine_cutoff field.

Return type:

List[NeighborsListOptions]

forward(systems: List[System], options: ModelEvaluationOptions, check_consistency: bool) → Dict[str, TensorMap]

Run the wrapped model and return the corresponding outputs.

Before running the model, this will convert the system data from the engine unit to the model unit, including all neighbors lists distances.

After running the model, this will convert all the outputs from the model units to the engine units.

Parameters:

• system – input system on which we should run the model. The system should already contain all neighbors lists corresponding to the options in requested_neighbors_lists().

• options (ModelEvaluationOptions) – options for this run of the model

• check_consistency (bool) – Should we run additional check that everything is consistent? This should be set to True when verifying a model, and to False once you are sure everything is running fine.

• systems (List[System]) –

Returns:

A dictionary containing all the model outputs

Return type:

Dict[str, TensorMap]

export(file: str, collect_extensions: str | None = None)

Export this model to a file that can then be loaded by simulation engine.

Parameters:

• file (str) – where to save the model. This can be a path or a file-like object.

• collect_extensions (str | None) – if not None, all currently loaded PyTorch extension will be collected in this directory. If this directory already exists, it is removed and re-created.

metatensor.torch.atomistic.check_atomistic_model(path: str)

Check that the file at path contains an exported metatensor atomistic model, and that this model can be loaded in the current process.

This function should be called before torch.jit.load() when loading an existing model.

Parameters:

path (str) –

Model-related metadata

class metatensor.torch.atomistic.ModelOutput(quantity: str = '', unit: str = '', per_atom: bool = False, explicit_gradients: List[str] = [])

Description of one of the quantity a model can compute.

Parameters:

• quantity (str) –

• unit (str) –

• per_atom (bool) –

• explicit_gradients (List[str]) –

quantity: str

Quantity of the output (e.g. energy, dipole, …). If this is an empty string, no unit conversion will be performed.

unit: str

Unit of the output. If this is an empty string, no unit conversion will be performed.

per_atom: bool

Is the output defined per-atom or for the overall structure

explicit_gradients: List[str]

Which gradients should be computed eagerly and stored inside the output TensorMap.

class metatensor.torch.atomistic.ModelCapabilities(length_unit: str = '', species: List[int] = [], outputs: Dict[str, ModelOutput] = {})

Description of a model capabilities, i.e. everything a model can do.

Parameters:

• length_unit (str) –

• species (List[int]) –

• outputs (Dict[str, ModelOutput]) –

length_unit: str

unit of lengths the model expects as input

species: List[int]

which atomic species the model can handle

outputs: Dict[str, ModelOutput]

All possible outputs from this model and corresponding settings.

During a specific run, a model might be asked to only compute a subset of these outputs.

class metatensor.torch.atomistic.ModelEvaluationOptions(length_unit: str = '', outputs: Dict[str, ModelOutput] = {}, selected_atoms: Labels | None = None)

Options requested by the simulation engine/evaluation code when doing a single model evaluation.

Parameters:

• length_unit (str) –

• outputs (Dict[str, ModelOutput]) –

• selected_atoms (Labels | None) –

length_unit: str

unit of lengths the engine uses for the model input

outputs: Dict[str, ModelOutput]

requested outputs for this run and corresponding settings

selected_atoms: Labels | None

Only run the calculation for a selected subset of atoms.

If this is set to None, run the calculation on all atoms. If this is a set of metatensor.torch.Labels, it will have two dimensions named "system" and "atom", containing the 0-based indices of all the atoms in the selected subset.