festifun/backend/festival4d/geometry.py

"""Coordinate conversions, quaternion math, ray casting, triangulation.

The single most important thing in this file is :func:`colmap_to_threejs` — the
COLMAP world->camera pose to Three.js camera conversion (spec M5). It is a **FROZEN
CONTRACT**: it ships in foundation with a passing unit test (``test_geometry.py``), and
the JS mirror lives in ``frontend/src/lib/pose.js`` with the *same* embedded test
vectors. Lanes B and C **consume** it; they never reimplement or modify it.

The remaining functions (:func:`slerp_pose`, :func:`ray_from_pixel`,
:func:`triangulate_rays`, :func:`nearest_point_on_ray`) are stubs with frozen
signatures for lane B (spec M2 + M8). Lane B fills the bodies and adds their tests.

Conventions
-----------
- Quaternions are COLMAP order **[w, x, y, z]** (scalar first), unit norm.
- A pose ``(q, t)`` is world->camera: ``x_cam = R(q) @ x_world + t``.
- COLMAP camera axes: +x right, +y down, +z forward (into the scene).
- Three.js cameras look down -z with +y up; hence the ``diag(1, -1, -1)`` flip.
"""

from __future__ import annotations

import numpy as np
from numpy.typing import ArrayLike, NDArray

# Camera-axis flip that takes COLMAP camera-local axes (x right, y down, z forward)
# to Three.js camera-local axes (x right, y up, z backward). Part of the frozen contract.
_FLIP_YZ = np.diag([1.0, -1.0, -1.0])


def quat_to_mat(q: ArrayLike) -> NDArray[np.float64]:
    """Convert a unit quaternion ``[w, x, y, z]`` to a 3x3 rotation matrix.

    Hamilton convention, right-handed, active rotation: the returned ``R`` is COLMAP's
    world->camera matrix when ``q`` is a COLMAP pose quaternion. FROZEN (used by the
    conversion contract).
    """
    w, x, y, z = (float(v) for v in np.asarray(q, dtype=np.float64).reshape(4))
    n = w * w + x * x + y * y + z * z
    if n < 1e-12:
        raise ValueError("quaternion has near-zero norm")
    s = 2.0 / n
    wx, wy, wz = s * w * x, s * w * y, s * w * z
    xx, xy, xz = s * x * x, s * x * y, s * x * z
    yy, yz, zz = s * y * y, s * y * z, s * z * z
    return np.array(
        [
            [1.0 - (yy + zz), xy - wz, xz + wy],
            [xy + wz, 1.0 - (xx + zz), yz - wx],
            [xz - wy, yz + wx, 1.0 - (xx + yy)],
        ],
        dtype=np.float64,
    )


def mat_to_quat(R: ArrayLike) -> NDArray[np.float64]:
    """Convert a 3x3 rotation matrix to a unit quaternion ``[w, x, y, z]`` (w >= 0).

    Inverse of :func:`quat_to_mat`. Used by tests and the pose export path.
    """
    m = np.asarray(R, dtype=np.float64).reshape(3, 3)
    trace = m[0, 0] + m[1, 1] + m[2, 2]
    if trace > 0.0:
        s = np.sqrt(trace + 1.0) * 2.0
        w = 0.25 * s
        x = (m[2, 1] - m[1, 2]) / s
        y = (m[0, 2] - m[2, 0]) / s
        z = (m[1, 0] - m[0, 1]) / s
    elif m[0, 0] > m[1, 1] and m[0, 0] > m[2, 2]:
        s = np.sqrt(1.0 + m[0, 0] - m[1, 1] - m[2, 2]) * 2.0
        w = (m[2, 1] - m[1, 2]) / s
        x = 0.25 * s
        y = (m[0, 1] + m[1, 0]) / s
        z = (m[0, 2] + m[2, 0]) / s
    elif m[1, 1] > m[2, 2]:
        s = np.sqrt(1.0 + m[1, 1] - m[0, 0] - m[2, 2]) * 2.0
        w = (m[0, 2] - m[2, 0]) / s
        x = (m[0, 1] + m[1, 0]) / s
        y = 0.25 * s
        z = (m[1, 2] + m[2, 1]) / s
    else:
        s = np.sqrt(1.0 + m[2, 2] - m[0, 0] - m[1, 1]) * 2.0
        w = (m[1, 0] - m[0, 1]) / s
        x = (m[0, 2] + m[2, 0]) / s
        y = (m[1, 2] + m[2, 1]) / s
        z = 0.25 * s
    q = np.array([w, x, y, z], dtype=np.float64)
    q /= np.linalg.norm(q)
    if q[0] < 0:  # canonical sign: non-negative scalar part
        q = -q
    return q


