nanomanifold 0.1.1__py3-none-any.whl

nanomanifold/SE3/__init__.py

@@ -0,0 +1,21 @@
+from .canonicalize import canonicalize
+from .conversions.matrix import from_matrix, to_matrix
+from .conversions.rt import from_rt, to_rt
+from .exp import exp
+from .inverse import inverse
+from .log import log
+from .multiply import multiply
+from .transform_points import transform_points
+
+__all__ = [
+    "from_matrix",
+    "to_matrix",
+    "from_rt",
+    "to_rt",
+    "canonicalize",
+    "multiply",
+    "inverse",
+    "transform_points",
+    "log",
+    "exp",
+]

nanomanifold/SE3/canonicalize.py

@@ -0,0 +1,25 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3.canonicalize import canonicalize as canonicalize_quat
+
+
+def canonicalize(se3: Float[Any, "... 7"]) -> Float[Any, "... 7"]:
+    """Canonicalize SE(3) representation by canonicalizing the quaternion part.
+
+    Args:
+        se3: SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+
+    Returns:
+        Canonicalized SE(3) representation with quaternion w >= 0
+    """
+    xp = get_namespace(se3)
+
+    quat = se3[..., :4]
+    translation = se3[..., 4:7]
+
+    quat_canonical = canonicalize_quat(quat)
+
+    return xp.concatenate([quat_canonical, translation], axis=-1)

nanomanifold/SE3/conversions/__init__.py

@@ -0,0 +1,9 @@
+from .matrix import from_matrix, to_matrix
+from .rt import from_rt, to_rt
+
+__all__ = [
+    "from_matrix",
+    "to_matrix",
+    "from_rt",
+    "to_rt",
+]

nanomanifold/SE3/conversions/matrix.py

@@ -0,0 +1,57 @@
+"""Matrix conversions for SE(3) transformations."""
+
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold import SO3
+from nanomanifold.common import get_namespace
+
+from ..canonicalize import canonicalize
+
+
+def to_matrix(se3: Float[Any, "... 7"]) -> Float[Any, "... 4 4"]:
+    """Convert SE(3) representation to 4x4 transformation matrix.
+
+    Args:
+        se3: SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+
+    Returns:
+        4x4 transformation matrix (..., 4, 4)
+    """
+    xp = get_namespace(se3)
+
+    quat = se3[..., :4]
+    translation = se3[..., 4:7]
+
+    R = SO3.to_matrix(quat)
+
+    translation_column = translation[..., None]
+    top_block = xp.concatenate([R, translation_column], axis=-1)
+
+    zeros = xp.zeros(top_block.shape[:-2] + (1, 3), dtype=se3.dtype)
+    ones = xp.ones(top_block.shape[:-2] + (1, 1), dtype=se3.dtype)
+    bottom_row = xp.concatenate([zeros, ones], axis=-1)
+
+    return xp.concatenate([top_block, bottom_row], axis=-2)
+
+
+def from_matrix(matrix: Float[Any, "... 4 4"]) -> Float[Any, "... 7"]:
+    """Convert 4x4 transformation matrix to SE(3) representation.
+
+    Args:
+        matrix: 4x4 transformation matrix (..., 4, 4)
+
+    Returns:
+        SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+    """
+    xp = get_namespace(matrix)
+
+    R = matrix[..., :3, :3]
+
+    quat = SO3.from_matrix(R)
+
+    translation = matrix[..., :3, 3]
+
+    se3 = xp.concatenate([quat, translation], axis=-1)
+    return canonicalize(se3)

nanomanifold/SE3/conversions/rt.py

