lars
68fb99bed8
Adds model.conditioning = "physical" | "embedding": physical mode routes particle mass/charge and material Z_eff/A_eff/density/X0/lambda_int through small MLPs to replace the learned PDG/material embedding tables, so the surrogate generalizes to PDG codes/materials outside the training vocab instead of memorizing it. "embedding" stays available as the comparison baseline (old checkpoints without the key default to it). Stage 2 now regresses a secondary's mass/charge directly against a fixed physics-derived target instead of a learned/snapped embedding, and uses no snapping at inference — the model's raw predicted (mass, charge) is the secondary's physical identity, including for its own further rollout steps. A separate reporting-only nearest-known-PDG lookup (never fed back into the model) populates output pdg columns / the embedding-mode rollout fallback. giant/materials.py's table is populated with Geant4's own built-in NIST constants (Z_eff, A_eff, density, X0, lambda_int), extracted directly from the Geant4 11.4.1 build vendored in minicalosim via G4NistManager rather than hand-typed literature values. G4_LYSO is left unfilled: confirmed (both by runtime lookup and by searching minicalosim's history) that it's never actually a constructed Geant4 material there, only documentation/UI color-map text. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
290 lines
9.3 KiB
Python
290 lines
9.3 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(conditioning="embedding"):
|
|
s1 = DenoisingMLP(
|
|
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, conditioning=conditioning
|
|
)
|
|
s2 = SecondaryDecoder(
|
|
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, conditioning=conditioning
|
|
)
|
|
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_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, 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()
|