giant/tests/test_rollout.py

"""Tests for the autoregressive shower rollout driver."""

from collections import Counter
from pathlib import Path
from unittest.mock import patch

import numpy as np
import pytest
import torch

from giant.config import ConditioningAxisConfig, ParticleTypeConfig
from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX
from giant.data.loader import TopNMap
from giant.data.transforms import Normalizer
from giant.model.network import (
    Stage1Model,
    Stage2Autoregressive,
    Stage2OneShot,
    stage2_trunk_sec_dim,
)
from giant.rollout import L1DistCollector, make_seed_frontier, rollout

pytest.importorskip("sklearn")
from giant import geometry as g  # noqa: E402

PDG_MAP = {22: 0, 11: 1, -11: 2}
MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1}


def _models(conditioning="embedding"):
    particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1)
    material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=16, n_layers=1)
    s1 = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        n_sec_head_k_max=K_MAX,
    )
    s2 = Stage2OneShot(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        generator="flow",
        time_dim=16,
    )
    return s1.eval(), s2.eval()


def _norms():
    rng = np.random.default_rng(0)
    cond = Normalizer().fit(rng.standard_normal((1000, 15)).astype(np.float32))
    tgt = Normalizer().fit(rng.standard_normal((1000, 9)).astype(np.float32))
    sec_phys = Normalizer().fit(rng.standard_normal((1000, 2)).astype(np.float32))
    return cond, tgt, sec_phys


def _oracle():
    rng = np.random.default_rng(0)
    pos = rng.uniform(-200, 200, (20000, 3)).astype(np.float32)
    inside = (np.abs(pos) < 100).all(axis=1)
    mat = np.where(inside, "G4_PbWO4", "G4_AIR").astype(object)
    lay = np.where(inside, 0, -1).astype(np.int64)
    with patch.object(g, "_iter_point_batches", lambda p: iter([(pos, mat, lay)])):
        # Pinned to "knn" explicitly: this test's escape-threshold semantics
        # (tiny threshold -> escape even at a valid interior point, because no
        # training point is that close) are KNN-specific, and the fixture's
        # box geometry isn't a layer stack the "slab" method could fit anyway.
        return g.build_geometry_oracle([Path("x")], method="knn", subsample=20000)


def _seeds(n=6):
    return {
        "event_id": np.arange(n, dtype=np.int64),
        "pdg": np.full(n, 11, dtype=np.int64),
        "pre_pos": np.zeros((n, 3)),
        "pre_E": np.linspace(30.0, 90.0, n),
        "pre_dir": np.tile([0.0, 0.0, 1.0], (n, 1)),
    }


def _run(
    escape_threshold=1e9,
    energy_cutoff=1.0,
    max_steps=30,
    max_tracks_per_event=300,
    seeds=None,
    conditioning="embedding",
):
    torch.manual_seed(0)
    np.random.seed(0)
    s1, s2 = _models(conditioning)
    cond, tgt, sec_phys = _norms()
    return rollout(
        s1,
        s2,
        _oracle(),
        seeds or _seeds(),
        cond,
        tgt,
        sec_phys,
        PDG_MAP,
        MAT_MAP,
        energy_cutoff=energy_cutoff,
        max_steps=max_steps,
        steps=4,
        batch_size=128,
        max_tracks_per_event=max_tracks_per_event,
        escape_threshold=escape_threshold,
        particle_conditioning=conditioning,
        material_conditioning=conditioning,
    )


@pytest.fixture
def fake_material_props(monkeypatch):
    import giant.materials as gm

    fake = {
        "G4_AIR": gm.MaterialProperties(z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5),
        "G4_PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7),
    }
    monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
    return fake


def test_rollout_physical_conditioning_end_to_end(fake_material_props):
    """Physical-mode rollout runs to completion; spawned secondaries carry
    mass/charge forward (no snapping) and the output pdg column is populated
    via the reporting-only nearest-known-PDG label."""
    rec = _run(conditioning="physical")
    assert len(rec["event_id"]) > 0
    assert set(np.unique(rec["pdg"]).tolist()) <= set(PDG_MAP.keys())


def test_rollout_physical_conditioning_generalizes_to_out_of_vocab_pdg(
    fake_material_props,
):
    """A real, giant.particles-resolvable species outside the training PDG
    vocab (muon, 13) must run through physical-property conditioning rather
    than terminate via TERM_UNKNOWN_PDG — that generalization is the entire
    point of "physical" mode (see build_cond_features(strict=...))."""
    seeds = _seeds(6)
    seeds["pdg"] = np.full(6, 13, dtype=np.int64)
    assert 13 not in PDG_MAP
    rec = _run(seeds=seeds, conditioning="physical")
    assert len(rec["event_id"]) > 0
    assert TERM_UNKNOWN_PDG not in set(rec["termination_reason"].tolist())


