giant/tests/test_sample.py

"""Tests for giant/sample.py's v0.3.0 stage-model sampling — the AR loop
(`sample_secondaries_ar`) and non-"physical" `particle_type.target` coverage
for the one-shot samplers."""

import pytest
import torch

from giant.config import ConditioningAxisConfig, ParticleTypeConfig
from giant.constants import COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, X_DIM
from giant.model.network import (
    Stage1Model,
    Stage2Autoregressive,
    Stage2OneShot,
    stage2_trunk_sec_dim,
)
from giant.sample import (
    resolve_n_sec,
    sample_flow,
    sample_secondaries,
    sample_secondaries_ar,
    sample_secondaries_wgan,
    sample_wgan,
)

_PHYS_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)


def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[ConditioningAxisConfig, ConditioningAxisConfig]:
    cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
    return cfg, cfg


def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
    cond_cont = torch.randn(B, COND_DIM)
    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
    return cond_cont, cond_cat


def _conditioning_for(target: str) -> str:
    # target="embedding" regresses against the conditioning's own embedding
    # table — only meaningful when the
    # conditioning axis is itself "embedding".
    return "embedding" if target == "embedding" else "physical"


def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot:
    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
    particle_type_cfg = ParticleTypeConfig(target=target)
    # build_models (giant/model/network.py) computes sec_dim this same way
    # before constructing Stage2OneShot — its own default (SEC_DIM, the
    # "physical" width) is only correct for target="physical".
    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, K_MAX, emb_dim)
    return Stage2OneShot(
        pdg_vocab=pdg,
        mat_vocab=mat,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        generator=generator,
        time_dim=16,
        noise_dim=8,
        sec_dim=sec_dim,
        particle_type_cfg=particle_type_cfg,
    ).eval()


def _stage2_ar(
    target: str,
    generator: str,
    emb_dim: int = 6,
    pdg: int = 3,
    mat: int = 2,
    k_max: int = 5,
    history: str = "markov",
    n_sec_sampling: str = "greedy",
) -> Stage2Autoregressive:
    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
    return Stage2Autoregressive(
        pdg_vocab=pdg,
        mat_vocab=mat,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        generator=generator,
        time_dim=16,
        noise_dim=8,
        k_max=k_max,
        particle_type_cfg=ParticleTypeConfig(target=target),
        history=history,
        attn_n_heads=2,
        attn_n_layers=1,
        n_sec_sampling=n_sec_sampling,
    ).eval()


def _expected_type_dim(target: str, emb_dim: int) -> int:
    return PARTICLE_PHYS_DIM if target == "physical" else emb_dim


def _stage2_ar_stop_token(
    target: str,
    generator: str,
    n_sec_sampling: str = "greedy",
    emb_dim: int = 6,
    pdg: int = 3,
    mat: int = 2,
    k_max: int = 5,
) -> Stage2Autoregressive:
    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
    return Stage2Autoregressive(
        pdg_vocab=pdg,
        mat_vocab=mat,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        generator=generator,
        time_dim=16,
        noise_dim=8,
        k_max=k_max,
        particle_type_cfg=ParticleTypeConfig(target=target),
        build_n_sec_head=False,
        build_stop_head=True,
        n_sec_sampling=n_sec_sampling,
    ).eval()


def _force_stop_head_logit(decoder: Stage2Autoregressive, logit: float) -> None:
    """Zeroes stop_head's weights and pins its bias, so predict_stop returns
    `logit` for every row/slot regardless of conditioning — makes the AR
    loop's stop decision deterministic for testing."""
    assert decoder.stop_head is not None
    last_linear = decoder.stop_head[-1]
    with torch.no_grad():
        last_linear.weight.zero_()
        last_linear.bias.fill_(logit)


# ── Stage-1 n_sec ownership ──────────────────────────────────────────────────


def test_sample_flow_returns_none_n_sec_when_stage1_owns_no_head():
    model = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PHYS_CFG,
        material_cfg=_PHYS_CFG,
        hidden_dim=16,
        n_res_blocks=1,
    )
    cond_cont, cond_cat = _cond(4)
    sample, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
    assert sample.shape == (4, X_DIM)
    assert n_sec is None


def test_sample_wgan_returns_none_n_sec_when_stage1_owns_no_head():
    model = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PHYS_CFG,
        material_cfg=_PHYS_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        generator="wgan",
        noise_dim=8,
    )
    cond_cont, cond_cat = _cond(4)
    sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
    assert sample.shape == (4, X_DIM)
    assert n_sec is None


