SegmentedTensorProduct#

class cuequivariance.SegmentedTensorProduct( operands_and_subscripts: Sequence[tuple[SegmentedOperand | None, str]] | None = None, coefficient_subscripts: str = '', *, paths: Sequence[Path] | None = None, )#

Irreps-agnostic and dataflow-agnostic descriptor of a segmented tensor product

Parameters:

operands (list of operands) – The operands of the tensor product. To each operand corresponds subscripts and a list of segments.
paths (list of paths) – Each path contains coefficients and a list of indices. The indices are the indices of the segments of the operands.
coefficient_subscripts (str) – The subscripts of the coefficients.

We typically use the from_subscripts method to create a descriptor and then add segments and paths one by one.

Methods

classmethod from_subscripts( subscripts: Subscripts, ) → SegmentedTensorProduct#

Create a descriptor from a subscripts string.

Examples

>>> d = cue.SegmentedTensorProduct.from_subscripts("uv,ui,vj+ij")
>>> i0 = d.add_segment(0, (2, 3))
>>> i1 = d.add_segment(1, (2, 5))
>>> i2 = d.add_segment(2, (3, 4))
>>> d.add_path(i0, i1, i2, c=np.ones((5, 4)))
0
>>> print(d)
uv,ui,vj+ij operands=[(2, 3)],[(2, 5)],[(3, 4)] paths=[op0[0]*op1[0]*op2[0]*c c.shape=(5, 4) c.nnz=20]

classmethod empty_segments( num_segments: list[int], ) → SegmentedTensorProduct#

Create a descriptor with a simple structure.

Examples

>>> cue.SegmentedTensorProduct.empty_segments([2, 3, 4])
,, sizes=2,3,4 num_segments=2,3,4 num_paths=0

classmethod from_dict( data: dict[str, Any], ) → SegmentedTensorProduct#: Create a descriptor from a dictionary.

classmethod from_json( data: str, ) → SegmentedTensorProduct#: Create a descriptor from a JSON string.

classmethod from_bytes( data: bytes, ) → SegmentedTensorProduct#: Create a descriptor from compressed binary data.

classmethod from_base64( data: str, ) → SegmentedTensorProduct#: Create a descriptor from a base64 string.

property num_operands: int#: Number of operands.

property num_paths: int#: Number of paths.

property subscripts: Subscripts#: Subscripts of the tensor product.

property indices: ndarray#: Indices of the paths.

property coefficients_are_stackable: bool#: Check if the coefficients are stackable.

property stacked_coefficients: ndarray#: Stacked coefficients of the paths in a single array.

to_text( coefficient_formatter=<function SegmentedTensorProduct.<lambda>>, ) → str#

Human-readable text representation of the descriptor.

Parameters:: coefficient_formatter (callable, optional) – A function to format the coefficients.

Examples

>>> ((_, d),) = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1")
... ).polynomial.operations
>>> d = d.flatten_coefficient_modes()
>>> print(d.to_text())
uvw,u,v,w sizes=320,16,16,16 num_segments=5,4,4,4 num_paths=16 u=4 v=4 w=4
operand #0 subscripts=uvw
  | u: [4] * 5
  | v: [4] * 5
  | w: [4] * 5
operand #1 subscripts=u
  | u: [4] * 4
operand #2 subscripts=v
  | v: [4] * 4
operand #3 subscripts=w
  | w: [4] * 4
Flop cost: 0->1344 1->2368 2->2368 3->2368
Memory cost: 368
Path indices: 0 0 0 0, 1 0 1 1, 1 0 2 2, 1 0 3 3, 2 1 0 1, 2 2 0 2, ...
Path coefficients: [0.17...]

to_dict( extended: bool = False, ) → dict[str, Any]#: Dictionary representation of the descriptor.

to_json(extended: bool = False) → str#: JSON representation of the descriptor.

to_bytes(extended: bool = False) → bytes#: Compressed binary representation of the descriptor.

