Source code for metatensor.operations.slice

from . import _dispatch
from ._classes import (
    Labels,
    TensorBlock,
    TensorMap,
    check_isinstance,
    torch_jit_is_scripting,
)


[docs] def slice(tensor: TensorMap, axis: str, labels: Labels) -> TensorMap: """ Slice a :py:class:`TensorMap` along either the "samples" or "properties" `axis`. `labels` is a :py:class:`Labels` objects that specifies the samples/properties (respectively) names and indices that should be sliced, i.e. kept in the output tensor. For axis, either "samples" or "properties" should be specified. .. code-block:: python samples = Labels( names=["structure", "center"], values=np.array([[0, 1], [0, 6], [1, 6], [3, 16]]), # must be a 2D-array ) properties = Labels( names=["n"], # radial channel values=np.array([[3], [4], [5]]), ) sliced_tensor_samples = slice( tensor, axis="samples", labels=samples, ) sliced_tensor_properties = slice( tensor, axis="properties", labels=properties, ) Also note that this function will return a :py:class:`TensorMap` whose blocks are of equal or smaller dimensions (due to slicing) than those of the input. However, the returned :py:class:`TensorMap` will be returned with the same number of blocks and the corresponding keys as the input. If any block upon slicing is reduced to nothing, i.e. in the case that it has none of the specified `labels` along the "samples" or "properties" `axis`, an empty block will be returned but will still be accessible by its key. User warnings will be issued if any blocks are sliced to contain no values. For the empty blocks that may be returned, although there will be no actual values in its ``TensorBlock.values`` array, the shape of this array will be non-zero in the dimensions that haven't been sliced. This allows the slicing of dimensions to be tracked. For example, if a TensorBlock of shape (52, 1, 5) is passed, and only some samples are specified to be sliced but none of these appear in the input :py:class:`TensorBlock`, the returned block values array will be empty, but its shape will be (0, 1, 5) - i.e. the samples dimension has been sliced to zero but the components and properties dimensions remain in-tact. The same logic applies to any Gradient TensorBlocks the input TensorBlock may have associated with it. See the documentation for the :py:func:`slice_block` function to see how an individual :py:class:`TensorBlock` is sliced. :param tensor: the input :py:class:`TensorMap` to be sliced. :param axis: a :py:class:`str` indicating the axis along which slicing should occur. Should be either "samples" or "properties". :param labels: a :py:class:`Labels` object containing the names and indices of the "samples" or "properties" to keep in each of the sliced :py:class:`TensorBlock` of the output :py:class:`TensorMap`. :return: a :py:class:`TensorMap` that corresponds to the sliced input tensor. """ # Check input args if not torch_jit_is_scripting(): if not check_isinstance(tensor, TensorMap): raise TypeError( f"`tensor` must be a metatensor TensorMap, not {type(tensor)}" ) _check_args(tensor.block(0), axis=axis, labels=labels) return TensorMap( keys=tensor.keys, blocks=[ _slice_block(tensor[tensor.keys.entry(i)], axis, labels) for i in range(len(tensor.keys)) ], )
[docs] def slice_block(block: TensorBlock, axis: str, labels: Labels) -> TensorBlock: """ Slices an input :py:class:`TensorBlock` along either the "samples" or "properties" `axis`. `labels` is a :py:class:`Labels` objects that specify the sample/property names and indices that should be sliced, i.e. kept in the output :py:class:`TensorBlock`. Example: take an input :py:class:`TensorBlock` of shape (100, 1, 6), where there are 100 'samples', 1 'components', and 6 'properties'. Say we want to slice this tensor along the samples and properties dimensions. As in the code-block below, we can specify, for example, 4 samples and 3 properties indices to keep. The returned :py:class:`TensorBlock` will have shape (4, 1, 3). .. code-block:: python samples = Labels( names=["structure", "center"], values=np.array([[0, 1], [0, 6], [1, 6], [3, 16]]), # must be a 2D-array ) properties = Labels( names=["n"], # radial channel values=np.array([[3], [4], [5]]), ) sliced_block_samples = slice_block( block, axis="samples", labels=samples, ) sliced_block_properties = slice_block( block, axis="properties", labels=properties, ) For the empty blocks that may be returned, although there will be no actual values in its TensorBlock.values tensor, the shape of this tensor will be non-zero in the dimensions that haven't been sliced. This is created by slicing the input TensorBlock, as opposed to just returning an artificially-created empty one (with no shape or dimensions), and is intentional. It allows the slicing of dimensions to be tracked. For instance, if a TensorBlock of shape (52, 1, 5) is passed, and only some samples are specified to be sliced but none of these appear in the input TensorBlock, the returned TensorBlock values array will be empty, but its shape will be (0, 1, 5) - i.e. the samples dimension has been sliced to zero but the components and properties dimensions remain in-tact. The same logic applies to any Gradient TensorBlocks the input TensorBlock may have associated with it. :param block: the input :py:class:`TensorBlock` to be sliced. :param axis: a :py:class:`str` indicating the axis along which slicing should occur. Should be either "samples" or "properties". :param labels: a :py:class:`Labels` object containing the names and indices of the "samples" or "properties" to keep in the sliced output :py:class:`TensorBlock`. :return new_block: a :py:class:`TensorBlock` that corresponds to the sliced input. """ if not torch_jit_is_scripting(): if not check_isinstance(block, TensorBlock): raise TypeError( f"`block` must be a metatensor TensorBlock, not {type(block)}" ) _check_args(block, axis=axis, labels=labels) return _slice_block( block, axis=axis, labels=labels, )
def _slice_block(block: TensorBlock, axis: str, labels: Labels) -> TensorBlock: """ Slices an input :py:class:`TensorBlock` along either the "samples" or "properties" `axis`. `labels` is :py:class:`Labels` object that specifies the sample/property names and indices that should be sliced, i.e. kept in the output :py:class:`TensorBlock`. :param block: the input :py:class:`TensorBlock` to be sliced. :param axis: a :py:class:`str` object containing `samples` or `properties` indicating the direction of slicing. :param labels: a :py:class:`Labels` object containing the names and indices of samples/properties to keep in the sliced output :py:class:`TensorBlock`. :return new_block: a :py:class:`TensorBlock` that corresponds to the sliced input. """ if axis == "samples": # only keep the same names as `labels` all_samples = block.samples.view(labels.names) # create an arrays of bools indicating which samples indices to keep samples_mask = _dispatch.bool_array_like( [all_samples.entry(i) in labels for i in range(len(all_samples))], block.values, ) new_values = _dispatch.mask(block.values, 0, samples_mask) new_samples = Labels( block.samples.names, _dispatch.mask(block.samples.values, 0, samples_mask), ) new_block = TensorBlock( values=new_values, samples=new_samples, components=block.components, properties=block.properties, ) # Create a map from the previous samples indexes to the new sample indexes # to update the gradient samples # sample_map contains at position old_sample the index of the # corresponding new sample sample_map = _dispatch.int_array_like( int_list=[-1] * len(samples_mask), like=samples_mask, ) last = 0 for i, picked in enumerate(samples_mask): if picked: sample_map[i] = last last += 1 for parameter, gradient in block.gradients(): if len(gradient.gradients_list()) != 0: raise NotImplementedError("gradients of gradients are not supported") # Create a samples filter for the Gradient TensorBlock grad_samples_mask = samples_mask[gradient.samples.column("sample")] new_grad_samples_values = _dispatch.mask( gradient.samples.values, 0, grad_samples_mask ) if new_grad_samples_values.shape[0] != 0: # update the "sample" column of the gradient samples # to refer to the new samples new_grad_samples_values[:, 0] = sample_map[ new_grad_samples_values[:, 0] ] new_grad_samples = Labels( names=gradient.samples.names, values=new_grad_samples_values, ) else: new_grad_samples = Labels( names=gradient.samples.names, values=_dispatch.empty_like( gradient.samples.values, [0, gradient.samples.values.shape[1]] ), ) new_grad_values = _dispatch.mask(gradient.values, 0, grad_samples_mask) # Add sliced gradient to the TensorBlock new_block.add_gradient( parameter=parameter, gradient=TensorBlock( values=new_grad_values, samples=new_grad_samples, components=gradient.components, properties=new_block.properties, ), ) else: assert axis == "properties" # only keep the same names as `labels` all_properties = block.properties.view(list(labels.names)) # create an arrays of bools indicating which samples indices to keep properties_mask = _dispatch.bool_array_like( [all_properties.entry(i) in labels for i in range(len(all_properties))], block.values, ) new_values = _dispatch.mask( block.values, len(block.values.shape) - 1, properties_mask ) new_properties = Labels( block.properties.names, _dispatch.mask(block.properties.values, 0, properties_mask), ) new_block = TensorBlock( values=new_values, samples=block.samples, components=block.components, properties=new_properties, ) # Slice each Gradient TensorBlock and add to the new_block. for parameter, gradient in block.gradients(): if len(gradient.gradients_list()) != 0: raise NotImplementedError("gradients of gradients are not supported") assert axis == "properties" new_grad_values = _dispatch.mask( gradient.values, len(gradient.values.shape) - 1, properties_mask ) new_grad_samples = gradient.samples # Add sliced gradient to the TensorBlock new_block.add_gradient( parameter=parameter, gradient=TensorBlock( values=new_grad_values, samples=new_grad_samples, components=gradient.components, properties=new_block.properties, ), ) return new_block def _check_args( block: TensorBlock, axis: str, labels: Labels, ): """ Checks the arguments passed to :py:func:`slice` and :py:func:`slice_block`. """ # check axis if axis not in ["samples", "properties"]: raise ValueError( f"``axis``: {axis} is not known as a slicing axis. Please use" "'samples' or 'properties'" ) if not torch_jit_is_scripting(): if not check_isinstance(labels, Labels): raise TypeError(f"`labels` must be metatensor Labels, not {type(labels)}") if axis == "samples": s_names = block.samples.names for name in labels.names: if name not in s_names: raise ValueError( f"invalid sample name '{name}' which is not part of the input" ) else: assert axis == "properties" p_names = block.properties.names for name in labels.names: if name not in p_names: raise ValueError( f"invalid property name '{name}' which is not part of the input" )