src/fpgm/physics/inference.py

"""Turn per-episode log-likelihoods into a posterior, an ESS, and an identifiability report.

This is where the stage's central arithmetic actually runs. Everything upstream
(:mod:`fpgm.physics.likelihood`) produces one raw ``(N,)`` vector per
episode-object; everything here does is:

    log w_i = sum_e logL_e(theta_i)          w = softmax(log w)

and then reports what ``w`` says about ``theta``. There is no step in this file
that looks at an episode individually once its ``loglik`` has been produced --
by the time :func:`accumulate` runs, "which episodes were informative" is no
longer a question anyone needs to ask, because an uninformative one already
canceled itself out in the sum above. Concretely: :func:`accumulate` computes
``log_weights`` as one ``np.sum`` over the stacked per-episode vectors and does
nothing else to them before softmax-ing -- no per-episode weighting, no
per-episode threshold, no episode ever dropped from the sum. That is the
property the null test in ``tests/test_physics_inference.py`` checks directly by
constructing a constant-``loglik`` episode and asserting the pooled posterior is
*bit-identical* with and without it.

**Four things this module reports about the resulting posterior, and what each
one is (and is not) for:**

* **ESS** (``1 / sum(w_i^2)``) -- how many of the ``N`` prior particles are
  effectively still doing any work. A low ESS means the prior's particles
  mostly missed the region the data favours, so the weighted mean/std below are
  being estimated from a handful of surviving particles and should not be
  trusted at face value. This is what motivates the refinement split below.
* **``ParamPosterior.contraction``** (per axis-aligned parameter) -- ``1 -
  Var_post/Var_prior``, in ``[0, 1]``. 0 = this parameter's marginal posterior
  is the same width as its prior (the data said nothing about it *alone*).
  ``learned`` is a label derived from a fixed threshold on this number, purely
  for a human reading a report table -- see the constant's docstring for why it
  can never gate anything upstream.
* **``sensitivity_values``/``sensitivity_directions``** -- the same
  "how much narrower than the prior" question, but for arbitrary linear
  combinations of parameters, not just the coordinate axes. See
  :func:`_sensitivity` for the construction and why it is not simply "SVD of
  the posterior covariance" (that would rank the wrong end first).
* **``info_nats``** -- a single scalar, an importance-sampling estimate of
  KL(posterior || prior) computed from the weights alone. See
  :func:`_snis_kl_estimate` for what it does and does not measure.

**The degeneracy-guard split, and why it is a split.** A naive spec for this
module would be "if ESS is too low, resample and re-simulate until it isn't."
That loop needs a simulator, and simulation is deliberately not something
:mod:`fpgm.physics` does -- it lives across a conda-env subprocess boundary
(:class:`~fpgm.physics.types.SimSpec`/``SimResult``), which this module has no
business depending on. So :func:`accumulate` never resamples or re-scores; it
only *reports*, in ``PosteriorResult.timings``, that ESS was low enough to
warrant another round. :func:`needs_refinement` and :func:`resample_particles`
are the two halves of the loop this module *can* offer a caller: "should I
refine?" and "give me the next particle set to score." Gluing those two calls
to a simulator and a second :func:`accumulate` call is an orchestrator's job,
not this module's -- and ``PosteriorResult.refinement_rounds`` is left at ``0``
by both :func:`accumulate` and :func:`pool` precisely so a result never claims
work that did not happen. A caller that does complete a refinement round is
expected to construct the final :class:`~fpgm.physics.types.PosteriorResult`
itself (or ask for this module to grow a thin wrapper, if that turns out to be
worth it) rather than have this module fake it.

**Why ``timings`` never stores ``timer.summary()``.** ``StepTimer.summary()``
freezes ``total_seconds`` on first call (see its own docstring: "idempotent";
``summary``/``report``/``total_seconds`` all trigger the freeze). Whether the
caller passes one fresh :class:`~fpgm.utils.timing.StepTimer` per
:func:`accumulate` call or shares one across an entire batch run is an
orchestration choice this module cannot see -- so calling ``summary()`` here
would either be redundant (fresh-timer case) or would silently freeze the
shared timer's total the first time any object finished inference, corrupting
every later step's fraction-of-total in the eventual batch report. Both
:func:`accumulate` and :func:`pool` only ever call ``timer.step(...)``
(never ``.summary()``/``.report()``/``.total_seconds``), and ``timings`` on the
returned :class:`~fpgm.physics.types.PosteriorResult` holds a small,
locally-computed dict instead (particle/episode counts, the refinement
verdict). The full step-by-step table is available from the shared timer itself
whenever the caller is ready to log or persist it.
"""

