nanomanifold 0.1.1__py3-none-any.whl

nanomanifold/SO3/conversions/euler.py

@@ -0,0 +1,149 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+from ..canonicalize import canonicalize
+from . import matrix
+
+
+def to_euler(q: Float[Any, "... 4"], convention: str = "ZYX") -> Float[Any, "... 3"]:
+    q = canonicalize(q)
+
+    R = matrix.to_matrix(q)
+
+    return _matrix_to_euler(R, convention)
+
+
+def from_euler(euler: Float[Any, "... 3"], convention: str = "ZYX") -> Float[Any, "... 4"]:
+    R = _euler_to_matrix(euler, convention)
+
+    return matrix.from_matrix(R)
+
+
+def _euler_to_matrix(euler, convention):
+    """Convert Euler angles to rotation matrix."""
+    xp = get_namespace(euler)
+
+    eye = xp.eye(3, dtype=euler.dtype)
+    output_shape = euler.shape[:-1] + (3, 3)
+    R = xp.broadcast_to(eye, output_shape)
+
+    is_extrinsic = convention.islower()
+    conv = convention.lower()
+
+    for i, axis in enumerate(conv):
+        angle = euler[..., i]
+        R_axis = _rotation_matrix(angle, axis)
+
+        if is_extrinsic:
+            R = xp.matmul(R_axis, R)
+        else:
+            R = xp.matmul(R, R_axis)
+
+    return R
+
+
+def _rotation_matrix(angle, axis):
+    """Create rotation matrix for given angle and axis."""
+    xp = get_namespace(angle)
+
+    cos_a = xp.cos(angle)
+    sin_a = xp.sin(angle)
+    zero = xp.zeros_like(cos_a)
+    ones = xp.ones_like(cos_a)
+
+    if axis == "x":
+        mat = xp.stack(
+            [xp.stack([ones, zero, zero], axis=-1), xp.stack([zero, cos_a, -sin_a], axis=-1), xp.stack([zero, sin_a, cos_a], axis=-1)],
+            axis=-2,
+        )
+    elif axis == "y":
+        mat = xp.stack(
+            [xp.stack([cos_a, zero, sin_a], axis=-1), xp.stack([zero, ones, zero], axis=-1), xp.stack([-sin_a, zero, cos_a], axis=-1)],
+            axis=-2,
+        )
+    elif axis == "z":
+        mat = xp.stack(
+            [xp.stack([cos_a, -sin_a, zero], axis=-1), xp.stack([sin_a, cos_a, zero], axis=-1), xp.stack([zero, zero, ones], axis=-1)],
+            axis=-2,
+        )
+    else:
+        raise ValueError(f"Invalid axis: {axis}")
+
+    return mat
+
+
+def _matrix_to_euler(matrix, convention):
+    """Convert rotation matrix to Euler angles."""
+    xp = get_namespace(matrix)
+
+    is_extrinsic = convention.islower()
+
+    if is_extrinsic:
+        convention = convention.upper()
+        convention = convention[::-1]
+
+    eulers = _matrix_to_euler_angles(matrix, convention)
+
+    if is_extrinsic:
+        return xp.stack([eulers[..., 2], eulers[..., 1], eulers[..., 0]], axis=-1)
+
+    return eulers
+
+
+def _matrix_to_euler_angles(matrix, convention):
+    """Extract Euler angles from rotation matrix using systematic approach."""
+    xp = get_namespace(matrix)
+
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+
+    if tait_bryan:
+        sign = -1.0 if i0 - i2 in [-1, 2] else 1.0
+        x = matrix[..., i0, i2] * sign
+
+        one = xp.ones_like(x)
+        eps = xp.finfo(x.dtype).eps * one
+        central_angle = xp.arcsin(xp.clip(x, -one + eps, one - eps))
+    else:
+        central_angle = xp.arccos(xp.clip(matrix[..., i0, i0], -1, 1))
+
+    first_angle = _angle_from_tan(convention[0], convention[1], matrix[..., i2], False, tait_bryan, xp)
+    third_angle = _angle_from_tan(convention[2], convention[1], matrix[..., i0, :], True, tait_bryan, xp)
+
+    return xp.stack([first_angle, central_angle, third_angle], axis=-1)
+
+
+def _angle_from_tan(axis, other_axis, data, horizontal, tait_bryan, xp):
+    """Compute angle from tangent using systematic indexing."""
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return xp.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return xp.atan2(-data[..., i2], data[..., i1])
+    return xp.atan2(data[..., i2], -data[..., i1])
+
+
+def _index_from_letter(letter):
+    """Convert axis letter to index."""
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+    raise ValueError("letter must be either X, Y or Z.")

