holonomy-lib 0.8.0__py3-none-any.whl

holonomy_lib/lie/so3.py

@@ -0,0 +1,385 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""SO(3) rotation-group primitives: batched, GPU-native.
+
+The 3D rotation group SO(3) has a 3-dim Lie algebra so(3) ≅ R³ via the
+hat map  ω ↦ [ω]_×, where [ω]_× is the 3×3 skew-symmetric matrix
+
+    [[ 0,    −ω_z,  ω_y],
+     [ ω_z,  0,    −ω_x],
+     [−ω_y,  ω_x,   0  ]].
+
+The Lie-exp map `so(3) → SO(3)` is the matrix exponential, which has
+the closed-form Rodrigues formula:
+
+    exp([ω]_×) = I + (sin θ / θ) [ω]_× + ((1 − cos θ) / θ²) [ω]_ײ,
+    θ = ‖ω‖.
+
+Its inverse, the Lie log, is smooth on rotations with angle below π.
+At π the axis has a sign ambiguity; see `so3_log` for the
+representative it returns.
+
+References:
+  Hall (2015), §3.1: SO(3) and the Rodrigues formula.
+  Shoemake (1992): uniform-random rotations via quaternions.
+  Murray, Li, Sastry (1994). A Mathematical Introduction to Robotic
+    Manipulation, §2.3: the hat/vee maps and the SO(3) exp/log
+    closed forms used here.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+# Mathematical dimension of SO(3): the group lives in R^{3×3} and its
+# Lie algebra is R^3. **Scale of validity**: hard-coded by the
+# definition of "3D rotation"; not tunable. Cataloged as `SO3_DIM`.
+SO3_DIM: int = 3
+
+# Quaternion dimension. **Scale of validity**: hard-coded by the
+# definition of "unit quaternion"; not tunable. Cataloged as
+# `QUATERNION_DIM`.
+QUATERNION_DIM: int = 4
+
+# The near-pi crossover was calibrated in float64 against Rodrigues round
+# trips. so3_log uses the larger of this and sqrt(eps) for the input dtype,
+# so float32 (sqrt(eps) about 3.5e-4) gets a wider near-pi band.
+# Cataloged as SO3_LOG_NEAR_PI_RAD; this is an angle in radians.
+SO3_LOG_NEAR_PI_RAD: float = 1e-7
+
+# For cos(theta) > 0, theta = arcsin(s) with s = sin(theta), so
+# theta / sin(theta) = arcsin(s) / s = 1 + s² / 3! + O(s⁴).
+# The small-angle neighborhood comes from dtype epsilon: with s² below
+# sqrt(eps), the dropped s⁴ term is of order eps, which is roundoff. No
+# separate angle tolerance is needed.
+SO3_LOG_SERIES_SECOND_ORDER: float = 1.0 / math.factorial(SO3_DIM)
+
+# Taylor coefficients of sin(theta)/theta and (1-cos(theta))/theta².
+# They are reciprocal factorials from Rodrigues' formula, not fitted values.
+SO3_EXP_SINC_LINEAR: float = 1.0 / 6.0
+SO3_EXP_SINC_QUADRATIC: float = 1.0 / 120.0
+SO3_EXP_COSC_LINEAR: float = 1.0 / 24.0
+SO3_EXP_COSC_QUADRATIC: float = 1.0 / 720.0
+
+
+def axis_angle_to_matrix(
+    axis: torch.Tensor, angle: torch.Tensor,
+) -> torch.Tensor:
+    """Rodrigues formula: convert (unit axis, angle) → rotation matrix.
+
+        R = I + sin(θ)·K + (1 − cos(θ))·K²,   K = [axis]_×.
+
+    The axis is normalized to unit length internally so callers do
+    not need to pre-normalize.
+
+    Args:
+      axis:  `(..., 3)` axis vectors; will be normalized.
+      angle: `(...,)` rotation angle in radians.
+
+    Returns:
+      `(..., 3, 3)` rotation matrices in SO(3).
+
+    References:
+      Hall (2015), Theorem 3.3.
+    """
+    if axis.ndim < 1 or axis.shape[-1] != SO3_DIM:
+        raise ValueError(
+            f"axis must end with dim {SO3_DIM}; got axis.shape={tuple(axis.shape)}"
+        )
+    if axis.shape[:-1] != angle.shape:
+        raise ValueError(
+            f"axis batch shape {tuple(axis.shape[:-1])} must match "
+            f"angle shape {tuple(angle.shape)}"
+        )
+    # A zero axis would divide by zero. The clamp keeps the division finite,
+    # and a zero axis then gives the identity.
+    floor = torch.finfo(axis.dtype).tiny
+    norm = torch.linalg.norm(axis, dim=-1, keepdim=True).clamp(min=floor)
+    axis_unit = axis / norm                                # (..., 3)
+
+    K = _hat(axis_unit)                                     # (..., 3, 3)
+    s = torch.sin(angle).unsqueeze(dim=-1).unsqueeze(dim=-1)  # (..., 1, 1)
+    c = torch.cos(angle).unsqueeze(dim=-1).unsqueeze(dim=-1)
+    eye = torch.eye(SO3_DIM, dtype=axis.dtype, device=axis.device)
+    eye = eye.expand(*axis.shape[:-1], SO3_DIM, SO3_DIM)
+    return eye + s * K + (1.0 - c) * (K @ K)
+
+
+def matrix_to_axis_angle(R: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    """Inverse Rodrigues: convert SO(3) matrix to (unit axis, angle).
+
+    Uses the Lie logarithm, then normalizes its rotation vector. At the
+    identity the angle is zero and the axis is chosen as (1, 0, 0).
+    The axis-angle representation is not differentiable at identity.
+
+    Args:
+      R: `(..., 3, 3)` rotation matrix.
+
+    Returns:
+      `(axis, angle)`:
+        axis  `(..., 3)` unit vector.
+        angle `(...,)` in `[0, π]`.
+
+    References:
+      Murray-Li-Sastry (1994), §2.3.3: the dual-branch logarithm.
+    """
+    if R.shape[-2:] != (SO3_DIM, SO3_DIM):
+        raise ValueError(
+            f"R must end with ({SO3_DIM}, {SO3_DIM}); got R.shape={tuple(R.shape)}"
+        )
+    omega = so3_log(R)
+    angle = torch.linalg.vector_norm(omega, dim=-1)
+    # The axis is undefined at the identity, where the rotation vector is
+    # zero. Only this axis-angle API has to pick one, so it uses (1, 0, 0).
+    nonzero = angle > 0
+    divisor = torch.where(nonzero, angle, torch.ones_like(angle))
+    axis = omega / divisor.unsqueeze(-1)
+    canonical = torch.zeros_like(axis)
+    canonical[..., 0] = 1.0
+    return torch.where(nonzero.unsqueeze(-1), axis, canonical), angle
+
+
+def so3_exp(omega: torch.Tensor) -> torch.Tensor:
+    """Matrix exponential on so(3): ω ∈ R³ ↦ R ∈ SO(3).
+
+    Evaluates Rodrigues' formula. Near zero it uses Taylor polynomials for
+    the two coefficients, so the result and its gradient stay finite at
+    omega = 0.
+
+    Args:
+      omega: `(..., 3)` Lie algebra elements.
+
+    Returns:
+      `(..., 3, 3)` rotation matrices.
+
+    References:
+      Hall (2015). Lie Groups, Lie Algebras, and Representations, §3.1.
+      Murray, Li, Sastry (1994). A Mathematical Introduction to Robotic
+        Manipulation, equation 2.14.
+    """
+    if omega.ndim < 1 or omega.shape[-1] != SO3_DIM:
+        raise ValueError(
+            f"omega must end with dim {SO3_DIM}; got omega.shape={tuple(omega.shape)}"
+        )
+    theta_squared = omega.square().sum(dim=-1)
+    small = theta_squared <= math.sqrt(torch.finfo(omega.dtype).eps)
+    # torch.where computes both branches everywhere, and a NaN or inf in the
+    # unused one would still turn into a NaN gradient. So the sqrt and the
+    # division get 1 wherever the small branch is active, and the polynomial
+    # sees 0 wherever it is not.
+    safe_squared = torch.where(small, torch.ones_like(theta_squared), theta_squared)
+    theta = torch.sqrt(safe_squared)
+    series_squared = torch.where(small, theta_squared, torch.zeros_like(theta_squared))
+    sinc = torch.where(
+        small,
+        1.0 - SO3_EXP_SINC_LINEAR * series_squared
+        + SO3_EXP_SINC_QUADRATIC * series_squared.square(),
+        torch.sin(theta) / theta,
+    )
+    cosc = torch.where(
+        small,
+        0.5 - SO3_EXP_COSC_LINEAR * series_squared
+        + SO3_EXP_COSC_QUADRATIC * series_squared.square(),
+        (1.0 - torch.cos(theta)) / safe_squared,
+    )
+    generator = _hat(omega)
+    eye = torch.eye(SO3_DIM, dtype=omega.dtype, device=omega.device)
+    return (
+        eye + sinc[..., None, None] * generator
+        + cosc[..., None, None] * (generator @ generator)
+    )
+
+
+def so3_log(R: torch.Tensor) -> torch.Tensor:
+    """Principal rotation vector of an SO(3) matrix, with norm in [0, π].
+
+    Away from π, the skew part is sin(theta) times the rotation axis, so
+    the rotation vector is that part scaled by theta / sin(theta). Near
+    zero that quotient is replaced by its series 1 + sin²(theta) / 6.
+    Near π the sine is too small to carry the axis, so the symmetric part
+    determines it instead. The sign at exactly π is intrinsically
+    ambiguous, so the logarithm is not smooth there.
+
+    Args:
+      R: `(..., 3, 3)` rotation matrices.
+
+    Returns:
+      `(..., 3)` rotation vectors.
+
+    References:
+      Murray, Li, Sastry (1994). A Mathematical Introduction to Robotic
+        Manipulation, §2.3.3, the logarithm of a rotation matrix.
+    """
+    if R.shape[-2:] != (SO3_DIM, SO3_DIM):
+        raise ValueError(
+            f"R must end with ({SO3_DIM}, {SO3_DIM}); got R.shape={tuple(R.shape)}"
+        )
+    skew_vector = _vee(0.5 * (R - R.mT))
+    sin_angle = torch.linalg.vector_norm(skew_vector, dim=-1)
+    sin_squared = (skew_vector * skew_vector).sum(dim=-1)
+    cos_angle = ((R.diagonal(dim1=-2, dim2=-1).sum(dim=-1) - 1.0) * 0.5)
+    angle = torch.atan2(sin_angle, cos_angle)
+    small = (cos_angle > 0) & (sin_squared < math.sqrt(torch.finfo(R.dtype).eps))
+    # The skew part only carries the axis when sin(theta) is not tiny, so
+    # this threshold switches to the symmetric part close to π.
+    near_pi = angle > math.pi - max(
+        SO3_LOG_NEAR_PI_RAD, math.sqrt(torch.finfo(R.dtype).eps),
+    )
+    # The angle / sine quotient is computed everywhere, so its sine is swapped
+    # for 1 in the small and near-π branches, where that quotient is not
+    # used. This matters in backward too: torch.where sends zero gradient
+    # into the unused branch, and zero times an infinite derivative is NaN.
+    safe_sine = torch.where(small | near_pi, torch.ones_like(sin_angle), sin_angle)
+    scale = torch.where(
+        small, 1.0 + SO3_LOG_SERIES_SECOND_ORDER * sin_squared, angle / safe_sine,
+    )
+    omega = scale.unsqueeze(-1) * skew_vector
+
+    # Sym(R) - cos(theta) I = (1 - cos(theta)) axis axis^T. Its largest
+    # diagonal selects a row with a nonzero axis component even at π.
+    eye = torch.eye(SO3_DIM, dtype=R.dtype, device=R.device)
+    symmetric = 0.5 * (R + R.mT) - cos_angle[..., None, None] * eye
+    # Away from π this part's axis is not used. At the identity, though,
+    # Sym(R) - cos(theta) I is exactly zero, so its row would be a 0 / 0
+    # normalization. Swap in the identity wherever this is not used. A NaN
+    # in an unused branch would still poison the gradient.
+    symmetric = torch.where(near_pi[..., None, None], symmetric, eye)
+    pivot = symmetric.diagonal(dim1=-2, dim2=-1).argmax(dim=-1)
+    row = symmetric.gather(
+        -2, pivot[..., None, None].expand(*R.shape[:-2], 1, SO3_DIM),
+    ).squeeze(-2)
+    axis = row / torch.linalg.vector_norm(row, dim=-1, keepdim=True)
+    alignment = (axis * skew_vector).sum(dim=-1, keepdim=True)
+    axis = torch.where(alignment < 0, -axis, axis)
+    return torch.where(near_pi.unsqueeze(-1), angle.unsqueeze(-1) * axis, omega)
+
+
+def compose(R1: torch.Tensor, R2: torch.Tensor) -> torch.Tensor:
+    """SO(3) group product `R1 · R2`, a plain matmul.
+
+    It lives in this module so the group operations are together.
+
+    Args:
+      R1, R2: broadcastable `(..., 3, 3)` rotation matrices.
+
+    Returns:
+      `(..., 3, 3)` rotation matrices, applying R2 before R1.
+
+    References:
+      Hall (2015). Lie Groups, Lie Algebras, and Representations, §1.2.
+    """
+    if R1.shape[-2:] != (SO3_DIM, SO3_DIM) or R2.shape[-2:] != (SO3_DIM, SO3_DIM):
+        raise ValueError(
+            f"both inputs must end with ({SO3_DIM}, {SO3_DIM}); got "
+            f"R1.shape={tuple(R1.shape)}, R2.shape={tuple(R2.shape)}"
+        )
+    return R1 @ R2
+
+
+def random_so3(
+    batch_size: int,
+    generator: torch.Generator | None = None,
+    device: str | torch.device = "cpu",
+    dtype: torch.dtype = torch.float64,
+) -> torch.Tensor:
+    """Haar-uniform random rotation matrices.
+
+    Construction: sample a unit quaternion uniformly on S³ (the
+    Shoemake 1992 method), then convert to a rotation matrix. The
+    quaternion-Haar measure pushes through to the SO(3)-Haar measure
+    under the standard 2-to-1 cover.
+
+    Args:
+      batch_size: number of independent rotations.
+      generator: torch.Generator for reproducibility.
+      device, dtype: tensor placement.
+
+    Returns:
+      `(batch_size, 3, 3)` SO(3) matrices.
+
+    References:
+      Shoemake (1992).
+    """
+    if batch_size < 0:
+        raise ValueError(f"batch_size must be >= 0, got {batch_size}")
+    # Shoemake: three uniform [0, 1] samples u1, u2, u3 → unit
+    # quaternion (w, x, y, z) with q on S³ uniformly. The sampling
+    # and trig run on `device` so a `device="cuda"` call doesn't
+    # silently fall back to CPU for the per-sample scalar work. The
+    # generator must be on the matching device: if the caller wants
+    # GPU sampling, they must pass a CUDA generator.
+    target_device = torch.device(device)
+    # A CUDA generator reports `cuda` with no index, so compare the device
+    # type, and the index only when both sides name one.
+    generator_device = None if generator is None else generator.device
+    if generator_device is not None and (
+        generator_device.type != target_device.type
+        or (generator_device.index is not None
+            and target_device.index is not None
+            and generator_device.index != target_device.index)
+    ):
+        raise ValueError(
+            f"generator is on {generator_device} but device={target_device} "
+            f"was requested; pass a generator on the target device or omit "
+            f"`generator` to use the default"
+        )
+    u = torch.rand(
+        batch_size, SO3_DIM,
+        generator=generator, dtype=dtype, device=target_device,
+    )
+    u1, u2, u3 = u[..., 0], u[..., 1], u[..., 2]
+    sqrt_u1 = torch.sqrt(u1)
+    sqrt_1_minus_u1 = torch.sqrt(1.0 - u1)
+    two_pi_u2 = 2.0 * math.pi * u2
+    two_pi_u3 = 2.0 * math.pi * u3
+    w = sqrt_1_minus_u1 * torch.sin(two_pi_u2)
+    x = sqrt_1_minus_u1 * torch.cos(two_pi_u2)
+    y = sqrt_u1 * torch.sin(two_pi_u3)
+    z = sqrt_u1 * torch.cos(two_pi_u3)
+    q = torch.stack([w, x, y, z], dim=-1)                   # (B, 4)
+    return _quaternion_to_matrix(q)
+
+
+# ============================================================
+# Internal helpers
+# ============================================================
+
+
+def _hat(v: torch.Tensor) -> torch.Tensor:
+    """Hat map: R³ → so(3). v ↦ [v]_× as a (..., 3, 3) skew-symmetric."""
+    zero = torch.zeros_like(v[..., 0])
+    x, y, z = v[..., 0], v[..., 1], v[..., 2]
+    row0 = torch.stack([zero, -z, y], dim=-1)
+    row1 = torch.stack([z, zero, -x], dim=-1)
+    row2 = torch.stack([-y, x, zero], dim=-1)
+    return torch.stack([row0, row1, row2], dim=-2)
+
+
+def _vee(K: torch.Tensor) -> torch.Tensor:
+    """Vee map: skew(3) → R³. The inverse of `_hat`."""
+    return torch.stack(
+        [K[..., 2, 1], K[..., 0, 2], K[..., 1, 0]], dim=-1,
+    )
+
+
+def _quaternion_to_matrix(q: torch.Tensor) -> torch.Tensor:
+    """Convert unit quaternions `(w, x, y, z)` to rotation matrices.
+
+    Standard formula; works for arbitrary batch shape ending in 4.
+    """
+    # Components are read by index, in the (w, x, y, z) layout. The last
+    # index is written as QUATERNION_DIM - 1 so it is tied to the constant.
+    # The others are plain 0, 1, 2.
+    last = QUATERNION_DIM - 1
+    w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., last]
+    # Each product shows up in more than one entry, so compute it once.
+    wx, wy, wz = w * x, w * y, w * z
+    xx, yy, zz = x * x, y * y, z * z
+    xy, xz, yz = x * y, x * z, y * z
+    row0 = torch.stack([1 - 2 * (yy + zz), 2 * (xy - wz), 2 * (xz + wy)], dim=-1)
+    row1 = torch.stack([2 * (xy + wz), 1 - 2 * (xx + zz), 2 * (yz - wx)], dim=-1)
+    row2 = torch.stack([2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (xx + yy)], dim=-1)
+    return torch.stack([row0, row1, row2], dim=-2)