def test_sample_flow_returns_n_sec_for_legacy_stage1():
    model = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=_PHYS_CFG,
        material_cfg=_PHYS_CFG,
        hidden_dim=16,
        n_res_blocks=1,
        n_sec_head_k_max=K_MAX,
    )
    cond_cont, cond_cat = _cond(5)
    _, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
    assert n_sec is not None and n_sec.shape == (5,)


# ── Stage2OneShot: non-"physical" particle_type.target ──────────────────────


@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_flow_shapes_by_target(target):
    B, emb_dim = 5, 6
    decoder = _stage2_oneshot(target, "flow", emb_dim=emb_dim)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
    sec_cont, sec_type, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
    assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
    assert sec_valid.shape == (B, K_MAX)
    assert torch.isfinite(sec_cont).all()
    assert torch.isfinite(sec_type).all()


@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_wgan_shapes_by_target(target):
    B, emb_dim = 5, 6
    decoder = _stage2_oneshot(target, "wgan", emb_dim=emb_dim)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
    sec_cont, sec_type, sec_valid = sample_secondaries_wgan(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
    assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
    assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
    assert sec_valid.shape == (B, K_MAX)


# ── Stage2Autoregressive ─────────────────────────────────────────────────────


@pytest.mark.parametrize("history", ["markov", "attention"])
@pytest.mark.parametrize("generator", ["flow", "wgan"])
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_ar_shapes(target, generator, history):
    B, k_max, emb_dim = 4, 5, 6
    decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max, history=history)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.randint(0, k_max + 1, (B,))
    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, k_max, CONT_SLOT_DIM)
    assert sec_type.shape == (B, k_max, _expected_type_dim(target, emb_dim))
    assert sec_valid.shape == (B, k_max)
    assert torch.isfinite(sec_cont).all()
    assert torch.isfinite(sec_type).all()


@pytest.mark.parametrize("generator", ["flow", "wgan"])
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
    B, k_max, emb_dim = 3, 5, 6
    decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 2, k_max])
    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
    for i, n in enumerate(n_sec_pred.tolist()):
        assert sec_valid[i, :n].all()
        assert not sec_valid[i, n:].any()


def test_sample_secondaries_ar_first_slot_has_no_history():
    """Slot 0 always has has_prev=False internally — nothing to assert on
    the public API directly, but a k_max=1 run should not crash on the
    "previous token" path at all (has_prev never true)."""
    B, emb_dim = 3, 6
    decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=1)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 1, 1])
    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
    assert sec_valid.tolist() == [[False], [True], [True]]


# ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────


@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_sampling):
    """A stop_head pinned to a large positive logit fires at slot 0 for
    every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
    a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
    generating any token."""
    B, k_max = 4, 5
    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
    _force_stop_head_logit(decoder, 50.0)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
    assert sec_valid.shape == (B, k_max)
    assert not sec_valid.any()


@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_sampling):
    """A stop_head pinned to a large negative logit never fires under either
    policy, so every row is capped at k_max (the safety cap, not a modeling
    ceiling)."""
    B, k_max = 4, 5
    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
    _force_stop_head_logit(decoder, -50.0)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
    assert sec_valid.all()


@pytest.mark.parametrize("generator", ["flow", "wgan"])
def test_sample_secondaries_ar_stop_token_valid_mask_is_always_a_prefix(generator):
    """Without forcing the stop head, per-row stop timing varies — but
    sec_valid must always be a contiguous prefix (slot k valid implies every
    slot < k is also valid), matching the "head"/"truth" contract."""
    B, k_max = 6, 5
    decoder = _stage2_ar_stop_token("physical", generator, k_max=k_max)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
    n = sec_valid.sum(dim=-1)
    expected = torch.arange(k_max).unsqueeze(0) < n.unsqueeze(1)
    assert torch.equal(sec_valid, expected)


def test_sample_secondaries_ar_stop_token_explicit_n_sec_pred_ignores_stop_head():
    """The scheduled-sampling training contract: passing n_sec_pred
    explicitly (as _assemble_stage2_ar_inputs_scheduled's self-sample call
    does, with ground-truth n_sec) must run the full k_max loop and mask by
    the given count, even though the decoder owns a stop_head that would
    otherwise stop early."""
    B, k_max = 3, 5
    decoder = _stage2_ar_stop_token("physical", "flow", k_max=k_max)
    _force_stop_head_logit(decoder, 50.0)  # would stop immediately if consulted
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 2, k_max])
    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
    for i, n in enumerate(n_sec_pred.tolist()):
        assert sec_valid[i, :n].all()
        assert not sec_valid[i, n:].any()


