src/fpgm/physics/likelihood.py

"""Per-episode log-likelihood: how well one particle's rollout matches the observed track.

This module produces exactly one number per particle per episode-object
(:class:`~fpgm.physics.types.EpisodeLogLik`), and the arithmetic contract on that
number is set by :mod:`fpgm.physics.types`, not by anything here: particles are
pooled by *summing* ``logL_e(theta_i)`` across episodes, then softmax-ing. That
means every choice below is judged by one question -- "does this preserve
`constant-per-episode => softmax no-op`?" -- not by whether it looks like a
normal likelihood function.

**Why the returned vector is never normalised, and never has its max
subtracted.** A per-episode ``logL - logsumexp(logL)`` (turning it into a proper
normalised log-probability over particles) is the textbook thing to do with a
log-likelihood, and it is exactly wrong here. ``logsumexp`` is *not* an additive
constant -- it depends on the full shape of ``logL`` across particles, so two
episodes with different informativeness would get shifted by different amounts
before summing, and the sum would no longer equal ``sum_e logL_e(theta_i)``
evaluated on the raw densities. An uninformative episode (flat ``logL``) would
still normalise to something informative-looking (a near-uniform but not exactly
flat vector, once floating point is involved) and stop being a true no-op. So
:func:`log_likelihood` returns a raw sum of per-frame log-densities and nothing
more; the only place softmax happens is once, in
:mod:`fpgm.physics.inference`, over the fully-summed weights.

**Why Student-t, not Gaussian.** ``ObservedTrack.sigma_trans_m`` /
``sigma_rot_rad`` are measured noise floors from frames the object provably did
not move (:mod:`fpgm.datagen.static_span`) -- but they describe the *typical*
frame, not the worst one. A PnP pose estimate occasionally fails badly on a
single frame (motion blur, a fleeting occlusion) without that meaning anything
about theta. Under a Gaussian, one such frame contributes a residual-squared
term that grows without bound and can dominate the whole episode's likelihood,
i.e. one bad frame vetoes an otherwise-correct particle. A Student-t's tails
decay polynomially instead of like ``exp(-x^2)``, so a single outlier frame's
log-density floors out instead of diverging -- it still counts against a
particle, just not catastrophically. ``nu=4`` is a fixed, moderate choice (not
fit per episode): heavy enough tails to matter, not so heavy that genuinely
informative frames stop contributing. This is a design default, not a measured
value.

**What "residual" means here.** :func:`se3_residuals` reduces each frame to two
*scalar* magnitudes -- a Euclidean translation distance (metres, >= 0) and a
geodesic rotation angle (radians, >= 0) -- rather than keeping the full 3-vector
translation error or a per-axis rotation error. The Student-t density below is
then evaluated on that scalar residual directly, symmetric around zero, with the
residual's actual sign discarded (it has none: a distance is not signed). This is
an approximation: the true sampling distribution of a Euclidean-norm residual
under isotropic per-axis noise is a (generalised) chi-like distribution, not a
Student-t. It is used anyway because (a) ``ObjectPoseDebug.pose_noise_mm`` is
itself only a scalar spread, not a full 3x3 covariance, so a per-axis model would
be inventing structure the upstream measurement does not support, and (b) the
qualitative behaviour that actually matters for the importance weighting --
small residual -> high density, residual >> sigma -> density decays slowly, not
catastrophically -- holds regardless. Treated as a calibrated, heavy-tailed
penalty on "how many noise-floors away was this frame", not as a rigorously
derived sampling density.

**Quaternion convention and why these helpers are local numpy, not
``scipy.spatial.transform``.** ``predicted_poses`` is MuJoCo's own wire format
straight out of :class:`~fpgm.physics.types.SimResult`: ``(w, x, y, z)``.
``scipy.spatial.transform.Rotation`` uses ``(x, y, z, w)``, so every call site
touching it would need a manual axis reorder -- precisely the kind of "array
looks right but is quietly transposed" bug
:mod:`fpgm.physics.types`'s own docstring calls out as the reason this stage's
wire formats are shape-checked in the first place. :mod:`fpgm.geometry.transforms`
already sets the precedent of implementing SE3 conversions as plain numpy for
this reason; :func:`_matrix_to_quat_wxyz` and :func:`_quat_geodesic_angle`
follow it, fully vectorised over arbitrary leading dimensions so ``(T, 3, 3)``
and ``(N, T, 4)`` arrays are converted/compared in one call each, no Python loop
over frames or particles. ``scipy.special.gammaln`` *is* used in
:func:`_student_t_logpdf`, deliberately: it is a numerical special function with
no convention to get backwards, unlike a rotation representation.

**Why ``observed.n_valid == 0`` returns exact zeros, not an error or ``-inf``.**
Zero valid frames means literally no observation was made of this object in this
episode -- e.g. it was occluded the whole clip, or every PnP estimate failed
S6's visibility gate. That is the purest case of "this episode carries no
information": the correct contribution to ``sum_e logL_e(theta_i)`` is a
constant (here, exactly zero) added to every particle, which is a no-op under
softmax by the same arithmetic the whole stage is built on. Raising would turn a
routine, expected situation (S6 does not promise every object is visible in
every episode) into a pipeline failure; returning ``-inf`` would incorrectly
veto every particle for an episode that said nothing at all.
"""