holonomy_lib/lie/spherical_harmonics.py

@@ -0,0 +1,189 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""Real spherical harmonics Y_lm on the unit 2-sphere, closed form
+for l ≤ 4.
+
+Real spherical harmonics are the orthonormal basis for square-
+integrable functions on `S² ⊂ R³` that diagonalizes the Laplace-
+Beltrami operator on the sphere. On the sphere, each `Y_lm` agrees with
+a homogeneous polynomial of degree l in `(x, y, z)`. There are `2l + 1`
+of them at each `l`, indexed by `m ∈ {−l, …, l}`. The total at
+truncation `l_max` is `(l_max + 1)²` components.
+
+In equivariant deep learning these serve as the spherical features that
+rotate in a known way under SO(3). e3nn and SE(3)-Transformer use them
+as steerable filters, and Clebsch-Gordan coefficients are what combine
+different degrees into equivariant products.
+
+The basis is orthonormal over the sphere and uses the Condon-Shortley-free
+convention: every prefactor is positive, with no `(−1)^m` phase. Inputs
+are direction vectors `(..., 3)`; the function normalizes to unit norm
+internally, so callers can pass unnormalized vectors.
+
+References:
+  Wikipedia, "Real spherical harmonics": the explicit formulas
+    transcribed here are the standard tesseral/sectoral basis.
+  Edmonds (1957), §2.5.
+  Geiger, M., Smidt, T. (2022). e3nn: Euclidean neural networks.
+    arXiv:2207.09453. The modern implementation in geometric
+    deep learning. To compare with e3nn.o3.spherical_harmonics, pass
+    directions[..., [1, 2, 0]], normalize=True, normalization="integral".
+    e3nn uses Cartesian order (x, y, z) at degree one; this basis uses
+    (y, z, x). See https://docs.e3nn.org/en/stable/api/o3/o3_sh.html.
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+
+_SUPPORTED_L_MAX: int = 4
+
+
+def real_spherical_harmonics(
+    directions: torch.Tensor, l_max: int = 2,
+) -> torch.Tensor:
+    """Real spherical harmonics `Y_lm` at the given direction vectors.
+
+    For each input direction `(x, y, z)` (normalized internally),
+    returns the `(l_max + 1)²` real spherical harmonics in the
+    standard ordering:
+
+        [Y_{0,0},
+         Y_{1,-1}, Y_{1,0}, Y_{1,1},
+         Y_{2,-2}, Y_{2,-1}, Y_{2,0}, Y_{2,1}, Y_{2,2},
+         …,
+         Y_{l_max, -l_max}, …, Y_{l_max, l_max}]
+
+    Args:
+      directions: `(..., 3)` direction vectors. Need not be unit; the
+        function normalizes internally.
+      l_max: highest degree to compute. Must be in `[0, 4]`. Higher
+        degrees need the associated-Legendre recurrence, which is not
+        implemented.
+
+    Returns:
+      `(..., (l_max + 1)²)` real spherical harmonic values.
+
+    References:
+      Wikipedia, "Table of spherical harmonics": real form.
+      Edmonds (1957), §2.5.
+    """
+    if directions.ndim < 1 or directions.shape[-1] != 3:
+        raise ValueError(
+            f"directions must end with dim 3; got shape="
+            f"{tuple(directions.shape)}"
+        )
+    if not (0 <= l_max <= _SUPPORTED_L_MAX):
+        raise ValueError(
+            f"l_max must be in [0, {_SUPPORTED_L_MAX}]; got l_max={l_max}. "
+            f"Higher degrees need the associated-Legendre recurrence "
+            f"(not implemented)."
+        )
+
+    # The clamp keeps a zero direction from dividing by zero. It then stays
+    # the zero vector, and the polynomials evaluated there are not meaningful
+    # SH values.
+    floor = torch.finfo(directions.dtype).tiny
+    norm = torch.linalg.norm(directions, dim=-1, keepdim=True).clamp(min=floor)
+    unit = directions / norm
+    x, y, z = unit[..., 0], unit[..., 1], unit[..., 2]
+
+    components: list[torch.Tensor] = []
+
+    # ----- l = 0 -----
+    # Y_{0,0} = 0.5 · sqrt(1/π). full_like keeps the batch shape, dtype and
+    # device of x, so the constant stacks with the other components.
+    y00 = torch.full_like(x, 0.5 * math.sqrt(1.0 / math.pi))
+    components.append(y00)
+
+    # Each block below adds one degree. Returning early for a smaller l_max
+    # means the higher blocks never run.
+    if l_max < 1:
+        return torch.stack(components, dim=-1)
+
+    # ----- l = 1 -----
+    # Y_{1,m} = sqrt(3/(4π)) · {y, z, x} for m = -1, 0, 1.
+    c1 = math.sqrt(3.0 / (4.0 * math.pi))
+    components.extend([c1 * y, c1 * z, c1 * x])
+
+    if l_max < 2:
+        return torch.stack(components, dim=-1)
+
+    # ----- l = 2 -----
+    # Y_{2,-2} = 0.5 · sqrt(15/π) · xy
+    # Y_{2,-1} = 0.5 · sqrt(15/π) · yz
+    # Y_{2,0}  = 0.25 · sqrt(5/π) · (3z² − 1)
+    # Y_{2,1}  = 0.5 · sqrt(15/π) · xz
+    # Y_{2,2}  = 0.25 · sqrt(15/π) · (x² − y²)
+    c2a = 0.5 * math.sqrt(15.0 / math.pi)
+    c20 = 0.25 * math.sqrt(5.0 / math.pi)
+    c22 = 0.25 * math.sqrt(15.0 / math.pi)
+    components.extend([
+        c2a * x * y,
+        c2a * y * z,
+        c20 * (3.0 * z * z - 1.0),
+        c2a * x * z,
+        c22 * (x * x - y * y),
+    ])
+
+    if l_max < 3:
+        return torch.stack(components, dim=-1)
+
+    # ----- l = 3 -----
+    # Y_{3,-3} = 0.25 · sqrt(35/(2π)) · y(3x² − y²)
+    # Y_{3,-2} = 0.5  · sqrt(105/π) · xyz
+    # Y_{3,-1} = 0.25 · sqrt(21/(2π)) · y(5z² − 1)
+    # Y_{3,0}  = 0.25 · sqrt(7/π) · (5z³ − 3z)
+    # Y_{3,1}  = 0.25 · sqrt(21/(2π)) · x(5z² − 1)
+    # Y_{3,2}  = 0.25 · sqrt(105/π) · z(x² − y²)
+    # Y_{3,3}  = 0.25 · sqrt(35/(2π)) · x(x² − 3y²)
+    c3_a = 0.25 * math.sqrt(35.0 / (2.0 * math.pi))
+    c3_b = 0.5 * math.sqrt(105.0 / math.pi)
+    c3_c = 0.25 * math.sqrt(21.0 / (2.0 * math.pi))
+    c3_0 = 0.25 * math.sqrt(7.0 / math.pi)
+    c3_d = 0.25 * math.sqrt(105.0 / math.pi)
+    components.extend([
+        c3_a * y * (3.0 * x * x - y * y),
+        c3_b * x * y * z,
+        c3_c * y * (5.0 * z * z - 1.0),
+        c3_0 * (5.0 * z ** 3 - 3.0 * z),
+        c3_c * x * (5.0 * z * z - 1.0),
+        c3_d * z * (x * x - y * y),
+        c3_a * x * (x * x - 3.0 * y * y),
+    ])
+
+    if l_max < 4:
+        return torch.stack(components, dim=-1)
+
+    # ----- l = 4 -----
+    # Y_{4,-4} = 0.75 · sqrt(35/π) · xy(x² − y²)
+    # Y_{4,-3} = 0.75 · sqrt(35/(2π)) · yz(3x² − y²)
+    # Y_{4,-2} = 0.75 · sqrt(5/π) · xy(7z² − 1)
+    # Y_{4,-1} = 0.75 · sqrt(5/(2π)) · yz(7z² − 3)
+    # Y_{4,0}  = 3/16 · sqrt(1/π) · (35z⁴ − 30z² + 3)
+    # Y_{4,1}  = 0.75 · sqrt(5/(2π)) · xz(7z² − 3)
+    # Y_{4,2}  = 0.375 · sqrt(5/π) · (x² − y²)(7z² − 1)
+    # Y_{4,3}  = 0.75 · sqrt(35/(2π)) · xz(x² − 3y²)
+    # Y_{4,4}  = 3/16 · sqrt(35/π) · (x⁴ − 6x²y² + y⁴)
+    c4_a = 0.75 * math.sqrt(35.0 / math.pi)
+    c4_b = 0.75 * math.sqrt(35.0 / (2.0 * math.pi))
+    c4_c = 0.75 * math.sqrt(5.0 / math.pi)
+    c4_d = 0.75 * math.sqrt(5.0 / (2.0 * math.pi))
+    c4_0 = (3.0 / 16.0) * math.sqrt(1.0 / math.pi)
+    c4_e = 0.375 * math.sqrt(5.0 / math.pi)
+    c4_f = (3.0 / 16.0) * math.sqrt(35.0 / math.pi)
+    components.extend([
+        c4_a * x * y * (x * x - y * y),
+        c4_b * y * z * (3.0 * x * x - y * y),
+        c4_c * x * y * (7.0 * z * z - 1.0),
+        c4_d * y * z * (7.0 * z * z - 3.0),
+        c4_0 * (35.0 * z ** 4 - 30.0 * z * z + 3.0),
+        c4_d * x * z * (7.0 * z * z - 3.0),
+        c4_e * (x * x - y * y) * (7.0 * z * z - 1.0),
+        c4_b * x * z * (x * x - 3.0 * y * y),
+        c4_f * (x ** 4 - 6.0 * x * x * y * y + y ** 4),
+    ])
+    return torch.stack(components, dim=-1)

