lars
bf3271f09e
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m6s
CI / Format (ruff format) (push) Successful in 1m6s
CI / Type check (ty) (push) Successful in 1m12s
CI / Tests (push) Successful in 3m15s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 28s
CI / Update README badges (version, test count) (push) Successful in 1m21s
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
296 lines
13 KiB
Markdown
296 lines
13 KiB
Markdown
<div align="center">
|
||
|
||
# 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)
|
||
|
||
</div>
|
||
|
||
---
|
||
|
||
## 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 <run_dir> --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.
|
||
|
||
---
|
||
|
||
<details>
|
||
<summary><h2 style="display:inline">Architecture</h2></summary>
|
||
|
||
**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` |
|
||
|
||
</details>
|
||
|
||
<details>
|
||
<summary><h2 style="display:inline">Configuration</h2></summary>
|
||
|
||
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.
|
||
|
||
</details>
|
||
|
||
<details>
|
||
<summary><h2 style="display:inline">Data</h2></summary>
|
||
|
||
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`.
|
||
|
||
</details>
|
||
|
||
<details>
|
||
<summary><h2 style="display:inline">CLI reference</h2></summary>
|
||
|
||
**`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 <run_dir> --gallery` | merge chunks → styled PDFs + HTML gallery (local, needs LaTeX) |
|
||
| `analyze metrics <train_run_dir>` | 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`.
|
||
|
||
</details>
|
||
|
||
<details>
|
||
<summary><h2 style="display:inline">Rollout & analysis</h2></summary>
|
||
|
||
`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`.
|
||
|
||
</details>
|
||
|
||
<details>
|
||
<summary><h2 style="display:inline">Install</h2></summary>
|
||
|
||
| 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`.
|
||
|
||
</details>
|
||
|
||
<details>
|
||
<summary><h2 style="display:inline">Development</h2></summary>
|
||
|
||
```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.
|
||
|
||
</details>
|
||
|
||
## License
|
||
|
||
Not yet decided — treat this repository as all-rights-reserved until a
|
||
`LICENSE` file is added.
|