from __future__ import annotations

from contextlib import nullcontext
from typing import TYPE_CHECKING

import numpy as np
from scipy.special import gammaln

from fpgm.physics.types import ObservedTrack, PhysicsError

if TYPE_CHECKING:
    from fpgm.utils.timing import StepTimer

__all__ = ["se3_residuals", "log_likelihood"]


def _step(timer: StepTimer | None, label: str, *, n: int | None = None):
    """``timer.step(...)`` if a timer was given, else a no-op context manager.

    Keeps every function below able to declare ``timer: StepTimer | None = None``
    without an ``if timer is not None: ... else: ...`` fork at every call site.
    """
    if timer is None:
        return nullcontext()
    return timer.step(label, n=n)


# --------------------------------------------------------------------------- #
# Quaternion helpers (local numpy -- see module docstring for why)
# --------------------------------------------------------------------------- #


def _matrix_to_quat_wxyz(rot: np.ndarray) -> np.ndarray:
    """Vectorised ``(..., 3, 3)`` rotation matrix -> ``(..., 4)`` unit quaternion, ``(w, x, y, z)``.

    Shepperd's method: pick the numerically stable branch (largest of ``trace``
    and the three diagonal entries) per-element via boolean masks, rather than
    the classic scalar if/elif chain, so this runs as four masked vector ops
    over the whole batch instead of a Python loop over ``T`` or ``N*T`` matrices.
    """
    rot = np.asarray(rot, dtype=np.float64)
    lead_shape = rot.shape[:-2]
    m = rot.reshape(-1, 3, 3)
    n = m.shape[0]
    q = np.zeros((n, 4), dtype=np.float64)

    trace = m[:, 0, 0] + m[:, 1, 1] + m[:, 2, 2]
    d00, d11, d22 = m[:, 0, 0], m[:, 1, 1], m[:, 2, 2]

    case_trace = trace > 0
    case_x = ~case_trace & (d00 >= d11) & (d00 >= d22)
    case_y = ~case_trace & (d11 > d00) & (d11 >= d22)
    case_z = ~case_trace & (d22 > d00) & (d22 > d11)

    if np.any(case_trace):
        c = case_trace
        s = np.sqrt(trace[c] + 1.0) * 2.0  # s = 4w
        q[c, 0] = 0.25 * s
        q[c, 1] = (m[c, 2, 1] - m[c, 1, 2]) / s
        q[c, 2] = (m[c, 0, 2] - m[c, 2, 0]) / s
        q[c, 3] = (m[c, 1, 0] - m[c, 0, 1]) / s

    if np.any(case_x):
        c = case_x
        s = np.sqrt(1.0 + d00[c] - d11[c] - d22[c]) * 2.0  # s = 4x
        q[c, 0] = (m[c, 2, 1] - m[c, 1, 2]) / s
        q[c, 1] = 0.25 * s
        q[c, 2] = (m[c, 0, 1] + m[c, 1, 0]) / s
        q[c, 3] = (m[c, 0, 2] + m[c, 2, 0]) / s

    if np.any(case_y):
        c = case_y
        s = np.sqrt(1.0 + d11[c] - d00[c] - d22[c]) * 2.0  # s = 4y
        q[c, 0] = (m[c, 0, 2] - m[c, 2, 0]) / s
        q[c, 1] = (m[c, 0, 1] + m[c, 1, 0]) / s
        q[c, 2] = 0.25 * s
        q[c, 3] = (m[c, 1, 2] + m[c, 2, 1]) / s

    if np.any(case_z):
        c = case_z
        s = np.sqrt(1.0 + d22[c] - d00[c] - d11[c]) * 2.0  # s = 4z
        q[c, 0] = (m[c, 1, 0] - m[c, 0, 1]) / s
        q[c, 1] = (m[c, 0, 2] + m[c, 2, 0]) / s
        q[c, 2] = (m[c, 1, 2] + m[c, 2, 1]) / s
        q[c, 3] = 0.25 * s

    return q.reshape(*lead_shape, 4)


