sort

metatensor.sort(tensor: TensorMap, axes: str | List[str] = 'all', descending: bool = False) → TensorMap

Sort the tensor according to the key values and the blocks for each specified axis in axes according to the label values along these axes.

Each block is sorted separately, see sort_block() for more information

Note: This function duplicates metadata on the CPU for the purpose of sorting.

Parameters:

• axes (str | List[str]) – axes to sort. The labels entries along these axes will be sorted in lexicographic order, and the arrays values will be reordered accordingly. Possible values are 'keys', 'samples', 'components', 'properties' and 'all' to sort everything.

• descending (bool) – if false, the order is ascending

• tensor (TensorMap) –

Returns:

sorted tensor map

Return type:

TensorMap

>>> import numpy as np
>>> import metatensor
>>> from metatensor import TensorBlock, TensorMap, Labels
>>> block_1 = TensorBlock(
...     values=np.arange(9).reshape(3, 3),
...     samples=Labels(["system", "atom"], np.array([[0, 3], [0, 1], [0, 2]])),
...     components=[],
...     properties=Labels(["n", "l"], np.array([[1, 0], [2, 0], [0, 0]])),
... )
>>> block_2 = TensorBlock(
...     values=np.arange(3).reshape(1, 3),
...     samples=Labels(["system", "atom"], np.array([[0, 0]])),
...     components=[],
...     properties=Labels(["n", "l"], np.array([[1, 0], [2, 0], [0, 0]])),
... )
>>> tensor = TensorMap(
...     keys=Labels(["types"], np.array([[1], [0]])), blocks=[block_2, block_1]
... )
>>> metatensor.sort(tensor, axes="keys")
TensorMap with 2 blocks
keys: types
        0
        1

metatensor.sort_block(block: TensorBlock, axes: str | List[str] = 'all', descending: bool = False) → TensorBlock

Rearrange the values of a block according to the order given by the sorted metadata of the given axes.

This function creates copies of the metadata on the CPU to sort the metadata.

Parameters:

• axes (str | List[str]) – axes to sort. The labels entries along these axes will be sorted in lexicographic order, and the arrays values will be reordered accordingly. Possible values are 'samples', 'components', 'properties' and 'all' to sort everything.

• descending (bool) – if false, the order is ascending

• block (TensorBlock) –

Returns:

sorted tensor block

Return type:

TensorBlock

>>> import numpy as np
>>> import metatensor
>>> from metatensor import TensorBlock, TensorMap, Labels
>>> block = TensorBlock(
...     values=np.arange(9).reshape(3, 3),
...     samples=Labels(["system", "atom"], np.array([[0, 3], [0, 1], [0, 2]])),
...     components=[],
...     properties=Labels(["n", "l"], np.array([[2, 0], [3, 0], [1, 0]])),
... )
>>> print(block)
TensorBlock
    samples (3): ['system', 'atom']
    components (): []
    properties (3): ['n', 'l']
    gradients: None
>>> # sorting axes one by one
>>> block_sorted_stepwise = metatensor.sort_block(block, axes=["properties"])
>>> # properties (last dimension of the array) are sorted
>>> block_sorted_stepwise.values
array([[2, 0, 1],
       [5, 3, 4],
       [8, 6, 7]])
>>> block_sorted_stepwise = metatensor.sort_block(
...     block_sorted_stepwise, axes=["samples"]
... )
>>> # samples (first dimension of the array) are sorted
>>> block_sorted_stepwise.values
array([[5, 3, 4],
       [8, 6, 7],
       [2, 0, 1]])
>>> # sorting both samples and properties at the same time
>>> sorted_block = metatensor.sort_block(block)
>>> np.all(sorted_block.values == block_sorted_stepwise.values)
True
>>> # This function can also sort gradients:
>>> sorted_block.add_gradient(
...     parameter="g",
...     gradient=TensorBlock(
...         values=np.arange(18).reshape(3, 2, 3),
...         samples=Labels(["sample"], np.array([[1], [2], [0]])),
...         components=[Labels.range("direction", 2)],
...         properties=sorted_block.properties,
...     ),
... )
>>> sorted_grad_block = metatensor.sort_block(sorted_block)
>>> sorted_grad_block.gradient("g").samples == Labels.range("sample", 3)
True
>>> sorted_grad_block.gradient("g").properties == sorted_block.properties
True
>>> # the components (middle dimensions) are also sorted:
>>> sorted_grad_block.gradient("g").values
array([[[12, 13, 14],
        [15, 16, 17]],

       [[ 0,  1,  2],
        [ 3,  4,  5]],

       [[ 6,  7,  8],
        [ 9, 10, 11]]])