nanomanifold/SO3/conversions/matrix.py

@@ -0,0 +1,79 @@
+"""Matrix conversions for SO(3) rotations."""
+
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+from ..canonicalize import canonicalize
+
+
+def to_matrix(q: Float[Any, "... 4"]) -> Float[Any, "... 3 3"]:
+    xp = get_namespace(q)
+    q = canonicalize(q)
+    w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
+
+    R = xp.stack(
+        [
+            xp.stack([1 - 2 * (y * y + z * z), 2 * (x * y - w * z), 2 * (x * z + w * y)], axis=-1),
+            xp.stack([2 * (x * y + w * z), 1 - 2 * (x * x + z * z), 2 * (y * z - w * x)], axis=-1),
+            xp.stack([2 * (x * z - w * y), 2 * (y * z + w * x), 1 - 2 * (x * x + y * y)], axis=-1),
+        ],
+        axis=-2,
+    )
+
+    return R
+
+
+def from_matrix(R: Float[Any, "... 3 3"]) -> Float[Any, "... 4"]:
+    xp = get_namespace(R)
+
+    trace = R[..., 0, 0] + R[..., 1, 1] + R[..., 2, 2]
+
+    zero = trace * 0
+    one = zero + 1
+    eps = one * 1e-10
+    quarter = one * 0.25
+    two = one * 2
+
+    s1 = xp.sqrt(xp.maximum(zero, trace + one)) * two  # s = 4 * w
+    s1_safe = xp.where(s1 < eps, eps, s1)  # Avoid division by zero
+    w1 = quarter * s1
+    x1 = (R[..., 2, 1] - R[..., 1, 2]) / s1_safe
+    y1 = (R[..., 0, 2] - R[..., 2, 0]) / s1_safe
+    z1 = (R[..., 1, 0] - R[..., 0, 1]) / s1_safe
+
+    s2 = xp.sqrt(xp.maximum(zero, one + R[..., 0, 0] - R[..., 1, 1] - R[..., 2, 2])) * two  # s = 4 * x
+    s2_safe = xp.where(s2 < eps, eps, s2)  # Avoid division by zero
+    w2 = (R[..., 2, 1] - R[..., 1, 2]) / s2_safe
+    x2 = quarter * s2
+    y2 = (R[..., 0, 1] + R[..., 1, 0]) / s2_safe
+    z2 = (R[..., 0, 2] + R[..., 2, 0]) / s2_safe
+
+    s3 = xp.sqrt(xp.maximum(zero, one + R[..., 1, 1] - R[..., 0, 0] - R[..., 2, 2])) * two  # s = 4 * y
+    s3_safe = xp.where(s3 < eps, eps, s3)  # Avoid division by zero
+    w3 = (R[..., 0, 2] - R[..., 2, 0]) / s3_safe
+    x3 = (R[..., 0, 1] + R[..., 1, 0]) / s3_safe
+    y3 = quarter * s3
+    z3 = (R[..., 1, 2] + R[..., 2, 1]) / s3_safe
+
+    s4 = xp.sqrt(xp.maximum(zero, one + R[..., 2, 2] - R[..., 0, 0] - R[..., 1, 1])) * two  # s = 4 * z
+    s4_safe = xp.where(s4 < eps, eps, s4)  # Avoid division by zero
+    w4 = (R[..., 1, 0] - R[..., 0, 1]) / s4_safe
+    x4 = (R[..., 0, 2] + R[..., 2, 0]) / s4_safe
+    y4 = (R[..., 1, 2] + R[..., 2, 1]) / s4_safe
+    z4 = quarter * s4
+
+    cond1 = trace > 0
+    cond2 = (R[..., 0, 0] > R[..., 1, 1]) & (R[..., 0, 0] > R[..., 2, 2])
+    cond3 = R[..., 1, 1] > R[..., 2, 2]
+
+    w = xp.where(cond1, w1, xp.where(cond2, w2, xp.where(cond3, w3, w4)))
+    x = xp.where(cond1, x1, xp.where(cond2, x2, xp.where(cond3, x3, x4)))
+    y = xp.where(cond1, y1, xp.where(cond2, y2, xp.where(cond3, y3, y4)))
+    z = xp.where(cond1, z1, xp.where(cond2, z2, xp.where(cond3, z3, z4)))
+
+    q = xp.stack([w, x, y, z], axis=-1)
+
+    return canonicalize(q)