def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
    decoder = _stage2_ar("physical", "flow", k_max=5)  # head mode: no stop_head
    cond_cont, cond_cat = _cond(3)
    stage1_out = torch.randn(3, X_DIM)
    with pytest.raises(AssertionError):
        sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)


# ── Row compaction: full_length=False (default, inference) must agree with
# full_length=True (the pre-compaction behaviour, still exercised by
# _assemble_stage2_ar_inputs_scheduled's training-time self-sample) ────────


def _zero_randn(*size, **kwargs):
    """Drop-in replacement for `torch.randn` that returns zeros of the same
    shape — makes the ODE/WGAN noise deterministic so a compacted run and a
    full_length run can be compared row-for-row regardless of how many
    `torch.randn` calls each makes (compaction changes the batch size, and
    therefore the RNG stream position, at every slot)."""
    device = kwargs.get("device")
    dtype = kwargs.get("dtype")
    return torch.zeros(*size, device=device, dtype=dtype)


@pytest.mark.parametrize("history", ["markov", "attention"])
@pytest.mark.parametrize("generator", ["flow", "wgan"])
def test_sample_secondaries_ar_compaction_matches_full_length_head_mode(generator, history, monkeypatch):
    B, k_max, emb_dim = 4, 5, 6
    decoder = _stage2_ar("physical", generator, emb_dim=emb_dim, k_max=k_max, history=history)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 1, 3, k_max])

    monkeypatch.setattr(torch, "randn", _zero_randn)
    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=False
    )
    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
    )

    assert torch.equal(sec_valid_c, sec_valid_f)
    assert torch.equal(sec_valid_c, torch.arange(k_max).unsqueeze(0) < n_sec_pred.unsqueeze(1))
    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)


@pytest.mark.parametrize("generator", ["flow", "wgan"])
def test_sample_secondaries_ar_compaction_matches_full_length_stop_token(generator, monkeypatch):
    """`n_sec_sampling="greedy"` keeps the stop decision itself deterministic
    (no `torch.rand` draw), so only `torch.randn` needs zeroing."""
    B, k_max = 6, 5
    decoder = _stage2_ar_stop_token("physical", generator, n_sec_sampling="greedy", k_max=k_max)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)

    monkeypatch.setattr(torch, "randn", _zero_randn)
    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=False
    )
    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=True
    )

    assert torch.equal(sec_valid_c, sec_valid_f)
    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)


def test_sample_secondaries_ar_full_length_ignores_n_sec_pred_zero_rows():
    """A row with n_sec_pred == 0 would be dropped from the active set at
    slot 0 under compaction (full_length=False) — full_length=True must
    still run the model for it at every slot (only masked by sec_valid at
    the end), matching _assemble_stage2_ar_inputs_scheduled's contract."""
    B, k_max, emb_dim = 3, 4, 6
    decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=k_max)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 0, 0])
    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
    )
    assert not sec_valid.any()
    # every slot still ran the model (not left at the zero-init default) —
    # a real flow ODE output from randn-initialized noise is essentially
    # never exactly zero.
    assert not torch.allclose(sec_cont, torch.zeros_like(sec_cont))


# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────


def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
    """Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
    returns `bias` (broadcast over the batch) as logits regardless of
    conditioning — mirrors `_force_stop_head_logit`."""
    assert decoder.n_sec_head is not None
    last_linear = decoder.n_sec_head[-1]
    with torch.no_grad():
        last_linear.weight.zero_()
        last_linear.bias.copy_(bias)


@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
    """A logit vector overwhelmingly favoring one class gives the same
    answer under both policies — greedy because it's the argmax, sample
    because softmax puts ~all mass on it."""
    B, k_max = 8, 5
    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
    bias = torch.full((k_max + 1,), -50.0)
    bias[2] = 50.0
    _force_n_sec_head_bias(decoder, bias)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
    assert n_sec is not None
    assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))


def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
    B, k_max = 32, 5
    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
    assert n_sec is not None
    assert n_sec.unique().numel() == 1


def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
    """Under a flat logit vector, a categorical draw across a large batch
    should hit more than one class — the whole point of gitea #86: greedy
    always collapses to one, sample should not."""
    torch.manual_seed(0)
    B, k_max = 256, 5
    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
    assert n_sec is not None
    assert n_sec.unique().numel() > 1