Systems

class metatensor.torch.atomistic.System(types: Tensor, positions: Tensor, cell: Tensor)

A System contains all the information about an atomistic system; and should be used as the input of metatensor atomistic models.

You can create a System with types, positions and cell tensors, or convert data from other libraries.

Converting data to metatensor System

We provide a way to convert ase.Atoms instances to System using the systems_to_torch() function.

In addition, some external packages provide ways to create System using data from other libraries:

• rascaline has the rascaline.torch.systems_to_torch() function that can convert from ASE, chemfiles and PySCF.

Parameters:

• types (Tensor) – 1D tensor of integer representing the particles identity. For atoms, this is typically their atomic numbers.

• positions (Tensor) – 2D tensor of shape (len(types), 3) containing the Cartesian positions of all particles in the system.

• cell (Tensor) – 2D tensor of shape (3, 3), describing the bounding box/unit cell of the system. Each row should be one of the bounding box vector; and columns should contain the x, y, and z components of these vectors (i.e. the cell should be given in row-major order). Systems are assumed to obey periodic boundary conditions, non-periodic systems should set the cell to 0.

property types: Tensor

Tensor of 32-bit integers representing the particles identity

property positions: Tensor

Tensor of floating point values containing the particles cartesian coordinates

property cell: Tensor

Tensor of floating point values containing bounding box/cell of the system

property device: device

get the device of all the arrays stored inside this System

property dtype: dtype

get the dtype of all the arrays stored inside this System

Warning

Due to limitations in TorchScript C++ extensions, the dtype is returned as an integer, which can not be compared with torch.dtype instances. See TensorBlock.dtype for more information.

to(dtype: dtype | None = None, device: device | None = None) → System

Move all the arrays in this system to the given dtype and device.

Parameters:

• dtype (dtype | None) – new dtype to use for all arrays. The dtype stays the same if this is set to None.

• device (device | None) – new device to use for all arrays. The device stays the same if this is set to None.

Return type:

System

add_neighbor_list(options: NeighborListOptions, neighbors: TensorBlock)

Add a new neighbors list in this system corresponding to the given options.

The neighbors list should have the following samples: "first_atom", "second_atom", "cell_shift_a", "cell_shift_b", "cell_shift_c", containing the index of the first and second atoms (matching the “atom” sample in the positions); and the number of cell vector a/b/c to add to the positions difference to get the pair vector.

The neighbors should also have a single component "xyz" with values [0, 1, 2]; and a single property "distance" with value 0.

The neighbors values must contain the distance vector from the first to the second atom, i.e. positions[second_atom] - positions[first_atom] + cell_shift_a * cell_a + cell_shift_b * cell_b + cell_shift_c * cell_c.

Parameters:

• options (NeighborListOptions) – options of the neighbors list

• neighbors (TensorBlock) – list of neighbors stored in a TensorBlock

get_neighbor_list(options: NeighborListOptions) → TensorBlock

Retrieve a previously stored neighbors list with the given options, or throw an error if no such neighbors list exists.

Parameters:

options (NeighborListOptions) – options of the neighbors list to retrieve

Return type:

TensorBlock

known_neighbor_lists() → List[NeighborListOptions]

Get all the neighbors lists options registered with this System

Return type:

List[NeighborListOptions]

add_data(name: str, data: TensorBlock, override: bool = False)

Add custom data to this system, stored as TensorBlock.

This is intended for experimentation with models that need more data as input, and moved into a field of System later.

Parameters:

• name (str) – name of the custom data

• data (TensorBlock) – values of the custom data

• override (bool) – if True, allow this function to override existing data with the same name

get_data(name: str) → TensorBlock

Retrieve custom data stored in this System with the given name, or throw an error if no data can be found.

Parameters:

name (str) – name of the custom data to retrieve

Return type:

TensorBlock

known_data() → List[str]

Get the name of all the custom data registered with this System

Return type:

List[str]

class metatensor.torch.atomistic.NeighborListOptions(cutoff: float, full_list: bool, strict: bool, requestor: str = '')

Options for the calculation of a neighbors list

Parameters:

• cutoff (float) – spherical cutoff radius for the neighbors list, in the model units

• full_list (bool) – should the list be a full or half neighbors list

• strict (bool) – whether the list guarantee to have no pairs farther than cutoff

• requestor (str) – who requested this neighbors list, you can add additional requestors later using add_requestor()

property cutoff: float

Spherical cutoff radius for this neighbors list in model units

property length_unit: str

The unit of length used for the cutoff.

This is typically set by MetatensorAtomisticModel when collecting all neighbors list requests.

The list of possible units is available here.

engine_cutoff(engine_length_unit: str) → float

Spherical cutoff radius for this neighbors list in engine units.

The engine must provide the unit it uses for lengths, and the cutoff will automatically be converted.

Parameters:

engine_length_unit (str)

Return type:

float

property full_list: bool

Should the list be a full neighbors list (contains both the pair i->j and j->i) or a half neighbors list (contains only the pair i->j)

property strict: bool

Does the list guarantee to have no pairs beyond the cutoff (strict) or can it also have pairs that are farther apart (non strict)

requestors() → List[str]

Get the list of modules requesting this neighbors list

Return type:

List[str]

add_requestor(requestor: str)

Add another requestor to the list of modules requesting this neighbors list

Parameters:

requestor (str)

__eq__(other: NeighborListOptions) → bool

Return self==value.

Parameters:

other (NeighborListOptions)

Return type:

bool

__ne__(other: NeighborListOptions) → bool

Return self!=value.

Parameters:

other (NeighborListOptions)

Return type:

bool

metatensor.torch.atomistic.systems_to_torch(systems: IntoSystem | List[IntoSystem], dtype: dtype | None = None, device: device | None = None, positions_requires_grad: bool = False, cell_requires_grad: bool = False) → System | List[System]

Converts a system or a list of systems into a metatensor.torch.atomistic.System or a list of such objects.

Param:

systems: The system or list of systems to convert.

Param:

dtype: The dtype of the output tensors. If None, the default dtype is used.

Param:

device: The device of the output tensors. If None, the default device is used.

Param:

positions_requires_grad: Whether the positions tensors of the outputs should require gradients.

Param:

cell_requires_grad: Whether the cell tensors of the outputs should require gradients.

Returns:

The converted system or list of systems.

Parameters:

• systems (IntoSystem | List[IntoSystem])

• dtype (dtype | None)

• device (device | None)

• positions_requires_grad (bool)

• cell_requires_grad (bool)

Return type:

System | List[System]

metatensor.torch.atomistic.register_autograd_neighbors(system: System, neighbors: TensorBlock, check_consistency: bool)

Register a new torch autograd node going from (system.positions, system.cell) to the neighbors distance vectors.

This does not recompute the distance vectors, but work as-if all the data in neighbors.values was computed directly from system.positions and system.cell, allowing downstream models to use it directly with full autograd integration.

Parameters:

• system (System) – system containing the positions and cell used to compute the neighbors list

• neighbors (TensorBlock) – neighbors list, following the same format as System.add_neighbor_list()

• check_consistency (bool) – can be set to True to run additional checks in case the data in neighbors does not follow what’s expected.