giant

Geant4 Inference via Autoregressive Neural sTep surrogate — a play on Geant4 and the step function being the computationally heaviest part of the simulation.

Conditional generative surrogate for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the variable-length list of secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained generative model. A trained checkpoint autoregressively rolls out full showers, stepping each primary and pushing secondaries as new tracks.

Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.

Architecture

A two-stage model, both stages checkpointed together, with a choice of generative mode per stage (--mode):

• flow (default) — conditional flow matching (Lipman et al. 2022): an MLP learns a vector field mapping noise → step outcomes, sampled via ODE integration in ~10 steps.
• ddpm — a standard denoising diffusion baseline for comparison (giant/model/schedule.py:CosineSchedule).
• wgan — a single-pass Wasserstein-GAN-GP generator/critic (giant/model/wgan.py), trading iterative sampling for one forward pass; implemented, not yet validated against the flow-matching baseline.

Stage 1 — primary (9D, giant/constants.py:LOCAL_TARGET_NAMES):

Index	Variable	Encoding
0	step_length [mm]	log
1–2	edep_logit, sec_logit	ALR coords of the deposit/secondary/post-energy simplex
3–5	post_dir in local frame	unit vector
6–8	travel_dir (post_pos − pre_pos) in local frame	unit vector

The two energy logits decode via softmax over [edep_logit, sec_logit, 0] × pre_E, so edep + e_sec + post_E == pre_E exactly — energy conservation is built into the parametrization, not left to the loss (energy_simplex_decode). Stage 1 also has a classifier head (predict_n_sec) predicting the number of secondaries n_sec ∈ {0..K_MAX} (K_MAX = 15) from the conditioning alone, no diffusion noise involved.

Both post_dir and travel_dir are expressed in the coordinate frame where pre_dir = ẑ, making the scattering distribution nearly azimuthally symmetric. post_pos itself is not a raw target — it's reconstructed at inference as pre_pos + step_length * world_frame(travel_dir), so the two stay consistent by construction instead of being learned (and potentially diverging) independently.

Stage 2 — secondaries (SecondaryDecoder): conditioned on the pre-step state and the Stage-1 outcome, a second net generates all K_MAX secondary slots at once — (stick-breaking energy logit, local-frame direction, log-mass, charge) per slot, ordered by descending energy; slots beyond the predicted n_sec are masked. Secondary energies are a stick-breaking partition of the e_sec budget from Stage 1, so the whole chain conserves energy. A secondary's mass/charge are regressed directly against its ground-truth PDG code's physical values (giant.particles.particle_mass_charge) and used as-is at inference — including for its own conditioning if it takes further steps in a rollout. No snapping to a known PDG code happens in the model path; giant.particles.nearest_known_pdg is a reporting-only lookup used to populate a nominal pdg label on output rows.

Conditioning (--conditioning, per-checkpoint): pre-step position, log(pre-energy), pre-step direction, layer ID, plus particle/material physical properties — mass/charge (giant/particles.py) and Z_eff/A_eff/density/X0/λ_int (giant/materials.py). Two mutually exclusive modes:

• physical (default) — the physical-property columns are routed through small MLPs, computable for any PDG code / material, letting the surrogate generalize to species/materials outside the training menu.
• embedding — the original design: a learned nn.Embedding per PDG code / material, kept as a generalization-comparison baseline (memorizes the training menu).

n_sec and e_sec are model outputs, not conditioning inputs — a rollout is self-contained and never injects ground truth.

Mixture-of-experts routing (--router, opt-in): giant/model/network.py also implements a pluggable Router contract (ROUTER_REGISTRY: energy, pdg, process, plus a composed router combining several axes) that splits DenoisingMLP/SecondaryDecoder into per-expert trunks, soft-gated in training and top-1 dispatched at eval. Implemented; first rollout benchmark needs a retrain with a load-balancing loss and better-seeded router centers (see Roadmap). See --router-type/--n-experts/--router-axis on giant train/giant new-run.

Roadmap

Phase 1 (done): n_sec and total secondary energy e_sec were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).

Phase 2 (implemented — baseline): the two-stage model above predicts n_sec and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.

Physical-property conditioning (implemented): replaces learned PDG/material embeddings with physical-property MLPs (see above); Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. Not yet done: the held-out-material/species generalization comparison against the embedding baseline — the natural dataset for that is the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes).

Faster-eval architectures (implemented, validation in progress): both target a ~10× native-Geant4 eval budget. WGAN-GP (--mode wgan) has no rollout-vs-reference analysis run against it yet. The MoE router (--router) had its first rollout benchmark diverge from Geant4 despite matching bulk deposited energy — the experts weren't specializing (near-uniform gating), traced to a missing load-balance loss and a center-init that didn't match the real energy distribution; both are now fixable via lambda_balance > 0 and quantile-seeded router centers, but a re-run to confirm hasn't happened yet.

