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