TensorBlock

class metatensor.TensorBlock(values: Array, samples: Labels, components: Sequence[Labels], properties: Labels)

Basic building block for a TensorMap.

A single block contains a n-dimensional metatensor.data.Array, and n sets of Labels (one for each dimension). The first dimension is the samples dimension, the last dimension is the properties dimension. Any intermediate dimension is called a component dimension.

Samples should be used to describe what we are representing, while properties should contain information about how we are representing it. Finally, components should be used to describe vectorial or tensorial components of the data.

A block can also contain gradients of the values with respect to a variety of parameters. In this case, each gradient is a TensorBlock with a separate set of samples and possibly components, but which shares the same property labels as the original TensorBlock.

>>> import numpy as np
>>> block = TensorBlock(
...     values=np.array(
...         [
...             [1, 2, 4],
...             [3, 5, 6],
...         ]
...     ),
...     samples=Labels("samples", np.array([[4], [2]])),
...     components=[],
...     properties=Labels("properties", np.array([[0], [1], [2]])),
... )
>>> block
TensorBlock
    samples (2): ['samples']
    components (): []
    properties (3): ['properties']
    gradients: None
>>> block.samples
Labels(
    samples
       4
       2
)
>>> block.values[block.samples.position([2])]
array([3, 5, 6])

Parameters:

• values (Array) – array containing the values for this block

• samples (Labels) – labels describing the samples (first dimension of the array)

• components (Sequence[Labels]) – list of labels describing the components (intermediate dimensions of the array). This should be an empty list for scalar/invariant data.

• properties (Labels) – labels describing the properties (last dimension of the array)

property shape

Get the shape of the values array in this block.

copy() → TensorBlock

get a deep copy of this block, including all the data and metadata

Return type:

TensorBlock

property values: Array

Get the values for this block.

The array type depends on how the block was created. Currently, numpy ndarray and torch Tensor are supported.

property samples: Labels

Get the sample Labels for this block.

The entries in these labels describe the first dimension of the values array.

property components: List[Labels]

Get the component Labels for this block.

The entries in these labels describe intermediate dimensions of the values array.

property properties: Labels

Get the property Labels for this block.

The entries in these labels describe the last dimension of the values array. The properties are guaranteed to be the same for values and gradients in the same block.

gradient(parameter: str) → TensorBlock

Get the gradient of the block values with respect to the given parameter.

Parameters:

parameter (str) – check for gradients with respect to this parameter (e.g. positions, cell, …)

Return type:

TensorBlock

>>> import numpy as np
>>> from metatensor import Labels, TensorBlock
>>> block = TensorBlock(
...     values=np.full((3, 1, 5), 1.0),
...     samples=Labels(["system"], np.array([[0], [2], [4]])),
...     components=[Labels.range("component", 1)],
...     properties=Labels.range("property", 5),
... )
>>> positions_gradient = TensorBlock(
...     values=np.full((2, 3, 1, 5), 11.0),
...     samples=Labels(["sample", "atom"], np.array([[0, 2], [2, 3]])),
...     components=[
...         Labels.range("direction", 3),
...         Labels.range("component", 1),
...     ],
...     properties=Labels.range("property", 5),
... )

>>> block.add_gradient("positions", positions_gradient)
>>> cell_gradient = TensorBlock(
...     values=np.full((2, 3, 3, 1, 5), 15.0),
...     samples=Labels.range("sample", 2),
...     components=[
...         Labels.range("direction_1", 3),
...         Labels.range("direction_2", 3),
...         Labels.range("component", 1),
...     ],
...     properties=Labels.range("property", 5),
... )
>>> block.add_gradient("cell", cell_gradient)

>>> positions_gradient = block.gradient("positions")
>>> print(positions_gradient)
Gradient TensorBlock ('positions')
    samples (2): ['sample', 'atom']
    components (3, 1): ['direction', 'component']
    properties (5): ['property']
    gradients: None

>>> cell_gradient = block.gradient("cell")
>>> print(cell_gradient)
Gradient TensorBlock ('cell')
    samples (2): ['sample']
    components (3, 3, 1): ['direction_1', 'direction_2', 'component']
    properties (5): ['property']
    gradients: None

add_gradient(parameter: str, gradient: TensorBlock)

Add gradient with respect to parameter in this block.

Parameters:

• parameter (str) – add gradients with respect to this parameter (e.g. positions, cell, …)

• gradient (TensorBlock) –

  a TensorBlock whose values contain the gradients of this TensorBlock values with respect to parameter. The labels of the gradient TensorBlock should be organized as follows:

  • its samples must contain "sample" as the first dimension, with values containing the index of the corresponding samples in this TensorBlock, and arbitrary supplementary samples dimension;

  • its components must contain at least the same components as this TensorBlock, with any additional components coming before those;

  • its properties must match exactly those of this TensorBlock.

>>> import numpy as np
>>> from metatensor import Labels, TensorBlock
>>> block = TensorBlock(
...     values=np.full((3, 1, 1), 1.0),
...     samples=Labels(["system"], np.array([[0], [2], [4]])),
...     components=[Labels.range("component", 1)],
...     properties=Labels.range("property", 1),
... )
>>> gradient = TensorBlock(
...     values=np.full((2, 1, 1), 11.0),
...     samples=Labels(["sample", "parameter"], np.array([[0, -2], [2, 3]])),
...     components=[Labels.range("component", 1)],
...     properties=Labels.range("property", 1),
... )
>>> block.add_gradient("parameter", gradient)
>>> print(block)
TensorBlock
    samples (3): ['system']
    components (1): ['component']
    properties (1): ['property']
    gradients: ['parameter']

gradients_list() → List[str]

get a list of all gradients defined in this block

Return type:

List[str]

has_gradient(parameter: str) → bool

Check if this block contains gradient information with respect to the given parameter.

Parameters:

parameter (str) – check for gradients with respect to this parameter (e.g. positions, cell, …)

Return type:

bool

gradients() → Generator[Tuple[str, TensorBlock], None, None]

Get an iterator over all gradients defined in this block.

Return type:

Generator[Tuple[str, TensorBlock], None, None]

property dtype: dtype | dtype

Get the dtype of all the values and gradient arrays stored inside this TensorBlock.

property device: str | device

Get the device of all the values and gradient arrays stored inside this TensorBlock.

to(*args, **kwargs) → TensorBlock

Move all the arrays in this block (values and gradients) to the given dtype, device and arrays backend.

Parameters:

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

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

• arrays – new backend to use for the arrays. This can be either "numpy", "torch" or None (keeps the existing backend); and must be given as a keyword argument (arrays="numpy").

Return type:

TensorBlock