.. _core-classes-overview:
Overview
========
This page presents the core classes of metatensor from the ground-up in a
somewhat abstract way, without being tied to any programming language API
specifics. If you prefer to read concrete examples and tutorials, you should
start with our :ref:`first steps ` tutorial instead!
.. py:currentmodule:: metatensor
.. _concept-TensorMap:
TensorMap
^^^^^^^^^
The core type of metatensor is the :py:class:`TensorMap`: a high dimensional
block-sparse tensor containing both data and metadata. A TensorMap contains a
list of blocks (represented as :ref:`concept-TensorBlock`), each associated with
a key; and the set of all keys is stored in a :ref:`concept-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.
.. figure:: ../../static/images/TensorMap.*
:width: 600px
:align: center
Illustration of a metatensor TensorMap object, made of a set of keys and
associated :ref:`concept-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 types dimensions**: when using a one-hot encoding of different
atomic types (or atomic elements) in a structure, the resulting data is
sparse, containing implicit zeros if a given type is not present in a
structure. This is the case of the ``center_type`` and various
``neighbor_type`` 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 ``o3_lambda`` key dimension produced by
`rascaline`_: different blocks will contain the :math:`\lambda = 1` and
:math:`\lambda = 2` parts of an equivariant representation.
.. _rascaline: https://github.com/Luthaf/rascaline/
.. _concept-Labels:
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 :py:class:`TensorMap`, stored as an instance of
the :py:class:`Labels` class. This class is also used to store all other
metadata in metatensor, i.e. all the metadata associated with a given
:py:class:`TensorBlock`.
.. _fig-labels:
.. figure:: ../../static/images/Labels.*
:width: 600px
:align: center
Illustration of two different :py:class:`Labels` instances, corresponding to
potential *samples* (green, on the left) and *properties* (red, on the
right) of a :py:class:`TensorBlock`.
A set of :py:class:`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, :py:class:`Labels` entries and
dimensions' names can be accessed and manipulated in different ways, please
refer to the corresponding :ref:`API documentation ` for
more information.
.. _concept-TensorBlock:
TensorBlock
^^^^^^^^^^^
The final core object of metatensor is the :py:class:`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 :py:class:`Labels`;
- metadata describing the columns of this array, called **properties**, also
stored as a set of :py:class:`Labels`.
The samples store information about **what objects** the data represents, 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 systems in a TensorBlock, the
samples would contain only a single ``"system"`` dimension, and multiple
entries — one per structure — going from 0 to ``len(systems)``. 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 (:math:`(r_{ij})^k` for :math:`k \in [1, n]`), the samples
would contain the index of the ``"first_atom"`` (:math:`i`) and the
``"second_atom"`` (:math:`j`); while the properties would contain the value of
``"k"``. The data array would contain the values of :math:`(r_{ij})^k`.
- if we are storing an atom-centered machine learning representation, the
samples would contain the index of the atom ``"atom"`` and the index of the
corresponding ``"system"``; while the properties would contain information
about the e.g. the basis functions used to define the representation. The
:ref:`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``:superscript:`th` entry of the samples and the
``j``:superscript:`th` entry of the properties.
.. figure:: ../../static/images/TensorBlock-Basic.*
:width: 300px
:align: center
Illustration of the simplest possible :py:class:`TensorBlock`: a two
dimensional data array, and two :py:class:`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.
.. admonition:: 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``,
- A very bare-bone N-dimensional array in metatensor C++ API:
:cpp:class:`metatensor::SimpleDataArray`
It is possible to integrate new array types with metatensor, look into the
:py:func:`metatensor.data.register_external_data_wrapper` function in Python, the
:c:struct:`mts_array_t` struct in C, the :cpp:class:`metatensor::DataArrayBase`
abstract base class in C++, and the `metatensor::Array`_ trait in Rust.
.. _Numpy's ndarray: https://numpy.org/doc/stable/reference/arrays.ndarray.html
.. _PyTorch's Tensor: https://pytorch.org/docs/stable/tensors.html
.. _Rust's ndarray: https://docs.rs/ndarray/
.. _metatensor::Array: ../reference/rust/metatensor/trait.Array.html
Gradients
---------
In addition to storing data and metadata together, a :py:class:`TensorBlock` can
also store values and gradients together. The gradients are stored in another
:py:class:`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).
.. figure:: ../../static/images/TensorBlock-Gradients.*
:width: 550px
:align: center
Illustration of gradients stored inside a :py:class:`TensorBlock`.
.. TODO: explain how the gradient sample works in a separate tutorial
Components
----------
There is one more thing :py:class:`TensorBlock` can contain. When working with
vectorial data, we also handle vector **components** in both data and metadata.
In its most generic form, a :py:class:`TensorBlock` contains a
:math:`N`-dimensional data array (with :math:`N \geqslant 2`), and :math:`N` set
of :py:class:`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 :math:`x/y/z` direction of the gradient, associated with a
**component** :py:class:`Labels`. Getting back to the example where we store
energy in the :py:class:`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 :py:class:`Labels` allow to find the :math:`x/y/z`
component of the forces.
.. figure:: ../../static/images/TensorBlock-Components.*
:width: 400px
:align: center
Illustration of a :py:class:`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`_ :math:`SO(3)`), all the spherical harmonics
:math:`Y_\lambda^\mu` with the same angular momentum :math:`\lambda` and
corresponding :math:`\mu` should be considered simultaneously: the different
:math:`\mu` are **components** of a single irreducible representation.
.. _group of 3D rotations: https://en.wikipedia.org/wiki/3D_rotation_group
When handling the gradients of equivariant data, we quickly realize that we
might need more than one component in a given :py:class:`TensorBlock`. Gradients
with respect to positions of an equivariant representation based on spherical
harmonics will need both the gradient direction :math:`x/y/z` and the spherical
harmonics :math:`m` as components. This impacts metadata associated with
:py:class:`TensorBlock` in two ways:
- :py:class:`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 :py:class:`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.