O(3) symmetries

The metatomic.torch.o3 module provides tools to deal with O(3) symmetries. First it offers functions to rotate and invert System and TensorMap objects, for example to generate randomly rotated copies of a structure for data augmentation. Second, it offers tools to convert TensorMap between Cartesian and spherical forms.

Conventions for O(3) representations

To transform a TensorMap (or a TensorBlock), the O3Transformations.transform_tensormap() (resp. O3Transformations.transform_block()) need to know, for each component axis, whether it carries a Cartesian or a spherical tensor. This is inferred from the axis name: axis, whether it carries a Cartesian or a spherical tensor. This is inferred from the axis name:

  • Cartesian axes are named xyz, or xyz_1, xyz_2, … for blocks with several Cartesian axes (e.g. rank-2 Cartesian tensors). These are rotated directly with the (3, 3) transformation matrix \(R\) passed to O3Transformations, following the usual column-vector convention \(v' = R v\) (equivalently, since values are stored as rows, values_transformed = values @ R.T).

  • Spherical axes are named o3_mu, or o3_mu_1, o3_mu_2, … They are rotated with the real Wigner-D matrix for the angular momentum o3_lambda (respectively o3_lambda_1, o3_lambda_2, …) found in the block’s key. Real, rather than complex, spherical harmonics are used throughout, matching the convention used elsewhere in metatomic and metatensor for spherical targets.

Wigner-D matrices

Complex Wigner-D matrices are computed by the wigners package, using the convention

\[D^{\ell}_{m',m}(\alpha, \beta, \gamma) = \langle \ell, m' | e^{-i J_z \alpha} e^{-i J_y \beta} e^{-i J_z \gamma} | \ell, m \rangle,\]

with ZYZ Euler angles \((\alpha, \beta, \gamma)\) extracted from the proper part of \(R\), i.e. from \(R\) itself when \(\det R = 1\), or from \(-R\) when \(\det R = -1\), such that this proper part equals \(R_z(\alpha) R_y(\beta) R_z(\gamma)\).

These are then converted to real Wigner-D matrices through the unitary change of basis \(T\) mapping complex spherical harmonics \(Y_{\ell}^{m}\) to real ones:

\[\begin{split}Y_{\ell, m}^{\text{real}} = \begin{cases} \dfrac{1}{\sqrt{2}} \left( Y_{\ell}^{-m} + (-1)^m Y_{\ell}^{m} \right) & m > 0 \\[6pt] Y_{\ell}^{0} & m = 0 \\[6pt] \dfrac{i}{\sqrt{2}} \left( Y_{\ell}^{m} - (-1)^m Y_{\ell}^{-m} \right) & m < 0 \end{cases}\end{split}\]

so that the real Wigner-D matrix is \(D^{\ell}_{\text{real}} = T^{*} D^{\ell} T^{T}\). Spherical components follow the same column-vector convention as Cartesian ones, \(f' = D^{\ell}_{\text{real}} f\) (equivalently, since values are stored as rows, values_transformed = values @ D.T).

For an improper rotation (a rotation composed with an inversion), Cartesian axes are flipped as part of the transformation matrix itself, while spherical axes pick up an extra parity factor \((-1)^{\ell} \sigma\), where \(\ell\) is o3_lambda and \(\sigma\) is the block’s o3_sigma (its behavior under inversion, +1 or -1), also read from the key.

A TensorBlock in general contains values associated with several systems. This information is contained in the "system" samples dimension; each row is rotated with the transformation of the system it belongs to. Gradient blocks are routed the same way, via their parent block’s "system" column. When only one system is being transformed, the "system" column is optional and, if present, is ignored: every row is rotated with the (single) given transformation.


class metatomic.torch.o3.O3Transformations(matrices: Tensor, max_angular_momentum: int)[source]

A batch of one or more O(3) transformations, represented by (N, 3, 3) rotation or improper-rotation matrices.

The module stores the matrices, an improper-operations mask, and a packed buffer of Wigner-D matrices. It is constructed eagerly in __init__ (building the Wigner-D matrices requires the wigners package); the transform methods are TorchScript compatible and can be called from a scripted forward.

Transform methods accept an indices argument to select a sub-batch of operations, and a add_inversion flag to compose every selected operation with the inversion. Inverse transformations are available through the inverse_transform_* methods, which share their kernel with the forward ones and use transposed matrices / Wigner-D matrices.

Parameters:
  • matrices (Tensor) – (N, 3, 3) rotation or improper-rotation matrices

  • max_angular_momentum (int) – non-negative maximum angular momentum for which Wigner-D matrices are built

property matrices: Tensor

The (N, 3, 3) batch of rotation or improper-rotation matrices.

property max_angular_momentum: int

The maximum angular momentum with available Wigner-D matrices.

property improper_mask: Tensor

Boolean mask marking the improper operations in the batch.

property dtype: dtype

The dtype of the transformation matrices.

property device: device

The device of the transformation matrices.

wigner_D_matrices(ell: int) → Tensor[source]

Return the proper-part Wigner-D matrices for ell.

For improper operations, the inversion-parity factor sigma * (-1) ** ell is applied separately when transforming spherical values.

Parameters:

ell (int) – angular momentum in [0, max_angular_momentum]

Returns:

(N, 2*ell+1, 2*ell+1) stack of Wigner-D matrices

Return type:

Tensor

inverse_wigner_D_matrices(ell: int) → Tensor[source]

Return the transposed (inverse) proper-part Wigner-D matrices.

The inverse of a real Wigner-D matrix is its transpose. This is the accessor used by projection integrals that contract values against the inverse representation directly.

Parameters:

ell (int) – angular momentum in [0, max_angular_momentum]

Returns:

(N, 2*ell+1, 2*ell+1) stack of transposed Wigner-D matrices

Return type:

Tensor

transform_cartesian(vectors: Tensor, add_inversion: bool = False) → Tensor[source]

Apply the transformations to Cartesian vectors.

Parameters:
  • vectors (Tensor) – (..., 3) tensor of Cartesian vectors

  • add_inversion (bool) – compose every operation with the inversion

Returns:

transformed vectors, with the input shape for a single transformation or a leading batch axis ((N, ..., 3)) for a batch of more than one

Return type:

Tensor

inverse_transform_cartesian(vectors: Tensor, add_inversion: bool = False) → Tensor[source]

Apply the inverse transformations to Cartesian vectors.

The inverse of an orthogonal matrix is its transpose.

Parameters:
  • vectors (Tensor) – (..., 3) tensor of Cartesian vectors

  • add_inversion (bool) – compose every operation with the inversion before inverting

Returns:

transformed vectors, with the input shape for a single transformation or a leading batch axis for a batch of more than one

Return type:

Tensor

transform_spherical(values: Tensor, ell: int, sigma: int, add_inversion: bool = False) → Tensor[source]

Apply the transformations to spherical values.

Parameters:
  • values (Tensor) – (…, 2*ell+1) tensor of spherical values

  • ell (int) – angular momentum in [0, max_angular_momentum]

  • sigma (int) – +1 for a proper spherical representation or -1 for a pseudo one. Under an improper transformation, the representation acquires the factor sigma * (-1) ** ell.

  • add_inversion (bool) – compose every operation with the inversion

Returns:

transformed values, with the input shape for a single transformation or a leading batch axis for a batch of more than one

Return type:

Tensor

inverse_transform_spherical(values: Tensor, ell: int, sigma: int, add_inversion: bool = False) → Tensor[source]

Apply the inverse transformations to spherical values.

The inverse Wigner-D matrix is its transpose; the inversion-parity factor for an improper operation is unchanged by the transposition.

Parameters:
  • values (Tensor) – (…, 2*ell+1) tensor of spherical values

  • ell (int) – angular momentum in [0, max_angular_momentum]

  • sigma (int) – +1 for a proper spherical representation or -1 for a pseudo one

  • add_inversion (bool) – compose every operation with the inversion before inverting

Returns:

transformed values, with the input shape for a single transformation or a leading batch axis for a batch of more than one

Return type:

Tensor

transform_systems(systems: List[System], add_inversion: bool = False, start: int = 0, stop: int | None = None) → List[System][source]

Apply transformations to a list of systems.

systems[i] is transformed by operation i. To apply a batch of transformations to the same input system, pass that system repeatedly in the list.

Positions, cell vectors, neighbor-list displacements, and custom data following Conventions for O(3) representations are transformed. Atomic types and periodic-boundary flags are preserved.

Parameters:
  • systems (List[System]) – input systems, one per operation, matching the transformation matrices in dtype and device

  • add_inversion (bool) – compose every operation with the inversion

  • start (int) – first operation to apply

  • stop (int | None) – end of the range of operations to apply, defaulting to all of them

Returns:

one transformed System per input system

Return type:

List[System]

inverse_transform_systems(systems: List[System], add_inversion: bool = False, start: int = 0, stop: int | None = None) → List[System][source]

Apply the inverse transformations to a list of systems.

See transform_systems() for the per-system routing; the inverse uses transposed matrices and Wigner-D matrices.

Parameters:
  • systems (List[System]) – input systems, one per operation

  • add_inversion (bool) – compose every operation with the inversion before inverting

  • start (int) – first operation to invert

  • stop (int | None) – end of the range of operations to invert, defaulting to all of them

Returns:

one transformed System per input system

Return type:

List[System]

transform_tensormap(tensor: TensorMap, system_ids: Tensor | None = None, add_inversion: bool = False, start: int = 0, stop: int | None = None) → TensorMap[source]

Apply the transformations to a TensorMap and its gradients.

Scalar, Cartesian, and spherical data are identified by their component-axis names, following Conventions for O(3) representations. With a batch of more than one operation, the "system" sample label assigns each value row to an operation: when system_ids is None, the labels index the batch directly, and otherwise rows labelled system_ids[i] use operation i. Gradient rows use the operation of the value row referenced by their "sample" label. With a single operation, the "system" label is optional and ignored.

Parameters:
  • tensor (TensorMap) – TensorMap to transform, matching the transformation matrices in dtype and device

  • system_ids (Tensor | None) – optional one-dimensional tensor with one distinct "system" sample label per operation

  • add_inversion (bool) – compose every operation with the inversion

  • start (int) – first operation to apply

  • stop (int | None) – end of the range of operations to apply, defaulting to all of them

Returns:

transformed TensorMap with the same metadata and global information

Return type:

TensorMap

inverse_transform_tensormap(tensor: TensorMap, system_ids: Tensor | None = None, add_inversion: bool = False, start: int = 0, stop: int | None = None) → TensorMap[source]

Apply the inverse transformations to a TensorMap and its gradients.

The inverse of an orthogonal matrix is its transpose, and the inverse of a real Wigner-D matrix is its transpose, so this shares its kernel with transform_tensormap() using transposed matrices and Wigner-D matrices. See transform_tensormap() for row routing.

Parameters:
  • tensor (TensorMap) – TensorMap to transform, matching the transformation matrices in dtype and device

  • system_ids (Tensor | None) – optional one-dimensional tensor with one distinct "system" sample label per operation

  • add_inversion (bool) – compose every operation with the inversion before inverting

  • start (int) – first operation to invert

  • stop (int | None) – end of the range of operations to invert, defaulting to all of them

Returns:

transformed TensorMap with the same metadata and global information

Return type:

TensorMap

transform_block(key: LabelsEntry, block: TensorBlock, system_ids: Tensor | None = None, add_inversion: bool = False) → TensorBlock[source]

Apply the transformations to a TensorBlock and its gradients.

See transform_tensormap() for the row-routing conventions. The key must be the block’s key, carrying the o3_lambda / o3_sigma metadata needed for spherical components.

Parameters:
  • key (LabelsEntry) – the block’s key

  • block (TensorBlock) – TensorBlock to transform, matching the transformation matrices in dtype and device

  • system_ids (Tensor | None) – optional one-dimensional tensor with one distinct "system" sample label per operation

  • add_inversion (bool) – compose every operation with an inversion

Returns:

transformed TensorBlock with the same metadata

Return type:

TensorBlock

inverse_transform_block(key: LabelsEntry, block: TensorBlock, system_ids: Tensor | None = None, add_inversion: bool = False) → TensorBlock[source]

Apply the inverse transformations to a TensorBlock and its gradients.

See transform_block() for the parameters; the inverse uses transposed matrices and Wigner-D matrices.

Parameters:
  • key (LabelsEntry) – the block’s key

  • block (TensorBlock) – TensorBlock to transform, matching the transformation matrices in dtype and device

  • system_ids (Tensor | None) – optional one-dimensional tensor with one distinct "system" sample label per operation

  • add_inversion (bool) – compose every operation with an inversion before inverting

Returns:

transformed TensorBlock with the same metadata

Return type:

TensorBlock

metatomic.torch.o3.random_transformations(n: int, max_angular_momentum: int = 0, *, device: device, dtype: dtype, add_inversions: bool = False, generator: Generator | None = None) → O3Transformations[source]

Sample n transformations uniformly from SO(3), or from O(3) when inversions are included.

Rotations are sampled from the Haar measure on SO(3) via random unit quaternions. When add_inversions is True, each matrix is independently negated with probability 0.5, giving a uniform distribution over the full O(3) group.

Parameters:
  • n (int) – positive number of transformations to generate

  • max_angular_momentum (int) – non-negative maximum angular momentum for Wigner-D matrices

  • device (device) – target device for the output tensors

  • dtype (dtype) – target dtype for the output tensors; must be torch.float32 or torch.float64

  • add_inversions (bool) – if True, sample from O(3) instead of SO(3)

  • generator (Generator | None) – optional torch.Generator for reproducible sampling; when None the global RNG is used

Returns:

a single O3Transformations holding n operations

Return type:

O3Transformations

Cartesian and spherical tensors

cartesian_to_spherical() converts a Cartesian tensor of arbitrary rank, stored on a product basis with one xyz_1, xyz_2, … component axis per Cartesian index, into a set of blocks on a basis of real spherical harmonics (with a single o3_mu component axis, and o3_lambda / o3_sigma keys). This is done by coupling the Cartesian components two by two with Clebsch-Gordan coefficients, starting from the last ones. For tensors of rank 3 and above, there are multiple ways to get the same o3_lambda, and the coupling history is stored in additional l_{i} and k_{i} keys. spherical_to_cartesian() performs the inverse transformation. These functions can be used from TorchScript, in which case the Clebsch-Gordan coefficients have to be computed in advance with calculate_cg_coefficients().


metatomic.torch.o3.cartesian_to_spherical(tensor: TensorMap, components: List[str], keep_l_in_keys: bool | None = None, remove_blocks_threshold: float | None = 1e-09, cg_coefficients: TensorMap | None = None) → TensorMap[source]

Transform a tensor of arbitrary rank from Cartesian form to a spherical form.

Starting from a tensor on a basis of product of Cartesian coordinates, this function computes the same tensor using a basis of real spherical harmonics Y^M_L. For example, a rank 1 tensor with a single xyz component would be represented as a single L=1 spherical harmonic; while a rank 5 tensor using a product basis ℝ^3 ⊗ ℝ^3 ⊗ ℝ^3 ⊗ ℝ^3 ⊗ ℝ^3 would require multiple blocks up to L=5 spherical harmonics.

A single TensorBlock in the input might correspond to multiple TensorBlock in the output. The output keys will contain all the dimensions of the input keys, plus o3_lambda (indicating the spherical harmonics degree) and o3_sigma (indicating that this block is a proper or improper tensor with +1 and -1 respectively). If keep_l_in_keys is True or if the input tensor is a tensor of rank 3 or more, the keys will also contain multiple l_{i} and k_{i} dimensions, which indicate which angular momenta have been coupled together in which order to get this block.

components specifies which ones of the components of the input TensorMap should be transformed from Cartesian to spherical. All these components will be replaced in the output by a single o3_mu component, corresponding to the spherical harmonics M. If components is empty, the tensor is treated as a rank 0 (scalar) tensor, and the output contains a single o3_lambda=0, o3_sigma=1 block for each input block, with an additional o3_mu component placed after all the existing components.

If the input keys only contain the _ placeholder dimension (as created by metatensor.torch.Labels.single()), it is not included in the output keys.

By default, symmetric tensors will only contain blocks corresponding to o3_sigma=1. This is achieved by checking the norm of the blocks after the full calculation; and dropping any block with a norm below remove_blocks_threshold. To keep all blocks regardless of their norm, you can set remove_blocks_threshold=None.

Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.

Parameters:
  • tensor (TensorMap) – input tensor, using a Cartesian product basis

  • components (List[str]) – components of the input tensor to transform into spherical components, or an empty list for scalar tensors

  • keep_l_in_keys (bool | None) –

    should the output contains the values of angular momenta that were combined together? This defaults to False for rank 1 and 2 tensors, and True for all other tensors.

    Keys named l_{i} correspond to the input components, with l_1 being the last entry in components and l_N the first one. Keys named k_{i} correspond to intermediary spherical components created during the calculation, i.e. a k_{i} used to be o3_lambda.

  • remove_blocks_threshold (float | None) – Numerical tolerance to use when determining if a block’s norm is zero or not. Blocks with zero norm will be excluded from the output. Set this to None to keep all blocks in the output.

  • cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with calculate_cg_coefficients(), with lambda_max at least equal to the number of components. This is optional except when using this function from TorchScript.

Returns:

TensorMap containing spherical components instead of Cartesian components.

Return type:

TensorMap

metatomic.torch.o3.spherical_to_cartesian(tensor: TensorMap, new_component_names: List[str] | None = None, rank: int | None = None, cg_coefficients: TensorMap | None = None) → TensorMap[source]

Transform a tensor from spherical form back to Cartesian form.

This is the exact inverse of cartesian_to_spherical(). Starting from a tensor expressed on a basis of real spherical harmonics Y^M_L (as produced by cartesian_to_spherical()), this function reconstructs the corresponding Cartesian product-basis tensor. Blocks that were removed from the spherical tensor (because their norm was zero, see the remove_blocks_threshold parameter of cartesian_to_spherical()) are treated as blocks full of zeros.

The input keys must contain o3_lambda and o3_sigma, and the single o3_mu component will be replaced by a number of Cartesian components (according to the rank of the tensor, each with values [0, 1, 2]). For tensors of rank 3 and above the keys must also contain the k_{i} dimensions describing the coupling history (these are added automatically by cartesian_to_spherical()).

The rank of the tensor is given by the number of l_{i} dimensions in the keys. If these dimensions are missing (for example when using keep_l_in_keys=False in cartesian_to_spherical()), the rank must be given explicitly. It can not be determined from the o3_lambda in the keys, since some of the blocks might be missing: for example, an antisymmetric rank 2 tensor only contains a o3_lambda=1 block, like a rank 1 tensor.

Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.

Parameters:
  • tensor (TensorMap) – input tensor, using a spherical harmonics basis

  • new_component_names (List[str] | None) – names to give to the reconstructed Cartesian components. There should be as many of them as is the rank of the tensor. If None, the components are named xyz_1, xyz_2, …

  • rank (int | None) – rank of the Cartesian tensor. This is required if the keys do not contain the l_{i} dimensions, and must match the number of l_{i} dimensions otherwise. Use rank=0 to convert scalar tensors, removing the o3_mu component.

  • cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with calculate_cg_coefficients(), with lambda_max at least equal to the rank of the tensor. This is optional except when using this function from TorchScript to convert tensors of rank 2 and above.

Returns:

TensorMap containing Cartesian components instead of spherical components.

Return type:

TensorMap

metatomic.torch.o3.calculate_cg_coefficients(lambda_max: int, *, dtype: dtype, device: device) → TensorMap[source]

Calculate the Clebsch-Gordan coefficients coupling two real spherical harmonics of degrees l1 and l2 into a single real spherical harmonic of degree lambda, for all l1, l2 and lambda up to lambda_max.

These coefficients are used by cartesian_to_spherical() and spherical_to_cartesian(), and must be computed in advance and given to these functions explicitly when calling them from TorchScript. Converting a tensor of rank N requires lambda_max >= N.

The coefficients are stored in a TensorMap with l1, l2 and lambda keys. Each block contains a single sample (_), the m1, m2 and mu components (with values from -l to l for the corresponding degree), and a single property (cg_coefficient).

Parameters:
  • lambda_max (int) – non-negative maximal degree of the spherical harmonics

  • dtype (dtype) – dtype of the coefficients, either torch.float32 or torch.float64

  • device (device) – device on which the coefficients should be stored

Returns:

TensorMap containing the Clebsch-Gordan coefficients

Return type:

TensorMap

Coupled and uncoupled spherical tensors

A block with several spherical component axes o3_mu_1, o3_mu_2, … (with the corresponding o3_lambda_1 / o3_sigma_1, o3_lambda_2 / o3_sigma_2, … keys) transforms like a product of spherical harmonics. This is the uncoupled form, used for example for the blocks of a Hamiltonian expressed on a basis of atomic orbitals. uncoupled_to_coupled() uses Clebsch-Gordan coefficients to transform such a tensor into the coupled form, where each block has a single o3_mu component axis and transforms like a single spherical harmonic, with degree and parity given by the o3_lambda and o3_sigma keys. The degrees and parities of the uncoupled components are kept in the l_1 / s_1, l_2 / s_2, … keys of the coupled tensor, and coupled_to_uncoupled() uses them to perform the inverse transformation. As for the conversion between Cartesian and spherical tensors, the Clebsch-Gordan coefficients have to be computed in advance with calculate_cg_coefficients() to use these functions from TorchScript.


metatomic.torch.o3.uncoupled_to_coupled(tensor: TensorMap, components: List[str] | None = None, remove_blocks_threshold: float | None = None, cg_coefficients: TensorMap | None = None) → TensorMap[source]

Transform a tensor from an uncoupled spherical form to a coupled spherical form.

In the uncoupled form, the blocks contain multiple spherical components (named o3_mu_1, o3_mu_2, …), each transforming like a real spherical harmonic with degree and parity given by the corresponding o3_lambda_1/o3_sigma_1, o3_lambda_2/o3_sigma_2, … keys. Such a block transforms like a product of spherical harmonics Y^{m_1}_{l_1} ⊗ Y^{m_2}_{l_2} ⊗ .... In the coupled form, these components are replaced by a single o3_mu component and each block transforms like a single spherical harmonic Y^M_L, using Clebsch-Gordan coefficients to combine the spherical harmonics. For example, coupling two components with l_1 = 1 and l_2 = 2 gives three blocks, with L = 1, 2, 3.

The output keys contain o3_lambda and o3_sigma (the degree and parity of the coupled spherical harmonic), followed by the keys of the input. The o3_sigma of the output is given by the product of the o3_sigma of the coupled components, times (-1)^(l_1 + l_2 + ... + L). The degree and parity of the coupled components are needed to go back to the uncoupled form with coupled_to_uncoupled(), so the o3_lambda_{i} and o3_sigma_{i} keys are kept in the output, renamed to l_{i} and s_{i}. When coupling three or more components, the components are coupled two by two, starting from the last ones, and the degrees of the intermediate spherical harmonics are stored in k_{i} keys.

Depending on the symmetries of the input (for example when coupling a component with itself), some of the blocks will be zero everywhere. Any block with a squared norm below remove_blocks_threshold is not included in the output. To keep all blocks regardless of their norm, you can set remove_blocks_threshold=None.

Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.

Parameters:
  • tensor (TensorMap) – input tensor, in uncoupled form

  • components (List[str] | None) – names of the spherical components to couple. These must be consecutive components named o3_mu_{suffix}, with the corresponding o3_lambda_{suffix} and o3_sigma_{suffix} keys. If this is None, all the components named o3_mu_{suffix} are coupled, in the order they appear in the blocks.

  • remove_blocks_threshold (float | None) – Numerical tolerance to use when determining if a block’s norm is zero or not. Blocks with zero norm will be excluded from the output. Set this to None to keep all blocks in the output (the default).

  • cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with calculate_cg_coefficients(), with lambda_max at least equal to the sum of the degrees of the coupled components. This is optional except when using this function from TorchScript.

Returns:

TensorMap in coupled form

Return type:

TensorMap

metatomic.torch.o3.coupled_to_uncoupled(tensor: TensorMap, components: List[str] | None = None, cg_coefficients: TensorMap | None = None) → TensorMap[source]

Transform a tensor from a coupled spherical form back to an uncoupled spherical form.

This is the inverse of uncoupled_to_coupled(). The single o3_mu component of the input is replaced by the spherical components, using the degrees stored in the corresponding l_{suffix} keys. The l_{suffix} and s_{suffix} keys are renamed back to o3_lambda_{suffix} and o3_sigma_{suffix} in the output, to match the components. All the blocks that only differ by their o3_lambda and o3_sigma (and k_{i} when uncoupling three or more components) are combined into a single block, and any block missing from the input (for example because its norm was zero, see the remove_blocks_threshold parameter of uncoupled_to_coupled()) is treated as a block full of zeros.

Gradients are not supported yet, and this function raises an error if any of the input blocks has gradients.

Parameters:
  • tensor (TensorMap) – input tensor, in coupled form. The keys must contain o3_lambda, o3_sigma, the l_{suffix} and s_{suffix} for all components, and the k_{i} describing the intermediate couplings when uncoupling three or more components.

  • components (List[str] | None) – names of the spherical components to create, in the same order as they were given to uncoupled_to_coupled(). These must be named o3_mu_{suffix}. If this is None, one component is created for each l_{suffix} key, in the order they appear in the keys.

  • cg_coefficients (TensorMap | None) – Clebsch-Gordan coefficients computed with calculate_cg_coefficients(), with lambda_max at least equal to the sum of the degrees of the uncoupled components. This is optional except when using this function from TorchScript.

Returns:

TensorMap in uncoupled form

Return type:

TensorMap