to_base64(extended: bool = False) → str#

Base64 representation of the descriptor.

Examples

>>> d = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1"),
...     cue.Irreps("SO3", "4x0+4x1")
... ).polynomial.operations[0][1]
>>> print(d.to_base64())
eJytkstuwjAQRX/F8r...lTF2zlX91/fHyvj2Z4=

get_dimensions_dict() → dict[str, set[int]]#: Get the dimensions of the tensor product.

get_dims(m: str) → set[int]#

Get the dimensions of a specific mode.

Examples

>>> d = cue.descriptors.fully_connected_tensor_product(
...     cue.Irreps("SO3", "4x0+8x1"),
...     cue.Irreps("SO3", "3x0+3x1"),
...     cue.Irreps("SO3", "5x0+7x1")
... ).polynomial.operations[0][1]
>>> d.get_dims("u")
{8, 4}
>>> d.get_dims("v")
{3}
>>> d.get_dims("w")
{5, 7}

get_path_dimensions_dict( path: int | Path, *, returns_sets: bool = False, ) → dict[str, int | set[int]]#: Get the dimensions of a specific path.

get_path_dim( path: int | Path, m: str, *, returns_set=False, ) → int | set[int]#: Get the dimension of a specific mode in a specific path.

segment_slice( operand: int, path: int | Path, ) → slice#: Get the slice of the segment in the given operand selected by the given path.

get_segment_shape( operand: int, path: int | Path, ) → tuple[int, ...]#: Get the shape of the segment in the given operand selected by the given path.

all_segments_are_used() → bool#: Check if all segments are used in the tensor product.

compressed_path_segment( operand: int, ) → ndarray#

Starting indices of paths for the segments of the specified operand.

Note: This method requires that the paths are sorted by the specified operand.

Parameters:: operand (int) – The index of the operand for which to find path groups.
Returns:: An array of starting path indices for each segment in the specified operand.
Return type:: np.ndarray

Examples

indices[:, operand], operands[operand].num_segments -> compressed_path_segment(operand)
[0, 0, 1, 1, 1], 2 -> [0, 2, 5]
[0, 0, 1, 1, 1], 3 -> [0, 2, 5, 5]
[0, 0, 2, 2, 2], 3 -> [0, 2, 2, 5]
[1, 1], 2 -> [0, 0, 2]

operands_with_identical_segments() → frozenset[frozenset[int]]#: Groups of operands sharing the same segments.

symmetries() → list[tuple[int, ...]]#: List of permutations that leave the tensor product invariant.

coefficients_equal_one() → bool#: Check if all coefficients are equal to one.

flop( operand: int, batch_size: int = 1, algorithm: str = 'optimal', ) → int#

Compute the number of flops needed to compute the specified operand.

Parameters:

operand (int) – The operand for which to compute the flop cost.
batch_size (int, optional) – The batch size for the computation. Defaults to 1.
algorithm (str, optional) – The algorithm to use to compute the cost. Can be ‘optimal’ or ‘naive’.

Returns:

The number of flops needed to compute the specified operand.

Return type:

int

memory(batch_sizes: list[int]) → int#: Compute the memory usage of the tensor product.

insert_path( path_index, *segments: int | tuple[int, ...] | dict[str, int] | None, c: ndarray, dims: dict[str, int] | None = None, ) → int#: Insert a path at a specific index.

add_path( *segments: int | tuple[int, ...] | dict[str, int] | None, c: ndarray, dims: dict[str, int] | None = None, ) → int#

Add a path to the descriptor.

Parameters:

segments – Specifies the segments of the operands that are contracted in the path.
c (np.ndarray) – The coefficients of the path.
dims (dict[str, int], optional) – The extent of the modes.

Returns:

The index of the added path.

Return type:

int

Examples

>>> d = cue.SegmentedTensorProduct.from_subscripts("uv,ui,vj+ij")
>>> i1 = d.add_segment(1, (2, 3))
>>> i2 = d.add_segment(2, (2, 5))

