Metatomic interface

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

vesin.metatomic.neighbor_lists_for_model(system_length_unit: str, model: AtomisticModel | ModelInterface, model_length_unit: str | None = None, torchscript: bool = False, check_consistency: bool = False) → List[NeighborList]

Get a list of calculators, corresponding to the neighbor lists requested by the model .

from vesin.metatomic import neighbor_lists_for_model

model = MyModel()  # some model following metatomic API

# Get the neighbor lists calculators for the model
calculators = neighbor_lists_for_model(systems, "angstrom", model)

systems = [System(...), System(...)]  # some systems

# Compute and store the neighbor lists in the systems
for calculator in calculators:
    calculator.add_neighbor_list(systems)
Parameters:

  • system_length_unit – unit of length used for the data in the 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.

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

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, copy: bool = True) → 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.

  • copy – whether to copy the neighbor list values before returning. If False, the neighbor list will contain a pointer to memory allocated inside the calculator, saving memory and removing one copy. In this case, the user MUST ensure that the calculator is kept alive at least as long as the neighbor list is used.

add_neighbor_list(systems: System | List[System], copy: bool = True)

Compute the neighbor list for all the given systems and add it to them.

Parameters:

  • systems –

    a system or a list of systems for which to compute and add the neighbor list. If the positions or cell of these systems 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.

  • copy – whether to copy the neighbor list values before adding them to the system. If False, the neighbor list will contain a pointer to memory allocated inside the calculator, saving memory and removing one copy. In this case, the user MUST ensure that the calculator is kept alive at least as long as the neighbor list is used. copy=False is only supported when computing neighbor list for a single system.

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.

Warning

This function is deprecated, please use vesin.metatomic.neighbor_lists_for_model() instead.

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

vesin.metatomic.compute_requested_neighbors_from_options(systems: List[System], options: List[NeighborListOptions], system_length_unit: str, check_consistency: bool) → None

Compute all neighbors lists requested by the options, and store them inside all the systems.

Warning

This function is deprecated, please use vesin.metatomic.neighbor_lists_for_model() instead.

Parameters:

  • systems – list of systems for which we need to compute the neighbor lists that required by the options.

  • options – list of NeighborListOptions

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

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