from __future__ import annotations

from collections.abc import Sequence
from contextlib import nullcontext
from typing import TYPE_CHECKING

import numpy as np

from fpgm.physics.types import (
    EpisodeLogLik,
    GaussianPrior,
    ParamPosterior,
    PhysicsError,
    PosteriorResult,
)

if TYPE_CHECKING:
    from fpgm.utils.timing import StepTimer

__all__ = ["accumulate", "pool", "needs_refinement", "resample_particles"]

#: Contraction above which a parameter is reported ``learned=True`` in the log
#: table. A fixed, documented cutoff rather than a fit one -- and, as
#: ``ParamPosterior`` itself documents, a label for a human reading a report,
#: never an input to which episodes or particles get used. 0.2 is a mild bar
#: (noticeably narrower than the prior, not "pinned"); it is a reporting
#: default, not a measured threshold.
_LEARNED_CONTRACTION = 0.2

#: Roughening-jitter std in :func:`resample_particles`, as a fraction of the
#: fitted posterior std per parameter. Standard particle-filter practice to
#: keep post-resample duplicates from collapsing to point masses; 0.3 is a
#: conventional, undemanding default, not derived from this problem's data.
_JITTER_FRAC = 0.3


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


def _centred(loglik: np.ndarray) -> np.ndarray:
    """Shift one episode's log-likelihood so its finite maximum is exactly 0.

    Subtracting a per-episode constant is a mathematical no-op: ``softmax`` is
    invariant to any shift applied uniformly across particles, and this shift is
    the same scalar for every particle by construction. So the posterior is
    unchanged in exact arithmetic -- but *not* in float64, and that difference
    is the entire reason this function exists.

    Measured, on 512 particles with two informative episodes, appending a
    constant-valued (uninformative) episode and comparing the resulting weights
    against the run without it:

    ==================  =======================
    constant            max relative weight error
    ==================  =======================
    -1e3                7.6e-14
    -1e6                4.7e-11
    -1e9                8.0e-08
    ==================  =======================

    Summing raw log-likelihoods of wildly different magnitudes destroys the
    low-order bits of the informative terms -- ``-8000.0 + (-1e9)`` keeps about
    seven significant digits of the ``-8000``. The stage's central claim is that
    an uninformative episode changes the posterior *exactly* not *nearly*, and a
    claim that degrades with the magnitude of an arbitrary additive constant is
    not the claim worth making. After centring, a constant episode contributes a
    vector of exact zeros, so bit-identity holds for any constant and any
    real-valued log-likelihoods -- not only for the integer-valued ones that
    happen to add exactly.

    ``-inf`` entries (particles whose rollout diverged, i.e. genuinely
    inconsistent with the observation) stay ``-inf``: they are excluded from the
    maximum and unaffected by the shift. An episode where *every* particle is
    ``-inf`` is left alone rather than centred, because there is no finite
    maximum to centre on and the all-``-inf`` state is itself the meaningful
    result.
    """
    finite = loglik[np.isfinite(loglik)]
    if finite.size == 0:
        return loglik
    return loglik - finite.max()


# --------------------------------------------------------------------------- #
# Public entry points
# --------------------------------------------------------------------------- #