nanomanifold/SO3/distance.py

@@ -0,0 +1,52 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+
+def distance(q1: Float[Any, "... 4"], q2: Float[Any, "... 4"]) -> Float[Any, "..."]:
+    """Compute the angular distance between two quaternions representing SO(3) rotations.
+
+    The angular distance is the smallest angle needed to rotate from one orientation
+    to another, measured in radians. This is equivalent to the geodesic distance
+    on the SO(3) manifold.
+
+    Args:
+        q1: First quaternion in [w, x, y, z] format
+        q2: Second quaternion in [w, x, y, z] format
+
+    Returns:
+        Angular distance in radians, in range [0, π]
+    """
+    xp = get_namespace(q1)
+
+    norm1 = xp.sqrt(xp.sum(q1**2, axis=-1, keepdims=True))
+    norm2 = xp.sqrt(xp.sum(q2**2, axis=-1, keepdims=True))
+    q1_unit = q1 / norm1
+    q2_unit = q2 / norm2
+
+    # Flip sign of q2 when dot(q1, q2) < 0 so the relative rotation
+    # always measures the shorter geodesic on the double cover.
+    dot_keepdims = xp.sum(q1_unit * q2_unit, axis=-1, keepdims=True)
+    q2_unit = xp.where(dot_keepdims < 0, -q2_unit, q2_unit)
+
+    w1 = q1_unit[..., :1]
+    v1 = q1_unit[..., 1:]
+    w2 = q2_unit[..., :1]
+    v2 = q2_unit[..., 1:]
+
+    cross = xp.stack(
+        [
+            v1[..., 1] * v2[..., 2] - v1[..., 2] * v2[..., 1],
+            v1[..., 2] * v2[..., 0] - v1[..., 0] * v2[..., 2],
+            v1[..., 0] * v2[..., 1] - v1[..., 1] * v2[..., 0],
+        ],
+        axis=-1,
+    )
+
+    vec = w1 * v2 - w2 * v1 - cross
+    vec_norm = xp.sqrt(xp.sum(vec**2, axis=-1))
+    w = w1 * w2 + xp.sum(v1 * v2, axis=-1, keepdims=True)
+
+    return 2.0 * xp.atan2(vec_norm, w[..., 0])

nanomanifold/SO3/exp.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from .conversions.axis_angle import from_axis_angle
+
+
+def exp(tangent_vector: Float[Any, "... 3"]) -> Float[Any, "... 4"]:
+    """Compute the exponential map from so(3) tangent space to SO(3) manifold.
+
+    The exponential map takes a tangent vector in the Lie algebra so(3)
+    and returns the corresponding rotation quaternion. This is the inverse
+    operation of log().
+
+    The exponential map is mathematically equivalent to converting an axis-angle
+    representation to its corresponding quaternion.
+
+    Args:
+        tangent_vector: Tangent vector in so(3) (axis-angle representation)
+                       The magnitude is the rotation angle, direction is the rotation axis
+
+    Returns:
+        Quaternion in [w, x, y, z] format representing the rotation
+    """
+    return from_axis_angle(tangent_vector)

nanomanifold/SO3/hat.py

