mindscience.e3nn.so2_conv.SO3Rotation

class mindscience.e3nn.so2_conv.SO3Rotation(lmax, irreps_in, irreps_out)[source]

Class for handling SO(3) rotations of spherical-harmonic irreps.

Parameters

  • lmax (int) – Maximum angular momentum to be considered.

  • irreps_in (Union[str, Irreps]) – Input irreps specification.

  • irreps_out (Union[str, Irreps]) – Output irreps specification.

Examples

>>> from mindscience.e3nn.so2_conv import SO3Rotation
>>> rot = SO3Rotation(lmax=2, irreps_in="1x0e + 1x1o", irreps_out="1x1o")
>>> wigner, wigner_inv = rot.set_wigner(rot_mat3x3)
>>> rotated = rot.rotate(embedding, wigner)
static narrow(inputs, axis, start, length)[source]

Narrow (slice) a tensor along a specified axis.

Parameters

  • inputs (Tensor) – The tensor to be sliced.

  • axis (int) – The axis along which to perform the slice.

  • start (int) – The starting index of the slice.

  • length (int) – The number of elements to include in the slice.

Returns

Tensor, The sliced tensor.

rotate(embedding, wigner)[source]

Rotate an embedding tensor according to the provided Wigner-D matrices.

Parameters

  • embedding (Tensor) – Input tensor of shape (…, irreps_in.dim) containing the spherical-harmonic coefficients to be rotated.

  • wigner (tuple[Tensor]) – Tuple of Wigner-D matrices for l = 0 … lmax, each of shape (…, 2l+1, 2l+1).

Returns

tuple[Tensor], Tuple of rotated tensors, one per irrep in irreps_in, each of shape (…, mul, 2l+1).

rotate_inv(embedding, wigner_inv)[source]

Apply the inverse SO(3) rotation to an embedding tensor using the provided inverse Wigner-D matrices.

Parameters

  • embedding (tuple[Tensor]) – Tuple of tensors, one per irrep in irreps_out, each of shape (…, mul, 2l+1).

  • wigner_inv (tuple[Tensor]) – Tuple of inverse (transposed) Wigner-D matrices for l = 0 … lmax, each of shape (…, 2l+1, 2l+1).

Returns

Tensor, The rotated-back tensor of shape (…, irreps_out.dim) obtained by concatenating the inverse-rotated irreps.

static rotation_to_wigner_d_matrix(edge_rot_mat, start_lmax, end_lmax)[source]

Convert a batch of \(3 imes 3\) rotation matrices into Wigner-D matrices for the specified range of angular momenta.

Parameters

  • edge_rot_mat (Tensor) – Batch of SO(3) rotation matrices of shape (…, 3, 3).

  • start_lmax (int) – Minimum angular momentum to include.

  • end_lmax (int) – Maximum angular momentum to include.

Returns

list[Tensor], List of Wigner-D matrices for l = start_lmax … end_lmax, each of shape (…, 2l+1, 2l+1).

set_wigner(rot_mat3x3)[source]

Compute Wigner-D matrices and their inverses from a batch of \(3 imes 3\) rotation matrices.

Parameters

rot_mat3x3 (Tensor) – Batch of SO(3) rotation matrices of shape (…, 3, 3).

Returns

tuple[list[Tensor], list[Tensor]], A tuple containing two lists.

  • wigner: List of Wigner-D matrices for l = 0 … lmax, each of shape (…, 2l+1, 2l+1).

  • wigner_inv: List of transposed (inverse) Wigner-D matrices for l = 0 … lmax, same shapes.