def accumulate(
    prior: GaussianPrior,
    particles: np.ndarray,
    episodes: Sequence[EpisodeLogLik],
    *,
    timer: StepTimer | None = None,
    ess_min_frac: float = 0.05,
    rng: np.random.Generator | None = None,
    refine: bool = True,
) -> PosteriorResult:
    """Sum ``episodes``' raw log-likelihoods over ``particles`` and report the posterior.

    Args:
        prior: The Gaussian prior ``particles`` were drawn from. Stored on the
            returned result unchanged, so contraction/sensitivity can compare
            against it later without needing it passed around separately.
        particles: ``(N, D)``, ``D == prior.space.dim``. Must be the exact same
            array every episode's rollout was simulated from -- this function
            has no way to check that beyond the shape, which is why
            :func:`pool` additionally checks bit-identity when combining
            results built from (presumably) the same particle set.
        episodes: Every episode-object contribution to sum. Order does not
            matter (addition commutes); an empty sequence is valid and yields
            ``log_weights = 0`` for every particle, i.e. the posterior equals
            the prior, exactly.
        timer: Optional :class:`~fpgm.utils.timing.StepTimer`.
        ess_min_frac: See :func:`needs_refinement`. Only consulted to decide
            whether ``timings["refinement_warranted"]`` is set; never changes
            ``particles``, ``log_weights``, or any other returned value.
        rng: Accepted for interface symmetry with :func:`resample_particles`
            (whose own ``rng`` argument a caller building a refinement loop
            will need anyway) and for forward compatibility. Unused in this
            function's body: every statistic below is a deterministic function
            of ``(particles, log_weights, prior)`` and needs no randomness.
        refine: If ``False``, skip the ESS-vs-threshold check entirely (e.g. a
            caller that already knows it will not run a refinement round and
            would rather not have a stale-looking ``refinement_warranted:
            False`` in the report). If ``True`` (default), the check runs and
            its outcome is recorded in ``timings`` -- see the module docstring
            for why it is only ever recorded, never acted on here.

    Returns:
        A new :class:`~fpgm.physics.types.PosteriorResult`, always with
        ``refinement_rounds=0``.

    Raises:
        PhysicsError: shape mismatches, or every particle ended up with
            ``-inf`` total log-weight (no particle is consistent with the
            accumulated episodes; there is nothing to report a posterior over).
    """
    particles = np.asarray(particles, dtype=np.float64)
    if particles.ndim != 2 or particles.shape[1] != prior.space.dim:
        raise PhysicsError(f"particles must be (N, {prior.space.dim}), got {particles.shape}")
    n_particles = particles.shape[0]

    with _step(timer, "sum_episode_loglik", n=max(len(episodes), 1)):
        if episodes:
            for e in episodes:
                if e.loglik.shape[0] != n_particles:
                    raise PhysicsError(
                        f"episode {e.uuid}/{e.label}: loglik has {e.loglik.shape[0]} "
                        f"entries for {n_particles} particles"
                    )
            log_weights = np.sum(
                np.stack([_centred(e.loglik) for e in episodes]), axis=0
            )
        else:
            log_weights = np.zeros(n_particles, dtype=np.float64)

    return _finalize(
        prior,
        particles,
        log_weights,
        tuple(episodes),
        timer=timer,
        ess_min_frac=ess_min_frac,
        refine=refine,
    )


