split

metatensor.split(tensor: TensorMap, axis: str, selections: List[List[int] | ndarray | 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 length of the returned list is equal to the number of Labels objects passed in selections. Each returned 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 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 as mts
>>> 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 = mts.split(
...     tensor,
...     axis="samples",
...     selections=[
...         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 selections belongs. Each TensorBlock in each returned TensorMap could have a reduced dimension along this axis, but the other axes will remain the same size.

• selections (List[List[int] | ndarray | Labels]) – a list of selections specifying what should be included in the different outputs. Each entry in the list can either be a Labels selection, an array or a List[int] indicating the raw indices that should be kept. The list can mix different types of selections. When using Labels selection, only a subset of the corresponding dimension names can be specified, and any entry with matching values will be selected.

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 selections.

Return type:

List[TensorMap]

metatensor.split_block(block: TensorBlock, axis: str, selections: List[List[int] | ndarray | Labels]) → List[TensorBlock]

Splits an input TensorBlock into multiple TensorBlock objects based on some specified selections, along either the "samples" or "properties" axis. The length of the returned list is equal to the number of selections passed in selections. Each returned 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 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 as mts
>>> 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 = mts.split_block(
...     block,
...     axis="samples",
...     selections=[
...         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 selections belongs. Each TensorBlock returned could have a reduced dimension along this axis, but the other axes will remain the same size.

• selections (List[List[int] | ndarray | Labels]) – a list of selections specifying what should be included in the different outputs. Each entry in the list can either be a Labels selection, an array or a List[int] indicating the raw indices that should be kept. The list can mix different types of selections. When using Labels selection, only a subset of the corresponding dimension names can be specified, and any entry with matching values will be selected.

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 selections.

Return type:

List[TensorBlock]