TensorMap#

class metatensor.torch.TensorMap(keys: Labels, blocks: List[TensorBlock])[source]#

A TensorMap is the main user-facing class of this library, and can store any kind of data used in atomistic machine learning.

A tensor map contains a list of TensorBlock, each one associated with a key. It also provides functions to access blocks associated with a key or part of a key, functions to merge blocks together and in general to manipulate this collection of blocks.

See also

The pure Python version of this class metatensor.TensorMap, and the differences between TorchScript and Python API for metatensor.

Parameters:

  • keys (Labels) – keys associated with each block

  • blocks (List[TensorBlock]) – set of blocks containing the actual data

property keys: Labels#

the set of keys labeling the blocks in this TensorMap

__len__() → int[source]#

get the number of key/block pairs in this TensorMap

Return type:

int

__getitem__(selection: int | Labels | LabelsEntry | Dict[str, int]) → TensorBlock[source]#

Get a single block with indexing syntax. This calls TensorMap.block() directly.

Parameters:

selection (int | Labels | LabelsEntry | Dict[str, int]) –

Return type:

TensorBlock

copy() → TensorMap[source]#

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

Return type:

TensorMap

static load(path: str) → TensorMap[source]#

Load a serialized TensorMap from the file at path, this is equivalent to metatensor.torch.load().

Parameters:

path (str) – Path of the file containing a saved TensorMap

Return type:

TensorMap

Warning

PyTorch can execute static functions (like this one) coming from a TorchScript extension, but fails when trying to save code calling this function with torch.jit.save(), giving the following error:

Failed to downcast a Function to a GraphFunction

This issue is reported as PyTorch#115639. In the mean time, you should use metatensor.torch.load() instead of this function to save your code to TorchScript.

static load_buffer(buffer: Tensor) → TensorMap[source]#

Load a serialized TensorMap from an in-memory buffer, this is equivalent to metatensor.torch.load_buffer().

Parameters:

buffer (Tensor) – torch Tensor representing an in-memory buffer

Return type:

TensorMap

Warning

PyTorch can execute static functions (like this one) coming from a TorchScript extension, but fails when trying to save code calling this function with torch.jit.save(), giving the following error:

Failed to downcast a Function to a GraphFunction

This issue is reported as PyTorch#115639. In the mean time, you should use metatensor.torch.load_buffer() instead of this function to save your code to TorchScript.

save(path: str)[source]#

Save this TensorMap to a file, this is equivalent to metatensor.torch.save().

Parameters:

path (str) – Path of the file. If the file already exists, it will be overwritten

save_buffer() → Tensor[source]#

Save this TensorMap to an in-memory buffer, this is equivalent to metatensor.torch.save_buffer().

Return type:

Tensor

items() → List[Tuple[LabelsEntry, TensorBlock]][source]#

get an iterator over (key, block) pairs in this TensorMap

Return type:

List[Tuple[LabelsEntry, TensorBlock]]

keys_to_samples(keys_to_move: str | List[str] | Tuple[str, ...] | Labels, sort_samples: bool = True) → TensorMap[source]#

Merge blocks along the samples axis, adding keys_to_move to the end of the samples labels dimensions.

This function will remove keys_to_move from the keys, and find all blocks with the same remaining keys values. It will then merge these blocks along the samples direction (i.e. do a vertical concatenation), adding keys_to_move to the end of the samples labels dimensions. The values taken by keys_to_move in the new samples labels will be the values of these dimensions in the merged blocks’ keys.

If keys_to_move is a set of Labels, it must be empty (len(keys_to_move) == 0), and only the Labels.names will be used.

The order of the samples in the merged blocks is controlled by sort_samples. If sort_samples is True, samples are re-ordered to keep them lexicographically sorted. Otherwise they are kept in the order in which they appear in the blocks.

This function is only implemented when the blocks to merge have the same properties values.

Parameters:

  • keys_to_move (str | List[str] | Tuple[str, ...] | Labels) – description of the keys to move

  • sort_samples (bool) – whether to sort the merged samples or keep them in the order in which they appear in the original blocks

Returns:

a new TensorMap with merged blocks

Return type:

TensorMap

keys_to_properties(keys_to_move: str | List[str] | Tuple[str, ...] | Labels, sort_samples: bool = True) → TensorMap[source]#

Merge blocks along the properties direction, adding keys_to_move at the beginning of the properties labels dimensions.

This function will remove keys_to_move from the keys, and find all blocks with the same remaining keys values. Then it will merge these blocks along the properties direction (i.e. do an horizontal concatenation).

If keys_to_move is given as strings, then the new property labels will only contain entries from the existing blocks. For example, merging a block with key a=0 and properties p=1, 2 with a block with key a=2 and properties p=1, 3 will produce a block with properties a, p = (0, 1), (0, 2), (2, 1), (2, 3).

If keys_to_move is a set of Labels and it is empty (len(keys_to_move) == 0), the Labels.names will be used as if they where passed directly.