def pool(
    results: Sequence[PosteriorResult],
    *,
    timer: StepTimer | None = None,
    ess_min_frac: float = 0.05,
    refine: bool = True,
) -> PosteriorResult:
    """Pool posteriors that share a ``ParamSpace``/particle set/prior into one.

    Each input is presumed to already be an :func:`accumulate` (or
    :func:`pool`) output over *some* subset of episodes of the same object
    class -- e.g. one result per episode, later combined here, or one result
    per camera, or a partial batch checkpointed mid-run. Because each input's
    own ``log_weights`` is already ``sum`` over its own episodes, pooling is
    exactly one more elementwise sum of already-summed vectors -- the same
    associative operation the whole stage is built on, applied one level
    higher, so appending information stays a pure no-op for an uninformative
    input here too (a ``results`` entry whose ``log_weights`` is constant across
    particles pools away to nothing, by the same argument as an uninformative
    episode).

    Raises:
        PhysicsError: ``results`` is empty, or two inputs disagree on parameter
            names, prior mean/std, or the particle array itself. The particle
            check uses ``np.array_equal``, not ``np.allclose``: a genuinely
            shared particle draw is bit-identical by construction (the same
            ``np.ndarray``, or one round-tripped through storage unmodified);
            anything less than exact equality means two different draws were
            used, which :func:`accumulate`'s per-particle sum would silently
            misinterpret as one.
    """
    if not results:
        raise PhysicsError("pool() needs at least one PosteriorResult")

    first = results[0]
    for other in results[1:]:
        if other.space.names != first.space.names:
            raise PhysicsError(
                f"pool: ParamSpace mismatch {other.space.names} vs {first.space.names}"
            )
        if not np.array_equal(other.particles, first.particles):
            raise PhysicsError(
                "pool: particle arrays differ -- inputs are not the same "
                "importance-sampling run and cannot be summed"
            )
        if not np.array_equal(other.prior.mean, first.prior.mean) or not np.array_equal(
            other.prior.std, first.prior.std
        ):
            raise PhysicsError("pool: prior mean/std mismatch")

    with _step(timer, "pool_sum_weights", n=len(results)):
        # Centred for the same reason accumulate() centres -- see _centred. A
        # result whose own log_weights are constant (every episode in it was
        # uninformative) must pool as an exact no-op, not a nearly-one.
        log_weights = np.sum(np.stack([_centred(r.log_weights) for r in results]), axis=0)
        episodes = tuple(e for r in results for e in r.episodes)

    return _finalize(
        first.prior,
        first.particles,
        log_weights,
        episodes,
        timer=timer,
        ess_min_frac=ess_min_frac,
        refine=refine,
    )


def needs_refinement(result: PosteriorResult, *, ess_min_frac: float = 0.05) -> bool:
    """Whether ``result``'s ESS is low enough that another refinement round is warranted.

    A pure re-derivation from ``result.ess`` and the particle count -- does not
    depend on whatever ``ess_min_frac`` :func:`accumulate` used internally when
    it computed ``timings["refinement_warranted"]``, so a caller auditing an
    existing result with a different threshold does not need to re-run
    inference.
    """
    n_particles = result.particles.shape[0]
    return bool(result.ess < ess_min_frac * n_particles)


def resample_particles(result: PosteriorResult, rng: np.random.Generator, n: int) -> np.ndarray:
    """Draw ``n`` new particles: importance-resample ``result``, then add roughening jitter.

    Standard particle-filter degeneracy fix, in two steps:

    1. Multinomial resampling by ``result.weights`` concentrates the new
       particle set where the posterior mass actually is.
    2. Resampling alone produces many *exact duplicates* -- the same
       high-weight particle drawn repeatedly -- which would silently collapse
       the next round's particle diversity to a handful of point masses before
       it is even re-scored. The Gaussian jitter (std = ``_JITTER_FRAC`` times
       this result's own fitted posterior std, per parameter) is centred on
       each *resampled* particle, not on the posterior mean, so it perturbs
       particles locally instead of pulling them back toward the prior.

    This function does not simulate or re-score anything -- see the module
    docstring for why that boundary is where it is. The caller must build a
    fresh :class:`~fpgm.physics.types.SimSpec`/rollout and
    :class:`~fpgm.physics.types.EpisodeLogLik` per episode for the particles
    returned here, then call :func:`accumulate` again.
    """
    if n <= 0:
        raise PhysicsError(f"resample_particles: n must be > 0, got {n}")
    w = result.weights
    idx = rng.choice(w.shape[0], size=n, replace=True, p=w)
    base = result.particles[idx]
    post_std = np.array([p.post_std for p in result.params], dtype=np.float64)
    jitter = rng.standard_normal((n, post_std.shape[0])) * (_JITTER_FRAC * post_std)[None, :]
    return base + jitter


