tests/test_method_spec_gaps.py

"""Literature-anchored specification tests — step-8 gap closure.

Complements tests/test_method_spec.py with the surfaces the step-8 audit
found untested in this repo: the reverse-posterior/remasking/carry-over
chain, the final greedy cleanup, decode temperature, label smoothing and
the loss weight clip. Every expected value derives from a cited source
or a derivation written in the docstring — never from the current output
of this or the sibling repo. The minihack twin file carries the same
chain/cleanup/temperature/smoothing assertions with the same inputs and
tolerances (weight clip and empty-mask zero already exist there).

References as in tests/test_method_spec.py (MDLM arXiv:2406.07524;
ReMDM arXiv:2503.00307; Holtzman arXiv:1904.09751).
"""

from __future__ import annotations

import math

import jax
import jax.numpy as jnp
import numpy as np
import pytest

from src.diffusion.loss import compute_loss
from src.diffusion.sampling import _decode, sample_plan
from src.diffusion.schedules import SCHEDULE_MAP

ATOL = 1e-6


# ---------------------------------------------------------------------------
# Reverse posterior + remasking + carry-over, jointly, through sample_plan
# ---------------------------------------------------------------------------


def _time_coded_apply(num_actions: int, horizon: int):
    """Stub model whose argmax token encodes the decode time.

    token 0 for t > 0.9, token 1 for 0.5 < t <= 0.9, token 2 otherwise.
    Deterministic under temperature=0 (argmax decode).
    """

    def apply_fn(params, obs, z, t, rng):
        idx = jnp.where(t[0] > 0.9, 0, jnp.where(t[0] > 0.5, 1, 2))
        logits = jax.nn.one_hot(idx, num_actions) * 30.0
        return jnp.broadcast_to(logits, (obs.shape[0], horizon, num_actions))

    return apply_fn


@pytest.mark.parametrize("eta", [0.0, 0.5])
def test_posterior_remask_carryover_token_distribution(eta):
    """Final token distribution matches the ReMDM Alg 1 posterior chain.

    Source: ReMDM eq (6) masked branch (unmask probability
    (alpha_s - (1-sigma) alpha_t) / (1 - alpha_t)), eq (7) sigma_max,
    Sec 4.1 rescale sigma = eta * sigma_max, and MDLM Sec 3.2.3 carry-over.

    Derivation (linear schedule, K=3, grid t=1,2/3,1/3 with s=2/3,1/3,0;
    the stub decodes token 0 at t=1, token 1 at t=2/3, token 2 at
    t<=1/3):
      step 1 (t=1):   alpha_t=0, alpha_s=1/3, p_unmask=1/3 -> token 0.
      step 2 (t=2/3): alpha_t=1/3, alpha_s=2/3, sigma_max=1, sigma=eta;
                      committed remask w.p. eta; masked unmask w.p.
                      (2/3-(1-eta)/3)/(2/3) = (1+eta)/2 -> token 1.
      step 3 (t=1/3): alpha_s=1 -> sigma_max=0 (no remask), p_unmask=1
                      -> every remaining mask becomes token 2.
    Hence P(0) = (1-eta)/3, P(1) = (1+eta)/3, P(2) = 1/3.
    eta=0 additionally pins carry-over: a committed token is never
    re-decided when sigma=0.

    Statistical: N = 512*32 = 16384 independent per-position outcomes;
    max sigma = sqrt(0.25/16384) = 0.0039; bound 0.02 = 5.1 sigma.
    """
    B, H, V = 512, 32, 3
    fn, _ = SCHEDULE_MAP["linear"]
    plan = sample_plan(
        _time_coded_apply(V, H), None, jax.random.PRNGKey(0),
        jnp.zeros((B, 4)), V, H, num_steps=3, schedule_fn=fn,
        remask_strategy="rescale", eta=eta, use_loop=False,
        temperature=0.0, top_p=None,
    )
    freq = np.bincount(np.asarray(plan).ravel(), minlength=V) / (B * H)
    expected = np.array([(1 - eta) / 3, (1 + eta) / 3, 1 / 3])
    assert np.all(np.abs(freq - expected) < 0.02), (freq, expected)


def test_final_cleanup_commits_all_remaining_masks():
    """With zero denoising steps, the final greedy cleanup decodes every
    position at t=0 (argmax), leaving no MASK token.

    Source: spec-method 4.9 (final-step commit of remaining masks is the
    project safety net; craftax cleanup is unconditional). The stub
    decodes token 2 at t=0, so the output must be all-2.
    """
    B, H, V = 8, 16, 3
    fn, _ = SCHEDULE_MAP["linear"]
    plan = sample_plan(
        _time_coded_apply(V, H), None, jax.random.PRNGKey(0),
        jnp.zeros((B, 4)), V, H, num_steps=0, schedule_fn=fn,
        remask_strategy="rescale", eta=0.0, use_loop=False,
        temperature=0.0, top_p=None,
    )
    assert (np.asarray(plan) == 2).all()


# ---------------------------------------------------------------------------
# Decode temperature
# ---------------------------------------------------------------------------


