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.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX
from giant.data.transforms import Normalizer
from giant.model.network import Stage1Model, Stage2OneShot
from giant.rollout import make_seed_frontier, rollout

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

# giant/rollout.py isn't updated yet — it drives Stage1Model/Stage2OneShot
# through giant.sample's sample_flow/sample_secondaries, which still call
# models with the pre-refactor positional convention
# (model(x, t, cond_cont, cond_cat)) that no longer matches these classes'
# forward signatures. Deferred to docs/v0.3.0-design.md step 6/§10.
pytestmark = pytest.mark.xfail(reason="giant/rollout.py not updated for v0.3.0 network.py yet (step 6)", strict=False)

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


def _models(conditioning="embedding"):
    particle_cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
    material_cfg = {"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,
        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, 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()