# --------------------------------------------------------------------------- #
# Shared statistics (accumulate and pool both reduce to this once they have a
# single (particles, log_weights, prior, episodes) tuple)
# --------------------------------------------------------------------------- #


def _finalize(
    prior: GaussianPrior,
    particles: np.ndarray,
    log_weights: np.ndarray,
    episodes: tuple[EpisodeLogLik, ...],
    *,
    timer: StepTimer | None,
    ess_min_frac: float,
    refine: bool,
) -> PosteriorResult:
    n_particles = particles.shape[0]

    with _step(timer, "normalize_weights", n=n_particles):
        # Numerically-stable softmax: subtract the max before exponentiating.
        # This is *not* the per-episode normalisation the module docstring
        # warns against -- it is applied once, after every episode's raw
        # log-likelihood has already been summed, and subtracting the same
        # scalar from the whole finished sum changes no ratio between
        # particles (unlike subtracting a different amount per episode before
        # summing, which would).
        top = np.max(log_weights)
        if not np.isfinite(top):
            raise PhysicsError(
                "every particle has -inf total log-weight: no particle is "
                "consistent with the accumulated episodes"
            )
        unnorm = np.exp(log_weights - top)
        w = unnorm / unnorm.sum()

    with _step(timer, "posterior_stats", n=n_particles):
        ess = float(1.0 / np.sum(w * w))
        post_mean = w @ particles
        centered = particles - post_mean[None, :]
        post_var = w @ (centered * centered)
        post_std = np.sqrt(np.maximum(post_var, 0.0))

        prior_var = prior.std * prior.std
        contraction = np.clip(1.0 - post_var / prior_var, 0.0, 1.0)

        params = tuple(
            ParamPosterior(
                name=name,
                prior_mean=float(prior.mean[i]),
                prior_std=float(prior.std[i]),
                post_mean=float(post_mean[i]),
                post_std=float(post_std[i]),
                contraction=float(contraction[i]),
                learned=bool(contraction[i] > _LEARNED_CONTRACTION),
            )
            for i, name in enumerate(prior.space.names)
        )

    with _step(timer, "info_nats", n=n_particles):
        info_nats = _snis_kl_estimate(w, n_particles)

    with _step(timer, "sensitivity_svd", n=n_particles):
        sensitivity_values, sensitivity_directions = _sensitivity(
            particles, post_mean, prior.std, w
        )

    warranted = bool(refine and ess < ess_min_frac * n_particles)
    timings: dict = {
        "n_particles": n_particles,
        "n_episodes": len(episodes),
        "ess_min_frac": ess_min_frac,
        "refinement_warranted": warranted,
    }
    if warranted:
        timings["refinement_note"] = (
            "ESS below ess_min_frac * N; accumulate()/pool() cannot re-score "
            "particles themselves (simulation lives outside fpgm.physics). "
            "Call needs_refinement(result) to confirm, draw new particles via "
            "resample_particles(result, rng, n), re-simulate and re-score them "
            "into fresh EpisodeLogLik objects, and call accumulate() again."
        )

    return PosteriorResult(
        space=prior.space,
        particles=particles,
        log_weights=log_weights,
        prior=prior,
        params=params,
        ess=ess,
        info_nats=info_nats,
        sensitivity_values=sensitivity_values,
        sensitivity_directions=sensitivity_directions,
        episodes=episodes,
        refinement_rounds=0,
        timings=timings,
    )


