calibration/imu_lidar/geometry.py

"""SE(3)/SO(3) utilities for LiDAR–IMU calibration."""

from __future__ import annotations

import math

import numpy as np


def skew(vector: np.ndarray) -> np.ndarray:
    """Return the skew-symmetric matrix such that ``skew(v) @ w == v x w``."""

    x, y, z = np.asarray(vector, dtype=float).reshape(3)
    return np.array([[0.0, -z, y], [z, 0.0, -x], [-y, x, 0.0]], dtype=float)


def so3_exp(rotation_vector: np.ndarray) -> np.ndarray:
    """Map a rotation vector in radians onto SO(3)."""

    vector = np.asarray(rotation_vector, dtype=float).reshape(3)
    angle = float(np.linalg.norm(vector))
    if angle < 1e-12:
        return np.eye(3) + skew(vector)
    axis_cross = skew(vector / angle)
    return np.eye(3) + math.sin(angle) * axis_cross + (1.0 - math.cos(angle)) * axis_cross @ axis_cross


def so3_log(rotation: np.ndarray) -> np.ndarray:
    """Map an SO(3) matrix to a rotation vector in radians."""

    rotation = np.asarray(rotation, dtype=float).reshape(3, 3)
    cos_angle = float(np.clip((np.trace(rotation) - 1.0) * 0.5, -1.0, 1.0))
    angle = math.acos(cos_angle)
    if angle < 1e-12:
        return 0.5 * np.array(
            [
                rotation[2, 1] - rotation[1, 2],
                rotation[0, 2] - rotation[2, 0],
                rotation[1, 0] - rotation[0, 1],
            ],
            dtype=float,
        )
    if abs(angle - math.pi) < 1e-6:
        # Near 180°: use eigenvector of the +1 eigenvalue.
        eigvals, eigvecs = np.linalg.eigh(0.5 * (rotation + rotation.T))
        axis = eigvecs[:, int(np.argmax(eigvals))]
        return axis * angle
    return (
        0.5
        * angle
        / math.sin(angle)
        * np.array(
            [
                rotation[2, 1] - rotation[1, 2],
                rotation[0, 2] - rotation[2, 0],
                rotation[1, 0] - rotation[0, 1],
            ],
            dtype=float,
        )
    )


def rotation_angle_deg(rotation: np.ndarray) -> float:
    """Return the rotation angle in degrees."""

    return float(np.degrees(np.linalg.norm(so3_log(rotation))))


def inverse_transform(transform: np.ndarray) -> np.ndarray:
    """Return the inverse of a rigid 4x4 transform."""

    transform = np.asarray(transform, dtype=float)
    if transform.shape != (4, 4):
        raise ValueError("a rigid transform must have shape (4, 4)")
    result = np.eye(4)
    result[:3, :3] = transform[:3, :3].T
    result[:3, 3] = -result[:3, :3] @ transform[:3, 3]
    return result


def make_transform(translation_m: np.ndarray, rotation: np.ndarray) -> np.ndarray:
    """Build ``T_A_B`` from its translation and rotation components."""

    translation_m = np.asarray(translation_m, dtype=float).reshape(3)
    rotation = np.asarray(rotation, dtype=float)
    if rotation.shape != (3, 3):
        raise ValueError("a rotation matrix must have shape (3, 3)")
    result = np.eye(4)
    result[:3, :3] = rotation
    result[:3, 3] = translation_m
    return result


def transform_points(points: np.ndarray, transform: np.ndarray) -> np.ndarray:
    """Apply ``T_A_B`` to an ``(N, 3)`` point array expressed in frame B."""

    points = np.asarray(points, dtype=float)
    if points.ndim != 2 or points.shape[1] != 3:
        raise ValueError("points must have shape (N, 3)")
    return points @ transform[:3, :3].T + transform[:3, 3]


def orthonormalize_rotation(rotation: np.ndarray) -> np.ndarray:
    """Project a near-rotation matrix onto SO(3)."""

    u, _, vt = np.linalg.svd(np.asarray(rotation, dtype=float).reshape(3, 3))
    result = u @ vt
    if np.linalg.det(result) < 0:
        u[:, -1] *= -1
        result = u @ vt
    return result


