src/fpgm/geometry/transforms.py

"""SE3 rigid-transform helpers, all on float64.

Every routine here is deliberately dependency-free numpy: no scipy.spatial.transform,
so the exact rotation-vector convention (Rodrigues, right-handed, magnitude = angle in
radians) is pinned down in one place instead of inherited from a library default.
"""

from __future__ import annotations

import numpy as np

from fpgm.types import GeometryError


def rotvec_to_matrix(rotvec: np.ndarray) -> np.ndarray:
    """Convert an axis-angle rotation vector ``(3,)`` to a ``(3, 3)`` rotation matrix.

    Uses Rodrigues' formula directly rather than a library call so the exact
    convention (right-handed, magnitude = angle in radians) is explicit here.
    """
    rotvec = np.asarray(rotvec, dtype=np.float64).reshape(3)
    theta = np.linalg.norm(rotvec)
    if theta < 1e-12:
        return np.eye(3, dtype=np.float64)
    axis = rotvec / theta
    kx, ky, kz = axis
    k = np.array([[0.0, -kz, ky], [kz, 0.0, -kx], [-ky, kx, 0.0]], dtype=np.float64)
    return np.eye(3) + np.sin(theta) * k + (1.0 - np.cos(theta)) * (k @ k)


def matrix_to_rotvec(rot: np.ndarray) -> np.ndarray:
    """Convert a ``(3, 3)`` rotation matrix to an axis-angle rotation vector ``(3,)``."""
    rot = np.asarray(rot, dtype=np.float64)
    cos_theta = np.clip((np.trace(rot) - 1.0) / 2.0, -1.0, 1.0)
    theta = np.arccos(cos_theta)
    if theta < 1e-12:
        return np.zeros(3, dtype=np.float64)
    if np.pi - theta < 1e-6:
        # Near-pi rotations: off-diagonal antisymmetric part vanishes, recover the
        # axis from the symmetric part instead.
        sym = (rot + np.eye(3)) / 2.0
        axis = np.sqrt(np.clip(np.diag(sym), 0.0, None))
        signs = np.sign(
            [rot[2, 1] - rot[1, 2], rot[0, 2] - rot[2, 0], rot[1, 0] - rot[0, 1]]
        )
        axis = axis * np.where(signs == 0, 1.0, signs)
        axis = axis / (np.linalg.norm(axis) + 1e-15)
        return axis * theta
    axis = (
        np.array([rot[2, 1] - rot[1, 2], rot[0, 2] - rot[2, 0], rot[1, 0] - rot[0, 1]])
        / (2.0 * np.sin(theta))
    )
    return axis * theta


def pose6_to_matrix(position: np.ndarray, rotvec: np.ndarray) -> np.ndarray:
    """Build a ``(4, 4)`` SE3 matrix from a ``(3,)`` position and ``(3,)`` rotation vector."""
    mat = np.eye(4, dtype=np.float64)
    mat[:3, :3] = rotvec_to_matrix(rotvec)
    mat[:3, 3] = np.asarray(position, dtype=np.float64).reshape(3)
    return mat


def invert_se3(transform: np.ndarray) -> np.ndarray:
    """Invert a ``(4, 4)`` SE3 matrix using the ``[R^T, -R^T t]`` identity.

    Deliberately not ``np.linalg.inv``: that solves a general 4x4 linear system and
    can drift off the SE3 manifold under float error, whereas the closed-form
    identity is exact given an orthonormal rotation block.
    """
    transform = np.asarray(transform, dtype=np.float64)
    rot = transform[:3, :3]
    t = transform[:3, 3]
    out = np.eye(4, dtype=np.float64)
    out[:3, :3] = rot.T
    out[:3, 3] = -rot.T @ t
    return out


def transform_points(transform: np.ndarray, points: np.ndarray) -> np.ndarray:
    """Apply a ``(4, 4)`` SE3 transform to points of shape ``(..., 3)``.

    Vectorised over arbitrary leading dimensions; no Python loop over points.
    """
    transform = np.asarray(transform, dtype=np.float64)
    points = np.asarray(points, dtype=np.float64)
    rot = transform[:3, :3]
    t = transform[:3, 3]
    return points @ rot.T + t