def _snis_kl_estimate(w: np.ndarray, n_particles: int) -> float:
    """Self-normalised-importance-sampling estimate of KL(posterior || prior), in nats.

    ``sum_i w_i * log(w_i * N)`` over normalised weights ``w`` (``sum w_i =
    1``) for particles drawn from the prior. This is the standard SNIS
    estimator: it needs no normalising constant for the (unnormalised) target
    density, only the particles and the weights already computed for the
    posterior mean/std above -- so it comes essentially for free.

    What this does *not* measure: the KL of the true continuous posterior, only
    of the particle-supported approximation to it, and its bias grows as ESS
    shrinks -- with few effectively-distinct particles the estimate is
    dominated by whichever handful happen to carry weight. That bias is not
    quantified here; ``PosteriorResult.ess`` is reported alongside specifically
    so a reader can judge how much to trust a given ``info_nats`` value rather
    than this function silently correcting for something it cannot measure.

    Particles with ``w == 0`` (including every particle behind an ``-inf``
    log-weight) contribute nothing, matching the ``0 * log(0) := 0`` convention
    for the entropy-like sum this is built from -- computed by masking them out
    rather than relying on ``0 * -inf`` evaluating to ``0`` in IEEE float (it
    does not; it is ``nan``).
    """
    mask = w > 0
    if not np.any(mask):
        return 0.0
    return float(np.sum(w[mask] * np.log(w[mask] * n_particles)))


def _sensitivity(
    particles: np.ndarray,
    post_mean: np.ndarray,
    prior_std: np.ndarray,
    w: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Which linear combinations of theta the data constrained, ranked most-informative first.

    Builds the posterior covariance in prior-std units -- ``d = sqrt(w) *
    (particles - post_mean) / prior_std``, so ``d.T @ d`` is the weighted
    covariance of the particles with every parameter on a comparable scale (a
    log-density unit and a friction-log unit are not otherwise commensurable) --
    then returns the eigendecomposition of ``I - d.T @ d``, not of ``d.T @ d``
    itself.

    That "``I -``" is the entire trick, and is worth stating explicitly because
    the more obvious thing to compute is wrong for this purpose: the SVD of the
    raw weighted covariance ranks directions by *leftover posterior spread*,
    which puts the *least*-constrained combination first -- backwards from
    "which directions did the data constrain". In prior-std units, a direction
    the likelihood pins down tightly has posterior variance near 0; a direction
    the likelihood never touches keeps the full prior variance, which is
    exactly 1 in these units. ``I - d.T @ d`` is therefore ~1 along a
    fully-identified combination and ~0 along an untouched one -- the vector
    generalisation of ``ParamPosterior.contraction = 1 - Var_post/Var_prior``,
    extended from the coordinate axes to arbitrary linear combinations by one
    eigendecomposition. This is exactly the measurement
    :class:`~fpgm.physics.types.PosteriorResult`'s own docstring describes:
    contact identifiability analyses *predict* that mass/friction/normal-load
    only enter observable dynamics through certain combinations; this reports
    whichever combinations this data actually constrained, without assuming in
    advance which ones those are.

    Eigenvalues are clipped to ``[0, inf)`` before sorting: a direction with
    true contraction ~0 can come out a few ULPs negative from finite-particle
    sampling noise, and "a negative amount constrained" is not reportable.
    ``np.linalg.eigh`` (not ``svd``) is used because ``I - d.T @ d`` is
    symmetric by construction, so its eigendecomposition already *is* its SVD
    with non-negative singular values -- ``eigh`` is the cheaper, better
    conditioned solver for a symmetric matrix.

    Returns:
        ``(values, directions)``: ``values`` is ``(D,)`` descending;
        ``directions`` is ``(D, D)`` with each row a unit eigenvector,
        expressed in prior-std-normalised coordinates (component ``j`` of a row
        is "how many prior-sigmas of parameter ``j``", so e.g. a row
        ``(0.71, 0.71, 0, ..., 0)`` reads as "the sum of parameters 0 and 1, in
        prior-std units, is the identified combination").
    """
    dim = particles.shape[1]
    sqrt_w = np.sqrt(w)
    scaled = (particles - post_mean[None, :]) / prior_std[None, :]
    d = sqrt_w[:, None] * scaled
    post_cov_scaled = d.T @ d
    sensitivity_matrix = np.eye(dim) - post_cov_scaled

    eigvals, eigvecs = np.linalg.eigh(sensitivity_matrix)
    order = np.argsort(eigvals)[::-1]
    values = np.clip(eigvals[order], 0.0, None)
    directions = eigvecs[:, order].T
    return values, directions