lars
550dc679c7
_Recorder previously accumulated every generated step across all events/ tracks/steps in Python lists, materialised once at the end and written via a single pq.write_table — memory scaled with n_events * max_steps * avg_tracks_per_event. rollout() now takes an optional on_chunk callback that streams each non-empty batch immediately (fixed per-key dtypes via _RECORD_DTYPES keep every chunk's table schema identical, which pq.ParquetWriter requires across writes); giant rollout wires this to an incrementally-written ParquetWriter, mirroring the row-group streaming giant predict already does on its input side. Without on_chunk, rollout() keeps its old buffered return for existing callers/tests. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
256 lines
8.1 KiB
Python
256 lines
8.1 KiB
Python
"""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
|
|
from giant.data.transforms import Normalizer
|
|
from giant.model.network import DenoisingMLP, SecondaryDecoder
|
|
from giant.rollout import 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():
|
|
s1 = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
|
s2 = SecondaryDecoder(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
|
return s1.eval(), s2.eval()
|
|
|
|
|
|
def _norms():
|
|
rng = np.random.default_rng(0)
|
|
cond = Normalizer().fit(rng.standard_normal((1000, 8)).astype(np.float32))
|
|
tgt = Normalizer().fit(rng.standard_normal((1000, 9)).astype(np.float32))
|
|
return cond, tgt
|
|
|
|
|
|
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,
|
|
):
|
|
torch.manual_seed(0)
|
|
np.random.seed(0)
|
|
s1, s2 = _models()
|
|
cond, tgt = _norms()
|
|
return rollout(
|
|
s1,
|
|
s2,
|
|
_oracle(),
|
|
seeds or _seeds(),
|
|
cond,
|
|
tgt,
|
|
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,
|
|
)
|
|
|
|
|
|
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_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 = _norms()
|
|
seeds = kwargs.pop("seeds", None) or _seeds()
|
|
return rollout(
|
|
s1,
|
|
s2,
|
|
_oracle(),
|
|
seeds,
|
|
cond,
|
|
tgt,
|
|
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()
|