@@ -0,0 +1,36 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from ..common import get_namespace
+
+
+def hat(w: Float[Any, "... 3"]) -> Float[Any, "... 3 3"]:
+    """Map vector to skew-symmetric matrix (hat operator).
+
+    Args:
+        w: (..., 3) array representing tangent vectors in so(3)
+
+    Returns:
+        (..., 3, 3) skew-symmetric matrices
+    """
+    xp = get_namespace(w)
+
+    w1 = w[..., 0]
+    w2 = w[..., 1]
+    w3 = w[..., 2]
+
+    zero = w1 * 0
+
+    # Build skew-symmetric matrix:
+    # [[ 0, -w3,  w2],
+    #  [ w3,  0, -w1],
+    #  [-w2,  w1,  0]]
+
+    row1 = xp.stack([zero, -w3, w2], axis=-1)
+    row2 = xp.stack([w3, zero, -w1], axis=-1)
+    row3 = xp.stack([-w2, w1, zero], axis=-1)
+
+    result = xp.stack([row1, row2, row3], axis=-2)
+
+    return result

nanomanifold/SO3/inverse.py

@@ -0,0 +1,16 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+from .canonicalize import canonicalize
+
+
+def inverse(q: Float[Any, "... 4"]) -> Float[Any, "... 4"]:
+    xp = get_namespace(q)
+    q = canonicalize(q)
+    w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
+    q_inv = xp.stack([w, -x, -y, -z], axis=-1)
+
+    return canonicalize(q_inv)

nanomanifold/SO3/log.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from .conversions.axis_angle import to_axis_angle
+
+
+def log(q: Float[Any, "... 4"]) -> Float[Any, "... 3"]:
+    """Compute the logarithmic map of a quaternion on the SO(3) manifold.
+
+    The logarithmic map takes a rotation and returns the corresponding
+    tangent vector (axis-angle representation) in the Lie algebra so(3).
+    This is the inverse operation of exp().
+
+    The logarithmic map is mathematically equivalent to converting a quaternion
+    to its axis-angle representation.
+
+    Args:
+        q: Quaternion in [w, x, y, z] format representing a rotation
+
+    Returns:
+        Tangent vector in so(3) (axis-angle representation)
+        The magnitude is the rotation angle, direction is the rotation axis
+    """
+    return to_axis_angle(q)

nanomanifold/SO3/multiply.py

@@ -0,0 +1,40 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+from .canonicalize import canonicalize
+
+
+def multiply(q1: Float[Any, "... 4"], q2: Float[Any, "... 4"]) -> Float[Any, "... 4"]:
+    """Multiply two quaternions representing SO(3) rotations.
+
+    The multiplication order matches rotation matrix multiplication:
+    multiply(q1, q2) represents the same composition as to_matrix(q1) @ to_matrix(q2)
+
+    This means q2 is applied first, then q1.
+
+    Args:
+        q1: First quaternion in [w, x, y, z] format
+        q2: Second quaternion in [w, x, y, z] format
+
+    Returns:
+        Product quaternion representing the composed rotation
+    """
+    xp = get_namespace(q1)
+
+    q1 = canonicalize(q1)
+    q2 = canonicalize(q2)
+
+    w1, x1, y1, z1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3]
+    w2, x2, y2, z2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3]
+
+    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
+    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
+    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
+    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
+
+    result = xp.stack([w, x, y, z], axis=-1)
+
+    return canonicalize(result)

nanomanifold/SO3/rotate_points.py

@@ -0,0 +1,50 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+from .canonicalize import canonicalize
+
+
+def rotate_points(q: Float[Any, "... 4"], points: Float[Any, "... N 3"]) -> Float[Any, "... N 3"]:
+    """Rotate points using quaternion rotation.
+
+    Args:
+        q: Quaternion in [w, x, y, z] format of shape (..., 4)
+        points: Points to rotate of shape (..., N, 3)
+
+    Returns:
+        Rotated points of shape (..., N, 3)
+    """
+    xp = get_namespace(q)
+    q = canonicalize(q)
+
+    w, x, y, z = q[..., 0], q[..., 1], q[..., 2], q[..., 3]
+
+    w = w[..., None]
+    x = x[..., None]
+    y = y[..., None]
+    z = z[..., None]
+
+    px, py, pz = points[..., 0], points[..., 1], points[..., 2]
+
+    # Apply quaternion rotation using the formula:
+    # p' = q * p * q^(-1)
+    # This can be computed efficiently as:
+    # p' = p + 2 * cross(q_vec, cross(q_vec, p) + w * p)
+    # where q_vec = [x, y, z]
+
+    cross1_x = y * pz - z * py + w * px
+    cross1_y = z * px - x * pz + w * py
+    cross1_z = x * py - y * px + w * pz
+
+    cross2_x = y * cross1_z - z * cross1_y
+    cross2_y = z * cross1_x - x * cross1_z
+    cross2_z = x * cross1_y - y * cross1_x
+
+    result_x = px + 2 * cross2_x
+    result_y = py + 2 * cross2_y
+    result_z = pz + 2 * cross2_z
+
+    return xp.stack([result_x, result_y, result_z], axis=-1)

