giant

Geant4 Inference via Autoregressive Neural sTep surrogate.

A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.

Quick start

uv sync --extra cpu               # install deps (CPU torch; use --extra cuda for GPU)

giant new-run --hidden-dim 512 --lr 3e-4          # scaffold config.toml + run dir
giant model summary --config config.toml          # parameter counts + which config keys actually bite
giant train path/to/steps.parquet                 # train (flow + wgan by default)
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt

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

Every command takes --help for the full flag list, and --config config.toml for anything not exposed as a flag.

Architecture

A two-stage model, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (--stage1-generator/--stage2-generator, or --mode to set both at once): flow (conditional flow matching, ODE-sampled in ~10 steps), ddpm (denoising diffusion), or wgan (single-pass WGAN-GP generator/critic).

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 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

• Energy logits decode via softmax over [edep_logit, sec_logit, 0] × pre_E, so edep + e_sec + post_E == pre_E exactly — conservation is architectural, not learned.
• post_dir/travel_dir live in the frame where pre_dir = ẑ. post_pos isn't a target — it's reconstructed as pre_pos + step_length * world_frame(travel_dir).

Stage 2 — secondaries. Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (--stage2-decoder):

• autoregressive — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (markov: previous token only, or attention: causal self-attention, KV-cached at inference)
• one_shot — all K_MAX slots generated in a single forward pass, masked past the predicted n_sec

Either way, secondary energies stick-break the e_sec budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as onehot (categorical, top-N PDG codes + "other"), physical (continuous log-mass/charge), or embedding (nearest-neighbour lookup).

Conditioning. 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 particle identity above. The particle and material axes are configured independently (conditioning.particle.type / conditioning.material.type; --conditioning sets both at once) and may mix — the physical representation generalizes to species/materials outside the training menu since it's computed rather than looked up. n_sec/e_sec are always model outputs, never conditioning inputs.

MoE routing (--router): a pluggable Router (energy/pdg/process/composed axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk. The CLI flags configure Stage 1's router; Stage 2 has its own stage2_model.router block, config-file only.

Data

• Input: parquet files produced by miniCaloSim, or converted from ROOT via dwarf convert. One row = one Geant4 step.
• Conditioning (pre-step) columns: event_id, pdg, pre_x/pre_y/pre_z, pre_E, pre_dx/pre_dy/pre_dz (direction), material, layer_id.
• Primary outcome (post-step) columns: post_x/post_y/post_z, post_E, post_dx/post_dy/post_dz, step_length, edep (energy deposited in this step), e_sec (total energy carried off by secondaries), child_track_ids (its length gives n_sec).
• Secondary columns, one variable-length list per step: sec_pdg_list, sec_E_list, sec_dx_list/sec_dy_list/sec_dz_list — padded/truncated to K_MAX (15) slots on load, ordered by descending energy.
• Optional: process — the physics process that produced the step (e.g. compt, phot, eBrem); a post-step label used only as classifier supervision (ProcessRouter), never as conditioning.
• Train/val split is by event_id (--seed-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
• Loading a directory or .manifest of multiple parquet files (each one Geant4 job, event_id restarting from 0) offsets each file's event_ids by a fixed per-file stride so ids stay globally unique across files.

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)
│   │   └── setup_cache.py  # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
│   ├── model/
│   │   ├── models.py       # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
│   │   ├── builders.py     # build_models / build_critics — config dict → assembled stage models
│   │   ├── encoders.py     # ConditionEncoder (physical / embedding / onehot, per axis)
│   │   ├── layers.py       # ResBlock/AdaLNResBlock registry, SinusoidalEmbedding, MLP heads
│   │   ├── trunks.py       # trunk registry (resmlp, none) + RoutedTrunk (MoE expert bodies)
│   │   ├── routers.py      # Router registry: energy / pdg / process / composed / none
│   │   ├── history.py      # stage-2 AR history encoders: markov / attention (KV-cached) / none
│   │   ├── objectives.py   # flow / ddpm / wgan objective registry
│   │   ├── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
│   │   ├── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
│   │   ├── summary.py      # build-only introspection behind `giant model summary`
│   │   ├── _legacy.py      # v0.2 checkpoint model_config/state-dict migration
│   │   └── network.py      # re-export shim over all of the above
│   ├── constants.py        # output/conditioning dims, K_MAX, secondary slot layout, schema keys
│   ├── cond_layout.py      # single source of truth for the cond_cont/cond_cat column layout
│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
│   ├── 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)
│   ├── training/           # two-stage training: loop, per-stage trainers, metrics, checkpointing
│   │   ├── loop.py         #   epoch loop, graceful shutdown, best-checkpoint selection
│   │   ├── trainers.py     #   StageSpec + flow/ddpm and WGAN-GP per-stage trainers
│   │   ├── stage2_inputs.py#   ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
│   │   ├── metrics.py      #   MetricsCollector: metrics.csv columns, W&B logging, progress/summary
│   │   ├── amp.py          #   bf16 autocast (`train.precision`)
│   │   ├── plots.py        #   training-progress plots (`giant analyze metrics`)
│   │   └── checkpoint.py   #   checkpoint assembly/restore (format unchanged since v0.2)
│   ├── sample.py           # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
│   ├── checkpoint_io.py    # checkpoint → ready-to-run models/normalizers (predict + rollout)
│   ├── geometry.py         # GeometryOracle: position → (material, layer_id, escaped) for rollout
│   ├── rollout.py          # autoregressive shower rollout driver
│   ├── validate.py         # step-level marginal + KL-divergence validation
│   ├── _migration.py       # shared v0.2 → v0.3 facts used by both migration surfaces
│   ├── analysis/           # rollout-vs-reference analysis pipeline (see `giant analyze` below)
│   │   ├── sources.py      #   canonical LazyFrames + secondary view
│   │   ├── variables.py    #   per-step value expressions shared by range sizing and the catalog
│   │   ├── 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
│   │   ├── reduced.py      #   Partial/Reduced — the compact JSON a compute job emits
│   │   ├── catalog.py      #   declarative PlotSpec registry (`giant analyze list`)
│   │   ├── router_gating.py / type_embedding_distance.py   # checkpoint-bound diagnostics
│   │   ├── runtime_estimate.py  #   per-(plot, chunk) walltime estimates for submit
│   │   ├── condor.py       #   prep / compute-one / merge / submit-description plumbing
│   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
│   └── cli.py              # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze`
├── giant/tools/                        # 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
uv sync --extra cpu --extra analysis  # matplotlib/polars/plotstyle, for `giant analyze render`
uv sync --extra cpu --extra convert   # uproot/awkward/polars, for `dwarf convert`
uv sync --extra cpu --extra wandb     # W&B logging (`giant train --wandb`)

