giant/README.md

# giant

**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep 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

```bash
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](https://gitlab.etp.kit.edu/lbogner/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_id`s 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

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

```bash
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` `MetricSpec`s
- `--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):

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

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