@@ -0,0 +1,30 @@
+from nanomanifold.common import get_namespace
+
+
+def from_rt(quat, translation):
+    """Create SE(3) representation from rotation quaternion and translation.
+
+    Args:
+        quat: Rotation quaternion (..., 4) as [w, x, y, z]
+        translation: Translation vector (..., 3)
+
+    Returns:
+        SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+    """
+    xp = get_namespace(quat)
+    return xp.concatenate([quat, translation], axis=-1)
+
+
+def to_rt(se3):
+    """Extract rotation quaternion and translation from SE(3) representation.
+
+    Args:
+        se3: SE(3) representation (..., 7) as [w, x, y, z, tx, ty, tz]
+
+    Returns:
+        quat: Rotation quaternion (..., 4) as [w, x, y, z]
+        translation: Translation vector (..., 3)
+    """
+    quat = se3[..., :4]
+    translation = se3[..., 4:7]
+    return quat, translation

nanomanifold/SE3/exp.py

@@ -0,0 +1,73 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3 import exp as so3_exp
+from nanomanifold.SO3 import hat
+
+
+def exp(tangent_vector: Float[Any, "... 6"]) -> Float[Any, "... 7"]:
+    """Compute the exponential map from se(3) tangent space to SE(3) manifold.
+
+    The exponential map takes a tangent vector in the Lie algebra se(3)
+    and returns the corresponding SE(3) transformation. This is the inverse
+    operation of log().
+
+    The se(3) exponential map takes a 6-vector [ω, ρ] where:
+    - ω ∈ ℝ³ is the angular velocity (rotation part)
+    - ρ ∈ ℝ³ is the translational velocity
+
+    The formula involves:
+    - R = exp_SO3(ω) for the rotation quaternion
+    - t = V * ρ where V is the left Jacobian matrix
+
+    Args:
+        tangent_vector: Tangent vector in se(3) as [ω, ρ] of shape (..., 6)
+
+    Returns:
+        SE(3) transformation in [w, x, y, z, tx, ty, tz] format of shape (..., 7)
+    """
+    xp = get_namespace(tangent_vector)
+
+    omega = tangent_vector[..., :3]
+    rho = tangent_vector[..., 3:6]
+
+    q = so3_exp(omega)
+    omega_norm = xp.linalg.norm(omega, axis=-1, keepdims=True)
+
+    eps = xp.finfo(omega.dtype).eps
+    small_angle_threshold = xp.asarray(max(1e-6, float(eps)), dtype=omega.dtype)
+    small_angle_mask = omega_norm < small_angle_threshold
+
+    omega_cross = hat(omega)
+    omega_cross_sq = xp.matmul(omega_cross, omega_cross)
+
+    identity = xp.eye(3, dtype=omega.dtype)
+    identity = xp.broadcast_to(identity, omega.shape[:-1] + (3, 3))
+
+    V_small = identity + 0.5 * omega_cross + (1.0 / 12.0) * omega_cross_sq
+
+    cos_norm = xp.cos(omega_norm)
+    sin_norm = xp.sin(omega_norm)
+
+    safe_norm = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm)
+    safe_norm_sq = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm**2)
+    safe_norm_cub = safe_norm_sq * safe_norm
+
+    A = (1.0 - cos_norm) / safe_norm_sq
+    B = (safe_norm - sin_norm) / safe_norm_cub
+    A = xp.reshape(A, A.shape[:-1])[..., None, None]
+    B = xp.reshape(B, B.shape[:-1])[..., None, None]
+
+    V_large = identity + A * omega_cross + B * omega_cross_sq
+
+    mask = xp.reshape(small_angle_mask, small_angle_mask.shape[:-1])[..., None, None]
+    V = xp.where(mask, V_small, V_large)
+    V = xp.reshape(V, omega.shape[:-1] + (3, 3))
+
+    t = xp.matmul(V, rho[..., None])[..., 0]
+
+    se3 = xp.concatenate([q, t], axis=-1)
+
+    return se3

nanomanifold/SE3/inverse.py

