Skip to content

Frames and spherical transforms

Spherical features can be rotated into a local frame, reordered between component layouts, or evaluated on a spherical grid. These are distinct operations: local-frame transforms rotate the features, IrrepLayoutTransform changes their ordering, and S2GridTransform maps between spherical coefficients and grid values.

The local-frame functions use Wigner matrices from wigner_D. Their output layouts also connect to the SO(2) linear operations.

build_local_frame 

build_local_frame(z_direction: Tensor, x_reference: Tensor) -> Tensor

Build proper global-to-local frames from z directions and x references.

Each nonzero z_direction is normalized to form the local z axis. The corresponding x_reference is projected onto its orthogonal plane and normalized to form the local x axis; the local y axis completes a right-handed frame. A stable Cartesian reference is used when the supplied reference is nearly parallel to z.

Parameters:

  • z_direction (Tensor) –

    Nonzero Cartesian directions with shape (..., 3).

  • x_reference (Tensor) –

    Cartesian reference directions broadcastable with z_direction and on the same device with the same dtype. The projected direction fixes the remaining SO(2) gauge freedom around z.

Returns:

  • Tensor –

    Proper rotation matrices with shape (..., 3, 3), whose rows are the

  • Tensor –

    local x, y, and z axes in global coordinates. Therefore the matrices

  • Tensor –

    map global Cartesian coordinates into the constructed local frames.

to_local_frame 

to_local_frame(
    so3_tensor: Tensor,
    wigner_matrices: Tensor,
    l_max: int,
    m_max: int | None = None,
    *,
    so2_layout: Literal["l_major", "m_major"] = "l_major",
) -> Tensor

Rotate global SO(3) features into local frames and optionally truncate m.

Apply the Wigner rotation, retain components with \(|m|\leq m_{\max}\), and pack the local features in so2_layout. The complete global input is always l-major. Channels are not mixed by this operation.

Parameters:

  • so3_tensor (Tensor) –

    Complete global features shaped (..., (l_max+1)**2, channels) in l-major order.

  • wigner_matrices (Tensor) –

    Flattened matrices from wigner_D(l_max, frames), where frames maps global Cartesian coordinates to local ones, as returned by build_local_frame. Shape is (..., sum((2*l+1)**2 for l in range(l_max+1))); leading dimensions, dtype, and device must match so3_tensor.

  • l_max (int) –

    Highest included non-negative degree.

  • m_max (int | None, default: None ) –

    Highest retained absolute order; None retains all orders.

  • so2_layout (Literal['l_major', 'm_major'], default: 'l_major' ) –

    Local output layout. l_major orders by ascending l, then ascending m within each l block. m_major orders blocks as m=0,+1,-1,+2,-2,..., with ascending l inside each block.

Returns:

  • Tensor –

    Local so2_tensor shaped (..., K, channels) in so2_layout, where

  • Tensor –

    K = sum(2*min(l, M)+1 for l in range(l_max+1)) and

  • Tensor –

    M = l_max if m_max is None else m_max. Leading dimensions,

  • Tensor –

    channels, dtype, and device are preserved.

Use to_global_frame with the same matrices and settings to rotate back. If orders were truncated, their original values cannot be recovered.

to_global_frame 

to_global_frame(
    so2_tensor: Tensor,
    wigner_matrices: Tensor,
    l_max: int,
    m_max: int | None = None,
    *,
    so2_layout: Literal["l_major", "m_major"] = "l_major",
) -> Tensor

Zero-fill local features and rotate them into the global frame.

Restore l-major ordering, fill omitted m components with zeros, and apply the transpose Wigner rotation. This is the adjoint of to_local_frame: without truncation it undoes that transform, but with truncation it cannot recover discarded components. Channels are not mixed.

Parameters:

  • so2_tensor (Tensor) –

    Local features shaped (..., K, channels) in so2_layout, where K = sum(2*min(l, M)+1 for l in range(l_max+1)) and M = l_max if m_max is None else m_max.

  • wigner_matrices (Tensor) –

    The same global-to-local Wigner matrices used by to_local_frame; do not transpose them yourself. Shape is (..., sum((2*l+1)**2 for l in range(l_max+1))); leading dimensions, dtype, and device must match so2_tensor.

  • l_max (int) –

    Highest included non-negative degree.

  • m_max (int | None, default: None ) –

    Highest absolute order present in the local input; None means all orders are present.

  • so2_layout (Literal['l_major', 'm_major'], default: 'l_major' ) –

    Local input layout, as described by to_local_frame. This must match the layout used to construct so2_tensor.

Returns:

  • Tensor –

    Complete global so3_tensor shaped (..., (l_max+1)**2, channels)

  • Tensor –

    in l-major order. Leading dimensions, channels, dtype, and device

  • Tensor –

    are preserved.

