lars
630d8d3992
CI / Format (ruff format) (push) Successful in 27s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 34s
CI / Type check (ty) (push) Successful in 38s
CI / Format (ruff format) (pull_request) Successful in 48s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 48s
CI / Tests (pull_request) Successful in 2m40s
CI / Tests (push) Successful in 2m46s
The two-stage architecture description had drifted from the v0.3.0 stage-2-autoregressive redesign (8 commits, eb6dd27..da7cde3) — it still documented the old one-shot-only SecondaryDecoder and continuous mass/charge secondary target. Restructured for faster onboarding: a Quick start section up front, bullet-point Architecture and training-flag docs instead of dense paragraphs, and a Data section listing the actual parquet columns consumed by giant/data/loader.py. Dropped the Roadmap section (status/history, not architecture) and CLI-flag default callouts from Architecture, keeping it focused on net structure. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
167 lines
13 KiB
Markdown
167 lines
13 KiB
Markdown
# giant
|
||
|
||
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
|
||
|
||
A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.
|
||
|
||
## Quick start
|
||
|
||
```bash
|
||
uv sync --extra cpu # install deps (CPU torch; use --extra cuda for GPU)
|
||
|
||
giant new-run --hidden-dim 512 --lr 3e-4 # scaffold config.toml + run dir
|
||
giant train path/to/steps.parquet # train (flow + wgan by default)
|
||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||
|
||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl # needed for rollout
|
||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||
```
|
||
|
||
Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.
|
||
|
||
## Architecture
|
||
|
||
A **two-stage model**, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (`--stage1-generator`/`--stage2-generator`, or `--mode` to set both at once): `flow` (conditional flow matching, ODE-sampled in ~10 steps), `ddpm` (denoising diffusion), or `wgan` (single-pass WGAN-GP generator/critic).
|
||
|
||
**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
|
||
|
||
| 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 |
|
||
|
||
- Energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — conservation is architectural, not learned.
|
||
- `post_dir`/`travel_dir` live in the frame where `pre_dir = ẑ`. `post_pos` isn't a target — it's reconstructed as `pre_pos + step_length * world_frame(travel_dir)`.
|
||
|
||
**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (`--stage2-decoder`):
|
||
|
||
- `autoregressive` — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (`markov`: previous token only, or `attention`: causal self-attention, KV-cached at inference)
|
||
- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
|
||
|
||
Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).
|
||
|
||
**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above (`--conditioning`) — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
|
||
|
||
**MoE routing** (`--router`, either stage): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk.
|
||
|
||
## Data
|
||
|
||
- Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from ROOT via `dwarf convert`. One row = one Geant4 step.
|
||
- **Conditioning (pre-step) columns:** `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz` (direction), `material`, `layer_id`.
|
||
- **Primary outcome (post-step) columns:** `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep` (energy deposited in this step), `e_sec` (total energy carried off by secondaries), `child_track_ids` (its length gives `n_sec`).
|
||
- **Secondary columns**, one variable-length list per step: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX` (15) slots on load, ordered by descending energy.
|
||
- **Optional:** `process` — the physics process that produced the step (e.g. `compt`, `phot`, `eBrem`); a post-step label used only as classifier supervision (`ProcessRouter`), never as conditioning.
|
||
- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
|
||
- Loading a directory or `.manifest` of multiple parquet files (each one Geant4 job, `event_id` restarting from 0) offsets each file's `event_id`s by a fixed per-file stride so ids stay globally unique across files.
|
||
|
||
## 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, Stage1Model, Stage2OneShot/Stage2Autoregressive, Router/MoE, CriticModel
|
||
│ │ ├── 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; onehot/embedding secondary-identity decode
|
||
│ ├── 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 — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||
|
||
## Training, prediction, rollout
|
||
|
||
```bash
|
||
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir
|
||
giant train path/to/steps.parquet # train (flow stage 1 + wgan stage 2, default)
|
||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||
|
||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl # position → material/layer_id
|
||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||
```
|
||
|
||
Useful flags on `giant train`:
|
||
|
||
- `--mode {flow,ddpm,wgan}` sets both stages' objective at once; `--stage1-generator`/`--stage2-generator` override per stage
|
||
- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
|
||
- `--conditioning {physical,embedding,onehot}` — conditioning representation
|
||
- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
|
||
- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
|
||
- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
|
||
|
||
Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
|
||
|
||
`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.
|
||
|
||
## Validation and analysis
|
||
|
||
- `giant.validate.validate_marginals` — step-level marginal + KL-divergence checks during training (`--validate-every`)
|
||
- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
|
||
|
||
```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` needs 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
|
||
```
|