We can use None to add a new segment on the fly: >>> d.add_path(None, i1, i2, c=np.ones((3, 5))) 0

The descriptor has now a new segment (2, 2) in the first operand: >>> d uv,ui,vj+ij operands=[(2, 2)],[(2, 3)],[(2, 5)] paths=[op0[0]*op1[0]*op2[0]*c c.shape=(3, 5) c.nnz=15] >>> d.add_path(0, None, None, c=np.ones((10, 10))) 1 >>> d uv,ui,vj+ij sizes=4,26,30 num_segments=1,2,2 num_paths=2 i={3, 10} j={5, 10} u=2 v=2

When the dimensions of the modes cannot be inferred, we can provide them: >>> d.add_path(None, None, None, c=np.ones((2, 2)), dims={“u”: 2, “v”: 2}) 2 >>> d uv,ui,vj+ij sizes=8,30,34 num_segments=2,3,3 num_paths=3 i={2, 3, 10} j={2, 5, 10} u=2 v=2

insert_segments( operand: int, sid: int, segments: list[tuple[int, ...]], )#: Insert segments at a specific index.

add_segment( operand: int, segment: tuple[int, ...] | dict[str, int], ) → int#: Add a segment to the descriptor.

add_segments( operand: int, segments: list[tuple[int, ...] | dict[str, int]], )#: Add segments to the descriptor.

canonicalize_subscripts() → SegmentedTensorProduct#

Return a new descriptor with a canonical representation of the subscripts.

Examples

>>> d = cue.SegmentedTensorProduct.from_subscripts("ab,ax,by+yx")
>>> d.canonicalize_subscripts()
uv,ui,vj+ij sizes=0,0,0 num_segments=0,0,0 num_paths=0 i= j= u= v=

This is useful to identify equivalent descriptors.

add_or_rename_modes( subscripts: str, *, mapping: dict[str, str] | None = None, ) → SegmentedTensorProduct#

Return a new descriptor with the modes renamed according to the new subscripts.

Parameters:

subscripts (str) – The new subscripts that contains the new names of the modes. The new subscripts can also be a superset of the old subscripts.
mapping (dict of str to str, optional) – The mapping between the old and new modes.

Returns:

The new descriptor with the renamed modes.

Return type:

SegmentedTensorProduct

add_or_transpose_modes( subscripts: str, dims: dict[str, int] | None = None, ) → SegmentedTensorProduct#

Return a new descriptor with the modes transposed according to the new subscripts.

Parameters:

subscripts (str) – A new subscripts that contains a permutation of the modes of the current subscripts.
dims (dict of str to int, optional) – The dimensions of the new modes.

Returns:

The new descriptor with the transposed modes.

Return type:

SegmentedTensorProduct

append_modes_to_all_operands( modes: str, dims: dict[str, int], ) → SegmentedTensorProduct#

Return a new descriptor with the modes appended (to the right) to all operands.

Parameters:

modes (str) – The new segment modes to append to all operands.
dims (dict of str to int) – The dimensions of the new modes.

Returns:

The new descriptor with the appended modes.

Return type:

SegmentedTensorProduct

permute_operands( perm: tuple[int, ...], ) → SegmentedTensorProduct#: Permute the operands of the descriptor.

move_operand( operand: int, new_index: int, ) → SegmentedTensorProduct#: Move an operand to a new index.

move_operand_first( operand: int, ) → SegmentedTensorProduct#: Move an operand to the first position.

move_operand_last( operand: int, ) → SegmentedTensorProduct#: Move an operand to the last position.

permute_segments( operand: int, perm: tuple[int, ...], ) → SegmentedTensorProduct#: Permute the segments of an operand.

sort_paths( operands_ordering: int | Sequence[int] | None = None, ) → SegmentedTensorProduct#