IrrepLayoutTransform 

IrrepLayoutTransform(
    l_max: int, m_max: int | None = None, *, device: device | str | None = None
)

Bases: Module

Reorder SO(3) spherical tensors between l-major and m-major layouts.

Both layouts contain the same components with 0 <= l <= l_max and |m| <= min(l,m_max). The m_max cutoff does not remove any l block; it removes the components with |m| > m_max from every higher-l block. Thus with l_max=3 and m_max=2, the l=3 block retains m=-2,...,+2 but omits m=-3 and m=+3.

For this example, the l-major layout, written as (l,m), is:

l=0: [(0,0)] l=1: [(1,-1), (1,0), (1,+1)] l=2: [(2,-2), (2,-1), (2,0), (2,+1), (2,+2)] l=3: [(3,-2), (3,-1), (3,0), (3,+1), (3,+2)]

The corresponding m-major layout, written as (m,l), is:

m=0: [(0,0), (0,1), (0,2), (0,3)] m=+1: [(+1,1), (+1,2), (+1,3)] m=-1: [(-1,1), (-1,2), (-1,3)] m=+2: [(+2,2), (+2,3)] m=-2: [(-2,2), (-2,3)]

This transform does not truncate m components. It accepts tensors already truncated to |m| <= m_max and only reorders them. to_m_major and to_l_major are exact inverses.

Parameters:

  • l_max (int) –

    Highest non-negative l.

  • m_max (int | None, default: None ) –

    Highest retained |m|. Defaults to l_max.

  • device (device | str | None, default: None ) –

    Initial device of the permutation indices.

Shape
  • Input: [..., K, C]
  • Output: [..., K, C]

Here K = (m_max+1)**2+(l_max-m_max)*(2*m_max+1).

Attributes:

  • m_sizes –

    Number of retained l values for each non-negative m.

  • n_components –

    Number of retained spherical components.

to_m_major 

to_m_major(l_major_tensor: Tensor) -> Tensor

Parameters:

  • l_major_tensor (Tensor) –

    [..., K, C].

Returns:

  • so2_tensor ( Tensor ) –

    [..., K, C].

to_l_major 

to_l_major(so2_tensor: Tensor) -> Tensor

Parameters:

  • so2_tensor (Tensor) –

    [..., K, C].

Returns:

  • l_major_tensor ( Tensor ) –

    [..., K, C].

S2GridTransform 

S2GridTransform(
    l_max: int,
    m_max: int | None = None,
    grid_shape: tuple[int, int] | None = None,
    *,
    component_layout: Literal["l_major", "m_major"] = "l_major",
    device: device | str | None = None,
    dtype: dtype | None = None,
)

Bases: Module

Transform band-limited real spherical functions to and from an S2 grid.

The component basis is the normalized Wikipedia real spherical-harmonic basis used throughout ELFES. Setting m_max < l_max keeps only |m| <= m_max in every \(l\) block. Components use either l_major layout, with consecutive \(l\) blocks ordered by m = -l, ..., l, or m_major layout, with the \(m=0\) block followed by the positive and negative blocks for each \(|m|\) and \(l\) increasing within every block.

The equiangular grid uses polar angles beta_i = pi * (i + 1/2) / n_polar and azimuthal angles alpha_j = 2 * pi * j / n_azimuthal. Its Cartesian directions are (sin(beta) * cos(alpha), sin(beta) * sin(alpha), cos(beta)). The inverse uses the Kostelec–Rockmore quadrature for this grid.

Parameters:

  • l_max (int) –

    Highest non-negative \(l\).

  • m_max (int | None, default: None ) –

    Highest retained \(|m|\). Defaults to l_max.

  • grid_shape (tuple[int, int] | None, default: None ) –

    Number of polar and azimuthal samples. Defaults to (2 * (l_max + 1), 2 * m_max + 1).

  • component_layout (Literal['l_major', 'm_major'], default: 'l_major' ) –

    Layout of the spherical components.

  • device (device | str | None, default: None ) –

    Initial buffer device.

  • dtype (dtype | None, default: None ) –

    Initial floating dtype. Defaults to PyTorch's default dtype.

Shape
  • Components: (..., n_components, channels)
  • Grid values: (..., n_polar, n_azimuthal, channels)

Here n_components = sum(2 * min(l, m_max) + 1 for l in range(l_max + 1)).

forward 

forward(components: Tensor) -> Tensor

Alias for to_grid.

to_grid 

to_grid(components: Tensor) -> Tensor

Evaluate spherical components on the configured grid.

from_grid 

from_grid(values: Tensor) -> Tensor

Project values on the configured grid back to components.