O(3) transformations

The metatomic.torch.o3 module rotates and inverts System and TensorMap objects, for example to generate randomly rotated copies of a structure for data augmentation.

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.

Reference

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

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

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

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

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

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

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

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) → List[System]

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

Returns:

one transformed System per input system

Return type:

List[System]

inverse_transform_systems(systems: List[System], add_inversion: bool = False) → List[System]

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

Returns:

one transformed System per input system

Return type:

List[System]

transform_tensormap(tensor: TensorMap, system_ids: Tensor | None = None, add_inversion: bool = False) → TensorMap

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

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) → TensorMap

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

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

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

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

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