A multi-material sampling-calorimeter dataset is a planned future direction, not yet built.

Data

Input: parquet files produced by miniCaloSim, or converted from a ROOT file via dwarf convert. Each row is one Geant4 step. Train/val split is by event_id (not row shuffle, and --seed-controlled) to avoid leaking correlated steps from the same shower.

Project structure

giant/
├── giant/
│   ├── data/
│   │   ├── loader.py       # parquet → numpy arrays (incl. streaming/chunked reads)
│   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
│   │   └── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
│   ├── model/
│   │   ├── network.py      # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic
│   │   ├── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
│   │   └── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
│   ├── constants.py        # output/conditioning dims, K_MAX, secondary slot layout, schema keys
│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; nearest-known-PDG lookup
│   ├── materials.py        # material name → (Z_eff, A_eff, density, X0, λ_int)
│   ├── config.py           # default hyperparameters, TOML config merging, device autodetect
│   ├── pipeline.py         # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
│   ├── train.py            # two-stage training loop, checkpointing, graceful shutdown, W&B logging
│   ├── sample.py           # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
│   ├── geometry.py         # GeometryOracle: position → (material, layer_id, escaped) for rollout
│   ├── rollout.py          # autoregressive shower rollout driver
│   ├── validate.py         # step-level marginal + KL-divergence validation
│   ├── analysis/           # rollout-vs-reference analysis pipeline (see `giant analyze` below)
│   │   ├── sources.py      #   canonical LazyFrames + secondary view
│   │   ├── reduce.py       #   streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
│   │   ├── grouping.py     #   fixed bin edges + energy/pdg/material group sets
│   │   ├── context.py      #   resolves grouping into `shared.json` once per run
│   │   ├── catalog.py      #   declarative PlotSpec registry
│   │   ├── condor.py       #   prep / compute-one / submit-description plumbing
│   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
│   └── cli.py              # `giant train` / `new-run` / `predict` / `rollout` / `analyze` Typer app
├── scripts/                            # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
│   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
│   │                                    #   update-manifest, create-manifest, make-root,
│   │                                    #   build-geometry-oracle, warm-cache, hparam-scan
│   ├── steps_to_parquet.py             # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
│   ├── steps_to_parquet_parallel.py    # fan out conversion over several ROOT files — `dwarf convert --jobs N`
│   ├── migrate_geant_steps.py          # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
│   ├── bump_dataset_version.py         # cut a new raw gen or parquet schema, with a logged reason —
│   │                                    #   `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
│   ├── create_root_files.py            # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
│   ├── geometry_oracle.py              # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
│   ├── warm_setup_cache.py             # precompute `giant train`'s setup-stage sidecar — `dwarf warm-cache`
│   ├── hparam_scan.py                  # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
│   └── profile_analysis_costs.py       # profiling helper for the `giant analyze` reduction pipeline
└── tests/

Setup

uv sync --extra cpu               # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
uv sync --extra cpu --extra dev   # add dev tools (pytest, ruff, ty)
uv sync --extra cpu --extra geometry  # add scikit-learn, for `dwarf build-geometry-oracle` / rollout

cpu and cuda are mutually exclusive extras selecting the torch build (pinned to 2.3.x); plain uv sync installs no torch at all. See CLAUDE.md for details.

Training, prediction, and rollout

giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir for a new run
giant train path/to/steps.parquet --mode flow               # train (flow matching; also --mode ddpm / wgan)
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt

# Full-shower rollout needs a geometry oracle (position → material/layer_id):
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl

train/predict accept a TOML config file (--config) and CLI overrides for hyperparameters; see --help on any command for the full option list. giant train --wandb logs per-epoch metrics (the same ones written to metrics.csv) to Weights & Biases; requires uv sync --extra wandb. A repeat giant train against the same dataset (e.g. a hyperparameter sweep) reuses a cached setup-stage sidecar (vocab maps, event split, normalizer stats) unless --no-cache-setup/--rebuild-setup-cache; dwarf warm-cache precomputes it ahead of time. giant rollout seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up material/layer_id from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.

Validation and analysis

giant.validate.validate_marginals runs step-level marginal and KL-divergence checks during training (--validate-every).

For deeper rollout-vs-reference diagnostics — marginals stratified by energy/pdg/material, per-event totals, shower profiles, species share, leakage, and secondaries — giant analyze runs a streaming compute/render pipeline against a giant rollout YAML sidecar:

giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot (compute only)
giant analyze render <run_dir> --gallery                   # local: styled PDFs + HTML gallery (needs LaTeX)

<run_dir> is derived next to the rollout parquet (analyze prep/submit print it). Compute jobs are polars/numpy only; only render imports plotstyle/LaTeX, so it always runs locally.

Development

uv run pytest            # run tests
uv run ruff check .      # lint
uv run ruff format .     # format
uv run ty check .        # type check