The dev extra pulls in convert, analysis, geometry and wandb as well.

cpu and cuda are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain uv sync installs no torch at all. See CLAUDE.md for details.

Training, prediction, rollout

giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir
giant train path/to/steps.parquet                          # train (flow stage 1 + wgan stage 2, default)
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt

dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # position → material/layer_id
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl

Useful flags on giant train:

• --mode {flow,ddpm,wgan} sets both stages' objective at once; --stage1-generator/--stage2-generator override per stage
• --stage2-decoder {autoregressive,one_shot} — Stage 2 decoding strategy (see Architecture)
• --conditioning {physical,embedding,onehot} — conditioning representation
• --router / --router-type / --n-experts / --router-axis — MoE routing
• --stage2-stage1-context {truth,sampled} — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via stage2_model.ctx_p_start/ctx_p_end)
• --precision {fp32,bf16} — bf16 autocast in the training loop
• --wandb — log per-epoch metrics to Weights & Biases (needs uv sync --extra wandb); metric names are <stage>/<split>/<metric> plus an unprefixed run-level tail, all derived from giant/training/trainers.py MetricSpecs
• --no-cache-setup / --rebuild-setup-cache — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); dwarf warm-cache precomputes it
• --stage1-init-from/--stage2-init-from (checkpoint .pt) + --stage1-freeze/--stage2-freeze — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint

Config-file-only knobs (no CLI flag — use --config config.toml): stage2_model.autoregressive.teacher_forcing/.history, stage2_model.particle_type.target/.class_weighting, stage2_model.n_sec.mode/.owner, conditioning.share_stages, stage*_model.trunk.* and the finer router knobs (lambda_balance, gumbel, learn_width, …). configs/ holds kept reference configs. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).

giant rollout seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up material/layer_id from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.

Validation and analysis

• giant.validate.validate_marginals — step-level marginal + KL-divergence checks during training (--validate-every)
• giant analyze — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):

giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot × chunk (compute only)
giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan  # N rollouts vs one shared reference
giant analyze render <run_dir> --gallery                   # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)

giant analyze list                                         # every catalog plot id
giant analyze prep rollout.yaml --chunks 8                 # just the run directory, no submission
giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0   # what a condor job runs
giant analyze merge-one --id marginal_edep --run-dir <run_dir>               # merge one plot's chunks (debugging)

<run_dir> defaults to <cwd>/analysis_runs/analysis_<id> (--run-dir overrides it; prep/submit print it). Multiple rollout YAMLs must all name the same reference (dataset) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only render needs LaTeX, so it always runs locally.

Separately, giant analyze metrics <train_run_dir> renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's metrics.csv.

Development

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