lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Update CLAUDE.md and README for the implemented Phase 2 model

Browse Source 
Bring the docs in line with the current two-stage code: energy ALR
simplex output, 8D conditioning (n_sec/e_sec now predicted, not given),
the SecondaryDecoder stage, and the shower rollout + geometry oracle.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-08 13:42:53 +02:00
parent 3faa272562
commit af8dce53a7
2 changed files with 48 additions and 27 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+32 -20
View File
@@ -2,35 +2,38 @@
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
Proof-of-concept surrogate model for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
Proof-of-concept surrogate model for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
## Architecture
Conditional **flow matching** model (Lipman et al. 2022): a small MLP learns a vector field mapping noise → step outcomes in ~10 ODE steps per sample. Falls back to DDPM for comparison.
A **two-stage conditional flow matching** model (Lipman et al. 2022): a small MLP learns a vector field mapping noise → step outcomes in ~10 ODE steps per sample. Falls back to DDPM for comparison.
**Output space (9D, diffused):**
**Stage 1 — primary (9D, diffused):**
| Index | Variable | Transform |
|-------|----------|-----------|
| Index | Variable | Encoding |
|-------|----------|----------|
| 0 | `step_length` [mm] | log |
| 1 | `ΔE = pre_E − post_E` [MeV] | log |
| 2 | `edep` [MeV] | log |
| 1–2 | `edep_logit`, `sec_logit` | ALR coords of the deposit/secondary/post-energy simplex |
| 3–5 | `post_dir` in local frame | unit vector |
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss. Stage 1 also has a classifier head predicting the number of secondaries `n_sec ∈ {0..15}` from the conditioning alone.
Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries (Phase 1 only — see Roadmap below).
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second flow net generates all `K_MAX = 15` secondary slots at once. Each slot carries a stick-breaking energy fraction, a local-frame direction, and a continuous particle-type embedding (snapped to the nearest PDG at inference), ordered by descending energy; slots beyond the predicted `n_sec` are masked. The secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the full chain conserves energy. Each secondary's momentum is reconstructed afterward from `(energy, direction, species)` rather than predicted.
**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID. (`n_sec` / `e_sec` are outputs now, not inputs.)
## Roadmap
The model is developed in two phases:
**Phase 1 (done):** `n_sec` and total secondary energy `e_sec` were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
**Phase 1 (current):** The number of secondaries produced in each step is passed as a conditioning input. This makes training easier because the model has direct access to multiplicity information and can focus on learning the continuous post-step kinematics.
**Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
**Phase 2 (target):** The number of secondaries is not given — the model must predict it jointly with all secondary properties (energy, direction, species) for each step. This requires extending the output space and likely an autoregressive or set-based generative approach for the variable-length secondary list.
**Next directions:** faster-eval architectures against a ~10× native-Geant4 budget (Wasserstein-GAN, mixture-of-experts routing tree), a multi-material sampling-calorimeter dataset, and physical-property conditioning over learned embeddings.
## Data
@@ -43,18 +46,21 @@ giant/
├── giant/
│ ├── data/
│ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads)
│ │ ├── transforms.py # log transforms, local-frame rotation, normaliser
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
│ ├── model/
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP
│ │ ├── network.py # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP, SecondaryDecoder
│ │ └── schedule.py # CosineSchedule (DDPM) and flow matching utilities
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run
│ ├── train.py # training loop, checkpointing, graceful shutdown
│ ├── sample.py # DDPM / DDIM / flow matching samplers
│ ├── train.py # two-stage training loop, checkpointing, graceful shutdown
│ ├── sample.py # DDPM / DDIM / flow matching samplers + secondary sampling
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
│ ├── rollout.py # autoregressive shower rollout driver
│ ├── validate.py # step-level marginal + KL-divergence validation
│ ├── analysis.py # notebook diagnostics: marginals, correlations, constraint checks
│ └── cli.py # `giant train` / `giant predict` Typer app
│ ├── analysis.py # step- and shower-level diagnostics: marginals, correlations, rollout observables
│ └── cli.py # `giant train` / `predict` / `rollout` Typer app
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`uv run dwarf --help`)
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
│ │ # update-manifest, create-manifest, make-root, hparam-scan
@@ -64,6 +70,7 @@ giant/
│ ├── bump_dataset_version.py # cut a new raw gen or parquet schema, with a logged reason —
│ │ # `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
│ ├── create_root_files.py # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
│ ├── geometry_oracle.py # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
│ └── hparam_scan.py # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
└── tests/
```
@@ -77,18 +84,23 @@ uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
`cpu` and `cuda` are mutually exclusive extras selecting the torch build; plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
## Training
## Training, prediction, and rollout
```bash
giant train path/to/steps.parquet --mode flow
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
# Full-shower rollout needs a geometry oracle (position → material/layer_id):
uv sync --extra cpu --extra geometry
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
```
Both accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on either for the full option list.
`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
## Validation
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper, notebook-driven diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical constraint violations — see `giant.analysis`.
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical-constraint violations, and shower-level rollout observables (longitudinal/transverse profiles, PDG energy shares) — see `giant.analysis`.
## Development