lars/giant
1
0
Fork 0
Code Issues 3 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
59eccbb5cbe887fe6462729969b458f33ab8b3bd
giant/tests/test_rollout.py
T
lars c1c4957e2f Implement n_sec.mode = "stop_token" for the AR secondary decoder (gitea #40) 
Stage 2's autoregressive decoder still predicted multiplicity the v0.2 way:
a one-shot n_sec_head classifier over conditioning alone, run before any
secondary token existed, with the AR loop then always executing k_max slots
and discarding the tail. This adds a real per-slot EOS mechanism instead:

- Stage2Autoregressive gains a stop_head (build_stop_head=True) that predicts
  P(n_sec == k | prefix) at each slot, mutually exclusive with n_sec_head
  (n_sec.mode = "stop_token" builds no n_sec_head at all).
- sample_secondaries_ar accepts n_sec_pred=None to drive generation off the
  stop head instead of a pre-resolved count: each row stops the first slot
  its stop logit fires (stage2_model.n_sec.stop_sampling = "greedy" — the
  default, threshold at 0 — or "sample", a Bernoulli draw), and the whole
  batch loop breaks once every row has stopped, so cost scales with the
  realized n_sec instead of a fixed k_max. Passing n_sec_pred explicitly
  (the scheduled-sampling self-sample path) is unchanged.
- resolve_n_sec returns None for a stop-token decoder instead of raising;
  rollout.py/cli.py/validate.py now derive the realized count from
  sample_stage2's returned sec_valid (sec_valid.sum(-1)) after sampling,
  rather than resolving it up front — a no-op reordering under every other
  n_sec.mode, where sec_valid was already built from n_sec_pred.
- Training: _stop_target_and_mask (giant/training/stage2_inputs.py) builds
  the per-slot target/mask (one slot wider than the existing token-content
  sec_mask, since the stop slot itself needs supervision) and
  StageTrainer._stop_loss trains it with masked BCE, gated on stop_head
  exactly like _n_sec_loss gates on n_sec_head. Wired into both the
  flow/ddpm trainer and the WGAN trainer (whose skip_g_step now also checks
  stop_head), weighted by the existing stage2_model.n_sec.lambda — the stop
  head replaces n_sec_head under this mode, so no new weight key.
- validate_config now accepts stop_token (requires decoder="autoregressive"
  and n_sec.owner="stage2") instead of always rejecting it.

Decisions made during planning: stop_sampling defaults to "greedy" for
deterministic rollouts; the stop head reuses stage2_model.heads.n_sec's
HeadConfig shape and stage2_model.n_sec.lambda's weight rather than adding
new config keys, since the two heads never coexist.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-17 10:44:14 +02:00

784 lines
28 KiB
Python
Raw Blame History

"""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,
stop_token=False,
):
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 and not stop_token,
build_stop_head=stop_token,
)
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)
@pytest.mark.parametrize("generator2", ["flow", "wgan"])
def test_rollout_stop_token_end_to_end(fake_material_props, generator2):
"""gitea #40: n_sec.mode='stop_token' (no n_sec_head on either stage —
resolve_n_sec must return None and let sample_stage2's AR loop derive
the count from its own stop head) must still run to completion, produce
secondaries, and conserve energy exactly like the 'head' mode."""
s1, s2 = _models_v3(decoder="autoregressive", generator2=generator2, stop_token=True)
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
Reference in New Issue View Git Blame Copy Permalink