split

metatensor.split(tensor: TensorMap, axis: str, grouped_labels: List[Labels]) → List[TensorMap]

Split a TensorMap into multiple TensorMap.

The operation is based on some specified groups of indices, along either the “samples” or “properties” axis. The number of returned TensorMap`s is equal to the number of :py:class:`Labels objects passed in grouped_labels. Each returned :py:class`TensorMap` will have the same keys and number of blocks at the input tensor, but with the dimensions of the blocks reduced to only contain the specified indices for the corresponding group.

For example, to split a tensor along the "samples" axis, according to the "system" index, where system 0, 6, and 7 are in the first returned :py:class`TensorMap`; 2, 3, and 4 in the second; and 1, 5, 8, 9, and 10 in the third:

>>> import numpy as np
>>> from metatensor import Labels, TensorBlock, TensorMap
>>> import metatensor
>>> block = TensorBlock(
...     values=np.random.rand(11, 3),
...     samples=Labels(
...         names=["system"],
...         values=np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).reshape(-1, 1),
...     ),
...     components=[],
...     properties=Labels.range("properties", 3),
... )
>>> keys = Labels(names=["key"], values=np.array([[0]]))
>>> tensor = TensorMap(keys, [block])
>>> splitted = metatensor.split(
...     tensor,
...     axis="samples",
...     grouped_labels=[
...         Labels(names=["system"], values=np.array([[0], [6], [7]])),
...         Labels(names=["system"], values=np.array([[2], [3], [4]])),
...         Labels(names=["system"], values=np.array([[1], [5], [8], [10]])),
...     ],
... )
>>> len(splitted)
3
>>> splitted[0].block(0).samples
Labels(
    system
      0
      6
      7
)
>>> splitted[1].block(0).samples
Labels(
    system
      2
      3
      4
)
>>> splitted[2].block(0).samples
Labels(
    system
      1
      5
      8
      10
)

Parameters:

• tensor (TensorMap) – a TensorMap to be split

• axis (str) – a str, either “samples” or “properties”, that indicates the TensorBlock axis along which the named index (or indices) in grouped_labels belongs. Each TensorBlock in each returned TensorMap could have a reduced dimension along this axis, but the other axes will remain the same size.

• grouped_labels (List[Labels]) – a list of Labels containing the names and values of the indices along the specified axis which should be in each respective output TensorMap.

Returns:

a list of:py:class:TensorMap that corresponds to the split input tensor. Each tensor in the returned list contains only the named indices in the respective py:class:Labels object of grouped_labels.

Return type:

List[TensorMap]

metatensor.split_block(block: TensorBlock, axis: str, grouped_labels: List[Labels]) → List[TensorBlock]

Splits an input TensorBlock into multiple TensorBlock objects based on some specified grouped_labels, along either the "samples" or "properties" axis. The number of returned TensorBlock is equal to the number of Labels objects passed in grouped_labels. Each returned :py:class`TensorBlock` will have the same keys and number of blocks at the input tensor, but with the dimensions of the blocks reduced to only contain the specified indices for the corresponding group.

For example, to split a block along the "samples" axis, according to the "system" index, where system 0, 6, and 7 are in the first returned :py:class`TensorBlock`; 2, 3, and 4 in the second; and 1, 5, 8, 9, and 10 in the third:

>>> import numpy as np
>>> from metatensor import Labels, TensorBlock, TensorMap
>>> import metatensor
>>> block = TensorBlock(
...     values=np.random.rand(11, 3),
...     samples=Labels(
...         names=["system"],
...         values=np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10]).reshape(-1, 1),
...     ),
...     components=[],
...     properties=Labels.range("properties", 3),
... )
>>> splitted = metatensor.split_block(
...     block,
...     axis="samples",
...     grouped_labels=[
...         Labels(names=["system"], values=np.array([[0], [6], [7]])),
...         Labels(names=["system"], values=np.array([[2], [3], [4]])),
...         Labels(names=["system"], values=np.array([[1], [5], [8], [10]])),
...     ],
... )
>>> len(splitted)
3
>>> splitted[0].samples
Labels(
    system
      0
      6
      7
)
>>> splitted[1].samples
Labels(
    system
      2
      3
      4
)
>>> splitted[2].samples
Labels(
    system
      1
      5
      8
      10
)

Parameters:

• block (TensorBlock) – a TensorBlock to be split

• axis (str) – a str, either “samples” or “properties”, that indicates the TensorBlock axis along which the named index (or indices) in grouped_labels belongs. Each TensorBlock returned could have a reduced dimension along this axis, but the other axes will remain the same size.

• grouped_labels (List[Labels]) – a list of Labels containing the names and values of the indices along the specified axis which should be in each respective output TensorBlock.

Returns:

a list of:py:class:TensorBlock that corresponds to the split input block. Each block in the returned list contains only the named indices in the respective py:class:Labels object of grouped_labels.

Return type:

List[TensorBlock]