Metatomic interface

Vesin offers an interface to compute neighbor lists for metatomic’s atomistic machine learning models.

vesin.metatomic.compute_requested_neighbors(systems: List[System] | System, system_length_unit: str, model: AtomisticModel | ModelInterface, model_length_unit: str | None = None, check_consistency: bool = False)

Compute all neighbors lists requested by the model through requested_neighbor_lists() member functions, and store them inside all the systems.

Parameters:

  • systems – Single system or list of systems for which we need to compute the neighbor lists that the model requires.

  • system_length_unit – unit of length used by the data in systems

  • model – AtomisticModel or any torch.nn.Module following the ModelInterface

  • model_length_unit – unit of length used by the model, optional. This is only required when giving a raw model instead of a AtomisticModel.

  • check_consistency – whether to run additional checks on the neighbor lists validity

class vesin.metatomic.NeighborList(options: NeighborListOptions, length_unit: str, torchscript: bool = False, check_consistency: bool = False)

A neighbor list calculator that can be used with metatomic’s models.

The main difference with the other calculators is the automatic handling of different length unit between what the model expects and what the System are using.

See also

The vesin.metatomic.compute_requested_neighbors() function can be used to automatically compute and store all neighbor lists required by a given model.

Parameters:

  • options – metatomic.torch.NeighborListOptions defining the parameters of the neighbor list

  • length_unit – unit of length used for the systems data

  • torchscript – whether this function should be compatible with TorchScript or not. If True, this requires installing the vesin-torch package.

  • check_consistency – whether to run additional checks on the neighbor list validity

Example

>>> from vesin.metatomic import NeighborList
>>> from metatomic.torch import System, NeighborListOptions
>>> import torch
>>> system = System(
...     positions=torch.eye(3).requires_grad_(True),
...     cell=4 * torch.eye(3).requires_grad_(True),
...     types=torch.tensor([8, 1, 1]),
...     pbc=torch.ones(3, dtype=bool),
... )
>>> options = NeighborListOptions(cutoff=4.0, full_list=True, strict=False)
>>> calculator = NeighborList(options, length_unit="Angstrom")
>>> neighbors = calculator.compute(system)
>>> neighbors
TensorBlock
    samples (18): ['first_atom', 'second_atom', 'cell_shift_a', 'cell_shift_b', 'cell_shift_c']
    components (3): ['xyz']
    properties (1): ['distance']
    gradients: None

The returned TensorBlock can then be registered with the system

>>> system.add_neighbor_list(options, neighbors)
compute(system: System) → TensorBlock

Compute the neighbor list for the given metatomic.torch.System.

Parameters:

system –

a metatomic.torch.System containing data about a single structure. If the positions or cell of this system require gradients, the neighbors list values computational graph will be set accordingly.

The positions and cell need to be in the length unit defined for this NeighborList calculator.