Systems

class metatomic.torch.System(types: Tensor, positions: Tensor, cell: Tensor, pbc: Tensor)

A System contains all the information about an atomistic system; and should be used as the input of 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:

• featomic has the featomic.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 that are not periodic along one or more directions should set the corresponding cell vectors to 0.

• pbc (Tensor) – tensor containing 3 boolean values, indicating which dimensions are periodic along each axis, in the same order as the cell vectors.

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 pbc: Tensor

Tensor of boolean values indicating which dimensions are periodic

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, non_blocking: bool = False) → 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.

• non_blocking (bool) – If this is True, the function tries to move the data asynchronously. See torch.Tensor.to() for more information.

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, tensor: TensorMap, 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

• tensor (TensorMap) – the data to store

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

get_data(name: str) → TensorMap

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:

TensorMap

known_data() → List[str]

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

Return type:

List[str]

class metatomic.torch.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 AtomisticModel 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

metatomic.torch.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 metatomic.torch.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]

class metatomic.torch.systems_to_torch.IntoSystem

A type that can be converted into a metatomic.torch.System.

This is an abstract class that is used to indicate a class whose objects can be converted into a System. For the moment, the only supported type is ase.Atoms.

metatomic.torch.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.

Serialization

Below are functions to load and save metatomic systems to disk. The serialization format is based on numpy’s .npz.

metatomic.torch.save(file: str | Path | BinaryIO, system: System) → None

Save a System object to a file.

The provided System must contain float64 data and be on the CPU device.

The saved file will be a zip archive containing the following files:

• types.npy, containing the atomic types in numpy’s NPY format;

• positions.npy, containing the systems’ positions in numpy’s NPY format;

• cell.npy, containing the systems’ cell in numpy’s NPY format;

• pbc.npy, containing the periodic boundary conditions in numpy’s NPY format;

For each neighbor list in the System object, the following files will be saved (where {nl_idx} is the index of the neighbor list):

• pairs/{nl_idx}/options.json: the NeighborListOptions object converted to a JSON string.

• pairs/{nl_idx}/data.mts: the neighbor list TensorBlock object

For each extra data in the System object, the following file will be saved (where {name} is the name of the extra data):

• data/{name}.mts: The extra data TensorMap

Parameters:

• file (str | Path | BinaryIO) – The path (or file-like object) to save the System to.

• system (System) – The System object to save.

Return type:

None

metatomic.torch.load_system(file: str | Path | BinaryIO) → System

Load a System object from a file.

The loaded System object will be on the CPU device and contain float64 data.

Parameters:

file (str | Path | BinaryIO) – The path (or file-like object) to load the System object from.

Returns:

The System object.

Return type:

System

metatomic.torch.save_buffer(system: System) → Tensor

Save the given system to an in-memory buffer, represented as a 1-dimensional torch.Tensor with uint8 dtype.

Parameters:

system (System) – The System object to serialize and save.

Return type:

Tensor

metatomic.torch.load_system_buffer(buffer: Tensor) → System

Load a previously saved System from an in-memory buffer, stored inside a 1-dimensional torch.Tensor of uint8.

Parameters:

buffer (Tensor) – CPU tensor with uint8 dtype representing an in-memory buffer.

Return type:

System