def test_seed_frontier_track_ids():
    seeds = _seeds(3)
    fr, counts = make_seed_frontier(**seeds)
    assert (fr["track_id"] == [0, 0, 0]).all()  # one primary per event -> id 0
    assert (fr["parent_id"] == -1).all()
    assert (fr["generation"] == 0).all()
    assert all(counts[e] == 1 for e in range(3))
    # pre_dir is normalised.
    np.testing.assert_allclose(np.linalg.norm(fr["pre_dir"], axis=1), 1.0, atol=1e-6)


def test_seed_frontier_embedding_mode_skips_unresolvable_pdg_lookup():
    """ "embedding" mode must not call giant.particles at all, so a seed PDG
    code it can't resolve (e.g. a fabricated/garbage code) must not crash
    frontier construction — mass/charge are simply zero-filled, unused."""
    seeds = _seeds(3)
    seeds["pdg"] = np.full(3, 999999999, dtype=np.int64)
    fr, _counts = make_seed_frontier(**seeds, particle_conditioning="embedding")
    np.testing.assert_array_equal(fr["mass"], 0.0)
    np.testing.assert_array_equal(fr["charge"], 0.0)


def test_rollout_embedding_mode_unresolvable_pdg_terminates_gracefully():
    """A rollout seeded with a PDG code giant.particles can't resolve, and
    which isn't in the training vocabulary either, must terminate via the
    existing TERM_UNKNOWN_PDG gate rather than crash in make_seed_frontier —
    "embedding" mode has no dependency on giant.particles at all."""
    seeds = _seeds(3)
    seeds["pdg"] = np.full(3, 999999999, dtype=np.int64)
    rec = _run(seeds=seeds, conditioning="embedding")
    assert len(rec["event_id"]) > 0
    assert set(rec["termination_reason"].tolist()) == {TERM_UNKNOWN_PDG}


def test_rollout_terminates_and_has_rows():
    rec = _run()
    assert len(rec["event_id"]) > 0
    # Every seed event appears.
    assert set(rec["event_id"].tolist()) == set(range(6))


def test_max_steps_respected():
    # Disable the energy cutoff so tracks survive long enough to hit the step cap.
    rec = _run(max_steps=5, energy_cutoff=0.0)
    assert rec["step_no"].max() <= 5
    assert (rec["termination_reason"] == TERM_MAX_STEPS).any()


def test_energy_conserved_deposit_plus_leak():
    seeds = _seeds()
    rec = _run(seeds=seeds)
    for i, ev in enumerate(seeds["event_id"]):
        m = rec["event_id"] == ev
        dep = rec["edep"][m].sum()
        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)


def test_secondaries_have_valid_parents():
    rec = _run()
    orphans = 0
    for ev in np.unique(rec["event_id"]):
        m = rec["event_id"] == ev
        tids = set(rec["track_id"][m].tolist())
        for pid in rec["parent_id"][m]:
            if pid >= 0 and pid not in tids:
                orphans += 1
    assert orphans == 0
    # At least one secondary (generation > 0) is produced by the tiny model.
    assert (rec["generation"] > 0).any()


def test_escape_terminates_immediately():
    # A tight escape threshold makes even the seed position (origin) escape.
    rec = _run(escape_threshold=1e-3)
    assert (rec["termination_reason"] == TERM_ESCAPED).all()
    assert rec["step_no"].max() == 0


def test_output_schema_complete():
    from giant.rollout import _RECORD_KEYS

    rec = _run()
    assert set(rec.keys()) == set(_RECORD_KEYS)
    n = len(rec["event_id"])
    assert all(len(v) == n for v in rec.values())


def test_max_tracks_cap_conserves_energy():
    # A very small cap forces sub-cap secondaries to deposit in place; energy
    # must still balance.
    seeds = _seeds()
    rec = _run(seeds=seeds, max_tracks_per_event=3)
    for i, ev in enumerate(seeds["event_id"]):
        m = rec["event_id"] == ev
        dep = rec["edep"][m].sum()
        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
        assert len(np.unique(rec["track_id"][m])) <= 3


# ── Streaming output (on_chunk) ──────────────────────────────────────────────


