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
CLAUDE.md
+16 -7
View File
@@ -12,9 +12,12 @@ uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle
pytest # run tests
giant train path/to/steps.parquet --mode flow # train (flow matching)
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
giant predict path/to/steps.parquet --checkpoint ckpt/best.pt # per-step predictions
giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl # full showers
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
# bump-schema, status, update-manifest, create-manifest,
# make-root, hparam-scan (see scripts/dwarf.py)
# make-root, build-geometry-oracle, hparam-scan
# (see scripts/dwarf.py)
```
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.
@@ -31,15 +34,19 @@ Part of the `dev` extra. Run these periodically (not just at commit time) to cat
## Architecture
GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome.
GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — now including the variable-length list of secondary particles the step produces (Phase 2, see Roadmap).
**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`). Train/val split is by `event_id` to avoid leaking correlated steps from the same shower.
**Output space (9D):** `step_length` (log), `ΔE` (log), `edep` (log), `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). `post_pos` itself 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.
**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:** 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).
**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.
**Model** (`giant/model/`): `DenoisingMLP` built from `ResBlock`s, with a `SinusoidalEmbedding` for the diffusion/flow time variable and a `ConditionEncoder` that fuses all conditioning inputs. `schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
**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`).
`schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
@@ -49,6 +56,8 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
## Roadmap
Phase 1 (current): number of secondaries is a conditioning input — model predicts only 9D post-step kinematics (including derived post_pos).
**Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
Phase 2 (target): model must jointly predict the number of secondaries and all their properties (energy, direction, species), requiring an extended output space and likely a set-based or autoregressive generation scheme for the variable-length secondary list.
**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`).
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