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¶

To transform a TensorBlock, transform_block() and transform_tensor() need to know, for each component 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 O3Transformation, 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}\).

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.O3Transformation(matrix: Tensor, max_angular_momentum: int)[source]¶

A single O(3) transformation, represented by a (3, 3) rotation or improper-rotation matrix.

Parameters:

matrix (Tensor) – (3, 3) rotation or improper-rotation matrix
max_angular_momentum (int) – maximum angular momentum of any spherical representation to be transformed by this transformation; Wigner-D matrices will be precomputed for all ell <= max_angular_momentum

property matrix: Tensor¶: The (3, 3) rotation or improper-rotation matrix.

property dtype: dtype¶: The dtype of the transformation matrix.

property device: device¶: The device of the transformation matrix.

property is_inverted: bool¶: Whether this transformation is an improper rotation (det < 0).

transform_cartesian(vectors: Tensor) → Tensor[source]¶

Apply the transformation to Cartesian vectors.

Parameters:: vectors (Tensor) – (…, 3) tensor of Cartesian vectors
Returns:: (…, 3) tensor of transformed vectors
Return type:: Tensor

transform_spherical(values: Tensor, ell: int, sigma: int) → Tensor[source]¶

Apply the transformation to spherical values.

Parameters:

values (Tensor) – (…, 2*ell+1) tensor of spherical values
ell (int) – angular momentum of the spherical representation
sigma (int) – inversion parity of the spherical representation

Returns:

(…, 2*ell+1) tensor of transformed spherical values

Return type:

Tensor

wigner_D_matrix(ell: int)[source]¶

Return the Wigner-D matrix for this transformation and angular momentum ell.

Parameters:: ell (int) – angular momentum of the spherical representation
Returns:: (2*ell+1, 2*ell+1) Wigner-D matrix

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

Sample n uniformly distributed O(3) transformations.

Rotations are sampled from the Haar measure on SO(3) via random unit quaternions. When include_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) – number of transformations to generate
max_angular_momentum (int) – maximum angular momentum for Wigner-D matrices
device (device) – target device for the output tensors
dtype (dtype) – target dtype for the output tensors
include_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:

list of n orthogonal (3, 3) tensors

Return type:

list[O3Transformation]

metatomic.torch.o3.transform_system(system: System, transformation: O3Transformation) → System[source]¶

Apply an O(3) transformation to a single System.

This function will transform positions, cell vectors, neighbor-list displacement vectors, and any custom data.

Parameters:

system (System) – input system
transformation (O3Transformation) – O(3) transformation to apply

Returns:

new System with transformed geometry

Return type:

System

metatomic.torch.o3.transform_tensor(tensor: TensorMap, systems: list[System], transformations: list[O3Transformation], system_ids: list[int] | Tensor | None = None) → TensorMap[source]¶

Rotate every block (and its gradients) of a TensorMap.

The tensor character of each component axis is inferred from its name, so a single TensorMap may freely mix scalar, Cartesian and spherical blocks, and blocks may carry gradients.

One of the samples dimensions must be "system", which is used to route each row to the correct system and transformation.

Parameters:

tensor (TensorMap) – TensorMap to rotate
systems (list[System]) – input systems
transformations (list[O3Transformation]) – per-system O(3) transformation matrices
system_ids (list[int] | Tensor | None) – index of the systems used in the samples of tensor; defaults to list(range(len(systems)))

Returns:

new TensorMap with rotated values and gradients

Return type:

TensorMap

metatomic.torch.o3.transform_block(key: LabelsEntry, block: TensorBlock, systems: list[System], transformations: list[O3Transformation], system_ids: list[int] | Tensor | None = None) → TensorBlock[source]¶

Rotate one block and all of its gradients.

Parameters:

key (LabelsEntry) – the block’s key, containing o3_lambda/o3_sigma values for spherical blocks
block (TensorBlock) – the block to rotate
systems (list[System]) – list of systems, one per transformation
transformations (list[O3Transformation]) – per-system O(3) transformation matrices
system_ids (list[int] | Tensor | None) – index of the systems used in the samples of block; defaults to list(range(len(systems)))

Returns:

new block with rotated values and gradients

Return type:

TensorBlock