Metatensor interface

vesin.torch.metatensor.compute_requested_neighbors(systems: List[System] | System, system_length_unit: str, model: MetatensorAtomisticModel | ModelInterface, model_length_unit: str | None = None)

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 – MetatensorAtomisticModel 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 MetatensorAtomisticModel.

class vesin.torch.metatensor.NeighborList(options: NeighborListOptions, length_unit: str)

A neighbor list calculator that can be used with metatensor’s atomistic 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.torch.metatensor.compute_requested_neighbors() function can be used to automatically compute and store all neighbor lists required by a given model.

Parameters:

Example

>>> from vesin.torch.metatensor import NeighborList
>>> from metatensor.torch.atomistic 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)
>>> calculator = NeighborList(options, length_unit="Angstrom")
>>> neighbors = calculator.compute(system)
>>> neighbors
TensorBlock
    samples (54): ['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 metatensor.torch.atomistic.System.

Parameters:

system –

a metatensor.torch.atomistic.System containing the data about a 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 :py:class`NeighborList` calculator.