GIANT

Geant4 Inference via Autoregressive Neural sTep 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.

---

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

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

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.

---

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.

# 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, 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_ids 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 <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.

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

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

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.