@@ -0,0 +1,37 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3 import inverse as so3_inverse
+from nanomanifold.SO3 import rotate_points
+
+from .canonicalize import canonicalize
+
+
+def inverse(se3: Float[Any, "... 7"]) -> Float[Any, "... 7"]:
+    """Compute the inverse of SE(3) transformations.
+
+    For an SE(3) transformation T = [R, t] represented as [q, t],
+    the inverse is T^(-1) = [R^T, -R^T * t] represented as [q^(-1), -q^(-1) * t].
+
+    Args:
+        se3: SE(3) transformation in [w, x, y, z, tx, ty, tz] format
+
+    Returns:
+        Inverse SE(3) transformation
+    """
+    xp = get_namespace(se3)
+
+    se3 = canonicalize(se3)
+
+    q = se3[..., :4]
+    t = se3[..., 4:7]
+
+    q_inv = so3_inverse(q)
+
+    t_inv = -rotate_points(q_inv, t[..., None, :]).squeeze(-2)
+
+    result = xp.concatenate([q_inv, t_inv], axis=-1)
+
+    return canonicalize(result)

nanomanifold/SE3/log.py

@@ -0,0 +1,75 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3 import hat
+from nanomanifold.SO3 import log as so3_log
+
+from .canonicalize import canonicalize
+
+
+def log(se3: Float[Any, "... 7"]) -> Float[Any, "... 6"]:
+    """Compute the logarithmic map of SE(3) to its Lie algebra se(3).
+
+    The logarithmic map takes an SE(3) transformation and returns the corresponding
+    tangent vector in the Lie algebra se(3). This is the inverse operation of exp().
+
+    The SE(3) logarithmic map computes a 6-vector [ω, ρ] where:
+    - ω ∈ ℝ³ is the angular velocity (rotation part, same as SO(3) log)
+    - ρ ∈ ℝ³ is the transformed translation part
+
+    The formula involves:
+    - ω = log_SO3(R) where R is the rotation quaternion
+    - ρ = V^(-1) * t where V is the left Jacobian inverse and t is the translation
+
+    Args:
+        se3: SE(3) transformation in [w, x, y, z, tx, ty, tz] format of shape (..., 7)
+
+    Returns:
+        Tangent vector in se(3) as [ω, ρ] of shape (..., 6)
+    """
+    xp = get_namespace(se3)
+    se3 = canonicalize(se3)
+
+    q = se3[..., :4]
+    t = se3[..., 4:7]
+
+    omega = so3_log(q)
+    omega_norm = xp.linalg.norm(omega, axis=-1, keepdims=True)
+
+    eps = xp.finfo(omega.dtype).eps
+    small_angle_threshold = xp.asarray(max(1e-6, float(eps)), dtype=omega.dtype)
+    small_angle_mask = omega_norm < small_angle_threshold
+
+    omega_cross = hat(omega)
+    omega_cross_sq = xp.matmul(omega_cross, omega_cross)
+
+    identity = xp.eye(3, dtype=omega.dtype)
+    identity = xp.broadcast_to(identity, omega.shape[:-1] + (3, 3))
+
+    V_inv_small = identity - 0.5 * omega_cross + (1.0 / 12.0) * omega_cross_sq
+
+    half_norm = omega_norm / 2.0
+    cos_half = xp.cos(half_norm)
+    sin_half = xp.sin(half_norm)
+
+    safe_sin_half = xp.where(small_angle_mask, xp.ones_like(sin_half), sin_half)
+    cot_half = cos_half / safe_sin_half
+
+    safe_norm = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm)
+    safe_norm_sq = xp.where(small_angle_mask, xp.ones_like(omega_norm), omega_norm**2)
+    B = (1.0 - 0.5 * safe_norm * cot_half) / safe_norm_sq
+    B = xp.reshape(B, B.shape[:-1])[..., None, None]
+
+    V_inv_large = identity - 0.5 * omega_cross + B * omega_cross_sq
+
+    mask = xp.reshape(small_angle_mask, small_angle_mask.shape[:-1])[..., None, None]
+    V_inv = xp.where(mask, V_inv_small, V_inv_large)
+    V_inv = xp.reshape(V_inv, omega.shape[:-1] + (3, 3))
+
+    rho = xp.matmul(V_inv, t[..., None])[..., 0]
+
+    tangent = xp.concatenate([omega, rho], axis=-1)
+
+    return tangent

