gitea-actions
7cd22a77ce
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 4m18s
CI / Format (ruff format) (push) Successful in 4m14s
CI / Type check (ty) (push) Successful in 4m18s
CI / Tests (push) Successful in 5m44s
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 18s
CI / Update README badges (version, test count) (push) Successful in 1m35s
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.
|