def integrate_gyro_rotation(
    times_s: np.ndarray,
    gyro_rad_s: np.ndarray,
    t0: float,
    t1: float,
    bias_rad_s: np.ndarray | None = None,
) -> np.ndarray:
    """Integrate gyroscope samples on ``[t0, t1]`` and return ``R(t0<-t1)`` wait.

    Returns ``R_i_j`` that maps vectors from the IMU frame at ``t1`` into the
    IMU frame at ``t0`` using right-invariant discrete integration:

        R <- R @ Exp(omega * dt)
    """

    times_s = np.asarray(times_s, dtype=float).reshape(-1)
    gyro_rad_s = np.asarray(gyro_rad_s, dtype=float).reshape(-1, 3)
    if times_s.size < 2:
        return np.eye(3)
    bias = np.zeros(3) if bias_rad_s is None else np.asarray(bias_rad_s, dtype=float).reshape(3)

    if t1 < t0:
        raise ValueError("t1 must be >= t0")

    # Include one sample before t0 and after t1 when possible for interpolation.
    left = int(np.searchsorted(times_s, t0, side="left") - 1)
    right = int(np.searchsorted(times_s, t1, side="right"))
    left = max(left, 0)
    right = min(right, times_s.size - 1)
    if right <= left:
        return np.eye(3)

    rotation = np.eye(3)
    for index in range(left, right):
        t_a = float(times_s[index])
        t_b = float(times_s[index + 1])
        if t_b <= t0 or t_a >= t1:
            continue
        seg0 = max(t_a, t0)
        seg1 = min(t_b, t1)
        dt = seg1 - seg0
        if dt <= 0:
            continue
        omega = 0.5 * (gyro_rad_s[index] + gyro_rad_s[index + 1]) - bias
        rotation = rotation @ so3_exp(omega * dt)
    return orthonormalize_rotation(rotation)


def rotation_matrix_to_quaternion_xyzw(rotation: np.ndarray) -> np.ndarray:
    """Convert SO(3) to quaternion ``[x, y, z, w]``."""

    rotation = orthonormalize_rotation(rotation)
    trace = float(np.trace(rotation))
    if trace > 0:
        s = math.sqrt(trace + 1.0) * 2.0
        w = 0.25 * s
        x = (rotation[2, 1] - rotation[1, 2]) / s
        y = (rotation[0, 2] - rotation[2, 0]) / s
        z = (rotation[1, 0] - rotation[0, 1]) / s
    elif rotation[0, 0] > rotation[1, 1] and rotation[0, 0] > rotation[2, 2]:
        s = math.sqrt(1.0 + rotation[0, 0] - rotation[1, 1] - rotation[2, 2]) * 2.0
        w = (rotation[2, 1] - rotation[1, 2]) / s
        x = 0.25 * s
        y = (rotation[0, 1] + rotation[1, 0]) / s
        z = (rotation[0, 2] + rotation[2, 0]) / s
    elif rotation[1, 1] > rotation[2, 2]:
        s = math.sqrt(1.0 + rotation[1, 1] - rotation[0, 0] - rotation[2, 2]) * 2.0
        w = (rotation[0, 2] - rotation[2, 0]) / s
        x = (rotation[0, 1] + rotation[1, 0]) / s
        y = 0.25 * s
        z = (rotation[1, 2] + rotation[2, 1]) / s
    else:
        s = math.sqrt(1.0 + rotation[2, 2] - rotation[0, 0] - rotation[1, 1]) * 2.0
        w = (rotation[1, 0] - rotation[0, 1]) / s
        x = (rotation[0, 2] + rotation[2, 0]) / s
        y = (rotation[1, 2] + rotation[2, 1]) / s
        z = 0.25 * s
    return np.array([x, y, z, w], dtype=float)


def rpy_deg_xyz(rotation: np.ndarray) -> np.ndarray:
    """Intrinsic XYZ Euler angles in degrees from a rotation matrix."""

    rotation = orthonormalize_rotation(rotation)
    sy = math.sqrt(rotation[0, 0] ** 2 + rotation[1, 0] ** 2)
    if sy > 1e-8:
        roll = math.atan2(rotation[2, 1], rotation[2, 2])
        pitch = math.atan2(-rotation[2, 0], sy)
        yaw = math.atan2(rotation[1, 0], rotation[0, 0])
    else:
        roll = math.atan2(-rotation[1, 2], rotation[1, 1])
        pitch = math.atan2(-rotation[2, 0], sy)
        yaw = 0.0
    return np.degrees(np.array([roll, pitch, yaw], dtype=float))