nanomanifold/SO3/slerp.py

@@ -0,0 +1,69 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+from .canonicalize import canonicalize
+
+
+def slerp(q1: Float[Any, "... 4"], q2: Float[Any, "... 4"], t: Float[Any, "... N"]) -> Float[Any, "... N 4"]:
+    """Spherical linear interpolation between two quaternions representing SO(3) rotations.
+
+    Performs geodesic interpolation on the SO(3) manifold, taking the shortest
+    path between two rotations. The interpolation parameter array t must have
+    values in [0, 1].
+
+    Args:
+        q1: Start quaternion in [w, x, y, z] format
+        q2: End quaternion in [w, x, y, z] format
+        t: Array of interpolation parameters in [0, 1]. Last dimension N represents
+           the number of interpolation points. For a single point, use shape [..., 1].
+           t values: 0 returns q1, 1 returns q2
+
+    Returns:
+        Interpolated quaternions with shape [..., N, 4] where N is the number of
+        interpolation points from the last dimension of t
+
+    Raises:
+        ValueError: If any t values are outside the range [0, 1]
+    """
+    xp = get_namespace(q1)
+
+    q1 = canonicalize(q1)
+    q2 = canonicalize(q2)
+
+    q1_expanded = xp.expand_dims(q1, axis=-2)
+    q2_expanded = xp.expand_dims(q2, axis=-2)
+
+    dot_product = xp.sum(q1_expanded * q2_expanded, axis=-1, keepdims=True)
+
+    q2_corrected = xp.where(dot_product < 0, -q2_expanded, q2_expanded)
+    dot_product = xp.where(dot_product < 0, -dot_product, dot_product)
+
+    dot_product = xp.clip(dot_product, 0.0, 1.0)
+
+    threshold = 0.9995
+
+    t_expanded = xp.expand_dims(t, axis=-1)
+
+    use_linear = dot_product > threshold
+
+    linear_result = (1.0 - t_expanded) * q1_expanded + t_expanded * q2_corrected
+    linear_norm = xp.sqrt(xp.sum(linear_result**2, axis=-1, keepdims=True))
+    linear_result = linear_result / linear_norm
+
+    omega = xp.acos(dot_product)
+    sin_omega = xp.sin(omega)
+
+    eps = 1e-8
+    sin_omega_safe = xp.where(xp.abs(sin_omega) < eps, eps, sin_omega)
+
+    weight1 = xp.sin((1.0 - t_expanded) * omega) / sin_omega_safe
+    weight2 = xp.sin(t_expanded * omega) / sin_omega_safe
+
+    spherical_result = weight1 * q1_expanded + weight2 * q2_corrected
+
+    result = xp.where(use_linear, linear_result, spherical_result)
+
+    return canonicalize(result)

nanomanifold/SO3/vee.py

@@ -0,0 +1,28 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from ..common import get_namespace
+
+
+def vee(W: Float[Any, "... 3 3"]) -> Float[Any, "... 3"]:
+    """Map skew-symmetric matrix to vector (vee operator).
+
+    Args:
+        W: (..., 3, 3) skew-symmetric matrices
+
+    Returns:
+        (..., 3) tangent vectors in so(3)
+    """
+    xp = get_namespace(W)
+
+    # Extract components from skew-symmetric matrix:
+    # [[ 0, -w3,  w2],
+    #  [ w3,  0, -w1],
+    #  [-w2,  w1,  0]]
+
+    w1 = W[..., 2, 1]
+    w2 = W[..., 0, 2]
+    w3 = W[..., 1, 0]
+
+    return xp.stack([w1, w2, w3], axis=-1)