def colmap_to_threejs(q: ArrayLike, t: ArrayLike) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Convert a COLMAP world->camera pose to a Three.js camera pose (spec M5).

    **FROZEN CONTRACT.** Mirrored in ``frontend/src/lib/pose.js``; do not change the math
    without a change request and a synchronized update to both sides + their test vectors.

    Parameters
    ----------
    q : array_like, shape (4,)
        COLMAP world->camera quaternion ``[w, x, y, z]``.
    t : array_like, shape (3,)
        COLMAP world->camera translation ``[tx, ty, tz]``.

    Returns
    -------
    position : ndarray, shape (3,)
        Camera center in world coordinates, ``C = -R^T t``. Assign to ``camera.position``.
    rotation_matrix : ndarray, shape (3, 3)
        Three.js camera world rotation ``R_three = R^T @ diag(1, -1, -1)``. Assign via
        ``camera.setRotationFromMatrix(...)`` (a proper rotation, det = +1).
    """
    R = quat_to_mat(q)                      # world -> cam
    t_vec = np.asarray(t, dtype=np.float64).reshape(3)
    R_c2w = R.T                             # cam -> world
    position = -R_c2w @ t_vec              # camera center in world coords
    rotation_matrix = R_c2w @ _FLIP_YZ     # flip camera-local y,z for Three.js
    return position, rotation_matrix


# ---------------------------------------------------------------------------
# Lane B stubs (spec M2 + M8). Signatures FROZEN; bodies raise NotImplementedError.
# ---------------------------------------------------------------------------
def slerp_pose(
    q0: ArrayLike,
    t0: ArrayLike,
    q1: ArrayLike,
    t1: ArrayLike,
    alpha: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Interpolate between two world->camera poses (spec M2 pose interpolation).

    Spherically interpolate the rotation (slerp on the shorter arc, handling the
    double-cover sign) and linearly interpolate the translation, at fraction
    ``alpha in [0, 1]`` from pose 0 to pose 1.

    Parameters
    ----------
    q0, q1 : array_like, shape (4,)
        COLMAP quaternions ``[w, x, y, z]`` at the endpoints.
    t0, t1 : array_like, shape (3,)
        COLMAP translations at the endpoints.
    alpha : float
        Interpolation fraction; 0 returns pose 0, 1 returns pose 1.

    Returns
    -------
    (q, t) : the interpolated quaternion ``[w, x, y, z]`` and translation ``[x, y, z]``.
    """
    raise NotImplementedError("lane B (M2): implement slerp_pose")


def ray_from_pixel(
    q: ArrayLike,
    t: ArrayLike,
    fx: float,
    fy: float,
    cx: float,
    cy: float,
    px: float,
    py: float,
) -> tuple[NDArray[np.float64], NDArray[np.float64]]:
    """Unproject a pixel into a world-space ray (spec M8 annotation resolution).

    Given the world->camera pose ``(q, t)`` and pinhole intrinsics, build the ray that
    passes through image pixel ``(px, py)``.

    Returns
    -------
    origin : ndarray, shape (3,)
        Ray origin = camera center in world coords.
    direction : ndarray, shape (3,)
        Unit ray direction in world coords, pointing into the scene.
    """
    raise NotImplementedError("lane B (M8): implement ray_from_pixel")


def triangulate_rays(
    origin_a: ArrayLike,
    dir_a: ArrayLike,
    origin_b: ArrayLike,
    dir_b: ArrayLike,
) -> tuple[NDArray[np.float64], float]:
    """Closest point between two world-space rays (spec M8 two-view triangulation).

    Returns
    -------
    point : ndarray, shape (3,)
        Midpoint of the shortest segment connecting the two rays.
    gap : float
        Length of that shortest segment (the mutual-nearest-approach distance). Callers
        reject the triangulation when the rays are near-parallel or ``gap`` exceeds the
        spec threshold (0.5 scene units).
    """
    raise NotImplementedError("lane B (M8): implement triangulate_rays")


def nearest_point_on_ray(
    origin: ArrayLike,
    direction: ArrayLike,
    points: ArrayLike,
    radius: float = 0.3,
) -> NDArray[np.float64] | None:
    """Nearest point-cloud point to a ray, within a cylinder (spec M8 single-view fallback).

    Among ``points`` (shape ``(N, 3)``) find the one whose perpendicular distance to the
    ray is smallest, considering only points within ``radius`` of the ray and in front of
    the origin.

    Returns
    -------
    point : ndarray shape (3,) or None
        The selected point-cloud point, or ``None`` if none lie within ``radius``.
    """
    raise NotImplementedError("lane B (M8): implement nearest_point_on_ray")