From 68fb99bed89d92f4feda85d71bb3300eb04841c4 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 17 Jul 2026 15:12:54 +0200 Subject: [PATCH 1/2] Condition on material/particle physical properties instead of learned embeddings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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 --- CLAUDE.md | 12 ++- giant/cli.py | 53 +++++++---- giant/config.py | 7 ++ giant/constants.py | 49 +++++++--- giant/data/dataset.py | 39 ++++---- giant/data/transforms.py | 184 +++++++++++++++++++++++++++++--------- giant/materials.py | 145 ++++++++++++++++++++++++++++++ giant/model/network.py | 104 ++++++++++++++++++---- giant/model/schedule.py | 21 +++-- giant/particles.py | 102 +++++++++++++++++++++ giant/pipeline.py | 47 ++++++---- giant/rollout.py | 68 ++++++++++---- giant/sample.py | 32 +++---- giant/train.py | 40 ++------- giant/validate.py | 65 ++++++++------ pyproject.toml | 1 + tests/test_materials.py | 105 ++++++++++++++++++++++ tests/test_particles.py | 118 ++++++++++++++++++++++++ tests/test_phase2.py | 188 ++++++++++++++++++++++++++++----------- tests/test_rollout.py | 50 +++++++++-- tests/test_router.py | 6 +- tests/test_transforms.py | 87 +++++++++++++++++- uv.lock | 33 +++++++ 23 files changed, 1252 insertions(+), 304 deletions(-) create mode 100644 giant/materials.py create mode 100644 giant/particles.py create mode 100644 tests/test_materials.py create mode 100644 tests/test_particles.py diff --git a/CLAUDE.md b/CLAUDE.md index 4402bfb..6712701 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -40,11 +40,15 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep **Stage-1 output space (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):** `log_step_length`, two additive-log-ratio (ALR) coordinates `edep_logit`/`sec_logit` of a **deposit / secondary / post-energy simplex**, `post_dir` (post-scattering momentum direction, unit vector in the local frame), and `travel_dir` (direction of `post_pos - pre_pos`, unit vector in the local frame). The energy simplex decodes via softmax over `[edep_logit, sec_logit, 0]` × `pre_E` so `edep + e_sec + post_E == pre_E` holds by construction — energy conservation is architectural, not learned (see `energy_simplex_decode`). `post_pos` is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, since `step_length` already encodes that displacement's magnitude and duplicating it would let the two become inconsistent. -**Conditioning vector (8D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID — plus PDG code and material as embeddings. `n_sec` and `e_sec` are **no longer conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model now predicts them. +**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus, since particle/material physical-property conditioning (`model.conditioning`, see below), 7 more columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model predicts them. + +`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. **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`. -- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, continuous type embedding 16D)` = `SEC_SLOT_DIM=20`, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. The type embedding is trained against a detached PDG-embedding target (stops self-referential collapse) and snapped to the nearest PDG at inference (`snap_type_to_pdg_idx`). +- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, log-mass, charge)` = `SEC_SLOT_DIM=6`, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. A secondary's mass/charge are regressed directly against a fixed physics-derived target (its ground-truth PDG code's `giant.particles.particle_mass_charge`) — not a learned/moving embedding target, so nothing needs detaching. **No snapping at inference**: the predicted (mass, charge) are used as-is as the secondary's physical identity, including for its own future conditioning if it goes on to take further steps in a rollout. A separate, reporting-only nearest-known-PDG lookup (`giant.particles.nearest_known_pdg`) is used purely to populate a nominal `pdg` label for output rows / `"embedding"`-mode fallback conditioning — it never feeds back into the model. `schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching). @@ -60,4 +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). -**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; and preferring **material + particle physical properties** over learned embeddings for conditioning. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`). +**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. + +**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 935bc98..6b0ced5 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -42,9 +42,10 @@ from giant.data.transforms import ( ) from giant.geometry import GeometryOracle from giant.model.network import build_models +from giant.particles import nearest_known_pdg from giant.pipeline import run_train_job from giant.rollout import rollout as run_rollout -from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx +from giant.sample import sample_flow, sample_secondaries app = typer.Typer(no_args_is_help=True) @@ -178,6 +179,11 @@ class Mode(str, Enum): ddpm = "ddpm" +class Conditioning(str, Enum): + physical = "physical" + embedding = "embedding" + + class Coord(str, Enum): global_ = "global" local = "local" @@ -266,6 +272,15 @@ def train( "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)" ), ] = None, + conditioning: Annotated[ + Optional[Conditioning], + typer.Option( + "--conditioning", + help="Input conditioning: continuous physical properties " + "(mass/charge/Z_eff/A_eff/density/X0/lambda_int, default) or the " + "original learned PDG/material embeddings", + ), + ] = None, router: Annotated[ Optional[bool], typer.Option( @@ -390,6 +405,7 @@ def train( "n_blocks": n_blocks, "emb_dim": emb_dim, "dropout": dropout, + "conditioning": conditioning.value if conditioning is not None else None, }.items() if v is not None } @@ -432,6 +448,7 @@ def train( f"_h{m['hidden_dim']}" f"_b{m['n_blocks']}" f"_e{m['emb_dim']}" + f"_c{m['conditioning']}" f"_lr{t['lr']}" f"_bs{t['batch_size']}" ) @@ -573,11 +590,12 @@ def predict( assert batch_size_value is not None bs = batch_size_value + conditioning = model_cfg.get("conditioning", "embedding") pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()} - pdg_map_inv = {v: k for k, v in pdg_map.items()} mat_map = {str(k): v for k, v in ckpt["mat_map"].items()} cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) + sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"]) model, sec_decoder = build_models(model_cfg) _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) @@ -612,13 +630,13 @@ def predict( nonlocal writer, total if coord == Coord.local: - cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features( - piece, pdg_map, mat_map + cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features( + piece, pdg_map, mat_map, conditioning=conditioning ) cond_cont = cond_norm.transform(cond_cont) else: cond_cont, cond_cat = build_cond_features( - piece, pdg_map, mat_map, cond_norm + piece, pdg_map, mat_map, cond_norm, conditioning=conditioning ) cc = torch.from_numpy(cond_cont).float().to(_device) @@ -626,14 +644,10 @@ def predict( stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps) if coord == Coord.global_: - sec_cont, sec_type_emb, _sec_valid_pred = sample_secondaries( + sec_cont, sec_phys, _sec_valid_pred = sample_secondaries( sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps ) - sec_pdg_idx = snap_type_to_pdg_idx( - sec_type_emb, model.pdg_embedding_weight() - ) - sec_cont_np = sec_cont.cpu().numpy() - sec_pdg_idx_np = sec_pdg_idx.cpu().numpy() + sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy() n_sec_pred_np = n_sec_pred.cpu().numpy() pred = stage1_norm.cpu().numpy() # normalised @@ -692,14 +706,19 @@ def predict( piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local ) - sec_E, sec_dir_world, sec_pdg_code, _sec_valid = decode_secondaries( - sec_cont_np, - sec_pdg_idx_np, + sec_E, sec_dir_world, sec_mass, sec_charge, _sec_valid = decode_secondaries( + sec_full_np, n_sec_pred_np, e_sec_pred, piece["pre_dir"], - pdg_map_inv, + sec_phys_normalizer=sec_phys_norm, ) + # Reporting-only nearest-known-PDG label (never fed back into the + # model) for the sec_pdg_list output column — see + # giant/particles.py and the "no snapping at inference" design. + sec_pdg_code = nearest_known_pdg( + sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys() + ).reshape(sec_mass.shape) sec_pdg_list = [ sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np) ] @@ -931,10 +950,12 @@ def rollout( 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()} mat_map = {str(k): v for k, v in ckpt["mat_map"].items()} cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"]) + sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"]) model, sec_decoder = build_models(model_cfg) _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) @@ -981,6 +1002,7 @@ def rollout( seeds, cond_norm, tgt_norm, + sec_phys_norm, pdg_map, mat_map, energy_cutoff=energy_cutoff, @@ -991,6 +1013,7 @@ def rollout( max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, on_chunk=_write_chunk, + conditioning=conditioning, ) if writer is not None: writer.close() diff --git a/giant/config.py b/giant/config.py index 4834f3c..ab121eb 100644 --- a/giant/config.py +++ b/giant/config.py @@ -33,6 +33,13 @@ DEFAULT_CONFIG: dict = { "n_blocks": 6, "emb_dim": 16, "dropout": 0.1, + # "physical" conditions on material/particle physical properties via + # a small MLP (giant.model.network.ConditionEncoder); "embedding" + # keeps the original learned pdg/material embedding tables — kept + # available as the generalization-comparison baseline. Checkpoints + # from before this option existed have no "conditioning" key and + # load as "embedding" (see giant.model.network.build_models). + "conditioning": "physical", "router": { "enabled": False, "type": "energy", # selects the Router impl from ROUTER_REGISTRY diff --git a/giant/constants.py b/giant/constants.py index 01a16f7..5e26a2f 100644 --- a/giant/constants.py +++ b/giant/constants.py @@ -1,25 +1,50 @@ X_DIM = 9 -# Conditioning continuous-feature width (Phase 2): pre_pos(3), log(pre_E)(1), -# pre_dir(3), layer_id(1). n_sec and log(e_sec) are removed — they are now -# *outputs* predicted by Stage 1, not conditioning inputs. -COND_DIM = 8 +# Original Phase-2 continuous conditioning: pre_pos(3), log(pre_E)(1), +# pre_dir(3), layer_id(1). This is the slice ConditionEncoder's "embedding" +# mode reads from cond_cont (see giant/model/network.py); n_sec and +# log(e_sec) are not part of it — they are *outputs* predicted by Stage 1, +# not conditioning inputs. +COND_DIM_BASE = 8 + +# Particle physical-property conditioning: log(mass)(1), charge(1). See +# giant/particles.py. +PARTICLE_PHYS_DIM = 2 + +# Material physical-property conditioning: Z_eff(1), A_eff(1), +# log(density)(1), log(X0)(1), log(lambda_int)(1). See giant/materials.py. +MATERIAL_PHYS_DIM = 5 + +# Conditioning continuous-feature width. cond_cont is unconditionally this +# wide regardless of model_config["conditioning"]: "physical" mode computes +# the trailing PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM columns for real, +# "embedding" mode zero-fills them (and never reads them) — see +# giant.data.transforms.build_features/build_cond_features. Bumped 8->15 for +# physical-property conditioning, the same kind of breaking bump as Phase 1 +# (10) -> Phase 2 (8). +COND_DIM = COND_DIM_BASE + PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM # 15 # Maximum number of secondary slots. From data: max(n_sec)=14 in PbWO4 dataset; # K_MAX=15 covers it with one spare slot. K_MAX = 15 -# Per-slot secondary target dimension: 1 (stick-breaking logit) + 3 (local dir) + -# EMB_DIM (continuous type embedding). EMB_DIM must match DenoisingMLP.emb_dim. -# Default emb_dim=16 → SEC_SLOT_DIM=20. -SEC_SLOT_DIM = 20 # 1 + 3 + 16 -EMB_DIM = 16 # must match model emb_dim default +# Per-slot continuous width: stick-breaking logit(1) + local dir(3). +CONT_SLOT_DIM = 4 -# Per-slot continuous (non-embedding) width: stick-breaking logit + local dir. -CONT_SLOT_DIM = SEC_SLOT_DIM - EMB_DIM # 4 +# Per-slot secondary target dimension: CONT_SLOT_DIM (stick-breaking logit + +# local dir) + PARTICLE_PHYS_DIM (log(mass), charge — the secondary's +# predicted physical identity, regressed directly against real physics +# targets rather than a learned/snapped embedding). +SEC_SLOT_DIM = CONT_SLOT_DIM + PARTICLE_PHYS_DIM # 6 + +# ConditionEncoder's physical-property sub-MLP output width (see +# giant/model/network.py) and "embedding" mode's pdg_emb/mat_emb width. +# Independent of SEC_SLOT_DIM — unlike Phase 2, Stage 2's per-slot physical +# output width is fixed by PARTICLE_PHYS_DIM, not by this. +EMB_DIM = 16 # Flattened Stage-2 target dimension -SEC_DIM = K_MAX * SEC_SLOT_DIM # 15 * 20 = 300 +SEC_DIM = K_MAX * SEC_SLOT_DIM # 15 * 6 = 90 # Stage-1 9D target names (unchanged from energy-conservation PoC) LOCAL_TARGET_NAMES = [ diff --git a/giant/data/dataset.py b/giant/data/dataset.py index f319d93..c1d07d2 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -36,16 +36,17 @@ class StreamingStepsDataset(IterableDataset): numpy slicing instead of a per-row Python loop in the default collate. Each batch is a tuple: - (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) where: - cond_cont: (B, COND_DIM) float32 - cond_cat: (B, 2) int64 - target_s1: (B, 9) float32 — normalised Stage-1 primary target - n_sec: (B,) int64 — true secondary count per step - sec_cont: (B, K_MAX, 4) float32 — [stick_logit, local_dir] per slot - sec_pdg_idx: (B, K_MAX) int64 — PDG model-index per secondary slot - proc_idx: (B,) int64 — process-class label (ProcessRouter supervision - only; zeros when `proc_map` is None) + cond_cont: (B, COND_DIM) float32 + cond_cat: (B, 2) int64 + target_s1: (B, 9) float32 — normalised Stage-1 primary target + n_sec: (B,) int64 — true secondary count per step + sec_cont: (B, K_MAX, SEC_SLOT_DIM) float32 — [stick_logit, + local_dir, log_mass, charge] per slot (mass/charge + normalised iff `sec_phys_normalizer` was given) + proc_idx: (B,) int64 — process-class label (ProcessRouter supervision + only; zeros when `proc_map` is None) """ def __init__( @@ -60,6 +61,8 @@ class StreamingStepsDataset(IterableDataset): shuffle_buffer: int = 65536, shuffle: bool = True, proc_map: dict[str, int] | None = None, + conditioning: str = "embedding", + sec_phys_normalizer: Normalizer | None = None, ) -> None: self.files = list(files) self.split_events = split_events @@ -72,6 +75,8 @@ class StreamingStepsDataset(IterableDataset): self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle self.proc_map = proc_map + self.conditioning = conditioning + self.sec_phys_normalizer = sec_phys_normalizer def __iter__(self): worker_info = torch.utils.data.get_worker_info() @@ -88,7 +93,6 @@ class StreamingStepsDataset(IterableDataset): buf_tgt: list[np.ndarray] = [] buf_nsec: list[np.ndarray] = [] buf_sec: list[np.ndarray] = [] - buf_spdg: list[np.ndarray] = [] buf_proc: list[np.ndarray] = [] buf_n = 0 @@ -105,7 +109,6 @@ class StreamingStepsDataset(IterableDataset): target_s1, n_sec, sec_cont, - sec_pdg_idx, proc_idx, _, _, @@ -115,15 +118,16 @@ class StreamingStepsDataset(IterableDataset): self.mat_map, cond_normalizer=self.cond_normalizer, target_normalizer=self.target_normalizer, + sec_phys_normalizer=self.sec_phys_normalizer, proc_map=self.proc_map, require_secondaries=True, + conditioning=self.conditioning, ) buf_cont.append(cond_cont) buf_cat.append(cond_cat) buf_tgt.append(target_s1) buf_nsec.append(n_sec) buf_sec.append(sec_cont) - buf_spdg.append(sec_pdg_idx) buf_proc.append(proc_idx) buf_n += len(cond_cont) @@ -134,7 +138,6 @@ class StreamingStepsDataset(IterableDataset): buf_tgt, buf_nsec, buf_sec, - buf_spdg, buf_proc, buf_n, ) = yield from self._flush( @@ -143,7 +146,6 @@ class StreamingStepsDataset(IterableDataset): buf_tgt, buf_nsec, buf_sec, - buf_spdg, buf_proc, final=False, ) @@ -155,7 +157,6 @@ class StreamingStepsDataset(IterableDataset): buf_tgt, buf_nsec, buf_sec, - buf_spdg, buf_proc, final=True, ) @@ -167,7 +168,6 @@ class StreamingStepsDataset(IterableDataset): buf_tgt: list[np.ndarray], buf_nsec: list[np.ndarray], buf_sec: list[np.ndarray], - buf_spdg: list[np.ndarray], buf_proc: list[np.ndarray], final: bool, ): @@ -176,13 +176,12 @@ class StreamingStepsDataset(IterableDataset): tgt = np.concatenate(buf_tgt) nsec = np.concatenate(buf_nsec) sec = np.concatenate(buf_sec) - spdg = np.concatenate(buf_spdg) proc = np.concatenate(buf_proc) if self.shuffle: idx = np.random.permutation(len(cont)) cont, cat, tgt = cont[idx], cat[idx], tgt[idx] - nsec, sec, spdg, proc = nsec[idx], sec[idx], spdg[idx], proc[idx] + nsec, sec, proc = nsec[idx], sec[idx], proc[idx] bs = self.batch_size n = len(cont) @@ -195,12 +194,11 @@ class StreamingStepsDataset(IterableDataset): torch.from_numpy(tgt[start:end]).float(), torch.from_numpy(nsec[start:end]).long(), torch.from_numpy(sec[start:end]).float(), - torch.from_numpy(spdg[start:end]).long(), torch.from_numpy(proc[start:end]).long(), ) if final: - return [], [], [], [], [], [], [], 0 + return [], [], [], [], [], [], 0 rem = n_full * bs return ( [cont[rem:]], @@ -208,7 +206,6 @@ class StreamingStepsDataset(IterableDataset): [tgt[rem:]], [nsec[rem:]], [sec[rem:]], - [spdg[rem:]], [proc[rem:]], n - rem, ) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index f7d647e..9b74ea7 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -289,20 +289,28 @@ def encode_secondaries( sec_valid: np.ndarray, e_sec: np.ndarray, pre_dir: np.ndarray, + sec_pdg_list: np.ndarray | None = None, ) -> np.ndarray: """Encode per-secondary attributes into continuous per-slot targets. Secondaries must already be sorted descending by energy (as stored in the - parquet). Returns sec_cont of shape (N, K_MAX, 4): - slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z] + parquet). Returns sec_cont of shape (N, K_MAX, SEC_SLOT_DIM=6): + slot[i] = [stick_break_logit, local_dir_x, local_dir_y, local_dir_z, + log_mass, charge] Stick-breaking logit: for slot i, f_i = E_i / remaining_budget, where remaining_budget = e_sec - sum(E_0..E_{i-1}). The logit is log(f/(1-f)), clipped to ±_STICK_LOGIT_CLIP. The last valid slot gets +_STICK_LOGIT_CLIP (takes the full remaining budget). Padding slots get 0. - sec_pdg_idx (integer) is not processed here — kept separate so the loss - function can look up the embedding table at training time. + log_mass/charge are the secondary's real physical identity, looked up + from its ground-truth PDG code (`sec_pdg_list`) via + `giant.particles.particle_phys_array` — a fixed physics-derived + regression target, not a learned/moving one (unlike the embedding-table + target this replaced), so nothing needs to be detached at training time. + `sec_pdg_list` is optional so callers that only need the continuous + stick/dir block (e.g. inference-time re-encoding) can omit it; omitting + it zero-fills the last two columns, matching the padding-slot convention. """ N, K = sec_E_list.shape e_sec = np.asarray(e_sec, dtype=np.float64) @@ -336,36 +344,66 @@ def encode_secondaries( pre_dir[valid_mask], sec_dir_list[valid_mask, i] ) + if sec_pdg_list is not None: + from giant.particles import particle_phys_array + + # Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int), + # which isn't a resolvable particle — substitute a dummy resolvable + # code (22, photon) there, since the result is discarded below by + # the sec_valid mask regardless. + safe_pdg = np.where(sec_valid, sec_pdg_list, 22) + flat_mass_charge = particle_phys_array(safe_pdg.reshape(-1)) # (N*K, 2) + mass = flat_mass_charge[:, 0].reshape(N, K) + charge = flat_mass_charge[:, 1].reshape(N, K) + log_mass = np.where(sec_valid, log_transform(mass), 0.0).astype(np.float32) + charge = np.where(sec_valid, charge, 0.0).astype(np.float32) + else: + log_mass = np.zeros((N, K), dtype=np.float32) + charge = np.zeros((N, K), dtype=np.float32) + sec_cont = np.concatenate( - [stick_logits[:, :, None], dir_local], axis=-1 - ) # (N, K, 4) + [stick_logits[:, :, None], dir_local, log_mass[:, :, None], charge[:, :, None]], + axis=-1, + ) # (N, K, 6) return sec_cont.astype(np.float32) def decode_secondaries( sec_cont: np.ndarray, - sec_pdg_pred: np.ndarray, n_sec: np.ndarray, e_sec: np.ndarray, pre_dir: np.ndarray, - pdg_map_inv: dict[int, int], -) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + sec_phys_normalizer: "Normalizer | None" = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Inverse of encode_secondaries: continuous targets → physical secondary attrs. - sec_cont: (N, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z] - sec_pdg_pred: (N, K_MAX) integer PDG indices (from nearest-neighbor snap) + sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y, + local_dir_z, log_mass, charge] (log_mass/charge normalised iff + `sec_phys_normalizer` was applied when this was produced — e.g. a + raw model prediction; pass the same normalizer here to invert it) n_sec: (N,) integer secondary counts e_sec: (N,) total secondary energy budget [MeV] pre_dir: (N, 3) pre-step world-frame direction - pdg_map_inv: maps model index → PDG code - Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX). - The valid slots' energies (`sec_E[sec_valid]`, per row) always sum to - exactly `e_sec` — see the rescaling below. + Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each + shape (N, K_MAX). The valid slots' energies (`sec_E[sec_valid]`, per row) + always sum to exactly `e_sec` — see the rescaling below. mass/charge are + the model's raw predicted physical identity for each secondary, used + as-is (no snapping to a discrete PDG code) — see giant/particles.py for + the separate, reporting-only nearest-PDG lookup callers may apply on top + of this for display/bookkeeping purposes. """ + if sec_phys_normalizer is not None: + N_, K_, _ = sec_cont.shape + phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2)) + sec_cont = sec_cont.copy() + sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2) + N, K, _ = sec_cont.shape stick_logits = sec_cont[:, :, 0] # (N, K) - dir_local = sec_cont[:, :, 1:].copy() # (N, K, 3) + dir_local = sec_cont[:, :, 1:4].copy() # (N, K, 3) + log_mass = sec_cont[:, :, 4] # (N, K) + charge = sec_cont[:, :, 5] # (N, K) # Flow-matching output isn't guaranteed unit norm; normalise before the # rotation below, which preserves magnitude rather than fixing it up. @@ -409,15 +447,62 @@ def decode_secondaries( pre_dir[valid], dir_local[valid, i] ) - sec_pdg_code = np.array( - [ - [pdg_map_inv.get(int(sec_pdg_pred[n, i]), 0) for i in range(K)] - for n in range(N) - ], - dtype=np.int32, - ) + # mass is non-negative by construction (inv_log_transform of a real + # number is always > 0); clip to 0 for padded/invalid slots rather than + # leaving a spurious small positive floor from the log inverse. + sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32) + sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32) - return sec_E, sec_dir_world, sec_pdg_code, sec_valid + return sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid + + +def _physical_cond_columns( + data: dict[str, np.ndarray], conditioning: str +) -> np.ndarray: + """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns. + + "embedding" mode zero-fills (cheap, and ConditionEncoder never reads + these columns in that mode — so an unfilled giant.materials table can + never crash an "embedding"-mode run). "physical" mode computes them for + real: particle columns come from `data["mass"]`/`data["charge"]` when the + caller already knows them directly (rollout.py, for a track descended + from a model-predicted secondary — see giant/rollout.py's "no snapping" + design), else derived from `data["pdg"]` via giant.particles; material + columns always come from `data["material"]` via giant.materials, since + material is never itself a model prediction. + """ + from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM + + if conditioning == "embedding": + n = len(next(iter(data.values()))) + return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32) + if conditioning != "physical": + raise ValueError(f"unknown conditioning mode {conditioning!r}") + + from giant.materials import material_properties_array + from giant.particles import particle_phys_array + + if "mass" in data and "charge" in data: + mass = np.asarray(data["mass"], dtype=np.float32) + charge = np.asarray(data["charge"], dtype=np.float32) + else: + mass, charge = particle_phys_array(data["pdg"]).T + + z_eff, a_eff, density, x0, lambda_int = material_properties_array( + data["material"] + ).T + + return np.column_stack( + [ + log_transform(mass), + charge, + z_eff, + a_eff, + log_transform(density), + log_transform(x0), + log_transform(lambda_int), + ] + ).astype(np.float32) def build_cond_features( @@ -425,6 +510,7 @@ def build_cond_features( pdg_map: dict[int, int], mat_map: dict[str, int], cond_normalizer: "Normalizer | None" = None, + conditioning: str = "embedding", ) -> tuple[np.ndarray, np.ndarray]: """Build conditioning arrays only — no target, no post-step variables.""" cond_cont = np.column_stack( @@ -435,6 +521,9 @@ def build_cond_features( data["layer_id"].astype(np.float32), ] ).astype(np.float32) + cond_cont = np.column_stack( + [cond_cont, _physical_cond_columns(data, conditioning)] + ).astype(np.float32) pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) @@ -452,9 +541,11 @@ def build_features( mat_map: dict[str, int], cond_normalizer: Normalizer | None = None, target_normalizer: Normalizer | None = None, + sec_phys_normalizer: Normalizer | None = None, fit: bool = False, proc_map: dict[str, int] | None = None, require_secondaries: bool = False, + conditioning: str = "embedding", ) -> tuple[ np.ndarray, np.ndarray, @@ -462,17 +553,17 @@ def build_features( np.ndarray, np.ndarray, np.ndarray, - np.ndarray, Normalizer | None, Normalizer | None, ]: - """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx) arrays. + """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) arrays. target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1) n_sec: (N,) integer secondary counts (target for n_sec head) - sec_cont: (N, K_MAX, 4) continuous secondary targets [stick_logit, dir_local] - sec_pdg_idx: (N, K_MAX) integer PDG model-indices; used to look up embedding - targets in the training loop + sec_cont: (N, K_MAX, SEC_SLOT_DIM=6) continuous secondary targets + [stick_logit, dir_local, log_mass, charge] — mass/charge are + the secondary's real physical identity (from its ground-truth + PDG code), a fixed regression target, not a learned/snapped one. proc_idx: (N,) integer process-class label (ProcessRouter supervision only — never conditioning). Zeros when `proc_map` is None or the loaded data has no "process" column (e.g. pre-conversion parquet files). @@ -510,7 +601,10 @@ def build_features( data["pre_dir"], data["layer_id"].astype(np.float32), ] - ).astype(np.float32) # (N, COND_DIM=8) + ).astype(np.float32) # (N, COND_DIM_BASE=8) + cond_cont = np.column_stack( + [cond_cont, _physical_cond_columns(data, conditioning)] + ).astype(np.float32) # (N, COND_DIM=15) pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64) mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64) @@ -521,9 +615,9 @@ def build_features( ) # (N,) unclamped, for the valid-slot mask # Clamp the classification label to K_MAX: the head only has K_MAX+1 classes # (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already - # applied to sec_cont/sec_pdg_idx by the loader's list padding. Without this, - # a rare high-multiplicity step (real data goes up to ~37) hands - # cross_entropy an out-of-range target and CUDA asserts. + # applied to sec_cont by the loader's list padding. Without this, a rare + # high-multiplicity step (real data goes up to ~37) hands cross_entropy + # an out-of-range target and CUDA asserts. n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64) # (N,) # Secondary continuous targets @@ -534,14 +628,13 @@ def build_features( if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None: sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None] # (N, K_MAX) sec_cont = encode_secondaries( - sec_E_list, sec_dir_list, sec_valid, data["e_sec"], data["pre_dir"] - ) # (N, K_MAX, 4) - # Padding slots carry sentinel pdg 0 (see loader._pad_list_col_int), - # which is never a real PDG code, so `.get(..., 0)` naturally maps - # both real unknown codes and padding to the same masked-out index. - sec_pdg_idx = np.vectorize(lambda p: pdg_map.get(int(p), 0))( - sec_pdg_list - ).astype(np.int64) + sec_E_list, + sec_dir_list, + sec_valid, + data["e_sec"], + data["pre_dir"], + sec_pdg_list=sec_pdg_list, + ) # (N, K_MAX, 6) else: # Guard against silently training Stage 2 on zeroed targets: if any step # actually spawned secondaries (n_sec > 0, from child_track_ids) but the @@ -563,8 +656,7 @@ def build_features( "require_secondaries=False for Stage-1-only use." ) N = len(n_sec) - sec_cont = np.zeros((N, K_MAX, 4), dtype=np.float32) - sec_pdg_idx = np.zeros((N, K_MAX), dtype=np.int64) + sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32) if fit: cond_normalizer = Normalizer().fit(cond_cont) @@ -574,6 +666,11 @@ def build_features( cond_cont = cond_normalizer.transform(cond_cont) if target_normalizer is not None: target_s1 = target_normalizer.transform(target_s1) + if sec_phys_normalizer is not None: + N_, K_, _ = sec_cont.shape + phys = sec_phys_normalizer.transform(sec_cont[:, :, 4:6].reshape(-1, 2)) + sec_cont = sec_cont.copy() + sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2) process = data.get("process") if proc_map is not None and process is not None: @@ -587,7 +684,6 @@ def build_features( target_s1, n_sec, sec_cont, - sec_pdg_idx, proc_idx, cond_normalizer, target_normalizer, diff --git a/giant/materials.py b/giant/materials.py new file mode 100644 index 0000000..a15c8d2 --- /dev/null +++ b/giant/materials.py @@ -0,0 +1,145 @@ +"""Material physical-property table for the "physical" conditioning mode. + +Values are Geant4's own built-in NIST material constants, not hand-typed +literature numbers -- extracted directly from a Geant4 11.4.1 build (the one +vendored in /home/lars/Programming/minicalosim/lib/geant4, built at +minicalosim/build/geant4-install) via a small standalone C++ program linked +against that build (G4NistManager::FindOrBuildMaterial + G4Material:: +GetDensity/GetRadlen/GetNuclearInterLength + G4IonisParamMat::GetZeffective). +`a_eff` isn't directly exposed by Geant4, so it's computed with the same +atomic-number-density-weighted-average formula Geant4 itself uses for Zeff +(see G4IonisParamMat::BuildFluctModel in +lib/geant4/source/materials/src/G4IonisParamMat.cc), just applied to A +instead of Z -- for a single-element material this is exact; for a compound +it matches Geant4's own effective-Z convention rather than a different +weighting scheme. + +Never silently substitute a default for a material missing from this table +(see UnknownMaterialError/MaterialPropertiesNotFilledError below) -- a wrong +material property would corrupt a whole conditioning axis without any +visible symptom until deep into training. +""" + +from __future__ import annotations + +from typing import NamedTuple + +import numpy as np + + +class MaterialProperties(NamedTuple): + z_eff: float | None # effective atomic number + a_eff: float | None # effective atomic mass [g/mol] + density: float | None # [g/cm^3] + x0: float | None # radiation length [cm] + lambda_int: float | None # nuclear interaction length [cm] + + +class UnknownMaterialError(KeyError): + pass + + +class MaterialPropertiesNotFilledError(NotImplementedError): + pass + + +# Keys: every NIST material name seen in +# physics/detector-design/minicalosim-geometry.md, plus G4_AIR/G4_lAr which +# appear in dataset parquet files but not that doc. Values from Geant4's +# built-in NIST database (see module docstring) -- all present except +# G4_LYSO, which is not actually a stock Geant4 NIST material (confirmed: +# G4NistManager::FindOrBuildMaterial("G4_LYSO") fails to build in the +# vendored Geant4 11.4.1; it only appears as a plotting-color key in +# minicalosim/bind/G4Calo.py, never constructed in DetectorConstruction.cc) +# -- left unfilled until it's either built as a custom material (e.g. +# Lu1.8Y0.2SiO5:Ce) or dropped from the geometry menu. +MATERIAL_PROPERTIES: dict[str, MaterialProperties] = { + "G4_PbWO4": MaterialProperties( + z_eff=31.333333, + a_eff=75.843426, + density=8.28, + x0=0.892453, + lambda_int=20.739740, + ), + "G4_CESIUM_IODIDE": MaterialProperties( + z_eff=54.0, a_eff=129.904539, density=4.51, x0=1.860288, lambda_int=39.305990 + ), + "G4_Pb": MaterialProperties( + z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950 + ), + "G4_W": MaterialProperties( + z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580 + ), + "G4_Cu": MaterialProperties( + z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940 + ), + "G4_Fe": MaterialProperties( + z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300 + ), + "G4_BRASS": MaterialProperties( + z_eff=30.939130, + a_eff=68.500857, + density=8.52, + x0=1.367465, + lambda_int=16.947420, + ), + "G4_POLYSTYRENE": MaterialProperties( + z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880 + ), + "G4_PLASTIC_SC_VINYLTOLUENE": MaterialProperties( + z_eff=3.368421, + a_eff=6.219791, + density=1.032, + x0=42.544200, + lambda_int=69.969390, + ), + "G4_BGO": MaterialProperties( + z_eff=27.578947, + a_eff=65.565839, + density=7.13, + x0=1.118030, + lambda_int=22.710130, + ), + "G4_LYSO": MaterialProperties(None, None, None, None, None), + "G4_AIR": MaterialProperties( + z_eff=7.261982, + a_eff=14.547593, + density=1.204790e-3, + x0=30392.070000, + lambda_int=71009.500000, + ), + "G4_lAr": MaterialProperties( + z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400 + ), +} + + +def get_material_properties( + name: str, table: dict[str, MaterialProperties] | None = None +) -> MaterialProperties: + t = MATERIAL_PROPERTIES if table is None else table + if name not in t: + raise UnknownMaterialError( + f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES " + f"-- add it (known: {sorted(t)})" + ) + props = t[name] + if any(v is None for v in props): + raise MaterialPropertiesNotFilledError( + f"material {name!r} has un-filled physical properties in " + "giant/materials.py -- a physicist must populate real " + "z_eff/a_eff/density/x0/lambda_int values before " + "conditioning='physical' can be used with this material" + ) + return props + + +def material_properties_array( + names: np.ndarray, table: dict[str, MaterialProperties] | None = None +) -> np.ndarray: + """(N,) str material names -> (N, 5) float32 [z_eff, a_eff, density, x0, lambda_int].""" + out = np.array( + [get_material_properties(str(m), table) for m in np.asarray(names)], + dtype=np.float32, + ) + return out.reshape(-1, 5) diff --git a/giant/model/network.py b/giant/model/network.py index 57d25df..9646a0c 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -6,7 +6,16 @@ import torch import torch.nn as nn import torch.nn.functional as F -from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM +from giant.constants import ( + COND_DIM, + COND_DIM_BASE, + EMB_DIM, + K_MAX, + MATERIAL_PHYS_DIM, + PARTICLE_PHYS_DIM, + SEC_DIM, + X_DIM, +) class SinusoidalEmbedding(nn.Module): @@ -28,6 +37,23 @@ class SinusoidalEmbedding(nn.Module): class ConditionEncoder(nn.Module): + """Fuses continuous conditioning with particle/material identity. + + Two mutually exclusive ways to turn (pdg, material) identity into the + two `emb_dim`-wide vectors concatenated with the base continuous + conditioning before the fusion MLP: + - "embedding": a learned `nn.Embedding` lookup table per axis, indexed + by `cond_cat`'s dense training-vocab index. Memorizes the training + menu; the original Phase-2 design. + - "physical": a small MLP per axis, mapping the axis's raw physical + properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see + giant.data.transforms.build_features) to an `emb_dim`-wide vector — + a drop-in replacement for the embedding lookup, computable for any + PDG code / material name rather than only ones seen in training. + Both modes produce the same `in_dim = COND_DIM_BASE + 2*emb_dim` for the + fusion MLP, so only how the two vectors are produced differs. + """ + def __init__( self, pdg_vocab: int, @@ -35,11 +61,27 @@ class ConditionEncoder(nn.Module): cont_dim: int = COND_DIM, emb_dim: int = 16, out_dim: int = 128, + conditioning: str = "embedding", ) -> None: super().__init__() - self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) - self.mat_emb = nn.Embedding(mat_vocab, emb_dim) - in_dim = cont_dim + 2 * emb_dim + if conditioning not in ("embedding", "physical"): + raise ValueError(f"unknown conditioning mode {conditioning!r}") + self.conditioning = conditioning + if conditioning == "embedding": + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + self.mat_emb = nn.Embedding(mat_vocab, emb_dim) + else: + self.particle_mlp = nn.Sequential( + nn.Linear(PARTICLE_PHYS_DIM, emb_dim), + nn.SiLU(), + nn.Linear(emb_dim, emb_dim), + ) + self.material_mlp = nn.Sequential( + nn.Linear(MATERIAL_PHYS_DIM, emb_dim), + nn.SiLU(), + nn.Linear(emb_dim, emb_dim), + ) + in_dim = COND_DIM_BASE + 2 * emb_dim self.mlp = nn.Sequential( nn.Linear(in_dim, out_dim), nn.SiLU(), @@ -47,9 +89,17 @@ class ConditionEncoder(nn.Module): ) def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - pdg_e = self.pdg_emb(cond_cat[:, 0]) - mat_e = self.mat_emb(cond_cat[:, 1]) - x = torch.cat([cond_cont, pdg_e, mat_e], dim=-1) + if self.conditioning == "embedding": + pdg_e = self.pdg_emb(cond_cat[:, 0]) + mat_e = self.mat_emb(cond_cat[:, 1]) + else: + particle_phys = cond_cont[ + :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM + ] + material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :] + pdg_e = self.particle_mlp(particle_phys) + mat_e = self.material_mlp(material_phys) + x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1) return self.mlp(x) @@ -91,6 +141,7 @@ class DenoisingMLP(nn.Module): x_dim: int = X_DIM, dropout: float = 0.1, k_max: int = K_MAX, + conditioning: str = "embedding", ) -> None: super().__init__() self.time_emb = SinusoidalEmbedding(time_dim) @@ -99,6 +150,7 @@ class DenoisingMLP(nn.Module): mat_vocab=mat_vocab, emb_dim=emb_dim, out_dim=cond_out_dim, + conditioning=conditioning, ) merged_cond_dim = time_dim + cond_out_dim self.input_proj = nn.Linear(x_dim, hidden_dim) @@ -141,10 +193,6 @@ class DenoisingMLP(nn.Module): c_emb = self.cond_enc(cond_cont, cond_cat) return self.n_sec_head(c_emb) - def pdg_embedding_weight(self) -> torch.Tensor: - """Return the PDG embedding table weights for secondary type targets.""" - return self.cond_enc.pdg_emb.weight - class SecondaryConditionEncoder(nn.Module): """Encodes pre-step conditioning + Stage-1 output for the secondary decoder.""" @@ -158,6 +206,7 @@ class SecondaryConditionEncoder(nn.Module): stage1_dim: int = X_DIM, stage1_proj_dim: int = 64, out_dim: int = 128, + conditioning: str = "embedding", ) -> None: super().__init__() self.base = ConditionEncoder( @@ -165,6 +214,7 @@ class SecondaryConditionEncoder(nn.Module): mat_vocab=mat_vocab, emb_dim=emb_dim, out_dim=cond_out_dim, + conditioning=conditioning, ) self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim) fused_dim = cond_out_dim + stage1_proj_dim @@ -187,8 +237,12 @@ class SecondaryConditionEncoder(nn.Module): class SecondaryDecoder(nn.Module): """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously. - Each slot encodes (stick_break_logit, local_dir_3D, type_emb) for one - secondary ordered by descending energy. Padded slots are masked from loss. + Each slot encodes (stick_break_logit, local_dir_3D, log_mass, charge) for + one secondary ordered by descending energy — mass/charge are the + secondary's predicted physical identity, regressed directly against real + physics targets (see giant.data.transforms.encode_secondaries), used + as-is with no snapping to a discrete PDG code. Padded slots are masked + from loss. """ def __init__( @@ -203,6 +257,7 @@ class SecondaryDecoder(nn.Module): stage1_proj_dim: int = 64, sec_dim: int = SEC_DIM, dropout: float = 0.1, + conditioning: str = "embedding", ) -> None: super().__init__() self.time_emb = SinusoidalEmbedding(time_dim) @@ -213,6 +268,7 @@ class SecondaryDecoder(nn.Module): cond_out_dim=cond_out_dim, stage1_proj_dim=stage1_proj_dim, out_dim=cond_out_dim, + conditioning=conditioning, ) merged_cond_dim = time_dim + cond_out_dim self.input_proj = nn.Linear(sec_dim, hidden_dim) @@ -587,8 +643,8 @@ class RoutedDenoisingMLP(nn.Module): Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all tiny) across experts and routes only the trunk (where the FLOPs are). - Same `forward`/`predict_n_sec`/`pdg_embedding_weight` signatures as - `DenoisingMLP`, so sample.py/rollout.py/validate.py need no changes. + Same `forward`/`predict_n_sec` signatures as `DenoisingMLP`, so + sample.py/rollout.py/validate.py need no changes. """ def __init__( @@ -604,6 +660,7 @@ class RoutedDenoisingMLP(nn.Module): x_dim: int = X_DIM, dropout: float = 0.1, k_max: int = K_MAX, + conditioning: str = "embedding", ) -> None: super().__init__() self.router = router @@ -613,6 +670,7 @@ class RoutedDenoisingMLP(nn.Module): mat_vocab=mat_vocab, emb_dim=emb_dim, out_dim=cond_out_dim, + conditioning=conditioning, ) merged_cond_dim = time_dim + cond_out_dim self.experts = nn.ModuleList( @@ -656,10 +714,6 @@ class RoutedDenoisingMLP(nn.Module): c_emb = self.cond_enc(cond_cont, cond_cat) return self.n_sec_head(c_emb) - def pdg_embedding_weight(self) -> torch.Tensor: - """Return the PDG embedding table weights for secondary type targets.""" - return self.cond_enc.pdg_emb.weight - class RoutedSecondaryDecoder(nn.Module): """Routed drop-in for `SecondaryDecoder`. @@ -682,6 +736,7 @@ class RoutedSecondaryDecoder(nn.Module): stage1_proj_dim: int = 64, sec_dim: int = SEC_DIM, dropout: float = 0.1, + conditioning: str = "embedding", ) -> None: super().__init__() self.router = router @@ -693,6 +748,7 @@ class RoutedSecondaryDecoder(nn.Module): cond_out_dim=cond_out_dim, stage1_proj_dim=stage1_proj_dim, out_dim=cond_out_dim, + conditioning=conditioning, ) merged_cond_dim = time_dim + cond_out_dim self.experts = nn.ModuleList( @@ -732,6 +788,7 @@ _STAGE1_MODEL_KEYS = { "emb_dim", "dropout", "k_max", + "conditioning", } _SEC_DECODER_MODEL_KEYS = { "pdg_vocab", @@ -740,6 +797,7 @@ _SEC_DECODER_MODEL_KEYS = { "n_blocks", "emb_dim", "dropout", + "conditioning", } @@ -798,6 +856,13 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: is truthy; a missing/absent "router" key (pre-routing checkpoints) falls back to the monolithic pair unchanged, so this is a drop-in replacement for the ad-hoc constructions it replaces. + + `model_config.get("conditioning", "embedding")` — old checkpoints have no + "conditioning" key and must keep loading with their original embedding + tables, so the default here is "embedding", not the training-time + default (which is "physical" — see giant.config.DEFAULT_CONFIG). Read + once and passed to both stage1/sec_decoder, so they structurally always + share one mode. """ router_cfg = model_config.get("router") if router_cfg and router_cfg.get("enabled"): @@ -810,6 +875,7 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: expert_n_blocks=model_config.get("expert_n_blocks", 3), emb_dim=model_config.get("emb_dim", EMB_DIM), dropout=model_config.get("dropout", 0.1), + conditioning=model_config.get("conditioning", "embedding"), ) stage1 = RoutedDenoisingMLP( router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab), diff --git a/giant/model/schedule.py b/giant/model/schedule.py index a5a384c..05d1e18 100644 --- a/giant/model/schedule.py +++ b/giant/model/schedule.py @@ -81,19 +81,22 @@ def flow_matching_loss_secondary( ) -> torch.Tensor: """Flow matching loss for the secondary decoder with per-slot masking. - x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, type_emb) + x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, log_mass, charge) sec_mask: (B, K_MAX) bool — True for valid secondary slots Only valid-slot dimensions contribute to the loss; padded slots are zeroed before averaging, so the loss is not diluted by empty slots. Each slot packs CONT_SLOT_DIM continuous dims (stick_logit, dir) followed - by EMB_DIM type-embedding dims. A flat per-dimension mean would let the - 16 embedding dims outvote the 4 physically-interesting ones, so the two - blocks are each averaged over their own width first and then combined - with equal weight — this stays correct if EMB_DIM/CONT_SLOT_DIM change. + by PARTICLE_PHYS_DIM physical-identity dims (log_mass, charge) — the + secondary's predicted physical identity, a fixed regression target (see + giant.data.transforms.encode_secondaries). Even though the two blocks are + the same order of magnitude now (unlike the 16-wide learned embedding + block this replaced), they're still on different physical scales, so + they're each averaged over their own width first and then combined with + equal weight — this stays correct if PARTICLE_PHYS_DIM/CONT_SLOT_DIM change. """ - from giant.constants import CONT_SLOT_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM + from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM B = x1.size(0) t = torch.rand(B, device=x1.device) @@ -104,10 +107,10 @@ def flow_matching_loss_secondary( err = ((v_t - u_t) ** 2).view(B, K_MAX, SEC_SLOT_DIM) cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1) # (B, K_MAX) - emb_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + EMB_DIM].mean(dim=-1) + phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM].mean(dim=-1) mask = sec_mask.float() denom = mask.sum().clamp(min=1) cont_loss = (cont_err * mask).sum() / denom - emb_loss = (emb_err * mask).sum() / denom - return cont_loss + emb_loss + phys_loss = (phys_err * mask).sum() / denom + return cont_loss + phys_loss diff --git a/giant/particles.py b/giant/particles.py new file mode 100644 index 0000000..ef149ab --- /dev/null +++ b/giant/particles.py @@ -0,0 +1,102 @@ +"""Particle physical-property lookup (mass, charge) for "physical" conditioning. + +Uses the scikit-HEP `particle` package (PDG data tables) for standard particles +and ground-state nuclei; falls back to the Z/A decode formula (PDG's 10-digit +ion scheme `10LZZZAAAI`: Z and A decoded straight from the digits, no lookup +table involved) for isomer/excited nuclear codes the package's ground-state-only +nuclide table doesn't cover — confirmed necessary for ~32% of the nuclear codes +actually present in the multi-material dataset +(`0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet`). +""" + +from __future__ import annotations + +from functools import lru_cache + +import numpy as np +from particle import InvalidParticle, Particle, ParticleNotFound +from particle import pdgid as _pdgid + +# First-pass nuclear mass approximation (A * atomic mass unit); no +# binding-energy correction. Only used for codes missing from `particle`'s +# ground-state nuclide table -- ground-state codes get the package's real +# (binding-energy-corrected) mass. +_AMU_MEV = 931.494 + +# Nearest-neighbour distance weight for `nearest_known_pdg`: charge is a +# small conserved quantum number and should usually match exactly, so it's +# weighted far more heavily than the (already log-scaled) mass term. +_CHARGE_WEIGHT = 50.0 +_LOG_EPS = 1e-8 + + +@lru_cache(maxsize=None) +def particle_mass_charge(pdg: int) -> tuple[float, float]: + """Return (mass_MeV, charge_e) for a raw PDG code. + + Cached per unique code: the training vocabulary is typically O(100) + unique codes while a dataset can have O(1e8) rows, and each entry's + lookup (package query + possible ion decode) is nontrivial enough to be + worth memoizing rather than repeating per row. + """ + pdg = int(pdg) + try: + p = Particle.from_pdgid(pdg) + except (ParticleNotFound, InvalidParticle): + if _pdgid.is_nucleus(pdg): + z, a = _pdgid.Z(pdg), _pdgid.A(pdg) + if z is None or a is None: + raise ValueError( + f"PDG {pdg}: is_nucleus but Z/A decode failed" + ) from None + return float(a) * _AMU_MEV, float(z) + raise ValueError( + f"PDG code {pdg} could not be resolved via the `particle` package " + "and is not a nuclear/ion code (is_nucleus=False) -- no fallback " + "available; add explicit handling if this is a legitimate code" + ) from None + # Neutrinos have unmeasured mass in the PDG tables (Particle.mass is + # None) -- treat as exactly 0, same physical treatment as the photon. + mass = 0.0 if p.mass is None else float(p.mass) + charge = 0.0 if p.charge is None else float(p.charge) + return mass, charge + + +def particle_phys_array(pdg_codes: np.ndarray) -> np.ndarray: + """(N,) int PDG codes -> (N, 2) float32 [mass_MeV, charge_e].""" + out = np.array( + [particle_mass_charge(int(p)) for p in np.asarray(pdg_codes)], + dtype=np.float32, + ) + return out.reshape(-1, 2) + + +def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.ndarray: + """Reporting-only nearest-PDG label for predicted (mass, charge) pairs. + + Never used in the inference/training path -- a Stage-2 secondary's + physical identity is always its raw predicted (mass, charge). This is + only for populating an output row's nominal "pdg" column and as an + "embedding" conditioning-mode fallback lookup key for tracks with no + real PDG code (see giant/rollout.py). Nearest neighbour in + (log_mass, charge) space over `candidates` (an iterable of PDG codes, + typically a `pdg_map`'s keys — the training vocabulary), weighting + 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) + + mass = np.asarray(mass, dtype=np.float64) + charge = np.asarray(charge, dtype=np.float64) + query_log_mass = np.log(np.maximum(mass, 0.0) + _LOG_EPS) + + d2 = (query_log_mass[:, None] - table_log_mass[None, :]) ** 2 + _CHARGE_WEIGHT * ( + charge[:, None] - table_charge[None, :] + ) ** 2 + idx = d2.argmin(axis=1) + return codes[idx] diff --git a/giant/pipeline.py b/giant/pipeline.py index d6522c7..f3db978 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -5,7 +5,14 @@ import torch from torch.utils.data import DataLoader from giant import config -from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_SLOT_DIM, X_DIM +from giant.constants import ( + COND_DIM, + EMB_DIM, + K_MAX, + PARTICLE_PHYS_DIM, + SEC_SLOT_DIM, + X_DIM, +) from giant.data.loader import ( find_parquet_files, load_event_ids, @@ -67,27 +74,33 @@ def run_train_job( ) echo("fitting normalizer (streaming) …") + conditioning = m["conditioning"] cond_acc = _WelfordAccumulator(COND_DIM) tgt_acc = _WelfordAccumulator(X_DIM) + sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM) for path in files: for chunk in iter_file_chunks(path): mask = np.isin(chunk["event_id"], events_arr) if not mask.any(): continue chunk_tr = {k: v[mask] for k, v in chunk.items()} - cond_cont, _, target_s1, _n_sec, _sec_cont, _sec_pdg, _proc, _, _ = ( - build_features( - chunk_tr, - pdg_map, - mat_map, - proc_map=proc_map, - require_secondaries=True, - ) + cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features( + chunk_tr, + pdg_map, + mat_map, + proc_map=proc_map, + require_secondaries=True, + conditioning=conditioning, ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) + sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None] + sec_phys = sec_cont[:, :, 4:6][sec_valid] + if len(sec_phys) > 0: + sec_phys_acc.update(sec_phys) cond_norm = cond_acc.to_normalizer() tgt_norm = tgt_acc.to_normalizer() + sec_phys_norm = sec_phys_acc.to_normalizer() train_ds = StreamingStepsDataset( files=files, @@ -100,6 +113,8 @@ def run_train_job( shuffle_buffer=shuffle_buffer, shuffle=True, proc_map=proc_map, + conditioning=conditioning, + sec_phys_normalizer=sec_phys_norm, ) val_ds = StreamingStepsDataset( files=files, @@ -111,6 +126,8 @@ def run_train_job( batch_size=t["batch_size"], shuffle=False, proc_map=proc_map, + conditioning=conditioning, + sec_phys_normalizer=sec_phys_norm, ) pin = device.type == "cuda" @@ -128,11 +145,6 @@ def run_train_job( ) emb_dim = m.get("emb_dim", EMB_DIM) - # SEC_SLOT_DIM must match constants (1 stick + 3 dir + emb_dim) - assert SEC_SLOT_DIM == 1 + 3 + emb_dim, ( - f"SEC_SLOT_DIM={SEC_SLOT_DIM} must equal 1+3+emb_dim={1 + 3 + emb_dim}; " - "update giant/constants.py if emb_dim changed" - ) model_config = { "pdg_vocab": len(pdg_map), @@ -143,6 +155,7 @@ def run_train_job( "dropout": m["dropout"], "k_max": K_MAX, "sec_slot_dim": SEC_SLOT_DIM, + "conditioning": conditioning, "router": dict(router_cfg), "expert_hidden_dim": router_cfg["expert_hidden_dim"], "expert_n_blocks": router_cfg["expert_n_blocks"], @@ -183,7 +196,11 @@ def run_train_job( lambda_s2=t.get("lambda_s2", 1.0), lambda_balance=router_cfg.get("lambda_balance", 0.0), lambda_proc=router_cfg.get("lambda_proc", 0.0), - normalizer_dict={"cond": cond_norm.to_dict(), "target": tgt_norm.to_dict()}, + normalizer_dict={ + "cond": cond_norm.to_dict(), + "target": tgt_norm.to_dict(), + "sec_phys": sec_phys_norm.to_dict(), + }, pdg_map={str(k): v for k, v in pdg_map.items()}, mat_map={str(k): v for k, v in mat_map.items()}, proc_map=proc_map, diff --git a/giant/rollout.py b/giant/rollout.py index 8a9773d..2d7c03d 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -38,7 +38,8 @@ from giant.data.transforms import ( inv_log_transform, reconstruct_post_pos, ) -from giant.sample import sample_flow, sample_secondaries, snap_type_to_pdg_idx +from giant.particles import nearest_known_pdg, particle_phys_array +from giant.sample import sample_flow, sample_secondaries # Record columns produced per step / per terminal marker. _RECORD_KEYS = [ @@ -82,6 +83,8 @@ def _empty_frontier() -> dict[str, np.ndarray]: "pre_pos": np.empty((0, 3), dtype=np.float64), "pre_E": np.empty(0, dtype=np.float64), "pre_dir": np.empty((0, 3), dtype=np.float64), + "mass": np.empty(0, dtype=np.float64), + "charge": np.empty(0, dtype=np.float64), } @@ -215,16 +218,24 @@ def make_seed_frontier( dir_ = np.asarray(pre_dir, dtype=np.float64) 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 + frontier = { "event_id": event_id, "track_id": track_id, "parent_id": np.full(n, -1, dtype=np.int64), "generation": np.zeros(n, dtype=np.int64), "step_in_track": np.zeros(n, dtype=np.int64), - "pdg": np.asarray(pdg, dtype=np.int64), + "pdg": pdg_arr, "pre_pos": np.asarray(pre_pos, dtype=np.float64), "pre_E": np.asarray(pre_E, dtype=np.float64), "pre_dir": dir_, + "mass": mass.astype(np.float64), + "charge": charge.astype(np.float64), } return frontier, counts @@ -272,6 +283,7 @@ def rollout( seeds: dict[str, np.ndarray], cond_norm: Normalizer, tgt_norm: Normalizer, + sec_phys_norm: Normalizer, pdg_map: dict[int, int], mat_map: dict[str, int], *, @@ -283,6 +295,7 @@ def rollout( max_tracks_per_event: int | None = None, escape_threshold: float | None = None, on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None, + conditioning: str = "embedding", ) -> dict[str, np.ndarray] | RolloutSummary: """Run showers to completion. @@ -302,9 +315,6 @@ def rollout( if escape_threshold is not None: oracle.escape_threshold = float(escape_threshold) - pdg_map_inv = {v: k for k, v in pdg_map.items()} - pdg_emb_weight = stage1_model.pdg_embedding_weight() - frontier, counts = make_seed_frontier( seeds["event_id"], seeds["pdg"], @@ -327,10 +337,9 @@ def rollout( oracle, cond_norm, tgt_norm, + sec_phys_norm, pdg_map, mat_map, - pdg_map_inv, - pdg_emb_weight, rec, counts, energy_cutoff, @@ -338,6 +347,7 @@ def rollout( steps, device, max_tracks_per_event, + conditioning, ) ) frontier = _concat_frontiers(next_parts) @@ -357,10 +367,9 @@ def _step_chunk( oracle, cond_norm, tgt_norm, + sec_phys_norm, pdg_map, mat_map, - pdg_map_inv, - pdg_emb_weight, rec, counts, energy_cutoff, @@ -368,6 +377,7 @@ def _step_chunk( steps, device, max_tracks_per_event, + conditioning, ) -> dict[str, np.ndarray]: """Advance one chunk of tracks by a single step; return the next frontier.""" n = len(tr["event_id"]) @@ -421,6 +431,11 @@ def _step_chunk( layer_id = tr["_layer_id"] # --- Build conditioning and run the two stages --- + # "mass"/"charge" are the track's own already-resolved physical identity + # (real for a primary, the model's raw predicted values with no snapping + # for a track descended from a secondary — see _spawn_secondaries), used + # directly instead of re-deriving via a pdg lookup. "pdg" still flows + # through for cond_cat's embedding-mode index and the known_pdg gate. cond_dict = { "pre_pos": tr["pre_pos"], "pre_E": tr["pre_E"], @@ -428,8 +443,12 @@ def _step_chunk( "layer_id": layer_id, "material": material, "pdg": tr["pdg"], + "mass": tr["mass"], + "charge": tr["charge"], } - cond_cont, cond_cat = build_cond_features(cond_dict, pdg_map, mat_map, cond_norm) + cond_cont, cond_cat = build_cond_features( + cond_dict, pdg_map, mat_map, cond_norm, conditioning=conditioning + ) cc = torch.from_numpy(cond_cont).float().to(device) ck = torch.from_numpy(cond_cat).long().to(device) @@ -456,18 +475,24 @@ def _step_chunk( n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64) # --- Secondaries --- - sec_cont, sec_type_emb, _valid = sample_secondaries( + # No snapping: sec_mass/sec_charge are the model's raw predicted physical + # identity, used as-is for the spawned track's own future conditioning. + # sec_pdg_code below is a *separate*, reporting-only nearest-known-PDG + # label (never fed back into the model) — see giant/particles.py. + sec_cont, sec_phys, _valid = sample_secondaries( sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps ) - sec_pdg_idx = snap_type_to_pdg_idx(sec_type_emb, pdg_emb_weight) - sec_E, sec_dir_world, sec_pdg_code, sec_valid = decode_secondaries( - sec_cont.cpu().numpy(), - sec_pdg_idx.cpu().numpy(), + sec_full = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy() + sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries( + sec_full, n_sec_np, e_sec, tr["pre_dir"], - pdg_map_inv, + sec_phys_normalizer=sec_phys_norm, ) + sec_pdg_code = nearest_known_pdg( + sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys() + ).reshape(sec_mass.shape) edep = edep.astype(np.float64) post_E = post_E.astype(np.float64) @@ -480,6 +505,8 @@ def _step_chunk( sec_E, sec_dir_world, sec_pdg_code, + sec_mass, + sec_charge, counts, max_tracks_per_event, ) @@ -537,6 +564,8 @@ def _step_chunk( "pre_pos": post_pos[cont], "pre_E": post_E[cont], "pre_dir": post_dir_world[cont], + "mass": tr["mass"][cont], + "charge": tr["charge"][cont], } return _concat_frontiers([cont_frontier, new_tracks]) @@ -548,6 +577,8 @@ def _spawn_secondaries( sec_E, sec_dir_world, sec_pdg_code, + sec_mass, + sec_charge, counts, max_tracks_per_event, ) -> tuple[dict[str, np.ndarray], np.ndarray]: @@ -594,9 +625,14 @@ def _spawn_secondaries( "parent_id": tr["track_id"][pr_k], "generation": tr["generation"][pr_k] + 1, "step_in_track": np.zeros(len(pr_k), dtype=np.int64), + # Reporting-only nominal PDG (nearest-known-PDG label, never fed back + # into the model) — the track's actual physical identity going + # forward is "mass"/"charge" below, the model's raw prediction. "pdg": sec_pdg_code[pr_k, sl_k].astype(np.int64), "pre_pos": post_pos[pr_k], "pre_E": sec_E[pr_k, sl_k].astype(np.float64), "pre_dir": sec_dir_world[pr_k, sl_k].astype(np.float64), + "mass": sec_mass[pr_k, sl_k].astype(np.float64), + "charge": sec_charge[pr_k, sl_k].astype(np.float64), } return frontier, dropped_edep diff --git a/giant/sample.py b/giant/sample.py index 900b32f..e97fead 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -43,10 +43,15 @@ def sample_secondaries( n_sec_pred: (B,) int64 — number of valid secondaries per step - Returns (sec_cont, sec_type_emb, sec_valid): - sec_cont: (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z] - sec_type_emb: (B, K_MAX, emb_dim) — predicted type embedding per slot - sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred + Returns (sec_cont, sec_phys, sec_valid): + sec_cont: (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z] + sec_phys: (B, K_MAX, PARTICLE_PHYS_DIM) — predicted [log_mass, charge] + per slot (normalised iff the checkpoint's sec_phys + normalizer was applied at training time — denormalize + before treating as physical units; see + giant.data.transforms.decode_secondaries). Used as-is — + no snapping to a discrete PDG code. + sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred """ sec_decoder.eval() B = cond_cont.size(0) @@ -61,27 +66,12 @@ def sample_secondaries( x_slots = x.view(B, K_MAX, SEC_SLOT_DIM) sec_cont = x_slots[:, :, :4] - sec_type_emb = x_slots[:, :, 4:] + sec_phys = x_slots[:, :, 4:] sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze( 1 ) - return sec_cont, sec_type_emb, sec_valid - - -def snap_type_to_pdg_idx( - sec_type_emb: torch.Tensor, - pdg_emb_weight: torch.Tensor, -) -> torch.Tensor: - """Nearest-neighbour snap: predicted type embedding → PDG model-index. - - sec_type_emb: (B, K_MAX, emb_dim) - Returns (B, K_MAX) int64 with model-indices. - """ - B, K, D = sec_type_emb.shape - flat = sec_type_emb.reshape(-1, D) - dists = torch.cdist(flat.float(), pdg_emb_weight.float()) - return dists.argmin(dim=-1).reshape(B, K) + return sec_cont, sec_phys, sec_valid @torch.no_grad() diff --git a/giant/train.py b/giant/train.py index d796956..1849eaa 100644 --- a/giant/train.py +++ b/giant/train.py @@ -80,31 +80,6 @@ class _GracefulShutdown: ) -def _build_sec_x1( - sec_cont: torch.Tensor, - sec_pdg_idx: torch.Tensor, - pdg_emb_weight: torch.Tensor, -) -> torch.Tensor: - """Assemble the Stage-2 flow target by appending type embeddings. - - sec_cont: (B, K_MAX, 4) — [stick_logit, dir_local] - sec_pdg_idx: (B, K_MAX) — integer PDG model-indices - pdg_emb_weight: (pdg_vocab, emb_dim) — live embedding table weights - - Returns (B, SEC_DIM) = (B, K_MAX * (4 + emb_dim)). - - Detaches the looked-up rows: this tensor becomes x1 in the flow-matching - loss (u_t = x1 - x0), so without detaching, the Stage-2 loss could pull - the embedding table itself toward whatever the decoder already predicts - (a moving, self-referential regression target) instead of only pulling - the decoder toward the table. The table is still trained normally via - its Stage-1 conditioning role and `predict_n_sec`. - """ - type_emb = pdg_emb_weight[sec_pdg_idx].detach() # (B, K_MAX, emb_dim) - x1_s2 = torch.cat([sec_cont, type_emb], dim=-1) # (B, K_MAX, 4+emb_dim) - return x1_s2.flatten(1) # (B, SEC_DIM) - - @torch.no_grad() def _update_ema( ema_model: torch.nn.Module, model: torch.nn.Module, decay: float @@ -128,13 +103,12 @@ def _compute_losses( torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor ]: """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc) for one batch.""" - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx = batch + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) x1_s1 = x1_s1.to(device) n_sec = n_sec.to(device) sec_cont = sec_cont.to(device) - sec_pdg_idx = sec_pdg_idx.to(device) proc_idx = proc_idx.to(device) # Stage-1 flow loss @@ -150,14 +124,14 @@ def _compute_losses( # Stage-2 secondary flow loss # Use a noiseless Stage-1 target as context (detach to avoid back-prop - # coupling between the two flow paths through the same embedding table). - # The type-embedding lookup itself is also detached inside _build_sec_x1, - # so the shared PDG table is shaped only by its Stage-1 conditioning role - # and predict_n_sec, not by chasing the Stage-2 decoder's predictions. + # coupling between the two flow paths). sec_cont's log_mass/charge + # columns are already a fixed physics-derived regression target (see + # giant.data.transforms.encode_secondaries) rather than a learned/moving + # one, so — unlike the embedding-table target this replaced — no + # detaching is needed to keep the target from chasing the decoder. from giant.constants import K_MAX - pdg_emb_weight = stage1_model.pdg_embedding_weight() - x1_s2 = _build_sec_x1(sec_cont, sec_pdg_idx, pdg_emb_weight) + x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM) sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) l_s2 = flow_matching_loss_secondary( diff --git a/giant/validate.py b/giant/validate.py index 3372a06..3a82326 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -8,9 +8,10 @@ from giant.sample import ( sample_ddpm, sample_ddim, sample_secondaries, - snap_type_to_pdg_idx, ) +_SEC_PHYS_NAMES = ["log_mass", "charge"] + def _kw(steps: int | None) -> dict[str, int]: return {} if steps is None else {"steps": steps} @@ -63,12 +64,15 @@ def validate_marginals( mode, which always runs the full schedule. When `sec_decoder` is given, also validates Stage 2: n_sec distribution - (+ classification accuracy), secondary species distribution, and - per-slot energy-fraction marginals — restricted to each side's own valid - slots (real: `n_sec`; generated: the Stage-1 head's argmax), since the - two need not agree on how many slots are valid. Adds - {"n_sec_real", "n_sec_pred", "n_sec_accuracy", "species_real", - "species_generated", "energy_fraction_kl"} to the returned dict. + (+ classification accuracy), predicted secondary physical-identity + (log_mass, charge) marginals, and per-slot energy-fraction marginals — + restricted to each side's own valid slots (real: `n_sec`; generated: the + Stage-1 head's argmax), since the two need not agree on how many slots + are valid. Compared directly in normalised space (no denormalising — + KL estimated from a shared per-sample histogram is invariant to a shared + affine rescaling of both sides). Adds {"n_sec_real", "n_sec_pred", + "n_sec_accuracy", "phys_real", "phys_generated", "phys_kl", + "energy_fraction_kl"} to the returned dict. """ if device is None: device = next(model.parameters()).device @@ -78,15 +82,15 @@ def validate_marginals( all_real, all_gen = [], [] all_n_sec_real, all_n_sec_pred = [], [] - all_species_real, all_species_gen = [], [] + all_phys_real, all_phys_gen = [], [] all_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)] all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)] for i, batch in enumerate(val_loader): if n_batches is not None and i >= n_batches: break - # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, sec_pdg_idx, proc_idx). - cond_cont, cond_cat, x1, n_sec, sec_cont, sec_pdg_idx, _proc_idx = batch + # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx). + cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) @@ -112,9 +116,9 @@ def validate_marginals( real_valid = np.arange(K_MAX)[None, :] < n_sec_np[:, None] # (B, K_MAX) real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64))) - real_species = sec_pdg_idx.numpy() + real_phys = sec_cont[:, :, 4:6].numpy() # (B, K_MAX, 2) [log_mass, charge] - sec_cont_pred, sec_type_emb, sec_valid_pred = sample_secondaries( + sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries( sec_decoder, cond_cont, cond_cat, @@ -122,15 +126,14 @@ def validate_marginals( n_sec_pred, steps=steps if steps is not None else 10, ) - sec_pdg_pred = snap_type_to_pdg_idx(sec_type_emb, model.pdg_embedding_weight()) gen_frac = 1.0 / ( 1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)) ) - gen_species = sec_pdg_pred.cpu().numpy() + gen_phys = sec_phys_pred.cpu().numpy() gen_valid = sec_valid_pred.cpu().numpy() - all_species_real.append(real_species[real_valid]) - all_species_gen.append(gen_species[gen_valid]) + all_phys_real.append(real_phys[real_valid]) + all_phys_gen.append(gen_phys[gen_valid]) for j in range(K_MAX): all_frac_real[j].append(real_frac[real_valid[:, j], j]) all_frac_gen[j].append(gen_frac[gen_valid[:, j], j]) @@ -169,14 +172,12 @@ def validate_marginals( n_sec_real = np.concatenate(all_n_sec_real, axis=0) n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0) n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean()) - species_real = np.concatenate(all_species_real, axis=0) - species_gen = np.concatenate(all_species_gen, axis=0) + phys_real = np.concatenate(all_phys_real, axis=0) # (M, 2) + phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2) - n_classes = ( - max(int(species_real.max(initial=0)), int(species_gen.max(initial=0))) + 1 + phys_kl = np.array( + [_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)] ) - species_real_dist = _bincount_frac(species_real, n_classes) - species_gen_dist = _bincount_frac(species_gen, n_classes) energy_fraction_kl = np.full(K_MAX, np.nan) print( @@ -192,10 +193,17 @@ def validate_marginals( for v in range(max_n_sec): print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}") - print(f"\n{'pdg model-index':<20} {'real_frac':>10} {'gen_frac':>10}") - print("-" * 42) - for c in range(n_classes): - print(f"{c:<20} {species_real_dist[c]:>10.4f} {species_gen_dist[c]:>10.4f}") + print( + f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} " + f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}" + ) + print("-" * 68) + for j, name in enumerate(_SEC_PHYS_NAMES): + r, g = phys_real[:, j], phys_gen[:, j] + print( + f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} " + f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}" + ) print( f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} " @@ -219,8 +227,9 @@ def validate_marginals( "n_sec_real": n_sec_real, "n_sec_pred": n_sec_pred_all, "n_sec_accuracy": n_sec_accuracy, - "species_real": species_real, - "species_generated": species_gen, + "phys_real": phys_real, + "phys_generated": phys_gen, + "phys_kl": phys_kl, "energy_fraction_kl": energy_fraction_kl, } ) diff --git a/pyproject.toml b/pyproject.toml index cf06f84..bca28c4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -11,6 +11,7 @@ dependencies = [ "tqdm>=4.60,<5", "typer>=0.12,<1", "pyyaml>=6,<7", + "particle>=1.0,<2", ] [project.optional-dependencies] diff --git a/tests/test_materials.py b/tests/test_materials.py new file mode 100644 index 0000000..c0d0c8e --- /dev/null +++ b/tests/test_materials.py @@ -0,0 +1,105 @@ +import numpy as np +import pytest + +from giant.materials import ( + MaterialProperties, + MaterialPropertiesNotFilledError, + UnknownMaterialError, + get_material_properties, + material_properties_array, +) + + +def test_get_material_properties_unknown_name_raises(): + with pytest.raises(UnknownMaterialError): + get_material_properties("G4_Unobtainium") + + +def test_get_material_properties_unfilled_entry_raises(): + """G4_LYSO is not a stock Geant4 NIST material (confirmed against the + vendored Geant4 11.4.1 build) and is the one entry still shipped unfilled.""" + with pytest.raises(MaterialPropertiesNotFilledError): + get_material_properties("G4_LYSO") + + +def test_get_material_properties_returns_filled_entry_from_injected_table(): + table = { + "G4_Pb": MaterialProperties( + z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59 + ) + } + props = get_material_properties("G4_Pb", table) + assert props.z_eff == 82.0 + assert props.a_eff == 207.2 + assert props.density == 11.35 + assert props.x0 == 0.5612 + assert props.lambda_int == 17.59 + + +def test_material_properties_array_shape_and_values(): + table = { + "G4_Pb": MaterialProperties(82.0, 207.2, 11.35, 0.5612, 17.59), + "G4_W": MaterialProperties(74.0, 183.84, 19.3, 0.3504, 9.95), + } + names = np.array(["G4_Pb", "G4_W", "G4_Pb"], dtype=object) + arr = material_properties_array(names, table) + assert arr.shape == (3, 5) + assert arr.dtype == np.float32 + np.testing.assert_allclose(arr[0], [82.0, 207.2, 11.35, 0.5612, 17.59], rtol=1e-5) + np.testing.assert_allclose(arr[1], [74.0, 183.84, 19.3, 0.3504, 9.95], rtol=1e-5) + + +def test_all_known_materials_present_in_stub_table(): + """Every material referenced elsewhere in the repo must at least have a + stub entry (even if unfilled) -- an unknown name should never be the + failure mode a physicist hits when populating the table.""" + from giant.materials import MATERIAL_PROPERTIES + + expected = { + "G4_PbWO4", + "G4_CESIUM_IODIDE", + "G4_Pb", + "G4_W", + "G4_Cu", + "G4_Fe", + "G4_BRASS", + "G4_POLYSTYRENE", + "G4_PLASTIC_SC_VINYLTOLUENE", + "G4_BGO", + "G4_LYSO", + "G4_AIR", + "G4_lAr", + } + assert expected <= set(MATERIAL_PROPERTIES.keys()) + + +def test_all_materials_filled_except_lyso(): + """G4_LYSO is the sole intentionally-unfilled entry (not a stock Geant4 + NIST material); every other known material has real Geant4-derived + values -- see the module docstring for provenance.""" + from giant.materials import MATERIAL_PROPERTIES + + for name, props in MATERIAL_PROPERTIES.items(): + if name == "G4_LYSO": + assert all(v is None for v in props) + else: + assert all(v is not None for v in props), f"{name} unexpectedly unfilled" + + +def test_elemental_material_z_eff_matches_atomic_number(): + """Single-element materials' z_eff must equal the element's real Z.""" + pb = get_material_properties("G4_Pb") + assert pb.z_eff == pytest.approx(82.0) + w = get_material_properties("G4_W") + assert w.z_eff == pytest.approx(74.0) + fe = get_material_properties("G4_Fe") + assert fe.z_eff == pytest.approx(26.0) + + +def test_pbwo4_values_match_known_cms_ecal_reference(): + """PbWO4 (CMS ECAL crystal) has well-known reference values: X0~0.89cm, + density 8.28 g/cm^3 -- sanity check the Geant4-derived numbers land there.""" + pbwo4 = get_material_properties("G4_PbWO4") + assert pbwo4.density == pytest.approx(8.28) + assert pbwo4.x0 == pytest.approx(0.89, abs=0.01) + assert pbwo4.z_eff == pytest.approx(31.33, abs=0.01) diff --git a/tests/test_particles.py b/tests/test_particles.py new file mode 100644 index 0000000..af09127 --- /dev/null +++ b/tests/test_particles.py @@ -0,0 +1,118 @@ +import numpy as np +import pytest + +from giant.particles import ( + nearest_known_pdg, + particle_mass_charge, + particle_phys_array, +) + + +def test_photon_massless_neutral(): + mass, charge = particle_mass_charge(22) + assert mass == pytest.approx(0.0) + assert charge == pytest.approx(0.0) + + +def test_electron_mass_charge(): + mass, charge = particle_mass_charge(11) + assert mass == pytest.approx(0.51099895069, rel=1e-6) + assert charge == pytest.approx(-1.0) + + +def test_positron_is_charge_conjugate_of_electron(): + mass_e, charge_e = particle_mass_charge(11) + mass_p, charge_p = particle_mass_charge(-11) + assert mass_p == pytest.approx(mass_e) + assert charge_p == pytest.approx(-charge_e) + + +def test_proton_mass_charge(): + mass, charge = particle_mass_charge(2212) + assert mass == pytest.approx(938.27208943, rel=1e-6) + assert charge == pytest.approx(1.0) + + +def test_neutrino_unmeasured_mass_treated_as_zero(): + """PDG tables store an unmeasured neutrino mass as None -- must not + propagate a None/NaN into a physical conditioning feature.""" + mass, charge = particle_mass_charge(12) + assert mass == pytest.approx(0.0) + assert charge == pytest.approx(0.0) + + +def test_ground_state_nucleus_resolved_via_particle_package(): + """He-4 (Z=2, A=4) is a common nuclide in `particle`'s ground-state table.""" + mass, charge = particle_mass_charge(1000020040) + assert charge == pytest.approx(2.0) + assert mass == pytest.approx( + 4 * 931.494, rel=0.05 + ) # near A*amu, binding-energy-corrected + + +def test_nuclear_isomer_falls_back_to_z_a_decode(): + """An excited/isomer nuclear code (nonzero trailing digit) is absent from + `particle`'s ground-state-only nuclide table -- confirmed necessary for + ~32% of the nuclear codes in the multi-material dataset. Fe-56 isomer: + Z=26, A=56, isomer level 1 -> pdgid 1000260561.""" + pdg = 1000260561 + mass, charge = particle_mass_charge(pdg) + assert charge == pytest.approx(26.0) + assert mass == pytest.approx(56 * 931.494, rel=1e-6) + + +def test_invalid_pdg_code_raises(): + with pytest.raises(ValueError): + particle_mass_charge(999999999) + + +def test_particle_mass_charge_is_cached(): + particle_mass_charge.cache_clear() + particle_mass_charge(22) + particle_mass_charge(22) + info = particle_mass_charge.cache_info() + assert info.hits >= 1 + + +def test_particle_phys_array_shape_and_dtype(): + arr = particle_phys_array(np.array([22, 11, 2212])) + assert arr.shape == (3, 2) + assert arr.dtype == np.float32 + np.testing.assert_allclose(arr[0], [0.0, 0.0]) + np.testing.assert_allclose(arr[2], [938.27208943, 1.0], rtol=1e-5) + + +# ── nearest_known_pdg (reporting-only nearest-neighbour label) ────────────── + + +def test_nearest_known_pdg_exact_match(): + candidates = [22, 11, -11, 2212, 2112] + 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_prioritises_charge_match(): + """Charge is a small conserved quantum number and should usually match + exactly even when the queried mass is noisy/imperfect.""" + candidates = [22, 11, -11, 2212] + # Close to electron mass but not exact, positive charge like the positron. + result = nearest_known_pdg(np.array([0.6]), np.array([1.0]), candidates) + assert result[0] == -11 + + +def test_nearest_known_pdg_empty_candidates_raises(): + with pytest.raises(ValueError): + nearest_known_pdg(np.array([1.0]), np.array([0.0]), []) + + +def test_nearest_known_pdg_shape(): + candidates = [22, 11, -11, 2212, 2112] + n = 10 + result = nearest_known_pdg( + np.random.default_rng(0).uniform(0, 1000, n), + np.random.default_rng(1).uniform(-1, 1, n), + candidates, + ) + assert result.shape == (n,) + assert set(result.tolist()) <= set(candidates) diff --git a/tests/test_phase2.py b/tests/test_phase2.py index e3dec95..c2857d1 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -4,21 +4,33 @@ import numpy as np import pytest import torch -from giant.constants import COND_DIM, EMB_DIM, K_MAX, SEC_DIM, X_DIM +from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, X_DIM from giant.model.network import DenoisingMLP, SecondaryDecoder from giant.model.schedule import flow_matching_loss_secondary -from giant.sample import sample_secondaries, snap_type_to_pdg_idx +from giant.sample import sample_secondaries # ── helpers ────────────────────────────────────────────────────────────────── -def _stage1(pdg=3, mat=2): - return DenoisingMLP(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2) +def _stage1(pdg=3, mat=2, conditioning="embedding"): + return DenoisingMLP( + pdg_vocab=pdg, + mat_vocab=mat, + hidden_dim=32, + n_blocks=2, + conditioning=conditioning, + ) -def _sec_decoder(pdg=3, mat=2): - return SecondaryDecoder(pdg_vocab=pdg, mat_vocab=mat, hidden_dim=32, n_blocks=2) +def _sec_decoder(pdg=3, mat=2, conditioning="embedding"): + return SecondaryDecoder( + pdg_vocab=pdg, + mat_vocab=mat, + hidden_dim=32, + n_blocks=2, + conditioning=conditioning, + ) def _cond(B=8, pdg=3, mat=2): @@ -48,18 +60,35 @@ def test_predict_n_sec_no_nan(): assert torch.isfinite(logits).all() -def test_pdg_embedding_weight_shape(): - model = _stage1(pdg=5, mat=2) - w = model.pdg_embedding_weight() - assert w.shape == (5, EMB_DIM) +@pytest.mark.parametrize("conditioning", ["embedding", "physical"]) +def test_no_pdg_embedding_weight_method(conditioning): + """The Stage-2 species output no longer needs a shared embedding table.""" + model = _stage1(pdg=5, mat=2, conditioning=conditioning) + assert not hasattr(model, "pdg_embedding_weight") + + +def test_condition_encoder_physical_mode_has_no_embedding_tables(): + model = _stage1(pdg=5, mat=2, conditioning="physical") + assert not hasattr(model.cond_enc, "pdg_emb") + assert not hasattr(model.cond_enc, "mat_emb") + assert hasattr(model.cond_enc, "particle_mlp") + assert hasattr(model.cond_enc, "material_mlp") + + +def test_condition_encoder_embedding_mode_has_embedding_tables(): + model = _stage1(pdg=5, mat=2, conditioning="embedding") + assert hasattr(model.cond_enc, "pdg_emb") + assert hasattr(model.cond_enc, "mat_emb") + assert not hasattr(model.cond_enc, "particle_mlp") # ── SecondaryDecoder ────────────────────────────────────────────────────────── -def test_sec_decoder_output_shape(): +@pytest.mark.parametrize("conditioning", ["embedding", "physical"]) +def test_sec_decoder_output_shape(conditioning): B = 8 - decoder = _sec_decoder() + decoder = _sec_decoder(conditioning=conditioning) x_t = torch.randn(B, SEC_DIM) t = torch.rand(B) cond_cont, cond_cat = _cond(B) @@ -144,11 +173,11 @@ def test_sample_secondaries_shapes(): cond_cont, cond_cat = _cond(B, pdg, mat) stage1_out = torch.randn(B, X_DIM) n_sec_pred = torch.randint(0, K_MAX + 1, (B,)) - sec_cont, sec_type_emb, sec_valid = sample_secondaries( + sec_cont, sec_phys, sec_valid = sample_secondaries( decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3 ) assert sec_cont.shape == (B, K_MAX, 4) - assert sec_type_emb.shape == (B, K_MAX, EMB_DIM) + assert sec_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM) assert sec_valid.shape == (B, K_MAX) assert sec_valid.dtype == torch.bool @@ -167,16 +196,6 @@ def test_sample_secondaries_valid_mask_matches_n_sec(): assert not sec_valid[i, n:].any() -def test_snap_type_to_pdg_idx_shape(): - B, pdg_vocab = 4, 5 - emb_weight = torch.randn(pdg_vocab, EMB_DIM) - sec_type_emb = torch.randn(B, K_MAX, EMB_DIM) - idx = snap_type_to_pdg_idx(sec_type_emb, emb_weight) - assert idx.shape == (B, K_MAX) - assert idx.dtype == torch.int64 - assert (idx >= 0).all() and (idx < pdg_vocab).all() - - # ── encode_secondaries round-trip ───────────────────────────────────────────── @@ -192,6 +211,7 @@ def test_encode_secondaries_energy_conservation(): sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32) sec_dir_list[:, :, 2] = 1.0 + sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64) sec_valid = np.zeros((N, K_MAX), dtype=bool) for i in range(N): k = n_sec[i] @@ -199,12 +219,15 @@ def test_encode_secondaries_energy_conservation(): energies = np.sort(energies)[::-1] sec_E_list[i, :k] = energies.astype(np.float32) sec_valid[i, :k] = True + sec_pdg_list[i, :k] = 22 # photon — resolvable by giant.particles pre_dir = rng.standard_normal((N, 3)).astype(np.float32) pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True) - sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) - assert sec_cont.shape == (N, K_MAX, 4) + sec_cont = encode_secondaries( + sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list + ) + assert sec_cont.shape == (N, K_MAX, 6) assert np.isfinite(sec_cont).all() @@ -233,13 +256,54 @@ def test_encode_secondaries_direction_encoding(): np.testing.assert_allclose(norms_out, 1.0, atol=1e-5) +def test_encode_secondaries_physical_columns_without_pdg_list(): + """Omitting sec_pdg_list zero-fills the physical columns (no crash).""" + from giant.data.transforms import encode_secondaries + + N = 3 + e_sec = np.ones(N, dtype=np.float32) + sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) + sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32) + sec_dir_list[:, :, 2] = 1.0 + sec_valid = np.zeros((N, K_MAX), dtype=bool) + pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) + + sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir) + np.testing.assert_allclose(sec_cont[:, :, 4:6], 0.0) + + +def test_encode_secondaries_physical_columns_match_ground_truth_pdg(): + """log_mass/charge for a valid slot match giant.particles for that PDG.""" + from giant.data.transforms import encode_secondaries, log_transform + from giant.particles import particle_mass_charge + + N = 1 + e_sec = np.array([5.0], dtype=np.float32) + sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) + sec_E_list[0, 0] = 5.0 + sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32) + sec_dir_list[0, 0] = [0, 0, 1] + sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64) + sec_pdg_list[0, 0] = 11 # electron + sec_valid = np.zeros((N, K_MAX), dtype=bool) + sec_valid[0, 0] = True + pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32) + + sec_cont = encode_secondaries( + sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list + ) + mass, charge = particle_mass_charge(11) + assert sec_cont[0, 0, 4] == pytest.approx(log_transform(np.array([mass]))[0]) + assert sec_cont[0, 0, 5] == pytest.approx(charge) + + # ── decode_secondaries: exact energy conservation ──────────────────────────── def _random_sec_cont(rng, N, stick_logit_scale=1.0): - sec_cont = rng.standard_normal((N, K_MAX, 4)).astype(np.float32) + sec_cont = rng.standard_normal((N, K_MAX, 6)).astype(np.float32) sec_cont[:, :, 0] *= stick_logit_scale - dirs = sec_cont[:, :, 1:] + dirs = sec_cont[:, :, 1:4] dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True) return sec_cont @@ -257,13 +321,12 @@ def test_decode_secondaries_valid_slots_sum_to_e_sec(): rng = np.random.default_rng(0) N = 200 sec_cont = _random_sec_cont(rng, N) - sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) n_sec = rng.integers(0, K_MAX + 1, size=N) e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32) pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) - sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries( - sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22} + sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries( + sec_cont, n_sec, e_sec, pre_dir ) valid_sum = (sec_E * sec_valid).sum(axis=1) @@ -283,13 +346,12 @@ def test_decode_secondaries_zero_n_sec_has_zero_energy(): rng = np.random.default_rng(1) N = 10 sec_cont = _random_sec_cont(rng, N) - sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) n_sec = np.zeros(N, dtype=np.int64) e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32) pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) - sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries( - sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22} + sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries( + sec_cont, n_sec, e_sec, pre_dir ) assert not sec_valid.any() @@ -307,12 +369,11 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split(): n_sec = np.array([0, 1, 3, K_MAX]) for i, k in enumerate(n_sec): sec_cont[i, :k, 0] = -80.0 - sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32) pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) - sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries( - sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22} + sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries( + sec_cont, n_sec, e_sec, pre_dir ) for i, k in enumerate(n_sec): @@ -336,25 +397,48 @@ def test_decode_secondaries_rescale_preserves_relative_shares(): sec_cont = _random_sec_cont(rng, N) n_sec = np.array([4]) pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32) - sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) - sec_E_small, _, _, sec_valid = decode_secondaries( - sec_cont, - sec_pdg_pred, - n_sec, - np.array([5.0], dtype=np.float32), - pre_dir, - {0: 22}, + sec_E_small, _, _, _, sec_valid = decode_secondaries( + sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir ) - sec_E_large, _, _, _ = decode_secondaries( - sec_cont, - sec_pdg_pred, - n_sec, - np.array([50.0], dtype=np.float32), - pre_dir, - {0: 22}, + sec_E_large, _, _, _, _ = decode_secondaries( + sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir ) ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0] ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0] np.testing.assert_allclose(ratio_small, ratio_large, rtol=1e-4) + + +def test_decode_secondaries_mass_charge_round_trip_with_normalizer(): + from giant.data.transforms import Normalizer, decode_secondaries, encode_secondaries + + N = 1 + e_sec = np.array([5.0], dtype=np.float32) + sec_E_list = np.zeros((N, K_MAX), dtype=np.float32) + sec_E_list[0, 0] = 5.0 + sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32) + sec_dir_list[0, 0] = [0, 0, 1] + sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64) + sec_pdg_list[0, 0] = 2212 # proton + sec_valid = np.zeros((N, K_MAX), dtype=bool) + sec_valid[0, 0] = True + pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32) + + sec_cont = encode_secondaries( + sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list + ) + norm = Normalizer() + norm.mean = np.array([-2.0, 0.5], dtype=np.float32) + norm.std = np.array([3.0, 1.5], dtype=np.float32) + sec_cont_normed = sec_cont.copy() + sec_cont_normed[:, :, 4:6] = norm.transform( + sec_cont[:, :, 4:6].reshape(-1, 2) + ).reshape(N, K_MAX, 2) + + n_sec = np.array([1]) + _, _, sec_mass, sec_charge, _ = decode_secondaries( + sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm + ) + assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2) + assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 4d7c747..79057ae 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -20,17 +20,22 @@ 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) +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, 8)).astype(np.float32)) + cond = Normalizer().fit(rng.standard_normal((1000, 15)).astype(np.float32)) tgt = Normalizer().fit(rng.standard_normal((1000, 9)).astype(np.float32)) - return cond, tgt + sec_phys = Normalizer().fit(rng.standard_normal((1000, 2)).astype(np.float32)) + return cond, tgt, sec_phys def _oracle(): @@ -63,11 +68,12 @@ def _run( max_steps=30, max_tracks_per_event=300, seeds=None, + conditioning="embedding", ): torch.manual_seed(0) np.random.seed(0) - s1, s2 = _models() - cond, tgt = _norms() + s1, s2 = _models(conditioning) + cond, tgt, sec_phys = _norms() return rollout( s1, s2, @@ -75,6 +81,7 @@ def _run( seeds or _seeds(), cond, tgt, + sec_phys, PDG_MAP, MAT_MAP, energy_cutoff=energy_cutoff, @@ -83,9 +90,35 @@ def _run( 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) @@ -171,7 +204,7 @@ def _run_streaming(on_chunk, **kwargs): torch.manual_seed(0) np.random.seed(0) s1, s2 = _models() - cond, tgt = _norms() + cond, tgt, sec_phys = _norms() seeds = kwargs.pop("seeds", None) or _seeds() return rollout( s1, @@ -180,6 +213,7 @@ def _run_streaming(on_chunk, **kwargs): seeds, cond, tgt, + sec_phys, PDG_MAP, MAT_MAP, energy_cutoff=kwargs.pop("energy_cutoff", 1.0), diff --git a/tests/test_router.py b/tests/test_router.py index 31f6b54..8f3fa2a 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -576,11 +576,9 @@ def test_routed_denoising_mlp_predict_n_sec_shape(): assert logits.shape == (B, K_MAX + 1) -def test_routed_denoising_mlp_pdg_embedding_weight_shape(): +def test_routed_denoising_mlp_has_no_pdg_embedding_weight_method(): model = _routed_stage1(pdg=5, mat=2) - from giant.constants import EMB_DIM - - assert model.pdg_embedding_weight().shape == (5, EMB_DIM) + assert not hasattr(model, "pdg_embedding_weight") # ── RoutedSecondaryDecoder ─────────────────────────────────────────────────── diff --git a/tests/test_transforms.py b/tests/test_transforms.py index c95a36b..a7af975 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -1,7 +1,8 @@ import numpy as np import pytest -from giant.constants import K_MAX +from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX from giant.data.transforms import ( + build_cond_features, build_features, energy_simplex_decode, energy_simplex_encode, @@ -235,7 +236,7 @@ def test_build_features_clamps_n_sec_label_to_k_max(): pdg_map = {11: 0} mat_map = {"PbWO4": 0} - _, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map) + _, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map) assert n_sec.max() <= K_MAX np.testing.assert_array_equal(n_sec, [0, 5, K_MAX]) @@ -312,9 +313,87 @@ def test_build_features_require_secondaries_ok_when_no_secondaries(): data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32)) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} - _, _, _, _, sec_cont, sec_pdg_idx, *_ = build_features( + _, _, _, _, sec_cont, *_ = build_features( data, pdg_map, mat_map, require_secondaries=True ) assert not sec_cont.any() - assert not sec_pdg_idx.any() + + +# ── physical-property conditioning ──────────────────────────────────────────── + + +@pytest.fixture +def fake_material_props(monkeypatch): + """Inject a fully-populated fake materials table for "physical" mode + tests, independent of when the real giant/materials.py table is filled + in by the user (see giant.materials.MaterialPropertiesNotFilledError).""" + import giant.materials as gm + + fake = { + "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_build_features_embedding_mode_zero_fills_physical_columns(): + data = _minimal_step_data(3) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="embedding") + + assert cond_cont.shape[1] == COND_DIM + np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0) + + +def test_build_features_physical_mode_shape_and_values(fake_material_props): + from giant.particles import particle_mass_charge + + data = _minimal_step_data(3) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="physical") + + assert cond_cont.shape[1] == COND_DIM + mass, charge = particle_mass_charge(11) + expected_log_mass = log_transform(np.array([mass]))[0] + np.testing.assert_allclose( + cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5 + ) + np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], charge) + np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 2], 75.6) # z_eff + np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 3], 205.3) # a_eff + + +def test_build_features_physical_mode_unfilled_material_raises(): + """G4_LYSO is the one material giant/materials.py still ships unfilled + (not a stock Geant4 NIST material) — must fail loudly, not silently.""" + from giant.materials import MaterialPropertiesNotFilledError + + data = _minimal_step_data(2) + data["material"] = np.full(2, "G4_LYSO", dtype=object) + pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0} + + with pytest.raises(MaterialPropertiesNotFilledError): + build_features(data, pdg_map, mat_map, conditioning="physical") + + +def test_build_cond_features_mass_charge_override(fake_material_props): + """rollout.py's secondaries carry their own predicted mass/charge — when + present in `data`, these bypass the pdg-based lookup entirely (the "no + snapping" design: a track's own future conditioning must use its actual + predicted physical identity, not a value re-derived from a PDG code).""" + data = _minimal_step_data(2) + data["mass"] = np.array([123.0, 456.0], dtype=np.float32) + data["charge"] = np.array([2.0, -2.0], dtype=np.float32) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + cond_cont, _ = build_cond_features(data, pdg_map, mat_map, conditioning="physical") + + np.testing.assert_allclose( + cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])) + ) + np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0]) diff --git a/uv.lock b/uv.lock index a197032..ef69653 100644 --- a/uv.lock +++ b/uv.lock @@ -54,6 +54,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/d2/39/e7eaf1799466a4aef85b6a4fe7bd175ad2b1c6345066aa33f1f58d4b18d0/asttokens-3.0.1-py3-none-any.whl", hash = "sha256:15a3ebc0f43c2d0a50eeafea25e19046c68398e487b9f1f5b517f7c0f40f976a", size = 27047, upload-time = "2025-11-15T16:43:16.109Z" }, ] +[[package]] +name = "attrs" +version = "26.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/9a/8e/82a0fe20a541c03148528be8cac2408564a6c9a0cc7e9171802bc1d26985/attrs-26.1.0.tar.gz", hash = "sha256:d03ceb89cb322a8fd706d4fb91940737b6642aa36998fe130a9bc96c985eff32", size = 952055, upload-time = "2026-03-19T14:22:25.026Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/64/b4/17d4b0b2a2dc85a6df63d1157e028ed19f90d4cd97c36717afef2bc2f395/attrs-26.1.0-py3-none-any.whl", hash = "sha256:c647aa4a12dfbad9333ca4e71fe62ddc36f4e63b2d260a37a8b83d2f043ac309", size = 67548, upload-time = "2026-03-19T14:22:23.645Z" }, +] + [[package]] name = "awkward" version = "2.9.1" @@ -439,6 +448,7 @@ source = { editable = "." } dependencies = [ { name = "numpy" }, { name = "pandas" }, + { name = "particle" }, { name = "pyarrow" }, { name = "pyyaml" }, { name = "tqdm" }, @@ -486,6 +496,7 @@ requires-dist = [ { name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" }, { name = "numpy", specifier = ">=1.26,<3" }, { name = "pandas", specifier = ">=2.2,<4" }, + { name = "particle", specifier = ">=1.0,<2" }, { name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0,<2" }, { name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" }, { name = "pyarrow", specifier = ">=16,<25" }, @@ -502,6 +513,15 @@ requires-dist = [ ] provides-extras = ["cpu", "cuda", "dev", "geometry", "convert", "analysis"] +[[package]] +name = "hepunits" +version = "2.4.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e7/1f/c5f7525bf7e5d14d61750f91845f9f9350a8f33f5ac95672cec92f317038/hepunits-2.4.6.tar.gz", hash = "sha256:bca6ada937147166d66e9fa152566f2378868d798211cad6990080399560fa34", size = 17985, upload-time = "2026-06-16T09:23:36.685Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/85/10/7f9c58d1ec6a0b7f7783fe552f3593f39cda30c2e1d7a9d148ae711e748d/hepunits-2.4.6-py3-none-any.whl", hash = "sha256:089c52c3b84ef67a159b5e9ee9bdd50e1a442e3fd0c101303cc409c1e9011c4d", size = 17090, upload-time = "2026-06-16T09:23:35.35Z" }, +] + [[package]] name = "iniconfig" version = "2.3.0" @@ -1184,6 +1204,19 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" }, ] +[[package]] +name = "particle" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "attrs" }, + { name = "hepunits" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c3/66/09911bbb658fdffe960903c12edecab95f6cced40fef4909d1cc04bd288b/particle-1.0.0.tar.gz", hash = "sha256:49145dec1cb5044b07f3e8e902280fa050950fa845b058003e69de519bb50492", size = 285766, upload-time = "2026-06-25T14:48:26.893Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a8/92/05078b696cddbdd60963577c895d5b77e5bf829197b08677351b11761c2b/particle-1.0.0-py3-none-any.whl", hash = "sha256:fc2656f53e729be76e45430f56aa65dc20dea069565a393032016544425b64bf", size = 245760, upload-time = "2026-06-25T14:48:25.238Z" }, +] + [[package]] name = "pexpect" version = "4.9.0" From 06c9ad8e5fefb18906a62523f462b586cfdfde44 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 20 Jul 2026 10:36:25 +0200 Subject: [PATCH 2/2] Fix crashes in physical-property conditioning edge cases - make_seed_frontier only resolves particle mass/charge in "physical" mode, so "embedding"-mode rollouts no longer crash on a seed PDG code giant.particles can't resolve (the TERM_UNKNOWN_PDG gate now handles it). - nearest_known_pdg skips unresolvable candidate PDG codes instead of raising and killing the whole rollout/predict run. - predict/rollout fail with a clear message when a checkpoint predates the sec_phys normalizer, instead of a bare KeyError. - validate_marginals' phys_kl degrades to NaN (matching the energy_fraction_kl pattern) instead of crashing when a validated batch has zero secondaries on either side. - Correct CLAUDE.md's stale claim that the materials table is unfilled. Co-Authored-By: Claude Sonnet 5 --- CLAUDE.md | 4 ++-- giant/cli.py | 16 +++++++++++++ giant/particles.py | 23 +++++++++++++----- giant/rollout.py | 19 +++++++++++---- giant/validate.py | 12 +++++++--- tests/test_particles.py | 10 ++++++++ tests/test_rollout.py | 25 ++++++++++++++++++- tests/test_validate.py | 53 +++++++++++++++++++++++++++++++++++++++++ 8 files changed, 146 insertions(+), 16 deletions(-) create mode 100644 tests/test_validate.py 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()