diff --git a/CLAUDE.md b/CLAUDE.md index 6712701..726a99d 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -44,7 +44,7 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep `ConditionEncoder`/`SecondaryConditionEncoder` (`giant/model/network.py`) support two mutually exclusive `conditioning` modes, selected per-checkpoint (`model_config["conditioning"]`, defaulting to `"embedding"` for old checkpoints without the key, `"physical"` for new `giant train` runs — see `--conditioning`): - **`"embedding"`** (original Phase 2 design): a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab (`pdg_map`/`mat_map`). Memorizes the training menu. -- **`"physical"`** (default): the 7 physical-property columns above are each routed through a small MLP (`particle_mlp`/`material_mlp`) to the same `emb_dim` width the embedding tables would have produced — a drop-in replacement computable for any PDG code / material name, not just ones seen in training, which is what lets the surrogate generalize to a held-out material or species. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships as an intentionally-unfilled stub (`MaterialProperties(None, ...)` per material) that raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting — a physicist must populate real values before `"physical"` mode can train. +- **`"physical"`** (default): the 7 physical-property columns above are each routed through a small MLP (`particle_mlp`/`material_mlp`) to the same `emb_dim` width the embedding tables would have produced — a drop-in replacement computable for any PDG code / material name, not just ones seen in training, which is what lets the surrogate generalize to a held-out material or species. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships real Geant4-11.4.1-derived `z_eff`/`a_eff`/`density`/`x0`/`lambda_int` values for every material the detector geometry actually produces; the sole exception is `G4_LYSO` (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays `MaterialProperties(None, ...)` and raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting if it's ever requested. **Model** (`giant/model/network.py`): a two-stage model, both checkpointed together. - **Stage 1 — `DenoisingMLP`:** `ResBlock` stack with a `SinusoidalEmbedding` for the flow/diffusion time variable and a `ConditionEncoder` fusing the conditioning. Predicts the 9D primary vector field, plus an `n_sec_head` classifier over `{0..K_MAX}` (`K_MAX=15`) that runs on the condition encoding alone (no diffusion noise), callable via `predict_n_sec`. @@ -64,6 +64,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep **Phase 2 (implemented — baseline):** the two-stage model above jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected). This is the "get a baseline out" track agreed with Jan & Tobias (2026-07-07). -**Physical-property conditioning (implemented):** `model.conditioning = "physical" | "embedding"` (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. `"embedding"` stays available as the generalization-comparison baseline. **Not yet done:** `giant/materials.py`'s table needs real physicist-supplied values before `"physical"` mode can train (currently unfilled, fails loudly if used); once filled, the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment. +**Physical-property conditioning (implemented):** `model.conditioning = "physical" | "embedding"` (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. `"embedding"` stays available as the generalization-comparison baseline. `giant/materials.py`'s table is already filled with real values for every material the geometry produces. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment. **Next directions** (parallel, not yet built): faster-eval architectures measured against a ~10× native-Geant4 budget — a Wasserstein-GAN throwaway (single-pass eval) and a mixture-of-experts / routing tree of small nets selected per call (pdg / energy / process), with soft/differentiable gating on continuous routing axes; a sampling-calorimeter (multi-material) dataset. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`). diff --git a/giant/cli.py b/giant/cli.py index 6b0ced5..66d3566 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -568,6 +568,14 @@ def predict( ) raise typer.Exit(1) + if "sec_phys" not in ckpt.get("normalizer", {}): + typer.echo( + "error: checkpoint has no normalizer.sec_phys — retrain with the " + "current code", + err=True, + ) + raise typer.Exit(1) + model_cfg = ckpt["model_config"] if batch_size_auto: @@ -949,6 +957,14 @@ def rollout( ) raise typer.Exit(1) + if "sec_phys" not in ckpt.get("normalizer", {}): + typer.echo( + "error: checkpoint has no normalizer.sec_phys — retrain with the " + "current code", + err=True, + ) + raise typer.Exit(1) + model_cfg = ckpt["model_config"] conditioning = model_cfg.get("conditioning", "embedding") pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()} diff --git a/giant/particles.py b/giant/particles.py index ef149ab..4e5ea17 100644 --- a/giant/particles.py +++ b/giant/particles.py @@ -84,12 +84,23 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd charge heavily since it's a small conserved quantum number that should usually match exactly. """ - codes = np.array(sorted({int(c) for c in candidates}), dtype=np.int64) - if len(codes) == 0: - raise ValueError("nearest_known_pdg: candidates is empty") - table = particle_phys_array(codes) # (C, 2) - table_log_mass = np.log(table[:, 0].astype(np.float64) + _LOG_EPS) - table_charge = table[:, 1].astype(np.float64) + # Resolve each candidate individually and skip ones giant.particles can't + # resolve, rather than letting one bad code in the training vocabulary + # crash every rollout/predict run over this reporting-only lookup — an + # unresolvable code was never a valid label to begin with. + resolved = [] + for c in sorted({int(c) for c in candidates}): + try: + resolved.append((c, *particle_mass_charge(c))) + except ValueError: + continue + if len(resolved) == 0: + raise ValueError("nearest_known_pdg: no resolvable candidates") + codes = np.array([r[0] for r in resolved], dtype=np.int64) + table_log_mass = np.log( + np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS + ) + table_charge = np.array([r[2] for r in resolved], dtype=np.float64) mass = np.asarray(mass, dtype=np.float64) charge = np.asarray(charge, dtype=np.float64) diff --git a/giant/rollout.py b/giant/rollout.py index 2d7c03d..0acf148 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -200,6 +200,7 @@ def make_seed_frontier( pre_pos: np.ndarray, pre_E: np.ndarray, pre_dir: np.ndarray, + conditioning: str = "embedding", ) -> tuple[dict[str, np.ndarray], dict[int, int]]: """Build the initial frontier from primary entry states. @@ -219,10 +220,19 @@ def make_seed_frontier( dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None) pdg_arr = np.asarray(pdg, dtype=np.int64) - # Real primaries always have a genuine ground-truth PDG code, looked up - # once here and carried forward unchanged for the track's lifetime (its - # species never changes mid-track) — same lifecycle as "pdg" itself. - mass, charge = particle_phys_array(pdg_arr).T + if conditioning == "physical": + # Real primaries always have a genuine ground-truth PDG code, looked + # up once here and carried forward unchanged for the track's lifetime + # (its species never changes mid-track) — same lifecycle as "pdg" + # itself. + mass, charge = particle_phys_array(pdg_arr).T + else: + # "embedding" mode never reads mass/charge (see + # _physical_cond_columns), so resolving them here would only risk + # crashing an embedding-mode rollout on a PDG code giant.particles + # can't resolve, for a value that's never used. + mass = np.zeros(n, dtype=np.float64) + charge = np.zeros(n, dtype=np.float64) frontier = { "event_id": event_id, @@ -321,6 +331,7 @@ def rollout( seeds["pre_pos"], seeds["pre_E"], seeds["pre_dir"], + conditioning=conditioning, ) rec = _Recorder(sink=on_chunk) diff --git a/giant/validate.py b/giant/validate.py index 3a82326..ae69207 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -175,9 +175,15 @@ def validate_marginals( phys_real = np.concatenate(all_phys_real, axis=0) # (M, 2) phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2) - phys_kl = np.array( - [_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)] - ) + if len(phys_real) > 0 and len(phys_gen) > 0: + phys_kl = np.array( + [ + _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) + for j in range(2) + ] + ) + else: + phys_kl = np.full(2, np.nan) energy_fraction_kl = np.full(K_MAX, np.nan) print( diff --git a/tests/test_particles.py b/tests/test_particles.py index af09127..224e2c8 100644 --- a/tests/test_particles.py +++ b/tests/test_particles.py @@ -106,6 +106,16 @@ def test_nearest_known_pdg_empty_candidates_raises(): nearest_known_pdg(np.array([1.0]), np.array([0.0]), []) +def test_nearest_known_pdg_skips_unresolvable_candidate(): + """One unresolvable code in the candidate set (e.g. a training-vocab + entry giant.particles can't decode) must not crash the lookup -- it's + simply excluded from the nearest-neighbour candidate pool.""" + candidates = [22, 11, -11, 999999999] + mass_e, charge_e = particle_mass_charge(11) + result = nearest_known_pdg(np.array([mass_e]), np.array([charge_e]), candidates) + assert result[0] == 11 + + def test_nearest_known_pdg_shape(): candidates = [22, 11, -11, 2212, 2112] n = 10 diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 79057ae..ec96b17 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -8,7 +8,7 @@ import numpy as np import pytest import torch -from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS +from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG from giant.data.transforms import Normalizer from giant.model.network import DenoisingMLP, SecondaryDecoder from giant.rollout import make_seed_frontier, rollout @@ -130,6 +130,29 @@ def test_seed_frontier_track_ids(): 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 diff --git a/tests/test_validate.py b/tests/test_validate.py new file mode 100644 index 0000000..ae0025f --- /dev/null +++ b/tests/test_validate.py @@ -0,0 +1,53 @@ +import numpy as np +import torch + +from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM +from giant.model.network import DenoisingMLP, SecondaryDecoder +from giant.validate import validate_marginals + + +def _tiny_models(): + s1 = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=16, n_blocks=1) + s2 = SecondaryDecoder(pdg_vocab=3, mat_vocab=2, hidden_dim=16, n_blocks=1) + return s1.eval(), s2.eval() + + +def _zero_secondaries_loader(B=4, n_batches=2): + """A val_loader whose every batch has n_sec=0 (real side) — matches the + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) tuple shape + StreamingStepsDataset yields.""" + batches = [] + for _ in range(n_batches): + cond_cont = torch.randn(B, COND_DIM) + cond_cat = torch.zeros(B, 2, dtype=torch.long) + x1 = torch.randn(B, X_DIM) + n_sec = torch.zeros(B, dtype=torch.long) + sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM) + proc_idx = torch.zeros(B, dtype=torch.long) + batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx)) + return batches + + +def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch): + """If n_sec_pred collapses to 0 across the whole validated set (realistic + during early/unstable training), phys_kl must degrade to NaN instead of + crashing on the empty-array .min()/.max() reduction inside + _histogram_kl -- a regression the old species/bincount code this + replaced explicitly guarded against.""" + s1, s2 = _tiny_models() + loader = _zero_secondaries_loader() + + # Force the Stage-1 n_sec head's prediction to 0 for every sample too, so + # the generated side's valid-slot mask is also empty (real side is + # already all n_sec=0 by construction of the fake loader above). + def _fake_sample_flow(model, cond_cont, cond_cat, **kw): + B = cond_cont.size(0) + return torch.randn(B, X_DIM), torch.zeros(B, dtype=torch.long) + + monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow) + + result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2) + + assert np.asarray(result["phys_real"]).shape == (0, 2) + assert np.asarray(result["phys_generated"]).shape == (0, 2) + assert np.isnan(np.asarray(result["phys_kl"])).all()