Finally, if keys_to_move is a non empty set of Labels, the new properties labels will contains all of the entries of keys_to_move (regardless of the values taken by keys_to_move.names in the merged blocks’ keys) followed by the existing properties labels. For example, using a=2, 3 in keys_to_move, blocks with properties p=1, 2 will result in a, p = (2, 1), (2, 2), (3, 1), (3, 2). If there is no values (no block/missing sample) for a given property in the merged block, then the value will be set to zero.

When using a non empty Labels for keys_to_move, the properties labels of all the merged blocks must take the same values.

The order of the samples in the merged blocks is controlled by sort_samples. If sort_samples is True, samples are re-ordered to keep them lexicographically sorted. Otherwise they are kept in the order in which they appear in the blocks.

Parameters:

  • keys_to_move (str | List[str] | Tuple[str, ...] | Labels) – description of the keys to move

  • sort_samples (bool) – whether to sort the merged samples or keep them in the order in which they appear in the original blocks

Returns:

a new TensorMap with merged blocks

Return type:

TensorMap

components_to_properties(dimensions: str | List[str] | Tuple[str, ...]) → TensorMap[source]#

Move the given dimensions from the component labels to the property labels for each block.

Parameters:

dimensions (str | List[str] | Tuple[str, ...]) – name of the component dimensions to move to the properties

Return type:

TensorMap

blocks_matching(selection: Labels) → List[int][source]#

Get a (possibly empty) list of block indexes matching the selection.

This function finds all keys in this TensorMap with the same values as selection for the dimensions/names contained in the selection; and return the corresponding indexes.

The selection should contain a single entry.

Parameters:

selection (Labels) –

Return type:

List[int]

block_by_id(index: int) → TensorBlock[source]#

Get the block at index in this TensorMap.

Parameters:

index (int) – index of the block to retrieve

Return type:

TensorBlock

blocks_by_id(indices: List[int]) → List[TensorBlock][source]#

Get the blocks with the given indices in this TensorMap.

Parameters:

indices (List[int]) – indices of the block to retrieve

Return type:

List[TensorBlock]

block(selection: None | int | Labels | LabelsEntry | Dict[str, int] = None) → TensorBlock[source]#

Get the single block in this TensorMap matching the selection.

When selection is an int, this is equivalent to TensorMap.block_by_id().

When selection is an Labels, it should only contain a single entry, which will be used for the selection.

When selection is a Dict[str, int], it is converted into a single single LabelsEntry (the dict keys becoming the names and the dict values being joined together to form the LabelsEntry values), which is then used for the selection.

When selection is a LabelsEntry, this function finds the key in this TensorMap with the same values as selection for the dimensions/names contained in the selection; and return the corresponding indexes.

Parameters:

selection (None | int | Labels | LabelsEntry | Dict[str, int]) – description of the block to extract

Return type:

TensorBlock

blocks(selection: None | List[int] | int | Labels | LabelsEntry | Dict[str, int] = None) → List[TensorBlock][source]#

Get the blocks in this TensorMap matching the selection.

When selection is None (the default), all blocks are returned.

When selection is an int, this is equivalent to TensorMap.block_by_id(); and for a List[int] this is equivalent to TensorMap.blocks_by_id().

When selection is an Labels, it should only contain a single entry, which will be used for the selection.

When selection is a Dict[str, int], it is converted into a single single LabelsEntry (the dict keys becoming the names and the dict values being joined together to form the LabelsEntry values), which is then used for the selection.

When selection is a LabelsEntry, this function finds all keys in this TensorMap with the same values as selection for the dimensions/names contained in the selection; and return the corresponding blocks.

Parameters:

selection (None | List[int] | int | Labels | LabelsEntry | Dict[str, int]) – description of the blocks to extract

Return type:

List[TensorBlock]

property sample_names: List[str]#

names of the samples for all blocks in this tensor map

property component_names: List[str]#

names of the components for all blocks in this tensor map

property property_names: List[str]#

names of the properties for all blocks in this tensor map

print(max_keys: int) → str[source]#

Print this TensorMap to a string, including at most max_keys in the output.

Parameters:

max_keys (int) – how many keys to include in the output. Use -1 to include all keys.

Return type:

str

property device: device#

get the device of all the arrays stored inside this TensorMap

property dtype: dtype#

get the dtype of all the arrays stored inside this TensorMap

Warning

Due to limitations in TorchScript C++ extensions, the dtype is returned as an integer, which can not be compared with torch.dtype instances. See TensorBlock.dtype() for more information.

to(dtype: dtype | None = None, device: device | None = None, arrays: str | None = None) → TensorMap[source]#

Move all the data (keys and blocks) in this TensorMap to the given dtype, device and arrays backend.

Parameters:

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

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

  • arrays (str | None) – new backend to use for the arrays. This parameter is here for compatibility with the pure Python API, can only be set to "torch" or None and does nothing.

Return type:

TensorMap