def _run_streaming(on_chunk, **kwargs):
    torch.manual_seed(0)
    np.random.seed(0)
    s1, s2 = _models()
    cond, tgt, sec_phys = _norms()
    seeds = kwargs.pop("seeds", None) or _seeds()
    return rollout(
        s1,
        s2,
        _oracle(),
        seeds,
        cond,
        tgt,
        sec_phys,
        PDG_MAP,
        MAT_MAP,
        energy_cutoff=kwargs.pop("energy_cutoff", 1.0),
        max_steps=kwargs.pop("max_steps", 30),
        steps=4,
        batch_size=128,
        max_tracks_per_event=kwargs.pop("max_tracks_per_event", 300),
        escape_threshold=kwargs.pop("escape_threshold", 1e9),
        on_chunk=on_chunk,
    )


def test_on_chunk_receives_every_row_exactly_once():
    """Concatenating the streamed chunks must reproduce the buffered result."""
    from giant.rollout import _RECORD_KEYS

    buffered = _run()

    chunks: list[dict[str, np.ndarray]] = []
    summary = _run_streaming(chunks.append)

    streamed = {k: np.concatenate([c[k] for c in chunks]) for k in _RECORD_KEYS}
    assert summary["n_rows"] == len(buffered["event_id"])
    assert len(streamed["event_id"]) == len(buffered["event_id"])
    for k in _RECORD_KEYS:
        np.testing.assert_array_equal(streamed[k], buffered[k])


def test_on_chunk_summary_termination_reason_counts_match_buffered():
    buffered = _run()
    summary = _run_streaming(lambda row: None)

    expected = Counter(r for r in buffered["termination_reason"].tolist() if r)
    assert summary["termination_reason_counts"] == dict(expected)


def test_on_chunk_never_buffers_full_records():
    """Streaming mode must not accumulate rows for later to_dict() retrieval."""
    from giant.rollout import _Recorder

    rec = _Recorder(sink=lambda row: None)
    rec.add(
        event_id=np.array([0]),
        track_id=np.array([0]),
        parent_id=np.array([-1]),
        generation=np.array([0]),
        step_no=np.array([0]),
        pdg=np.array([11]),
        pre_x=np.array([0.0]),
        pre_y=np.array([0.0]),
        pre_z=np.array([0.0]),
        pre_E=np.array([1.0]),
        pre_dx=np.array([0.0]),
        pre_dy=np.array([0.0]),
        pre_dz=np.array([1.0]),
        post_x=np.array([0.0]),
        post_y=np.array([0.0]),
        post_z=np.array([1.0]),
        post_E=np.array([0.0]),
        post_dx=np.array([0.0]),
        post_dy=np.array([0.0]),
        post_dz=np.array([1.0]),
        edep=np.array([1.0]),
        step_length=np.array([1.0]),
        material=np.array(["G4_AIR"], dtype=object),
        layer_id=np.array([0]),
        n_sec_pred=np.array([0]),
        termination_reason=np.array(["natural_end"], dtype=object),
    )
    assert rec.n_rows == 1
    assert rec.termination_reason_counts == {"natural_end": 1}
    with pytest.raises(AssertionError):
        rec.to_dict()


# ── v0.3.0 step 6: per-stage generators, AR decoder, particle_type.target ───

# emb_dim=3: "other" (class idx 2) is shared by -11 and 13 (muon), matching
# the real shape build_pdg_topn_map_from_files produces — see
# decode_topn_class's docstring.
PDG_TOPN_MAP = TopNMap(
    class_map={22: 0, 11: 1, -11: 2, 13: 2},
    other_members={-11: 5, 13: 1},
)