nanomanifold/SO3/weighted_mean.py

@@ -0,0 +1,89 @@
+from typing import Any, Sequence
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+from .canonicalize import canonicalize
+
+
+def weighted_mean(quaternions: Sequence[Float[Any, "... 4"]], weights: Float[Any, "... N"]) -> Float[Any, "... 4"]:
+    """Compute the weighted mean of SO(3) rotations represented as quaternions.
+
+    This function implements the Riemannian mean on SO(3) by computing the weighted
+    average in quaternion space using the outer product method. The result is the
+    eigenvector corresponding to the largest eigenvalue of the weighted covariance matrix.
+
+    Args:
+        quaternions: Sequence of quaternions in [w, x, y, z] format. Each quaternion
+                    should have shape [..., 4] where the last dimension contains the
+                    quaternion components.
+        weights: Array of weights with shape [..., N] where N is the number of quaternions.
+                The weights are normalized internally.
+
+    Returns:
+        Weighted mean quaternion with shape [..., 4] in [w, x, y, z] format.
+        The result is canonicalized to ensure w >= 0.
+
+    Note:
+        This implementation follows the algorithm from:
+        "Averaging Quaternions" by F. Landis Markley et al.
+    """
+    xp = get_namespace(quaternions[0])
+    original_dtype = quaternions[0].dtype
+
+    quats = xp.stack([xp.asarray(q, dtype=original_dtype) for q in quaternions], axis=-2)
+    weights_array = xp.asarray(weights, dtype=original_dtype)
+
+    norms = xp.linalg.norm(quats, axis=-1, keepdims=True)
+    eps = xp.finfo(original_dtype).eps * 10
+    safe_norms = xp.where(norms < eps, eps, norms)
+    quats_normalized = quats / safe_norms
+
+    sign_mask = quats_normalized[..., 0:1] < 0
+    quats_canonical = xp.where(sign_mask, -quats_normalized, quats_normalized)
+
+    weights_normalized = weights_array / xp.sum(weights_array, axis=-1, keepdims=True)
+
+    weighted_quats = weights_normalized[..., :, None] * quats_canonical
+
+    M = xp.einsum("...nj,...nk->...jk", weighted_quats, quats_canonical)
+
+    if original_dtype == xp.float16:
+        M_compute = xp.asarray(M, dtype=xp.float32)
+        eigenvalues, eigenvectors = xp.linalg.eigh(M_compute)
+        eigenvectors = xp.asarray(eigenvectors, dtype=original_dtype)
+    else:
+        eigenvalues, eigenvectors = xp.linalg.eigh(M)
+
+    avg_quat = eigenvectors[..., :, -1]
+
+    avg_quat = avg_quat / xp.linalg.norm(avg_quat, axis=-1, keepdims=True)
+
+    return canonicalize(avg_quat)
+
+
+def mean(quaternions: Sequence[Float[Any, "... 4"]]) -> Float[Any, "... 4"]:
+    """Compute the mean of SO(3) rotations represented as quaternions.
+
+    This is equivalent to weighted_mean with uniform weights.
+
+    Args:
+        quaternions: Sequence of quaternions in [w, x, y, z] format. Each quaternion
+                    should have shape [..., 4] where the last dimension contains the
+                    quaternion components.
+
+    Returns:
+        Mean quaternion with shape [..., 4] in [w, x, y, z] format.
+        The result is canonicalized to ensure w >= 0.
+    """
+    if len(quaternions) == 0:
+        raise ValueError("Cannot compute mean of empty quaternion sequence")
+
+    xp = get_namespace(quaternions[0])
+
+    batch_shape = quaternions[0].shape[:-1]
+    num_quats = len(quaternions)
+    weights = xp.ones(batch_shape + (num_quats,), dtype=quaternions[0].dtype)
+
+    return weighted_mean(quaternions, weights)

nanomanifold/__init__.py

@@ -0,0 +1,3 @@
+from . import SE3, SO3
+
+__all__ = ["SO3", "SE3"]