# GIANT
### **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
calorimeter showers.
[](pyproject.toml)
[](pyproject.toml)
[](CHANGELOG.md)
[](tests/)
[](https://git.larsbogner.de/lars/giant/actions)
[](#license)
---
## The idea
Geant4's step function is the innermost loop of detector simulation — for every
particle, at every step, it stochastically samples where the particle goes next,
how much energy it deposits, and what secondaries it spawns. GIANT learns that
function instead of running it: given a pre-step particle state (position,
energy, direction, particle species, material), a two-stage model samples a
post-step outcome — including the variable-length list of secondaries — and
autoregressively rolls that out into whole showers. It trains entirely from
parquet dumps of a Geant4 steps tree; nothing downstream needs a Geant4 runtime.
Two guarantees are architectural, not learned:
- **Energy is conserved by construction.** Stage 1 decodes deposit / secondary /
post-step energy through a softmax simplex that sums to the pre-step energy
exactly; Stage 2's secondaries stick-break that same energy budget.
- **No shower leaks across the train/val split.** Steps are split by `event_id`,
never by row, so correlated steps from the same shower can't appear on both
sides.
## How a step becomes a shower
```mermaid
flowchart LR
A["pre-step state\nposition · energy · direction\nspecies · material"] --> B["ConditionEncoder\nphysical / embedding / onehot"]
B --> C["Stage 1\n9D post-step outcome"]
C --> D["Stage 2 (autoregressive)\nsecondaries, descending energy"]
D --> E["rollout step"]
E -->|"primary continues"| F["GeometryOracle\nposition to material, layer"]
E -->|"secondaries pushed"| G["track queue"]
F --> A
G --> A
E -->|"terminated"| H["deposited shower"]
```
## Quick start
```bash
uv sync --extra cpu # install deps (CPU torch; --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 stage 1 + wgan stage 2, by default)
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
giant analyze prep rollout.yaml && giant analyze render --gallery # rollout-vs-Geant4 diagnostics
```
Every command takes `--help` for its full flag list, and `--config config.toml`
for anything not exposed as a flag.
---
Architecture
**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 coordinates of the deposit / secondary / post-energy simplex |
| 3–5 | `post_dir` | unit vector, local frame (`pre_dir = ẑ`) |
| 6–8 | `travel_dir` (`post_pos − pre_pos`) | unit vector, local frame |
Energy logits decode via `softmax([edep_logit, sec_logit, 0]) × pre_E`, so
`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 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).
**Conditioning (15D).** 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 secondary species above, configured *independently* per axis
(`conditioning.particle.type` / `conditioning.material.type`). `physical`
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.
**Composable by design** — every stage assembles from small registries, so
swapping one axis doesn't touch the others:
| Registry | Choices |
|---|---|
| Objective | `flow` (matching, ~10-step ODE sample) · `ddpm` (denoising diffusion) · `wgan` (single-pass GAN) |
| Trunk | `resmlp` · `none`, optionally MoE-routed (`RoutedTrunk`) |
| Router | `energy` · `pdg` · `process` · `composed` · `none` — soft-mixed at train time, **top-1 dispatched at eval time**, which is the actual inference-speed win |
| History (stage 2 AR) | `markov` · `attention` · `none` |
Configuration
Every default lives in one place: frozen dataclasses in `giant/config.py`,
composed into `GiantConfig` (`conditioning` / `stage1_model` / `stage2_model` /
`train`). `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather
than hand-maintained, so the dataclasses can't drift from what actually gets
merged. TOML config keys are validated against that shape — an unknown key is
rejected with a did-you-mean suggestion. Precedence: CLI flag > `--config` file
> default.
```toml
# config.toml — resolved shape of the four blocks
[conditioning]
particle.type = "physical"
material.type = "physical"
[stage1_model]
generator = "flow"
[stage2_model]
generator = "wgan"
decoder = "autoregressive"
[train]
epochs = 100
batch_size = 4096
lr = 3e-4
```
Some knobs only exist in the config file, with no CLI flag:
`stage2_model.autoregressive.teacher_forcing`/`.history`,
`stage2_model.particle_type.target`/`.class_weighting`,
`stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`,
`stage2_model.router.*`, and the finer `router` knobs (`lambda_balance`,
`gumbel`, `learn_width`, …).
`configs/` holds kept reference configs — `baseline.toml` is the fixed
comparison point every experimental variant (routed trunk, WGAN, attention
history, embedding conditioning) is a single edit away from. v0.2 flat-schema
configs and checkpoints load and auto-migrate.
Data
Input is parquet — one row per Geant4 step — from
[miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from
ROOT via `dwarf convert`.
| Group | Columns |
|---|---|
| **Conditioning (pre-step)** | `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz`, `material`, `layer_id` |
| **Primary outcome (post-step)** | `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep`, `e_sec`, `child_track_ids` (length → `n_sec`) |
| **Secondaries** (variable-length lists) | `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX = 15` slots, descending energy |
| **Optional** | `process` — physics-process label, classifier supervision only (`ProcessRouter`), never conditioning |
Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so a
shower's correlated steps never straddle the split. Loading a directory or
`.manifest` of several parquet files offsets each file's `event_id`s by a
per-file stride so ids stay globally unique. The pre-epoch setup scan (vocab
maps, event split, normalizer stats) persists to a sidecar cache
(`--cache-setup`/`--rebuild-setup-cache`), precomputable ahead of time via
`dwarf warm-cache`.
CLI reference
**`giant`** — train, run, and analyze the surrogate:
| Command | Does |
|---|---|
| `new-run` | scaffold a `config.toml` + run directory from flags |
| `train DATA` | train the two-stage model |
| `model summary` | build-only parameter counts, without training |
| `predict DATA --checkpoint …` | per-step predictions from a checkpoint |
| `rollout DATA --checkpoint … --geometry …` | full autoregressive shower rollout |
| `analyze prep/submit` | build a run dir; `submit` also queues HTCondor compute jobs |
| `analyze compute-one` / `merge-one` | one plot × chunk reduction / merge (what a condor job runs) |
| `analyze render --gallery` | merge chunks → styled PDFs + HTML gallery (local, needs LaTeX) |
| `analyze metrics ` | training-progress plots from `metrics.csv` |
| `analyze list` | every catalog plot id |
**`dwarf`** — dataset/tooling CLI:
| Command | Does |
|---|---|
| `convert` | ROOT Steps tree → parquet (`--jobs N` fans out) |
| `migrate` | one-time move into the raw/processed/pools/derived layout |
| `bump-gen` / `bump-schema` / `status` | dataset versioning |
| `update-manifest` / `create-manifest` | point/build a manifest of parquet files |
| `make-root` | generate new ROOT shards via a minicalosim executable |
| `build-geometry-oracle` | fit position → (material, layer_id) for rollout |
| `warm-cache` | precompute `giant train`'s setup-stage sidecar |
| `hparam-scan` | grid-scan dropout × n_blocks × hidden_dim |
Worth knowing on `giant train` (full surface behind `--help`):
`--mode {flow,ddpm,wgan}` / `--stage1-generator` / `--stage2-generator`,
`--stage2-decoder {autoregressive,one_shot}`, `--conditioning
{physical,embedding,onehot}`, `--router` / `--router-type` / `--n-experts` /
`--router-axis`, `--stage{1,2}-init-from` + `--stage{1,2}-freeze` (retrain one
stage against a fixed other one), `--precision {fp32,bf16}`, `--wandb`.
Rollout & analysis
`giant rollout` seeds showers from each event's highest-energy entry step, then
autoregressively steps the model to completion — advancing all active tracks
breadth-first, batched — pushing secondaries as new tracks and looking up
`material`/`layer_id` from the geometry oracle each step. Tracks terminate on
one of six reasons (energy cutoff, max steps, detector escape, natural end,
unknown pdg, max tracks); every reason but escape deposits the remaining energy
locally, so showers conserve energy by construction — only `escaped` counts as
leakage.
`giant analyze` compares one or more rollouts against a single held-out
reference: `prep` resolves shared bin edges/groups once, `submit`/`compute-one`
run each (plot, `event_id`-disjoint chunk) pair as a polars/numpy-only HTCondor
job, `render` merges the chunks and produces the styled PDFs + HTML gallery
locally (the only step that needs LaTeX). Each rollout gets its own colored
series against one shared reference line. `giant analyze metrics` is a separate
entry point — training-progress plots straight from a run's `metrics.csv`.
Install
| Extra | Adds | For |
|---|---|---|
| `cpu` **or** `cuda` | torch 2.3.x | required — mutually exclusive, pick one |
| `geometry` | scikit-learn | `dwarf build-geometry-oracle`, rollout |
| `analysis` | matplotlib, plotstyle | `giant analyze render` |
| `convert` | uproot, awkward | `dwarf convert` |
| `wandb` | wandb | `giant train --wandb` |
| `dev` | pytest, ruff, ty, + all of the above | development |
```bash
uv sync --extra cpu --extra dev # everything needed to develop
```
Plain `uv sync` with no extra installs **no torch at all** — always include
`--extra cpu` or `--extra cuda`.
Development
```bash
uv run pytest # 964 tests
uv run ruff check . # lint
uv run ruff format . # format
uv run ty check . # type check
```
Gitea Actions (`.gitea/workflows/ci.yml`) runs lint + format-check + type-check
+ tests on every push and PR; merges to `master` auto-bump the patch version
and regenerate `CHANGELOG.md` — don't hand-edit either.
## License
Not yet decided — treat this repository as all-rights-reserved until a
`LICENSE` file is added.