holonomy_lib/manifolds/__init__.py

@@ -0,0 +1,53 @@
+# SPDX-License-Identifier: FSL-1.1-ALv2
+# Copyright 2026 Synoros
+
+"""holonomy_lib.manifolds; Riemannian manifold primitives, GPU-batched.
+
+Currently implemented:
+  fixed_rank.FixedRankManifold; M_r(m, n), the rank-r matrices in R^{m×n}
+  spd.SPDManifold             ; P(n), symmetric positive definite,
+                                 affine-invariant metric
+  lorentz.LorentzManifold     ; H^n_k, hyperboloid model of hyperbolic
+                                 space of curvature k < 0 (Nickel-Kiela 2018)
+  stereographic.KappaStereographicManifold; parametric κ ∈ R, interpolates
+                                 spherical (κ>0), Euclidean (κ=0),
+                                 hyperbolic (κ<0) via gyro-algebra
+                                 (Bachmann-Bécigneul-Ganea 2020)
+  lorentzian.LorentzianManifold; flat Minkowski spacetime R^{1, n-1}
+                                 with pseudo-Riemannian (1, n-1)
+                                 signature; causal structure (timelike/
+                                 null/spacelike) + proper time/distance
+                                 (MTW 1973; O'Neill 1983)
+  product.ProductManifold     ; Riemannian product M_1 × M_2 × ...
+                                 with the Pythagorean direct-sum
+                                 metric (Gu-Sala et al. 2019;
+                                 Skopek et al. 2019)
+Planned (port/wrap from geoopt, geomstats, pymanopt):
+  sphere, stiefel, grassmann
+"""
+
+from holonomy_lib.manifolds.comparison import (
+    model_anisotropic_flux,
+    model_ball_volume,
+    model_sphere_area,
+    sn_kappa,
+)
+from holonomy_lib.manifolds.fixed_rank import FixedRankManifold
+from holonomy_lib.manifolds.lorentz import LorentzManifold
+from holonomy_lib.manifolds.lorentzian import LorentzianManifold
+from holonomy_lib.manifolds.product import ProductManifold
+from holonomy_lib.manifolds.spd import SPDManifold
+from holonomy_lib.manifolds.stereographic import KappaStereographicManifold
+
+__all__ = [
+    "FixedRankManifold",
+    "KappaStereographicManifold",
+    "LorentzManifold",
+    "LorentzianManifold",
+    "ProductManifold",
+    "SPDManifold",
+    "model_anisotropic_flux",
+    "model_ball_volume",
+    "model_sphere_area",
+    "sn_kappa",
+]