def _quat_multiply(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
    """Hamilton product of two batches of ``(..., 4)`` ``(w, x, y, z)`` quaternions."""
    w1, x1, y1, z1 = q1[..., 0], q1[..., 1], q1[..., 2], q1[..., 3]
    w2, x2, y2, z2 = q2[..., 0], q2[..., 1], q2[..., 2], q2[..., 3]
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
    return np.stack([w, x, y, z], axis=-1)


def _quat_geodesic_angle(q1: np.ndarray, q2: np.ndarray) -> np.ndarray:
    """Angle in ``[0, pi]`` radians between two batches of unit quaternions, ``(..., 4)``.

    Computed as ``2 * atan2(||relative.xyz||, |relative.w|)``, where
    ``relative = conjugate(q1) * q2`` is the quaternion that rotates frame 1
    onto frame 2. The absolute value on ``relative.w`` accounts for ``q`` and
    ``-q`` representing the same rotation.

    This is deliberately *not* the more obvious ``2 * arccos(|dot(q1, q2)|)``:
    that formula is ill-conditioned exactly where it matters most for this
    module -- near-zero rotation error, which is the common case for a
    particle whose theta is roughly right. ``d/dx arccos(x)`` diverges as ``x
    -> 1``, so the ~1e-16 float noise ordinary quaternion arithmetic
    accumulates gets amplified to angle errors around 1e-7 rad, which is not
    negligible next to ``sigma_rot_rad`` (a few hundredths of a radian) --
    this was measured directly while writing this module's tests, not assumed.
    ``atan2(small, ~1)`` has no such singularity: its response to the same
    float noise stays at the noise's own scale. Near the opposite extreme
    (angle ~pi), ``arccos`` is equally ill-conditioned and ``atan2`` is
    equally fine, by the same argument applied to the ``w`` component instead.
    """
    q1 = np.asarray(q1, dtype=np.float64)
    q2 = np.asarray(q2, dtype=np.float64)
    q1 = q1 / (np.linalg.norm(q1, axis=-1, keepdims=True) + 1e-15)
    q2 = q2 / (np.linalg.norm(q2, axis=-1, keepdims=True) + 1e-15)
    q1_conj = q1 * np.array([1.0, -1.0, -1.0, -1.0])
    relative = _quat_multiply(q1_conj, q2)
    xyz_norm = np.linalg.norm(relative[..., 1:4], axis=-1)
    w = np.abs(relative[..., 0])
    return 2.0 * np.arctan2(xyz_norm, w)


# --------------------------------------------------------------------------- #
# Residuals
# --------------------------------------------------------------------------- #


def se3_residuals(
    observed: ObservedTrack,
    predicted_poses: np.ndarray,
    *,
    timer: StepTimer | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Per-frame, per-particle pose error between ``observed`` and ``predicted_poses``.

    Args:
        observed: The measured track. Its ``valid`` mask decides which frames
            participate; see the "invalid frames" note below.
        predicted_poses: ``(N, T, 7)`` ``[x, y, z, qw, qx, qy, qz]``, MuJoCo's
            quaternion order (matches :class:`~fpgm.physics.types.SimResult`).
        timer: Optional :class:`~fpgm.utils.timing.StepTimer`; a clean no-op
            when ``None``.

    Returns:
        ``(trans_residuals, rot_residuals)``, each ``(N, T)``:
        translation in metres (Euclidean distance between positions), rotation
        in radians (geodesic angle between orientations). Both are ``NaN`` at
        every frame where ``observed.valid`` is ``False`` -- never zero. Zero
        would silently claim "predicted and observed agreed here"; ``NaN`` is
        the honest "there is no observation to compare against", and it is what
        lets :func:`log_likelihood` exclude those frames via ``nansum`` instead
        of accidentally rewarding a rollout for matching a pose nobody measured.
    """
    predicted_poses = np.asarray(predicted_poses, dtype=np.float64)
    if predicted_poses.ndim != 3 or predicted_poses.shape[2] != 7:
        raise PhysicsError(f"predicted_poses must be (N, T, 7), got {predicted_poses.shape}")
    n_particles, n_frames, _ = predicted_poses.shape
    if n_frames != observed.T_world_obj.shape[0]:
        raise PhysicsError(
            f"{observed.label}: predicted_poses has {n_frames} frames, observed "
            f"track has {observed.T_world_obj.shape[0]}"
        )

    with _step(timer, "se3_residuals", n=n_particles * n_frames):
        observed_pos = observed.T_world_obj[:, :3, 3]  # (T, 3)
        observed_quat = _matrix_to_quat_wxyz(observed.T_world_obj[:, :3, :3])  # (T, 4)

        predicted_pos = predicted_poses[:, :, :3]  # (N, T, 3)
        predicted_quat = predicted_poses[:, :, 3:7]  # (N, T, 4)

        trans_res = np.linalg.norm(predicted_pos - observed_pos[None, :, :], axis=-1)
        rot_res = _quat_geodesic_angle(predicted_quat, observed_quat[None, :, :])

        invalid = ~observed.valid
        trans_res[:, invalid] = np.nan
        rot_res[:, invalid] = np.nan

    return trans_res, rot_res


# --------------------------------------------------------------------------- #
# Log-likelihood
# --------------------------------------------------------------------------- #


def _student_t_logpdf(x: np.ndarray, sigma: float, nu: float) -> np.ndarray:
    """Log-density of a location-0, scale-``sigma`` Student-t at ``x`` (``NaN`` in, ``NaN`` out)."""
    z = x / sigma
    return (
        gammaln((nu + 1.0) / 2.0)
        - gammaln(nu / 2.0)
        - 0.5 * np.log(nu * np.pi)
        - np.log(sigma)
        - 0.5 * (nu + 1.0) * np.log1p(z * z / nu)
    )


def log_likelihood(
    observed: ObservedTrack,
    predicted_poses: np.ndarray,
    ok: np.ndarray,
    *,
    nu: float = 4.0,
    timer: StepTimer | None = None,
) -> np.ndarray:
    """Raw per-particle log-likelihood of ``observed`` under each particle's rollout.

    Args:
        observed: The measured track.
        predicted_poses: ``(N, T, 7)``, see :func:`se3_residuals`.
        ok: ``(N,)`` bool, one per particle -- ``SimResult.ok``. ``False`` means
            the rollout diverged; that particle gets exactly ``-inf`` regardless
            of any residual computed from its (numerically meaningless) poses.
            ``-inf`` is a real statement ("this theta is inconsistent with
            physics itself, before even comparing to the observation"), not a
            dropped sample -- it still participates in the softmax denominator
            downstream in :mod:`fpgm.physics.inference`.
        nu: Student-t degrees of freedom; see the module docstring for why
            Student-t and why this default.
        timer: Optional :class:`~fpgm.utils.timing.StepTimer`.

    Returns:
        ``(N,)`` float64. Exactly ``np.zeros(N)`` if ``observed.n_valid == 0``
        (see module docstring). Otherwise the sum, over valid frames only, of
        the translation and rotation Student-t log-densities -- **not**
        normalised, **not** max-subtracted. That is not an oversight: see the
        module docstring's first section for why either operation would break
        the "uninformative episode is a softmax no-op" invariant this whole
        stage depends on.
    """
    predicted_poses = np.asarray(predicted_poses, dtype=np.float64)
    ok = np.asarray(ok, dtype=bool).reshape(-1)
    n_particles = predicted_poses.shape[0]
    if ok.shape[0] != n_particles:
        raise PhysicsError(
            f"{observed.label}: ok has {ok.shape[0]} entries for {n_particles} particles"
        )

    if observed.n_valid == 0:
        return np.zeros(n_particles, dtype=np.float64)

    trans_res, rot_res = se3_residuals(observed, predicted_poses, timer=timer)

    with _step(timer, "score_frames", n=n_particles):
        with np.errstate(over="ignore", invalid="ignore"):
            trans_ll = _student_t_logpdf(trans_res, observed.sigma_trans_m, nu)
            rot_ll = _student_t_logpdf(rot_res, observed.sigma_rot_rad, nu)
            loglik = np.nansum(trans_ll, axis=1) + np.nansum(rot_ll, axis=1)
        loglik = np.where(ok, loglik, -np.inf)

    return loglik