Rewrite README.md from scratch as a scannable landing page (hero, one
mermaid pipeline diagram, quick start, deep detail folded into
collapsible sections) instead of the old flat prose dump duplicating
CLAUDE.md.
Swap the static "CI" badge for a live Gitea Actions status badge, and
add an update-badges job to ci.yml that recomputes the version and
test-count badges on every push to master and pushes an update only
when they actually changed.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01JL1hFbhLv5uwjTXqkWTLnH
2026-09-02 15:58:02 +02:00
2 changed files with 299 additions and 161 deletions
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
# GIANT
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.
### **G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate
A conditional generative model that replaces the Geant4 step function — sample a
post-step outcome instead of simulating one, then roll that out into full
Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.
Every command takes `--help` for its 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:
**Stage 1 — primary step.** Predicts the 9D post-step outcome
(`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
| Index | Variable | Encoding |
| Index | Variable | Encoding |
|-------|----------|----------|
|-------|----------|----------|
| 0 | `step_length` [mm] | log |
| 0 | `step_length` [mm] | log |
| 1–2 | `edep_logit`, `sec_logit` | ALR coords of the deposit/secondary/post-energy simplex |
| 1–2 | `edep_logit`, `sec_logit` | ALR coordinates of the deposit / secondary / post-energy simplex |
| 3–5 | `post_dir`in local frame | unit vector |
| 3–5 | `post_dir`| unit vector, local frame (`pre_dir = ẑ`) |
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
| 6–8 | `travel_dir` (`post_pos − pre_pos`) | unit vector, local frame |
-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.
Energy logits decode via `softmax([edep_logit, sec_logit, 0])× pre_E`, so
-`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)`.
`edep + e_sec + post_E == pre_E` holds exactly. `post_pos` is not itself a
target — it's reconstructed as `pre_pos + step_length · world_frame(travel_dir)`,
since duplicating that magnitude in a second target would let the two drift out
of sync.
**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`):
**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's
outcome, it generates the variable-length secondary list, one token at a time in
descending-energy order (`autoregressive`, default) or all `K_MAX` slots in one
masked pass (`one_shot`). Autoregressive tokens condition on a running history —
`markov` (previous token only) or `attention` (causal self-attention, KV-cached
at inference). Either way, secondary energies stick-break the `e_sec` budget
handed down from Stage 1, so the whole chain conserves energy. A secondary's
species is represented `onehot` (categorical, top-N PDG codes + "other"),
`physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour
lookup).
-`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)
**Conditioning (15D).** Pre-step position / energy / direction / layer, plus
-`one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
particle mass/charge and material Z_eff/A_eff/density/X0/λ_int — encoded the
same three ways as secondary species above, configured *independently* per axis
computes rather than looks up, so it generalizes to species and materials
outside the training menu; that's the default. `n_sec`/`e_sec` are always model
outputs, never conditioning inputs.
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).
**Composable by design** — every stage assembles from small registries, so
swapping one axis doesn't touch the others:
**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. The particle and material axes are configured independently (`conditioning.particle.type` / `conditioning.material.type`; `--conditioning` sets both at once) and may mix — 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.
| Router | `energy` · `pdg` · `process` · `composed` · `none` — soft-mixed at train time, **top-1 dispatched at eval time**, which is the actual inference-speed win |
**MoE routing** (`--router`): 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. The CLI flags configure Stage 1's router; Stage 2 has its own `stage2_model.router` block, config-file only.
- 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.
Every default lives in one place: frozen dataclasses in `giant/config.py`,
composed into `GiantConfig` (`conditioning` / `stage1_model` / `stage2_model` /
- **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`).
`train`). `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather
- **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.
than hand-maintained, so the dataclasses can't drift from what actually gets
- **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.
merged. TOML config keys are validated against that shape — an unknown key is
- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
rejected with a did-you-mean suggestion. Precedence: CLI flag > `--config` file
- 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.
The `dev` extra pulls in `convert`, `analysis`, `geometry` and `wandb` as well.
Plain `uv sync` with no extra installs **no torch at all** — always include
`--extra cpu` or `--extra cuda`.
`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.
-`--stage2-stage1-context {truth,sampled}` — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via `stage2_model.ctx_p_start`/`ctx_p_end`)
-`--precision {fp32,bf16}` — bf16 autocast in the training loop
-`--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
-`--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`/`.class_weighting`, `stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`, `stage*_model.trunk.*` and the finer `router` knobs (`lambda_balance`, `gumbel`, `learn_width`, …). `configs/` holds kept reference configs. 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 × chunk (compute only)
giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan # N rollouts vs one shared reference
giant analyze render <run_dir> --gallery # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
giant analyze list # every catalog plot id
giant analyze prep rollout.yaml --chunks 8# just the run directory, no submission
giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0# what a condor job runs
`<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.
+ tests on every push and PR; merges to `master` auto-bump the patch version
and regenerate `CHANGELOG.md` — don't hand-edit either.
</details>
## License
Not yet decided — treat this repository as all-rights-reserved until a
`LICENSE` file is added.
Reference in New Issue
Block a user
Blocking a user prevents them from interacting with repositories, such as opening or commenting on pull requests or issues. Learn more about blocking a user.
lars