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

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

The constructor stores a copy of matrix.

Parameters:

  • matrix (Tensor) – (3, 3) rotation or improper-rotation matrix

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

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_improper: bool

Whether this transformation is improper, with negative determinant.

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 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.

Returns:

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

Return type:

Tensor

wigner_D_matrix(ell: int) → Tensor[source]

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

For an improper transformation, transform_spherical() applies the inversion-parity factor separately.

Parameters:

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

Returns:

(2*ell+1, 2*ell+1) Wigner-D matrix

Return type:

Tensor

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 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 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) – non-negative 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

  • 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 O3Transformation objects

Return type:

list[O3Transformation]

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

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

Positions, cell vectors, neighbor-list displacements, and custom data following Conventions are transformed. Atomic types and periodic-boundary flags are preserved.

Parameters:

  • system (System) – input system

  • transformation (O3Transformation) – O(3) transformation to apply, matching system.positions in dtype and device

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]

Apply per-system O(3) transformations to a TensorMap and its gradients.

Scalar, Cartesian, and spherical data are identified by their component-axis names, following Conventions; one TensorMap may contain all three kinds of data. At most ten component axes are supported in one value or gradient block.

With multiple systems, the "system" sample label assigns each value sample to a transformation: samples labelled system_ids[i] use transformations[i]. A block may contain samples for only some of the systems, but every "system" label present in the block must appear in system_ids. A gradient sample uses the same transformation as the parent value sample referenced by its "sample" label. With one system, the "system" label is optional and ignored.

Parameters:

  • tensor (TensorMap) – TensorMap to transform

  • systems (list[System]) – systems corresponding positionally to transformations

  • transformations (list[O3Transformation]) – one O(3) transformation per system, matching the tensor values in dtype and device when present

  • system_ids (list[int] | Tensor | None) – one distinct integer "system" sample label per system; entry i is paired with transformations[i]. A tensor argument must be one-dimensional and use the same device as the tensor values. Defaults to range(len(systems))

Returns:

transformed TensorMap with the same keys and global information; when systems is empty, the tensor is unchanged

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]

Apply per-system O(3) transformations to a block and its gradients.

With one system, the "system" sample label is optional and ignored, as in transform_tensor().

Parameters:

  • key (LabelsEntry) – parent block key, supplying the O(3) labels required by spherical component axes

  • block (TensorBlock) – block to transform

  • systems (list[System]) – systems corresponding positionally to transformations

  • transformations (list[O3Transformation]) – one O(3) transformation per system, matching block.values in dtype and device

  • system_ids (list[int] | Tensor | None) – one distinct integer "system" sample label per system; entry i is paired with transformations[i]. A tensor argument must be one-dimensional and use the same device as block.values. Defaults to range(len(systems))

Returns:

block with transformed values and gradients and unchanged labels; when systems is empty, the block is unchanged

Return type:

TensorBlock