Overview

TensorMap

The core type of metatensor is the TensorMap: a high dimensional block-sparse tensor containing both data and metadata. A TensorMap contains a list of blocks (represented as TensorBlock), each associated with a key; and the set of all keys is stored in a Labels object. Both these building block for TensorMap are explained in more details below.

The keys can contain multiple dimensions (in the illustration below we have two dimensions named key_1 and key_2), and each entry in the keys has one integer value for each dimension. Here for example, the first block is associated with key_1 = 0 and key_2 = 0, while the second block is associated with key_1 = 0 and key_2 = 1, and so on.

Illustration of a metatensor TensorMap object, made of a set of keys and associated TensorBlock.

Different key dimensions can have different purposes, but some typical keys dimensions you’ll encounter when working with atomistic data are the following:

• atomic species dimensions: when using a one-hot encoding of different atomic species (or atomic elements) in a structure, the resulting data is sparse, containing implicit zeros if a given species is not present in a structure. This is the case of the species_center and various species_neighbor key dimensions produced by rascaline.

• symmetry markers: Another use case for metatensor is to store and manipulate equivariant data, i.e. data that transforms in a known, specific way when the corresponding atomic structure is transformed. This is typically used to represent the symmetry property of some data with respect to rotations, by decomposing the properties of interest on a basis of spherical harmonics. When handling this kind of data, it is convenient to store and manipulate the data corresponding to different spherical harmonics (or generally different irreducible representations of the symmetry group) separately. This is the case of the spherical_harmonics_l key dimension produced by rascaline: different blocks will contain the \(\ell = 1\) and \(\ell = 2\) parts of an equivariant representation.

Labels

A fundamental part of metatensor is to carry simultaneously the data used in machine learning and the associated metadata. The first kind of metadata we encountered was the keys of a TensorMap, stored as an instance of the Labels class. This class is also used to store all other metadata in metatensor, i.e. all the metadata associated with a given TensorBlock.

Illustration of two different Labels instances, corresponding to potential samples (green, on the left) and properties (red, on the right) of a TensorBlock.

A set of Labels can be seen as a two dimensional array of integers, where each row corresponds to an entry in the data, and each column is a dimension, which is named. For example, in the illustration above, the set of Labels on the left has two dimensions (structure and center), and 10 entries (10 rows); while the Labels on the right has four dimensions and 8 entries. Depending on the language you use, Labels entries and dimensions’ names can be accessed and manipulated in different ways, please refer to the corresponding API documentation for more information.

TensorBlock

The final core object of metatensor is the TensorBlock, containing a dense array of data and metadata describing the different axes of this array. The simplest possible TensorBlock is represented below, and contains three things:

• a 2-dimensional data array;

• metadata describing the rows of this array, called samples and stored as a set of Labels;

• metadata describing the columns of this array, called properties, also stored as a set of Labels.

The samples store information about what objects does the data represent, while properties store information about how these objects are represented. Taking a couple of examples for clarity:

• if we are storing the energy of a list of structures in a TensorBlock, the samples would contain only a single "structure" dimension, and multiple entries — one per structure — going from 0 to len(structures). The properties would contain a single "energy" dimension with a single entry, which value does not carry information;

• if we are storing increasing powers of the bond lengths between pairs of atom in a structure (\((r_{ij})^k\) for \(k \in [1, n]\)), the samples would contain the index of the "first_atom" (\(i\)) and the "second_atom" (\(j\)); while the properties would contain the value of "k". The data array would contain the values of \((r_{ij})^k\).

• if we are storing an atom-centered machine learning representation, the samples would contain the index of the atom "center" and the index of the corresponding "structure"; while the properties would contain information about the e.g. the basis functions used to define the representation. The Labels figure above contains an example of samples and properties that one would find in machine learning representation.

In general, for a 2-dimensional data array, the value at index (i, j) is described by the i^th entry of the samples and the j^th entry of the properties.

Illustration of the simplest possible TensorBlock: a two dimensional data array, and two Labels describing these two axis. The metadata associated with the first axis (rows) describes samples, while the metadata associated with the second axis (columns) describes properties.

In addition to all this metadata, metatensor also carries around some data. This data can be stored in various arrays types, all integrated with metatensor. Metatensor then manipulate these arrays in an opaque way, without knowing what’s inside. This allows to integrate metatensor with multiple third-party libraries and ecosystems, for example having the data live on GPU, or using memory-mapped data arrays.

Advanced functionalities: integrating new array types with metatensor

Currently, the following array types are integrated with metatensor:

• Numpy’s ndarray from Python,

• PyTorch’s Tensor from Python and C++, including full support for autograd and different device (data living on CPU memory, GPU memory, …),

• Rust’s ndarray from Rust, more specifically ndarray::ArrayD<f64>,

• A very bare-bone N-dimensional array in metatensor C++ API: metatensor::SimpleDataArray

It is possible to integrate new array types with metatensor, look into the metatensor.data.register_external_data_wrapper() function in Python, the mts_array_t struct in C, the metatensor::DataArrayBase abstract base class in C++, and the metatensor::Array trait in Rust.

Gradients

In addition to storing data and metadata together, a TensorBlock can also store values and gradients together. The gradients are stored in another TensorBlock, associated with a parameter name, describing with respect to what the gradients are taken. Regarding metadata, the gradient properties always match the values properties; while the gradient sample are different from the value samples. The gradient samples contains both what a given row in the data is the gradient of, and with respect to what the gradient is taken. As illustrated below, multiple gradient rows can be gradients of the same values row, but with respect to different things (here the positions of different particles in the system).

Illustration of gradients stored inside a TensorBlock.

Components

There is one more thing TensorBlock can contain. When working with vectorial data, we also handle vector components in both data and metadata. In its most generic form, a TensorBlock contains a \(N\)-dimensional data array (with \(N \geqslant 2\)), and \(N\) set of Labels. The first Labels describe the samples, the last Labels describe the properties, and all the remaining Labels describe vectorial components (matching all remaining axes of the data array, from the second to one-before-last).

For example, gradients with respect to positions are actually a bit more complex than the illustration above. They always contain a supplementary axis in the data for the \(x/y/z\) direction of the gradient, associated with a component Labels. Getting back to the example where we store energy in the TensorBlock values, the gradient (i.e. the negative of the forces) samples describe with respect to which atom position we are taking gradient, and the component Labels allow to find the \(x/y/z\) component of the forces.

Illustration of a TensorBlock containing components as an extra set of metadata to describe supplementary axes of the data array.

Another use-case for components is the storage of equivariant data, where a given irreducible representation might have multiple elements. For example, when handling spherical harmonics (which are the irreducible representation of the group of 3D rotations \(SO(3)\)), all the spherical harmonics with the same \(\ell\) and corresponding \(m\) should be considered simultaneously: the different \(m\) are components of a single irreducible representation.

When handling the gradients of equivariant data, we quickly realize that we might need more than one component in a given TensorBlock. Gradients with respect to positions of an equivariant representation based on spherical harmonics will need both the gradient direction \(x/y/z\) and the spherical harmonics \(m\) as components. This impacts metadata associated with TensorBlock in two ways:

• TensorBlock can have an arbitrary number of components associated with the values, which will always occur “in between” samples and properties metadata;

• when values in a TensorBlock have components, and gradient with respect to some parameter would add more components, the resulting gradient components will contain first the new, gradient-specific components, and then all of the components already present in the values.