tests/test_physics_likelihood.py

"""Tests for :mod:`fpgm.physics.likelihood`: residuals and the Student-t log-likelihood.

GPU-free, no downloads, no mujoco: every observation and rollout here is
synthetic, built directly from :class:`~fpgm.physics.types.ObservedTrack` and
hand-constructed ``(N, T, 7)`` predicted-pose arrays.
"""

from __future__ import annotations

import numpy as np

from fpgm.physics.likelihood import log_likelihood, se3_residuals
from fpgm.physics.types import ObservedTrack

# --------------------------------------------------------------------------- #
# Local quaternion helpers, independent of anything in fpgm.physics.likelihood
# -- so a bug in the module under test can't hide itself from these tests.
# --------------------------------------------------------------------------- #


def _quat_to_matrix(q: np.ndarray) -> np.ndarray:
    w, x, y, z = q
    n = w * w + x * x + y * y + z * z
    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)],
        ]
    )


def _random_quat(rng: np.random.Generator) -> np.ndarray:
    axis = rng.normal(size=3)
    axis /= np.linalg.norm(axis)
    angle = rng.uniform(-0.6, 0.6)
    q = np.array([np.cos(angle / 2.0), *(axis * np.sin(angle / 2.0))])
    return q / np.linalg.norm(q)


def _make_track(n_frames=12, valid=None, seed=0, sigma_trans=0.005, sigma_rot=0.02):
    rng = np.random.default_rng(seed)
    T_world_obj = np.zeros((n_frames, 4, 4))
    quats = np.zeros((n_frames, 4))
    for t in range(n_frames):
        pos = np.array([0.3, 0.0, 0.1]) + rng.normal(size=3) * 0.05
        q = _random_quat(rng)
        quats[t] = q
        T = np.eye(4)
        T[:3, :3] = _quat_to_matrix(q)
        T[:3, 3] = pos
        T_world_obj[t] = T
    valid_arr = np.ones(n_frames, dtype=bool) if valid is None else np.asarray(valid, dtype=bool)
    track = ObservedTrack(
        uuid="ep0",
        camera_serial="cam0",
        label="obj",
        T_world_obj=T_world_obj,
        valid=valid_arr,
        sigma_trans_m=sigma_trans,
        sigma_rot_rad=sigma_rot,
    )
    return track, quats


def _predicted_from_track(
    track: ObservedTrack, quats: np.ndarray, n_particles: int = 1
) -> np.ndarray:
    pos = track.T_world_obj[:, :3, 3]
    base = np.concatenate([pos, quats], axis=-1)  # (T, 7)
    return np.tile(base[None, :, :], (n_particles, 1, 1))


def _gaussian_logpdf_sum(residuals: np.ndarray, sigma: float) -> float:
    """Independent re-implementation of a Gaussian log-density sum, for the outlier test.

    Deliberately not calling anything in ``fpgm.physics.likelihood`` -- this is
    the "what a Gaussian would have done" baseline the Student-t result is
    compared against.
    """
    z = residuals / sigma
    return float(np.sum(-0.5 * z * z))


# --------------------------------------------------------------------------- #
# Identity checks
# --------------------------------------------------------------------------- #


def test_identity_residuals_are_zero():
    track, quats = _make_track()
    predicted = _predicted_from_track(track, quats)
    trans_res, rot_res = se3_residuals(track, predicted)
    assert np.allclose(trans_res, 0.0, atol=1e-12)
    assert np.allclose(rot_res, 0.0, atol=1e-8)


def test_identity_likelihood_is_max_over_sweep():
    track, quats = _make_track()
    rng = np.random.default_rng(1)
    n_frames = track.T_world_obj.shape[0]
    n_perturb = 6
    predicted = np.zeros((1 + n_perturb, n_frames, 7))
    predicted[0] = _predicted_from_track(track, quats, n_particles=1)[0]
    for i in range(1, 1 + n_perturb):
        pos = track.T_world_obj[:, :3, 3] + rng.normal(size=(n_frames, 3)) * 0.01 * i
        perturbed_quats = quats + rng.normal(size=quats.shape) * 0.02 * i
        perturbed_quats /= np.linalg.norm(perturbed_quats, axis=-1, keepdims=True)
        predicted[i] = np.concatenate([pos, perturbed_quats], axis=-1)

    ok = np.ones(predicted.shape[0], dtype=bool)
    ll = log_likelihood(track, predicted, ok)
    assert np.argmax(ll) == 0
    assert np.all(ll[0] >= ll[1:])


# --------------------------------------------------------------------------- #
# Invalid frames
# --------------------------------------------------------------------------- #


def test_invalid_frames_never_contribute():
    n_frames = 12
    valid = np.ones(n_frames, dtype=bool)
    valid[3] = False
    valid[7] = False
    track, quats = _make_track(n_frames=n_frames, valid=valid)

    rng = np.random.default_rng(2)
    predicted = _predicted_from_track(track, quats, n_particles=4)
    predicted = predicted + rng.normal(size=predicted.shape) * 0.01
    predicted[..., 3:7] /= np.linalg.norm(predicted[..., 3:7], axis=-1, keepdims=True)
    ok = np.array([True, True, False, True])

    trans_res, rot_res = se3_residuals(track, predicted)
    ll = log_likelihood(track, predicted, ok)

    garbage = predicted.copy()
    garbage[:, 3, :3] = 1e6
    garbage[:, 3, 3:7] = np.array([1e6, 2e6, -3e6, 4e6])
    garbage[:, 7, :3] = -1e6
    garbage[:, 7, 3:7] = np.array([5e6, 0.0, 0.0, 0.0])

    trans_res2, rot_res2 = se3_residuals(track, garbage)
    ll2 = log_likelihood(track, garbage, ok)

    assert np.array_equal(trans_res, trans_res2, equal_nan=True)
    assert np.array_equal(rot_res, rot_res2, equal_nan=True)
    assert np.array_equal(ll, ll2)
    assert np.all(np.isnan(trans_res[:, 3])) and np.all(np.isnan(trans_res[:, 7]))
    assert np.all(np.isnan(rot_res[:, 3])) and np.all(np.isnan(rot_res[:, 7]))


