tests/test_physics_inference.py

"""Tests for :mod:`fpgm.physics.inference`.

GPU-free, no downloads, no mujoco: particles, priors and per-episode
log-likelihoods are all synthetic. The most important test in this file is
:func:`test_null_episode_is_bitidentical_alone_and_appended` -- it is the
executable proof of the stage's central invariant (an uninformative episode is
an exact softmax no-op), which is the entire justification for having no
episode filter anywhere in this stage.
"""

from __future__ import annotations

import numpy as np
import pytest

from fpgm.physics.inference import accumulate, needs_refinement, pool, resample_particles
from fpgm.physics.types import EpisodeLogLik, GaussianPrior, ParamSpace, PhysicsError


def _prior(dim: int = 4) -> GaussianPrior:
    space = ParamSpace(tuple(f"p{i}" for i in range(dim)))
    return GaussianPrior(space=space, mean=np.zeros(dim), std=np.ones(dim))


def _particles(prior: GaussianPrior, n: int, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    return prior.sample(rng, n)


def _episode(loglik: np.ndarray, *, uuid: str = "ep", label: str = "obj", n_obs_frames: int = 5):
    return EpisodeLogLik(
        uuid=uuid, camera_serial="cam0", label=label, loglik=loglik, n_obs_frames=n_obs_frames
    )


def _assert_bitidentical(a, b) -> None:
    assert np.array_equal(a.weights, b.weights)
    assert a.ess == b.ess
    assert a.info_nats == b.info_nats
    assert np.array_equal(a.sensitivity_values, b.sensitivity_values)
    assert np.array_equal(a.sensitivity_directions, b.sensitivity_directions)
    assert len(a.params) == len(b.params)
    for pa, pb in zip(a.params, b.params, strict=True):
        assert pa.name == pb.name
        assert pa.prior_mean == pb.prior_mean
        assert pa.prior_std == pb.prior_std
        assert pa.post_mean == pb.post_mean
        assert pa.post_std == pb.post_std
        assert pa.contraction == pb.contraction
        assert pa.learned == pb.learned


# --------------------------------------------------------------------------- #
# The null test -- the whole point of the stage
# --------------------------------------------------------------------------- #


def test_null_episode_is_bitidentical_alone_and_appended():
    prior = _prior(dim=4)
    particles = _particles(prior, n=5000, seed=1)
    n = particles.shape[0]

    # Integer-valued so every addition/subtraction below is exact in float64
    # (no rounding at all) -- this test is about the arithmetic invariant, not
    # about floating-point luck, so it is built to make luck irrelevant.
    informative = _episode(-(np.arange(n) % 401).astype(np.float64), uuid="informative")

    for const in (0.0, 500.0, -1_000_000.0):
        null_ep = _episode(np.full(n, const), uuid="null", n_obs_frames=0)

        base_empty = accumulate(prior, particles, [])
        with_null_only = accumulate(prior, particles, [null_ep])
        _assert_bitidentical(base_empty, with_null_only)

        base_informative = accumulate(prior, particles, [informative])
        with_null_appended = accumulate(prior, particles, [informative, null_ep])
        _assert_bitidentical(base_informative, with_null_appended)


def test_null_episode_is_bitidentical_with_real_valued_logliks():
    """The invariant on log-likelihoods that look like real ones, not integers.

    ``test_null_episode_is_bitidentical_alone_and_appended`` deliberately uses
    integer-valued log-likelihoods so that float64 addition is exact, isolating
    the *algebra*. That is a fair test of the algebra and a poor test of the
    implementation: with the ordinary irrational values a Student-t likelihood
    actually produces, summing raw log-likelihoods of very different magnitudes
    loses the low-order bits of the informative terms, and the pre-``_centred``
    implementation drifted by 4.7e-11 relative at a constant of -1e6 and 8.0e-08
    at -1e9 -- small, but "exactly unchanged" was the claim, and a claim that
    decays with the magnitude of an arbitrary constant is a different, weaker
    claim.

    Constants here span 12 orders of magnitude and include a deliberately
    irrational one, because the failure this guards against is precisely the
    one that hides behind values that happen to add exactly.
    """
    prior = _prior(dim=4)
    particles = _particles(prior, n=512, seed=11)

    # Two genuinely informative episodes: one constrains p1 on its own, the
    # other only the combination p0 + p1, so the posterior is non-trivial in
    # both location and shape and a drift would show up somewhere.
    informative = [
        _episode(-0.5 * ((particles[:, 1] - 0.7) / 0.2) ** 2, uuid="peaked"),
        _episode(-0.5 * ((particles[:, 0] + particles[:, 1] - 1.0) / 0.3) ** 2, uuid="ridge"),
    ]
    base = accumulate(prior, particles, informative)

    for const in (0.0, -1e3, -1e6, 3.7e4, -1e9, -1e15, float(np.pi) * 1e7):
        null_ep = _episode(np.full(particles.shape[0], const), uuid="static", n_obs_frames=0)
        with_null = accumulate(prior, particles, [informative[0], null_ep, informative[1]])
        _assert_bitidentical(base, with_null)


def test_pooling_an_entirely_uninformative_result_is_bitidentical():
    """``pool`` carries the same guarantee as ``accumulate``.

    A per-object posterior whose every episode was uninformative has constant
    ``log_weights``; pooling it with a real one must be an exact no-op, or the
    invariant would hold per-episode and then leak away at the pooling step --
    which is exactly where a multi-episode run spends its time.
    """
    prior = _prior(dim=4)
    particles = _particles(prior, n=512, seed=12)
    informative = accumulate(
        prior, particles, [_episode(-0.5 * ((particles[:, 2] - 0.4) / 0.25) ** 2)]
    )
    all_null = accumulate(
        prior, particles, [_episode(np.full(particles.shape[0], -4.2e8), n_obs_frames=0)]
    )

    _assert_bitidentical(pool([informative]), pool([informative, all_null]))


# --------------------------------------------------------------------------- #
# Contraction semantics
# --------------------------------------------------------------------------- #


def test_flat_likelihood_gives_near_zero_contraction_and_posterior_near_prior():
    prior = _prior(dim=3)
    particles = _particles(prior, n=20_000, seed=5)
    flat = _episode(np.zeros(particles.shape[0]), n_obs_frames=0)

    result = accumulate(prior, particles, [flat])
    for p in result.params:
        assert abs(p.contraction) < 0.05
        assert abs(p.post_mean - p.prior_mean) < 0.05
        assert abs(p.post_std - p.prior_std) < 0.05


def test_peaked_likelihood_gives_high_contraction_only_on_its_own_dimension():
    prior = _prior(dim=3)
    particles = _particles(prior, n=20_000, seed=6)
    peaked = _episode(-0.5 * (particles[:, 0] / 0.05) ** 2)

    result = accumulate(prior, particles, [peaked])
    assert result.params[0].contraction > 0.9
    assert result.params[0].post_std < 0.1
    assert result.params[1].contraction < 0.1
    assert result.params[2].contraction < 0.1


# --------------------------------------------------------------------------- #
# -inf handling
# --------------------------------------------------------------------------- #


def test_diverged_particles_get_zero_weight_without_nan():
    prior = _prior(dim=2)
    particles = _particles(prior, n=2000, seed=7)
    n = particles.shape[0]
    loglik = -0.5 * particles[:, 0] ** 2
    diverged = np.arange(0, n, 7)
    loglik[diverged] = -np.inf

    result = accumulate(prior, particles, [_episode(loglik)])
    assert not np.any(np.isnan(result.weights))
    assert np.all(np.isfinite(result.weights))
    assert np.all(result.weights[diverged] == 0.0)
    assert np.isfinite(result.ess)
    for p in result.params:
        assert np.isfinite(p.post_mean)
        assert np.isfinite(p.post_std)


def test_all_particles_diverged_raises():
    prior = _prior(dim=2)
    particles = _particles(prior, n=100, seed=8)
    ep = _episode(np.full(100, -np.inf))
    with pytest.raises(PhysicsError):
        accumulate(prior, particles, [ep])


# --------------------------------------------------------------------------- #
# Sensitivity: detecting a theta_a + theta_b degeneracy
# --------------------------------------------------------------------------- #


def test_sensitivity_detects_ab_combination():
    prior = _prior(dim=4)
    particles = _particles(prior, n=40_000, seed=9)
    combo = particles[:, 0] + particles[:, 1]
    peaked = _episode(-0.5 * ((combo - 1.0) / 0.05) ** 2)

    result = accumulate(prior, particles, [peaked])
    top_direction = result.sensitivity_directions[0]
    target = np.array([1.0, 1.0, 0.0, 0.0]) / np.sqrt(2.0)
    cosine = abs(float(np.dot(top_direction, target)))

    assert cosine > 0.9
    assert result.sensitivity_values[0] > result.sensitivity_values[1]
    assert result.sensitivity_values[0] > result.sensitivity_values[2]
    assert result.sensitivity_values[0] > result.sensitivity_values[3]


# --------------------------------------------------------------------------- #
# ESS bounds
# --------------------------------------------------------------------------- #


def test_ess_bounds():
    prior = _prior(dim=2)
    particles = _particles(prior, n=5000, seed=10)
    n = particles.shape[0]

    uniform = accumulate(prior, particles, [_episode(np.zeros(n), n_obs_frames=0)])
    assert uniform.ess == pytest.approx(float(n), rel=1e-9)

    concentrated_ll = np.full(n, -1e9)
    concentrated_ll[0] = 0.0
    concentrated = accumulate(prior, particles, [_episode(concentrated_ll)])
    assert concentrated.ess == pytest.approx(1.0, abs=1e-9)


# --------------------------------------------------------------------------- #
# Pooling
# --------------------------------------------------------------------------- #


def test_pool_matches_direct_accumulate_over_the_same_episodes():
    prior = _prior(dim=2)
    particles = _particles(prior, n=3000, seed=13)
    ep_a = _episode(-0.5 * particles[:, 0] ** 2, uuid="a")
    ep_b = _episode(-0.5 * (particles[:, 1] - 0.5) ** 2, uuid="b")

    result_a = accumulate(prior, particles, [ep_a])
    result_b = accumulate(prior, particles, [ep_b])
    pooled = pool([result_a, result_b])
    direct = accumulate(prior, particles, [ep_a, ep_b])

    assert np.array_equal(pooled.weights, direct.weights)
    assert pooled.ess == direct.ess
    assert len(pooled.episodes) == 2


def test_pool_raises_on_particle_mismatch():
    prior = _prior(dim=2)
    particles1 = _particles(prior, n=500, seed=11)
    particles2 = _particles(prior, n=500, seed=12)
    ep = _episode(np.zeros(500))

    r1 = accumulate(prior, particles1, [ep])
    r2 = accumulate(prior, particles2, [ep])
    with pytest.raises(PhysicsError):
        pool([r1, r2])


def test_pool_raises_on_empty_input():
    with pytest.raises(PhysicsError):
        pool([])


# --------------------------------------------------------------------------- #
# Degeneracy guard / refinement reporting split
# --------------------------------------------------------------------------- #


def test_low_ess_is_reported_but_not_silently_refined():
    prior = _prior(dim=3)
    particles = _particles(prior, n=2000, seed=14)
    n = particles.shape[0]
    loglik = np.full(n, -1e9)
    loglik[:5] = 0.0
    ep = _episode(loglik)

    result = accumulate(prior, particles, [ep], ess_min_frac=0.5)
    assert result.refinement_rounds == 0  # never claims work it structurally cannot do
    assert result.timings["refinement_warranted"] is True
    assert needs_refinement(result, ess_min_frac=0.5) is True

    rng = np.random.default_rng(0)
    new_particles = resample_particles(result, rng, 500)
    assert new_particles.shape == (500, 3)
    assert np.all(np.isfinite(new_particles))


def test_refine_false_reports_not_warranted_regardless_of_ess():
    prior = _prior(dim=2)
    particles = _particles(prior, n=1000, seed=15)
    n = particles.shape[0]
    loglik = np.full(n, -1e9)
    loglik[0] = 0.0
    ep = _episode(loglik)

    result = accumulate(prior, particles, [ep], refine=False)
    assert result.timings["refinement_warranted"] is False
    assert "refinement_note" not in result.timings
    # needs_refinement is a pure re-derivation and is unaffected by refine=False
    assert needs_refinement(result) is True