def _models_v3(
    conditioning="physical",
    decoder="one_shot",
    target="physical",
    generator1="flow",
    generator2="flow",
    k_max=6,
    emb_dim=4,
    stage2_has_n_sec_head=True,
):
    particle_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
    material_cfg = ConditioningAxisConfig(type=conditioning, emb_dim=emb_dim, n_layers=1)
    # A fresh v0.3.0 Stage1Model — no n_sec_head_k_max, unlike _models() above
    # (n_sec ownership moves to stage 2 by default).
    s1 = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        generator=generator1,
        noise_dim=8,
    )
    particle_type_cfg = ParticleTypeConfig(target=target)
    # Explicit kwargs rather than a shared **common dict: a dict() call whose
    # values have heterogeneous types (str/int/dict/bool) widens under static
    # analysis to dict[str, <big union>], which then makes every constructor
    # keyword not itself part of that union (router, cond_enc, ...) look like
    # a type mismatch to `ty` even though every actual value passed is fine.
    if decoder == "one_shot":
        sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator2, k_max, emb_dim)
        s2 = Stage2OneShot(
            pdg_vocab=3,
            mat_vocab=2,
            particle_cfg=particle_cfg,
            material_cfg=material_cfg,
            hidden_dim=32,
            n_res_blocks=2,
            generator=generator2,
            time_dim=16,
            noise_dim=8,
            k_max=k_max,
            particle_type_cfg=particle_type_cfg,
            build_n_sec_head=stage2_has_n_sec_head,
            sec_dim=sec_dim,
        )
    else:
        s2 = Stage2Autoregressive(
            pdg_vocab=3,
            mat_vocab=2,
            particle_cfg=particle_cfg,
            material_cfg=material_cfg,
            hidden_dim=32,
            n_res_blocks=2,
            generator=generator2,
            time_dim=16,
            noise_dim=8,
            k_max=k_max,
            particle_type_cfg=particle_type_cfg,
            build_n_sec_head=stage2_has_n_sec_head,
        )
    return s1.eval(), s2.eval()


def _run_v3(
    s1,
    s2,
    escape_threshold=1e9,
    energy_cutoff=1.0,
    max_steps=15,
    max_tracks_per_event=100,
    seeds=None,
    conditioning="physical",
    sec_type_topn_map=None,
    other_policy="sample",
    seed=0,
    stage1_ddpm_steps=1000,
    l1_dist_collector=None,
):
    torch.manual_seed(0)
    np.random.seed(0)
    cond, tgt, sec_phys = _norms()
    return rollout(
        s1,
        s2,
        _oracle(),
        seeds or _seeds(),
        cond,
        tgt,
        sec_phys,
        PDG_MAP,
        MAT_MAP,
        energy_cutoff=energy_cutoff,
        max_steps=max_steps,
        steps=3,
        batch_size=128,
        max_tracks_per_event=max_tracks_per_event,
        escape_threshold=escape_threshold,
        particle_conditioning=conditioning,
        material_conditioning=conditioning,
        sec_type_topn_map=sec_type_topn_map,
        other_policy=other_policy,
        seed=seed,
        stage1_ddpm_steps=stage1_ddpm_steps,
        l1_dist_collector=l1_dist_collector,
    )


def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props):
    """A fresh v0.3.0 Stage1Model has no n_sec_head — n_sec must come from
    Stage2's own head instead, and the run must still complete and
    conserve energy."""
    s1, s2 = _models_v3()
    rec = _run_v3(s1, s2)
    assert len(rec["event_id"]) > 0
    seeds = _seeds()
    for i, ev in enumerate(seeds["event_id"]):
        m = rec["event_id"] == ev
        dep = rec["edep"][m].sum()
        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)


def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
    """Neither stage owning n_sec_head only happens for a
    stage2_model.n_sec.mode other than "head" — not a valid rollout-capable
    checkpoint, and must fail with a clear error rather than crash deep
    inside predict_n_sec."""
    s1, s2 = _models_v3(stage2_has_n_sec_head=False)
    with pytest.raises(RuntimeError, match="n_sec_head"):
        _run_v3(s1, s2)


@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
@pytest.mark.parametrize("generator2", ["flow", "wgan"])
def test_rollout_physical_target_decoder_generator_matrix(fake_material_props, decoder, generator2):
    """Every (decoder, stage2 generator) combination under
    particle_type.target="physical" must run to completion and conserve
    energy."""
    s1, s2 = _models_v3(decoder=decoder, generator2=generator2)
    rec = _run_v3(s1, s2)
    assert len(rec["event_id"]) > 0
    seeds = _seeds()
    for i, ev in enumerate(seeds["event_id"]):
        m = rec["event_id"] == ev
        dep = rec["edep"][m].sum()
        leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
        assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)


def test_sample_stage1_dispatches_ddpm_by_generator_kind():
    """stage1_model.generator="ddpm" must be dispatched to sample_ddpm
    (previously _step_chunk silently fell through to the flow ODE sampler
    regardless of the checkpoint's actual generator — see giant.sample.sample_stage1).
    A short T avoids the reverse-diffusion numerical blowup an untrained,
    random-weight network produces over many steps; that instability is a
    property of sampling from an untrained net, not of the dispatch logic
    under test here, so a full oracle-driven rollout isn't needed."""
    from giant.sample import sample_stage1

    s1, _ = _models_v3(generator1="ddpm")
    cond_cont = torch.randn(6, 15)
    cond_cat = torch.zeros(6, 2, dtype=torch.long)
    sample, n_sec = sample_stage1(s1, cond_cont, cond_cat, steps=10, ddpm_steps=5)
    assert sample.shape == (6, 9)
    assert n_sec is None  # fresh v0.3.0 Stage1Model owns no n_sec_head


