lars
8019a80563
CI / Format (ruff format) (push) Successful in 27s
CI / Lint (ruff check) (push) Successful in 27s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 39s
CI / Type check (ty) (push) Successful in 43s
CI / Format (ruff format) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 31s
CI / Tests (push) Successful in 2m35s
CI / Tests (pull_request) Successful in 2m36s
Every metric name used to exist in four places: the dict keys each
StageTrainer returned, the hardcoded _metrics_fields() column list, the
~110-line metrics_row assembly in train(), and the tqdm/summary
formatting. The two had to be kept in exact correspondence by hand or
csv.DictWriter would raise.
Each metric is now declared once, as a MetricSpec on the trainer that
computes it. MetricsCollector derives the CSV header and W&B payload from
those declarations and owns all accumulation, so train() no longer carries
a running sum, and every isinstance(tr, WGANStageTrainer) branch is gone —
replaced by four trainer hooks (batch_loss, summary, val_objective,
supports_val_loss).
giant/train.py (1875 lines) becomes giant/training/:
trainers.py StageSpec + shared StageTrainer base + the two subclasses
metrics.py MetricSpec, MetricsCollector
stage2_inputs.py the pure AR/teacher-forcing tensor helpers, moved verbatim
loop.py train() (225 lines, was ~514) + graceful shutdown
checkpoint.py build/load, lifted out of train()'s closures
The trainers shared ~15 identical constructor arguments and copy-pasted
their cosine-warmup lambda, EMA setup, state_dict/load_state_dict,
resume_lr and train_mode/eval_mode. StageSpec resolves one stage's config
once (constructors go from 24 and 22 keyword arguments to (spec, model,
device)), the base class holds the rest, and build_stage_trainers drops
from ~100 lines to 15.
Metric columns are renamed to a uniform stage/split/metric scheme
(stage1/train/loss, stage2/train/d_loss, stage1/lr, stage1/router/entropy,
val/loss, ...). Old metrics.csv files and W&B history are not comparable.
The checkpoint format is unchanged.
BEHAVIOR CHANGE — WGAN best-checkpoint selection. The old code meant to
score a WGAN stage on its marginal KL, but the guard
`{n: kl for n in wgan_names if n not in val_loss_per_stage}` could never
fire: val_loss_per_stage was pre-seeded with 0.0 for every stage, so a
WGAN stage contributed a flat 0.0 and the KL was written to metrics.csv
without ever influencing best.pt. val_objective now returns it as
intended. On the test harness's default flow+wgan config val_loss went
from 2.182 (stage 1 only) to 15.137 (stage 1 + KL 12.954), and which epoch
won changed. Runs before this commit picked their best checkpoint on the
non-adversarial stages alone. Written up in docs/v0.3.0-followups.md.
Verified: 699 tests pass; ruff, ruff format and ty clean. Baseline-vs-
refactor metrics.csv compared across five configs (flow+wgan, AR+onehot,
routed, both-flow, AR-flow) — every comparable value bit-identical except
val/loss where the fix applies. Resume appends without a duplicate header
and reproduces a HEAD worktree's per-epoch losses and LRs exactly across
the resume boundary. A refactored last.pt loads through
cli.py:_load_model_weights in both raw and ema modes.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
156 lines
14 KiB
Markdown
156 lines
14 KiB
Markdown
# giant
|
||
|
||
**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.
|
||
|
||
Conditional generative surrogate 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 variable-length list of secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained generative model. A trained checkpoint autoregressively rolls out full showers, stepping each primary and pushing secondaries as new tracks.
|
||
|
||
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
|
||
|
||
## Architecture
|
||
|
||
A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`):
|
||
|
||
- **`flow`** (default) — conditional flow matching (Lipman et al. 2022): an MLP learns a vector field mapping noise → step outcomes, sampled via ODE integration in ~10 steps.
|
||
- **`ddpm`** — a standard denoising diffusion baseline for comparison (`giant/model/schedule.py:CosineSchedule`).
|
||
- **`wgan`** — a single-pass Wasserstein-GAN-GP generator/critic (`giant/model/wgan.py`), trading iterative sampling for one forward pass; implemented, not yet validated against the flow-matching baseline.
|
||
|
||
**Stage 1 — primary (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):**
|
||
|
||
| Index | Variable | Encoding |
|
||
|-------|----------|----------|
|
||
| 0 | `step_length` [mm] | 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 (`energy_simplex_decode`). Stage 1 also has a classifier head (`predict_n_sec`) predicting the number of secondaries `n_sec ∈ {0..K_MAX}` (`K_MAX = 15`) from the conditioning alone, no diffusion noise involved.
|
||
|
||
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.
|
||
|
||
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second net generates all `K_MAX` secondary slots at once — `(stick-breaking energy logit, local-frame direction, log-mass, charge)` per slot, 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, so the whole chain conserves energy. A secondary's mass/charge are regressed directly against its ground-truth PDG code's physical values (`giant.particles.particle_mass_charge`) and used as-is at inference — including for its own conditioning if it takes further steps in a rollout. No snapping to a known PDG code happens in the model path; `giant.particles.nearest_known_pdg` is a reporting-only lookup used to populate a nominal `pdg` label on output rows.
|
||
|
||
**Conditioning (`--conditioning`, per-checkpoint):** pre-step position, log(pre-energy), pre-step direction, layer ID, plus particle/material physical properties — mass/charge (`giant/particles.py`) and Z_eff/A_eff/density/X0/λ_int (`giant/materials.py`). Two mutually exclusive modes:
|
||
|
||
- **`physical`** (default) — the physical-property columns are routed through small MLPs, computable for any PDG code / material, letting the surrogate generalize to species/materials outside the training menu.
|
||
- **`embedding`** — the original design: a learned `nn.Embedding` per PDG code / material, kept as a generalization-comparison baseline (memorizes the training menu).
|
||
|
||
`n_sec` and `e_sec` are model outputs, not conditioning inputs — a rollout is self-contained and never injects ground truth.
|
||
|
||
**Mixture-of-experts routing (`--router`, opt-in):** `giant/model/network.py` also implements a pluggable `Router` contract (`ROUTER_REGISTRY`: `energy`, `pdg`, `process`, plus a `composed` router combining several axes) that splits `DenoisingMLP`/`SecondaryDecoder` into per-expert trunks, soft-gated in training and top-1 dispatched at eval. Implemented; first rollout benchmark needs a retrain with a load-balancing loss and better-seeded router centers (see Roadmap). See `--router-type`/`--n-experts`/`--router-axis` on `giant train`/`giant new-run`.
|
||
|
||
## Roadmap
|
||
|
||
**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 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.
|
||
|
||
**Physical-property conditioning (implemented):** replaces learned PDG/material embeddings with physical-property MLPs (see above); Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. Not yet done: the held-out-material/species generalization comparison against the `embedding` baseline — the natural dataset for that is the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes).
|
||
|
||
**Faster-eval architectures (implemented, validation in progress):** both target a ~10× native-Geant4 eval budget. WGAN-GP (`--mode wgan`) has no rollout-vs-reference analysis run against it yet. The MoE router (`--router`) had its first rollout benchmark diverge from Geant4 despite matching bulk deposited energy — the experts weren't specializing (near-uniform gating), traced to a missing load-balance loss and a center-init that didn't match the real energy distribution; both are now fixable via `lambda_balance > 0` and quantile-seeded router centers, but a re-run to confirm hasn't happened yet.
|
||
|
||
A multi-material sampling-calorimeter dataset is a planned future direction, not yet built.
|
||
|
||
## Data
|
||
|
||
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle, and `--seed`-controlled) to avoid leaking correlated steps from the same shower.
|
||
|
||
## Project structure
|
||
|
||
```
|
||
giant/
|
||
├── giant/
|
||
│ ├── data/
|
||
│ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads)
|
||
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
|
||
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
||
│ ├── model/
|
||
│ │ ├── network.py # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic
|
||
│ │ ├── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||
│ │ └── wgan.py # WGAN-GP gradient penalty / critic / generator losses
|
||
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
|
||
│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; nearest-known-PDG lookup
|
||
│ ├── materials.py # material name → (Z_eff, A_eff, density, X0, λ_int)
|
||
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
|
||
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
|
||
│ ├── training/ # two-stage training: loop, per-stage trainers, metrics, checkpointing
|
||
│ │ ├── loop.py # epoch loop, graceful shutdown, best-checkpoint selection
|
||
│ │ ├── trainers.py # StageSpec + flow/ddpm and WGAN-GP per-stage trainers
|
||
│ │ ├── stage2_inputs.py# ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
|
||
│ │ ├── metrics.py # MetricsCollector: metrics.csv columns, W&B logging, progress/summary
|
||
│ │ └── checkpoint.py # checkpoint assembly/restore (format unchanged since v0.2)
|
||
│ ├── sample.py # DDPM / DDIM / flow matching / WGAN 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/ # rollout-vs-reference analysis pipeline (see `giant analyze` below)
|
||
│ │ ├── sources.py # canonical LazyFrames + secondary view
|
||
│ │ ├── reduce.py # streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
|
||
│ │ ├── grouping.py # fixed bin edges + energy/pdg/material group sets
|
||
│ │ ├── context.py # resolves grouping into `shared.json` once per run
|
||
│ │ ├── catalog.py # declarative PlotSpec registry
|
||
│ │ ├── condor.py # prep / compute-one / submit-description plumbing
|
||
│ │ └── render.py # PDFs + HTML gallery (only module importing plotstyle/LaTeX)
|
||
│ └── cli.py # `giant train` / `new-run` / `predict` / `rollout` / `analyze` Typer app
|
||
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
|
||
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
|
||
│ │ # update-manifest, create-manifest, make-root,
|
||
│ │ # build-geometry-oracle, warm-cache, hparam-scan
|
||
│ ├── steps_to_parquet.py # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
|
||
│ ├── steps_to_parquet_parallel.py # fan out conversion over several ROOT files — `dwarf convert --jobs N`
|
||
│ ├── migrate_geant_steps.py # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
|
||
│ ├── 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`
|
||
│ ├── warm_setup_cache.py # precompute `giant train`'s setup-stage sidecar — `dwarf warm-cache`
|
||
│ ├── hparam_scan.py # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
|
||
│ └── profile_analysis_costs.py # profiling helper for the `giant analyze` reduction pipeline
|
||
└── tests/
|
||
```
|
||
|
||
## Setup
|
||
|
||
```bash
|
||
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
|
||
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
|
||
uv sync --extra cpu --extra geometry # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
|
||
```
|
||
|
||
`cpu` and `cuda` are mutually exclusive extras selecting the torch build (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||
|
||
## Training, prediction, and rollout
|
||
|
||
```bash
|
||
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir for a new run
|
||
giant train path/to/steps.parquet --mode flow # train (flow matching; also --mode ddpm / wgan)
|
||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||
|
||
# Full-shower rollout needs a geometry oracle (position → material/layer_id):
|
||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
|
||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||
```
|
||
|
||
`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant train --wandb` logs per-epoch metrics (the same ones written to `metrics.csv`) to Weights & Biases; requires `uv sync --extra wandb`. Metric names are `<stage>/<split>/<metric>` (e.g. `stage1/train/loss`, `stage2/train/d_loss`, `stage1/lr`), plus an unprefixed run-level tail (`val/loss`, `grad_norm`, `epoch_time_s`, ...); each is declared once on the `StageTrainer` that computes it (`giant/training/trainers.py`), so the CSV header and the W&B panel names are derived, never hand-maintained. A repeat `giant train` against the same dataset (e.g. a hyperparameter sweep) reuses a cached setup-stage sidecar (vocab maps, event split, normalizer stats) unless `--no-cache-setup`/`--rebuild-setup-cache`; `dwarf warm-cache` precomputes it ahead of time. `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 and analysis
|
||
|
||
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`).
|
||
|
||
For deeper rollout-vs-reference diagnostics — marginals stratified by energy/pdg/material, per-event totals, shower profiles, species share, leakage, and secondaries — `giant analyze` runs a streaming compute/render pipeline against a `giant rollout` YAML sidecar:
|
||
|
||
```bash
|
||
giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot (compute only)
|
||
giant analyze render <run_dir> --gallery # local: styled PDFs + HTML gallery (needs LaTeX)
|
||
```
|
||
|
||
`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` imports plotstyle/LaTeX, so it always runs locally.
|
||
|
||
## Development
|
||
|
||
```bash
|
||
uv run pytest # run tests
|
||
uv run ruff check . # lint
|
||
uv run ruff format . # format
|
||
uv run ty check . # type check
|
||
```
|