nanomanifold/SE3/multiply.py

@@ -0,0 +1,48 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+from nanomanifold.SO3 import multiply as so3_multiply
+from nanomanifold.SO3 import rotate_points
+
+from .canonicalize import canonicalize
+
+
+def multiply(se3_1: Float[Any, "... 7"], se3_2: Float[Any, "... 7"]) -> Float[Any, "... 7"]:
+    """Multiply two SE(3) transformations.
+
+    The multiplication order matches transformation matrix multiplication:
+    multiply(se3_1, se3_2) represents the same composition as to_matrix(se3_1) @ to_matrix(se3_2)
+
+    This means se3_2 is applied first, then se3_1.
+
+    For SE(3) transformations [q1, t1] and [q2, t2], the result is:
+    - Quaternion: q1 * q2 (quaternion multiplication)
+    - Translation: R1 * t2 + t1 (where R1 rotates by q1)
+
+    Args:
+        se3_1: First SE(3) transformation in [w, x, y, z, tx, ty, tz] format
+        se3_2: Second SE(3) transformation in [w, x, y, z, tx, ty, tz] format
+
+    Returns:
+        Product SE(3) transformation representing the composed transformation
+    """
+    xp = get_namespace(se3_1)
+
+    se3_1 = canonicalize(se3_1)
+    se3_2 = canonicalize(se3_2)
+
+    q1 = se3_1[..., :4]
+    t1 = se3_1[..., 4:7]
+    q2 = se3_2[..., :4]
+    t2 = se3_2[..., 4:7]
+
+    q_result = so3_multiply(q1, q2)
+
+    t2_rotated = rotate_points(q1, t2[..., None, :]).squeeze(-2)
+    t_result = t2_rotated + t1
+
+    result = xp.concatenate([q_result, t_result], axis=-1)
+
+    return canonicalize(result)

nanomanifold/SE3/transform_points.py

@@ -0,0 +1,34 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.SO3 import rotate_points
+
+from .canonicalize import canonicalize
+
+
+def transform_points(se3: Float[Any, "... 7"], points: Float[Any, "... N 3"]) -> Float[Any, "... N 3"]:
+    """Transform 3D points using SE(3) transformation.
+
+    Applies both rotation and translation: p' = R * p + t
+    where SE(3) = [q, t] with q being the quaternion and t being the translation.
+
+    Args:
+        se3: SE(3) transformation in [w, x, y, z, tx, ty, tz] format of shape (..., 7)
+        points: Points to transform of shape (..., N, 3)
+
+    Returns:
+        Transformed points of shape (..., N, 3)
+    """
+    se3 = canonicalize(se3)
+
+    q = se3[..., :4]
+    t = se3[..., 4:7]
+
+    rotated_points = rotate_points(q, points)
+
+    t_expanded = t[..., None, :]
+
+    transformed_points = rotated_points + t_expanded
+
+    return transformed_points

nanomanifold/SO3/__init__.py

@@ -0,0 +1,35 @@
+from .canonicalize import canonicalize
+from .conversions.axis_angle import from_axis_angle, to_axis_angle
+from .conversions.euler import from_euler, to_euler
+from .conversions.matrix import from_matrix, to_matrix
+from .distance import distance
+from .exp import exp
+from .hat import hat
+from .inverse import inverse
+from .log import log
+from .multiply import multiply
+from .rotate_points import rotate_points
+from .slerp import slerp
+from .vee import vee
+from .weighted_mean import mean, weighted_mean
+
+__all__ = [
+    "to_axis_angle",
+    "from_axis_angle",
+    "to_euler",
+    "from_euler",
+    "to_matrix",
+    "from_matrix",
+    "canonicalize",
+    "rotate_points",
+    "inverse",
+    "multiply",
+    "distance",
+    "log",
+    "exp",
+    "hat",
+    "vee",
+    "slerp",
+    "weighted_mean",
+    "mean",
+]