def pose_error(current: np.ndarray, target: np.ndarray) -> np.ndarray:
    """SE3 error ``(6,)`` = ``[position_error, base-frame rotation vector]``.

    ``e[:3]`` is ``target - current`` translation; ``e[3:]`` is
    ``matrix_to_rotvec(R_target @ R_current.T)`` -- the *base-frame* (spatial)
    rotation vector that rotates ``current``'s orientation onto ``target``'s,
    expressed in the frame both poses share (not ``current``'s own body frame).
    This is the convention the geometric Jacobian's angular rows (``J[3:] =
    z.T``, base-frame joint axes) consume; using the body-frame log instead
    would silently converge IK to the wrong attitude.

    Deliberately not the exact SE(3) logarithm (no left-Jacobian correction
    coupling rotation back into the translation error): that correction only
    perturbs the descent direction, not the fixed point the solver converges
    to, so it is skipped for one less dependency per damped-least-squares step.
    """
    current = np.asarray(current, dtype=np.float64)
    target = np.asarray(target, dtype=np.float64)
    err = np.empty(6, dtype=np.float64)
    err[:3] = target[:3, 3] - current[:3, 3]
    err[3:] = matrix_to_rotvec(target[:3, :3] @ current[:3, :3].T)
    return err


def rpy_to_matrix(rpy: np.ndarray) -> np.ndarray:
    """Convert URDF ``rpy="roll pitch yaw"`` to a ``(3, 3)`` rotation matrix.

    Pins the URDF/ROS convention explicitly (extrinsic XYZ, i.e. intrinsic
    ZYX): ``R = Rz(yaw) @ Ry(pitch) @ Rx(roll)``. No ``scipy.spatial.transform``
    -- the exact axis order is a frequent source of silent bugs when inherited
    from a library default instead of stated in code.
    """
    roll, pitch, yaw = (float(v) for v in np.asarray(rpy, dtype=np.float64).reshape(3))
    cr, sr = np.cos(roll), np.sin(roll)
    cp, sp = np.cos(pitch), np.sin(pitch)
    cy, sy = np.cos(yaw), np.sin(yaw)
    rz = np.array([[cy, -sy, 0.0], [sy, cy, 0.0], [0.0, 0.0, 1.0]])
    ry = np.array([[cp, 0.0, sp], [0.0, 1.0, 0.0], [-sp, 0.0, cp]])
    rx = np.array([[1.0, 0.0, 0.0], [0.0, cr, -sr], [0.0, sr, cr]])
    return rz @ ry @ rx


def matrix_to_rpy(rot: np.ndarray) -> np.ndarray:
    """Inverse of :func:`rpy_to_matrix`: ``(3, 3)`` rotation matrix -> ``(roll, pitch, yaw)``.

    Standard ZYX-Euler extraction, with the gimbal-lock case (``|R[2,0]| ~ 1``,
    pitch = +-pi/2) falling back to zero roll -- roll and yaw become degenerate
    (only their difference/sum is observable) at that singularity, matching the
    usual convention.
    """
    rot = np.asarray(rot, dtype=np.float64)
    sp = np.clip(-rot[2, 0], -1.0, 1.0)
    pitch = np.arcsin(sp)
    if np.abs(rot[2, 0]) > 1.0 - 1e-9:
        roll = 0.0
        yaw = np.arctan2(-rot[0, 1], rot[1, 1])
    else:
        roll = np.arctan2(rot[2, 1], rot[2, 2])
        yaw = np.arctan2(rot[1, 0], rot[0, 0])
    return np.array([roll, pitch, yaw], dtype=np.float64)


def assert_valid_se3(transform: np.ndarray, atol: float = 1e-4) -> None:
    """Raise :class:`~fpgm.types.GeometryError` unless ``transform`` is a valid SE3 matrix.

    Checks the bottom row is ``[0, 0, 0, 1]`` and the rotation block is orthonormal
    with determinant +1. This is the cheap early check that catches an accidental
    cam2world/world2cam swap before it silently corrupts every downstream position.
    """
    transform = np.asarray(transform, dtype=np.float64)
    if transform.shape != (4, 4):
        raise GeometryError(f"expected a (4, 4) SE3 matrix, got shape {transform.shape}")
    bottom = transform[3, :]
    if not np.allclose(bottom, [0.0, 0.0, 0.0, 1.0], atol=atol):
        raise GeometryError(f"bottom row of SE3 matrix is not [0,0,0,1]: {bottom}")
    rot = transform[:3, :3]
    should_be_identity = rot @ rot.T
    if not np.allclose(should_be_identity, np.eye(3), atol=atol):
        raise GeometryError("rotation block is not orthonormal (R @ R.T != I)")
    det = np.linalg.det(rot)
    if not np.isclose(det, 1.0, atol=atol):
        raise GeometryError(f"rotation block determinant is {det}, expected +1")