def test_decode_temperature_scales_logits_before_sampling():
    """Sampling frequencies follow softmax(logits / temperature).

    Source: spec-method 5.2 (softmax temperature before filtering;
    standard technique). Derivation: logits [0, 1] at temperature 0.5
    give softmax([0, 2]) = [1/(1+e^2), e^2/(1+e^2)] = [0.1192, 0.8808].
    Statistical: 8192 draws; sigma = sqrt(0.8808*0.1192/8192) = 0.00358;
    bound 0.0143 = 4 sigma.
    """
    logits = jnp.broadcast_to(jnp.array([0.0, 1.0]), (8192, 1, 2))
    tokens = np.asarray(
        _decode(jax.random.PRNGKey(0), logits, temperature=0.5, top_p=None)
    ).ravel()
    p1 = math.exp(2) / (1 + math.exp(2))
    assert abs(tokens.mean() - p1) < 0.0143


# ---------------------------------------------------------------------------
# Label smoothing
# ---------------------------------------------------------------------------


def test_label_smoothing_matches_closed_form():
    """Smoothed CE = -[(1-eps)+eps/V] log p_true - (eps/V) sum log p_other.

    Source: spec-method 3.7 (smoothing target (1-eps)*onehot + eps/V;
    eps=0 is the exact ELBO). Derivation: model probs [0.7,0.1,0.1,0.1],
    true class 0, V=4, eps=0.3: coefficient on -log 0.7 is
    (1-0.3)+0.3/4 = 0.775; each other class gets 0.3/4 = 0.075.
    t pinned to 1 (linear, w=1, everything masked) makes the loss equal
    that CE exactly. Same inputs and expectations as the minihack twin.
    """
    V, H, B = 4, 8, 4
    fn, deriv = SCHEDULE_MAP["linear"]
    probs = jnp.array([0.7, 0.1, 0.1, 0.1])

    def apply_fn(params, obs, z, t, rng):
        return jnp.broadcast_to(jnp.log(probs), (obs.shape[0], H, V))

    x0 = jnp.zeros((B, H), dtype=jnp.int32)
    obs = jnp.zeros((B, 3))
    expected = -0.775 * math.log(0.7) - 0.075 * 3 * math.log(0.1)
    loss, _ = compute_loss(
        apply_fn, None, jax.random.PRNGKey(0), x0, obs, jnp.ones(B), V,
        fn, deriv, label_smoothing=0.3, t_min=1.0, t_max=1.0,
    )
    assert abs(float(loss) - expected) < 1e-5
    loss0, _ = compute_loss(
        apply_fn, None, jax.random.PRNGKey(0), x0, obs, jnp.ones(B), V,
        fn, deriv, label_smoothing=0.0, t_min=1.0, t_max=1.0,
    )
    assert abs(float(loss0) - (-math.log(0.7))) < 1e-5


# ---------------------------------------------------------------------------
# Loss weight clip and the empty-mask edge case (minihack twins exist)
# ---------------------------------------------------------------------------


def test_loss_weight_clip_bound():
    """w(t) is clipped at 1000 (project numerics guard, spec-method 3.4;
    the same _MAX_WEIGHT the minihack twin pins via an explicit zt).

    compute_loss samples zt internally, so this is statistical: at
    t=1e-4 (linear) the raw weight is 1/t = 10^4, clipped to 10^3.
    E[loss] = 1000 * ln V * P(mask) = 1000 * ln 5 * 1e-4 = 0.1 ln 5.
    The unclipped hypothesis gives 1.0 ln 5 (10x larger). Masked count
    over N = 131072*8 positions is ~Poisson(104.9); sigma(loss) =
    1000 * lnV * sqrt(104.9)/N = 0.0098 lnV; bound 0.04 lnV = 4.1 sigma.
    """
    V, H, B = 5, 8, 131072
    fn, deriv = SCHEDULE_MAP["linear"]

    def apply_fn(params, obs, z, t, rng):
        return jnp.zeros((obs.shape[0], H, V))

    x0 = jnp.zeros((B, H), dtype=jnp.int32)
    loss, _ = compute_loss(
        apply_fn, None, jax.random.PRNGKey(0), x0, jnp.zeros((B, 1)),
        jnp.ones(B), V, fn, deriv, t_min=1e-4, t_max=1e-4,
    )
    assert abs(float(loss) - 0.1 * math.log(V)) < 0.04 * math.log(V)


def test_loss_zero_when_nothing_masked():
    """Loss is exactly 0.0 when the batch has no masked positions.

    Source: spec-method 3.4 (zero-on-empty project convention; the
    minihack twin asserts the same through an explicit unmasked zt).
    At t=0, alpha=1 and the forward process keeps every token
    (uniform draws in [0,1) are always < 1), so nothing is masked.
    """
    V, H, B = 5, 8, 4
    fn, deriv = SCHEDULE_MAP["linear"]

    def apply_fn(params, obs, z, t, rng):
        return jnp.zeros((obs.shape[0], H, V))

    x0 = jnp.zeros((B, H), dtype=jnp.int32)
    loss, _ = compute_loss(
        apply_fn, None, jax.random.PRNGKey(0), x0, jnp.zeros((B, 1)),
        jnp.ones(B), V, fn, deriv, t_min=0.0, t_max=0.0,
    )
    assert float(loss) == 0.0