nanomanifold/SO3/canonicalize.py

@@ -0,0 +1,17 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+
+def canonicalize(q: Float[Any, "... 4"]) -> Float[Any, "... 4"]:
+    xp = get_namespace(q)
+
+    norm = xp.sqrt(xp.sum(q**2, axis=-1, keepdims=True))
+    q_normalized = q / norm
+
+    mask = q_normalized[..., 0:1] < 0
+    q_canonical = xp.where(mask, -q_normalized, q_normalized)
+
+    return q_canonical

nanomanifold/SO3/conversions/__init__.py

@@ -0,0 +1,14 @@
+"""SO(3) conversion functions."""
+
+from .axis_angle import from_axis_angle, to_axis_angle
+from .euler import from_euler, to_euler
+from .matrix import from_matrix, to_matrix
+
+__all__ = [
+    "to_axis_angle",
+    "from_axis_angle",
+    "to_euler",
+    "from_euler",
+    "to_matrix",
+    "from_matrix",
+]

nanomanifold/SO3/conversions/axis_angle.py

@@ -0,0 +1,63 @@
+from typing import Any
+
+from jaxtyping import Float
+
+from nanomanifold.common import get_namespace
+
+
+def to_axis_angle(q: Float[Any, "... 4"]) -> Float[Any, "... 3"]:
+    xp = get_namespace(q)
+
+    w = q[..., 0]
+    xyz = q[..., 1:4]
+
+    norm_xyz = xp.linalg.norm(xyz, axis=-1, keepdims=True)
+
+    # For small rotations, use the direct approximation: axis_angle ≈ 2*xyz
+    # This avoids numerical issues with acos when w ≈ 1
+    small_angle_threshold = 1e-6
+    small_angle_mask = norm_xyz < small_angle_threshold
+
+    axis_angle_small = 2.0 * xyz
+
+    w_clipped = xp.clip(xp.abs(w), 0.0, 1.0)  # Use abs(w) to handle canonicalization
+    angle = 2 * xp.acos(w_clipped)
+
+    safe_norm = xp.where(norm_xyz < 1e-15, 1.0, norm_xyz)  # Avoid true division by zero
+    axis = xyz / safe_norm
+
+    axis_angle_large = angle[..., None] * axis
+
+    axis_angle = xp.where(small_angle_mask, axis_angle_small, axis_angle_large)
+
+    return axis_angle
+
+
+def from_axis_angle(axis_angle: Float[Any, "... 3"]) -> Float[Any, "... 4"]:
+    xp = get_namespace(axis_angle)
+
+    angle = xp.linalg.norm(axis_angle, axis=-1)
+
+    small_angle_mask = angle < 1e-12
+    safe_angle = xp.where(small_angle_mask, 1.0, angle)
+    axis = axis_angle / safe_angle[..., None]
+
+    half_angle = angle / 2
+
+    # For small angles, use Taylor series: cos(x) ≈ 1 - x²/2, sin(x) ≈ x
+    cos_half = xp.cos(half_angle)
+    sin_half = xp.sin(half_angle)
+
+    w = cos_half[..., None]
+
+    xyz_normal = sin_half[..., None] * axis
+    xyz_small = axis_angle / 2.0
+
+    xyz = xp.where(small_angle_mask[..., None], xyz_small, xyz_normal)
+
+    q = xp.concat([w, xyz], axis=-1)
+
+    mask = q[..., 0:1] < 0
+    q = xp.where(mask, -q, q)
+
+    return q