@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
    """particle_type.target="onehot" resolves a concrete PDG via
    decode_topn_class (argmax + other_policy), and that PDG's real physics
    (giant.particles.particle_phys_array) become the secondary's identity —
    unlike "physical", not just a reporting label."""
    s1, s2 = _models_v3(decoder=decoder, target="onehot", emb_dim=3)
    rec = _run_v3(s1, s2, sec_type_topn_map=PDG_TOPN_MAP, other_policy="modal")
    assert len(rec["event_id"]) > 0
    # Every spawned secondary's nominal pdg must be one decode_topn_class can
    # actually produce (the topn map's known classes + its "other" members).
    possible = set(PDG_TOPN_MAP.class_map.keys()) | set(PDG_TOPN_MAP.other_members.keys())
    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
    assert secondary_pdgs <= possible


def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
    s1, s2 = _models_v3(target="onehot", emb_dim=3)
    with pytest.raises(RuntimeError, match="sec_type_topn_map"):
        _run_v3(s1, s2, sec_type_topn_map=None)


# --- conditioning.{particle,material}.type = "onehot" — a separate axis from
# stage2_model.particle_type.target above: this is what feeds cond_cat's
# extra top-N columns for ConditionEncoder's own "onehot" mode, not the
# secondary-species decode. ---------------------------------------------

COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={})
COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={})


def _onehot_conditioning_models():
    particle_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(PDG_MAP), n_layers=1)
    material_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(MAT_MAP), n_layers=1)
    s1 = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
    )
    s2 = Stage2OneShot(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        sec_dim=stage2_trunk_sec_dim(ParticleTypeConfig(target="physical"), "flow", K_MAX, 3),
        generator="flow",
        time_dim=16,
    )
    return s1.eval(), s2.eval()


def _run_onehot_conditioning(pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP):
    s1, s2 = _onehot_conditioning_models()
    cond, tgt, sec_phys = _norms()
    return rollout(
        s1,
        s2,
        _oracle(),
        _seeds(),
        cond,
        tgt,
        sec_phys,
        PDG_MAP,
        MAT_MAP,
        energy_cutoff=1.0,
        max_steps=30,
        steps=4,
        batch_size=128,
        max_tracks_per_event=300,
        escape_threshold=1e9,
        particle_conditioning="onehot",
        material_conditioning="onehot",
        pdg_topn_map=pdg_topn_map,
        mat_topn_map=mat_topn_map,
    )


def test_rollout_conditioning_onehot_end_to_end(fake_material_props):
    rec = _run_onehot_conditioning()
    assert len(rec["event_id"]) > 0
    assert set(rec["event_id"].tolist()) == set(range(6))


def test_rollout_conditioning_onehot_particle_missing_topn_map_raises(
    fake_material_props,
):
    with pytest.raises(RuntimeError, match="pdg_topn_map"):
        _run_onehot_conditioning(pdg_topn_map=None)


def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
    fake_material_props,
):
    with pytest.raises(RuntimeError, match="mat_topn_map"):
        _run_onehot_conditioning(mat_topn_map=None)


SEC_TYPE_TOPN_MAP_DIFFERENT_N = TopNMap(class_map={22: 0, 11: 1, -11: 2, 13: 3}, other_members={2112: 3, 2212: 1})