Sort the paths by their indices in lexicographical order.

Parameters:: operands_ordering (int or sequence of int, optional) – The order in which to sort the paths. If an integer, sort by that operand. If a sequence, sort by the first operand in the sequence and then by the second operand if equal, etc.
Returns:: The sorted descriptor.
Return type:: SegmentedTensorProduct

squeeze_modes( modes: str | None = None, ) → SegmentedTensorProduct#

Squeeze the descriptor by removing dimensions that are always 1.

Parameters:: modes (str, optional) – The modes to squeeze. If None, squeeze all modes that are always 1.
Returns:: The squeezed descriptor.
Return type:: SegmentedTensorProduct

split_mode( mode: str, size: int, ) → SegmentedTensorProduct#

Split a mode into multiple modes of a given size.

Parameters:

mode (str) – The mode to split. All its dimensions must be divisible by the size.
size (int) – The size of the new modes.

Returns:

The new descriptor.

Return type:

SegmentedTensorProduct

all_same_segment_shape() → bool#: Check if all segments have the same shape.

normalize_paths_for_operand( operand: int, ) → SegmentedTensorProduct#

Normalize the paths for an operand.

Parameters:: operand (int) – The index of the operand to normalize.

Assuming that the input operand have unit variance, this method computes the variance of each segment in the selected operand and uniformly normalize the coefficients of the paths by the square root of the variance of the segment. This is useful to ensure that the output has unit variance.

remove_zero_paths() → SegmentedTensorProduct#: Remove paths with zero coefficients.

fuse_paths_with_same_indices() → SegmentedTensorProduct#: Fuse paths with the same indices.

consolidate_paths() → SegmentedTensorProduct#: Consolidate the paths by merging duplicates and removing zeros.

sort_indices_for_identical_operands( operands: Sequence[int], ) → SegmentedTensorProduct#: Reduce the number of paths by sorting the indices for identical operands.

symmetrize_operands( operands: Sequence[int], force: bool = False, ) → SegmentedTensorProduct#: Symmetrize the specified operands permuting the indices.

remove_empty_segments() → SegmentedTensorProduct#: Remove empty segments.

flatten_modes( modes: Sequence[str], *, skip_zeros=True, force=False, ) → SegmentedTensorProduct#

Remove the specified modes by subdividing segments and paths.

Parameters:

modes (Sequence of str) – The modes to remove, they must precede the modes to keep in each operand.
skip_zeros (bool, optional) – Whether to skip paths with zero coefficients. Default is True.
force (bool, optional) – Whether to force the flattening by flattening extra necessary modes. Default is False.

flatten_coefficient_modes( *, skip_zeros=True, force=False, ) → SegmentedTensorProduct#

Flatten the coefficients of the descriptor. Create new segments and paths to flatten the coefficients.

Parameters:

skip_zeros (bool, optional) – Whether to skip paths with zero coefficients. Default is True.
force (bool, optional) – Whether to force the flattening by flattening extra necessary modes. Default is False.

consolidate_modes( modes: str | None = None, ) → SegmentedTensorProduct#

Consolidate the descriptor by merging modes together.

Parameters:: modes (str, optional) – The modes to consolidate. If None, consolidate all modes that are repeated.

round_coefficients_to_rational( max_denominator: int, ) → SegmentedTensorProduct#

Round the coefficients to the nearest p / q number with a given maximum denominator.

Parameters:: max_denominator (int) – The maximum denominator, q < max_denominator.

round_coefficients_to_sqrt_rational( max_denominator: int, ) → SegmentedTensorProduct#

Round the coefficients to the nearest sqrt(p / q) number with a given maximum denominator.

Parameters:: max_denominator (int) – The maximum denominator, q < max_denominator.

modify_coefficients( f: Callable[[ndarray], ndarray], ) → SegmentedTensorProduct#

Modify the coefficients of the descriptor.

Parameters:: f (callable) – The function to apply to the coefficients.