def test_invalid_frames_are_nan_not_zero():
    n_frames = 5
    valid = np.ones(n_frames, dtype=bool)
    valid[2] = False
    track, quats = _make_track(n_frames=n_frames, valid=valid)
    predicted = _predicted_from_track(track, quats)
    trans_res, rot_res = se3_residuals(track, predicted)
    assert np.isnan(trans_res[0, 2])
    assert np.isnan(rot_res[0, 2])
    assert np.all(np.isfinite(trans_res[0, valid]))
    assert np.all(np.isfinite(rot_res[0, valid]))


# --------------------------------------------------------------------------- #
# Geodesic rotation angle
# --------------------------------------------------------------------------- #


def test_rotation_residual_is_true_geodesic_angle():
    T = np.eye(4)
    track = ObservedTrack(
        uuid="ep0",
        camera_serial="cam0",
        label="obj",
        T_world_obj=T[None, :, :],
        valid=np.array([True]),
        sigma_trans_m=0.01,
        sigma_rot_rad=0.01,
    )
    q90 = np.array([np.cos(np.pi / 4.0), 0.0, 0.0, np.sin(np.pi / 4.0)])
    q180 = np.array([0.0, 0.0, 0.0, 1.0])
    predicted = np.zeros((2, 1, 7))
    predicted[0, 0, 3:7] = q90
    predicted[1, 0, 3:7] = q180

    _, rot_res = se3_residuals(track, predicted)
    assert np.isclose(rot_res[0, 0], np.pi / 2.0, atol=1e-8)
    assert np.isclose(rot_res[1, 0], np.pi, atol=1e-8)


# --------------------------------------------------------------------------- #
# n_valid == 0
# --------------------------------------------------------------------------- #


def test_zero_valid_frames_returns_exact_zeros():
    track, quats = _make_track(valid=np.zeros(12, dtype=bool))
    predicted = _predicted_from_track(track, quats, n_particles=5)
    ok = np.ones(5, dtype=bool)
    ll = log_likelihood(track, predicted, ok)
    assert np.array_equal(ll, np.zeros(5))

    # Even a wildly different rollout changes nothing: with no valid frames,
    # log_likelihood returns before ever looking at predicted_poses's values.
    ll2 = log_likelihood(track, predicted + 100.0, ok)
    assert np.array_equal(ll2, np.zeros(5))


# --------------------------------------------------------------------------- #
# Student-t vs Gaussian on an outlier frame
# --------------------------------------------------------------------------- #


def test_student_t_downweights_outlier_relative_to_gaussian():
    """One outlier frame should hurt far less under our Student-t than under a Gaussian.

    Particle A matches 5 of 6 frames exactly and misses the 6th by 10 sigma
    (one bad PnP frame). Particle B misses every one of 6 frames by a modest
    1.2 sigma (a consistently slightly-off rollout). Under a Gaussian, A's
    single 10-sigma residual contributes -50 nats (before B's -4.32) --
    A is crushed by one frame. Under Student-t(nu=4), that same 10-sigma frame
    only costs about -8.1 nats, so the A-vs-B gap shrinks by roughly an order
    of magnitude -- an outlier frame no longer single-handedly vetoes a
    particle that fits everything else well.
    """
    n_frames = 6
    T_world_obj = np.tile(np.eye(4)[None, :, :], (n_frames, 1, 1))
    track = ObservedTrack(
        uuid="ep0",
        camera_serial="cam0",
        label="obj",
        T_world_obj=T_world_obj,
        valid=np.ones(n_frames, dtype=bool),
        sigma_trans_m=1.0,
        sigma_rot_rad=1.0,  # rotation residual is 0 for both particles; term cancels in the gap
    )

    delta_a = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 10.0])
    delta_b = np.full(n_frames, 1.2)

    predicted = np.zeros((2, n_frames, 7))
    predicted[:, :, 3] = 1.0  # identity quaternion (w=1) for both particles, all frames
    predicted[0, :, 0] = delta_a
    predicted[1, :, 0] = delta_b

    ok = np.ones(2, dtype=bool)
    trans_res, rot_res = se3_residuals(track, predicted)
    # rotation term is identical (zero) for A and B, so it drops out of the gap below
    assert np.allclose(rot_res, 0.0, atol=1e-12)

    ll = log_likelihood(track, predicted, ok, nu=4.0)
    t_gap = float(ll[1] - ll[0])  # B - A under Student-t

    gaussian_gap = _gaussian_logpdf_sum(delta_b, 1.0) - _gaussian_logpdf_sum(delta_a, 1.0)

    assert gaussian_gap > 40.0  # sanity: the Gaussian baseline really is dominated by the outlier
    assert t_gap > 0.0  # B is still favoured under Student-t too
    assert abs(t_gap) < 0.3 * abs(gaussian_gap)  # but by far less: the outlier's veto is softened