def _run_conditioning_and_type_onehot_different_n_classes():
    """Both conditioning.particle.type="onehot" and
    stage2_model.particle_type.target="onehot" active at once, with
    stage2_model.particle_type.n_classes deliberately different from
    conditioning.particle.emb_dim (gitea #29)."""
    cond_emb_dim = len(PDG_MAP)  # 3
    type_n_classes = 5  # deliberately different from cond_emb_dim
    particle_cfg = ConditioningAxisConfig(type="onehot", emb_dim=cond_emb_dim, n_layers=1)
    material_cfg = ConditioningAxisConfig(type="onehot", emb_dim=len(MAT_MAP), n_layers=1)
    particle_type_cfg = ParticleTypeConfig(target="onehot", n_classes=type_n_classes)

    s1 = Stage1Model(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
    ).eval()
    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", K_MAX, type_n_classes)
    s2 = Stage2OneShot(
        pdg_vocab=3,
        mat_vocab=2,
        particle_cfg=particle_cfg,
        material_cfg=material_cfg,
        hidden_dim=32,
        n_res_blocks=2,
        sec_dim=sec_dim,
        generator="flow",
        time_dim=16,
        k_max=K_MAX,
        particle_type_cfg=particle_type_cfg,
    ).eval()
    # Sanity: the model's own type_dim followed n_classes, not cond_emb_dim.
    assert s2.type_dim == type_n_classes

    cond, tgt, sec_phys = _norms()
    return rollout(
        s1,
        s2,
        _oracle(),
        _seeds(),
        cond,
        tgt,
        sec_phys,
        PDG_MAP,
        MAT_MAP,
        energy_cutoff=1.0,
        max_steps=15,
        steps=3,
        batch_size=128,
        max_tracks_per_event=100,
        escape_threshold=1e9,
        particle_conditioning="onehot",
        material_conditioning="onehot",
        pdg_topn_map=COND_PDG_TOPN_MAP,
        mat_topn_map=COND_MAT_TOPN_MAP,
        sec_type_topn_map=SEC_TYPE_TOPN_MAP_DIFFERENT_N,
        other_policy="modal",
    )


def test_rollout_conditioning_and_type_onehot_with_different_n_classes(fake_material_props):
    """gitea #29 end-to-end: conditioning.particle.type="onehot" and
    stage2_model.particle_type.target="onehot" now use independently sized
    top-N maps (stage2_model.particle_type.n_classes != conditioning.particle
    .emb_dim), and rollout must decode secondaries using the type-side map,
    not silently reuse the conditioning-side one (the pre-#29 bug)."""
    rec = _run_conditioning_and_type_onehot_different_n_classes()
    assert len(rec["event_id"]) > 0
    possible = set(SEC_TYPE_TOPN_MAP_DIFFERENT_N.class_map.keys()) | set(
        SEC_TYPE_TOPN_MAP_DIFFERENT_N.other_members.keys()
    )
    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
    assert secondary_pdgs <= possible


@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
def test_rollout_embedding_target_end_to_end(decoder):
    """particle_type.target="embedding" L1-snaps to the nearest row of the
    conditioning's own particle embedding table, so every resolved PDG must
    be a real member of the dense training vocab (pdg_map) — unlike
    "onehot", there is no "other" bucket to fall outside of."""
    s1, s2 = _models_v3(conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4)
    rec = _run_v3(s1, s2, conditioning="embedding")
    assert len(rec["event_id"]) > 0
    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
    assert secondary_pdgs <= set(PDG_MAP.keys())


def test_l1_dist_collector_populated_only_for_embedding_target():
    """The L1-distance diagnostic only makes sense under
    particle_type.target="embedding" — a physical-target run must leave the
    collector empty rather than silently accumulating garbage."""
    s1, s2 = _models_v3(target="physical")
    collector = L1DistCollector()
    _run_v3(s1, s2, l1_dist_collector=collector)
    assert collector.n == 0
    assert collector.summary() is None


def test_l1_dist_collector_accumulates_for_embedding_target():
    s1, s2 = _models_v3(conditioning="embedding", target="embedding", emb_dim=4)
    collector = L1DistCollector()
    rec = _run_v3(s1, s2, conditioning="embedding", l1_dist_collector=collector)
    n_secondaries = int((rec["generation"] > 0).sum())
    assert n_secondaries > 0  # sanity: the tiny model does spawn secondaries
    summary = collector.summary()
    assert summary is not None
    assert summary["n"] == collector.n > 0
    assert summary["min"] <= summary["mean"] <= summary["max"]
    assert summary["std"] >= 0.0
    assert len(summary["hist_edges"]) == len(summary["hist_counts"]) + 1
    assert sum(summary["hist_counts"]) <= summary["n"]  # some may fall outside [lo, hi)


def test_l1_dist_collector_add_ignores_invalid_slots():
    collector = L1DistCollector()
    dist = np.array([[1.0, 5.0, 9.0]])
    valid = np.array([[True, False, True]])
    collector.add(dist, valid)
    assert collector.n == 2
    assert collector.minimum == 1.0
    assert collector.maximum == 9.0


def test_l1_dist_collector_add_empty_is_noop():
    collector = L1DistCollector()
    collector.add(np.zeros((0, 3)), np.zeros((0, 3), dtype=bool))
    assert collector.n == 0
    assert collector.summary() is None