giant/CLAUDE.md

# CLAUDE.md

This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.

## Commands

```bash
uv sync --extra cpu                # install dependencies with CPU-only torch (standard/default)
uv sync --extra cuda               # install dependencies with CUDA 11.8 torch
uv sync --extra cpu --extra dev    # add dev extras (pytest, etc.)
uv sync --extra cpu --extra geometry  # add scikit-learn for the geometry oracle (giant rollout)
pytest               # run tests
giant train path/to/steps.parquet --mode flow   # train (flow matching)
giant train path/to/steps.parquet --mode ddpm   # train (DDPM baseline)
giant predict path/to/steps.parquet --checkpoint ckpt/best.pt   # per-step predictions
giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl  # full showers
dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
                                                  # bump-schema, status, update-manifest, create-manifest,
                                                  # make-root, build-geometry-oracle, hparam-scan
                                                  # (see scripts/dwarf.py)
```

`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.

### Lint and type checking

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

Part of the `dev` extra. Run these periodically (not just at commit time) to catch drift early.

## Architecture

GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — now including the variable-length list of secondary particles the step produces (Phase 2, see Roadmap).

**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`). Train/val split is by `event_id` to avoid leaking correlated steps from the same shower.

**Stage-1 output space (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):** `log_step_length`, two additive-log-ratio (ALR) coordinates `edep_logit`/`sec_logit` of a **deposit / secondary / post-energy simplex**, `post_dir` (post-scattering momentum direction, unit vector in the local frame), and `travel_dir` (direction of `post_pos - pre_pos`, unit vector in the local frame). The energy simplex decodes via softmax over `[edep_logit, sec_logit, 0]` × `pre_E` so `edep + e_sec + post_E == pre_E` holds by construction — energy conservation is architectural, not learned (see `energy_simplex_decode`). `post_pos` is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, since `step_length` already encodes that displacement's magnitude and duplicating it would let the two become inconsistent.

**Conditioning vector (8D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID — plus PDG code and material as embeddings. `n_sec` and `e_sec` are **no longer conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model now predicts them.

**Model** (`giant/model/network.py`): a two-stage model, both checkpointed together.
- **Stage 1 — `DenoisingMLP`:** `ResBlock` stack with a `SinusoidalEmbedding` for the flow/diffusion time variable and a `ConditionEncoder` fusing the conditioning. Predicts the 9D primary vector field, plus an `n_sec_head` classifier over `{0..K_MAX}` (`K_MAX=15`) that runs on the condition encoding alone (no diffusion noise), callable via `predict_n_sec`.
- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, continuous type embedding 16D)` = `SEC_SLOT_DIM=20`, 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 (they sum to it), so the whole chain conserves energy. The type embedding is trained against a detached PDG-embedding target (stops self-referential collapse) and snapped to the nearest PDG at inference (`snap_type_to_pdg_idx`).

`schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).

**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.

**Validation** (`giant/validate.py`): step-level marginal comparisons. `giant/analysis.py` is a fully-streaming (lazy polars) diagnostics module, sized for predict/rollout files larger than RAM, with no in-memory `SampleCollection` and no full-array materialization. It covers one-step-ahead `giant predict --coord local` output (`compute_event_observables_pl` + `plot_total_energy`/`plot_longitudinal_profile`/etc. for shower-level observables, plus the marginal/correlation/constraint tiers) and, via the `RolloutVsTruth` source type, a full autoregressive `giant rollout` shower compared against held-out truth data (`compute_rollout_vs_truth_observables_pl` for shower-level observables, reusing the same plot functions) — see the module docstring.

**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.

## Roadmap

**Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).

**Phase 2 (implemented — baseline):** the two-stage model above jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected). This is the "get a baseline out" track agreed with Jan & Tobias (2026-07-07).

**Next directions** (parallel, not yet built): faster-eval architectures measured against a ~10× native-Geant4 budget — a Wasserstein-GAN throwaway (single-pass eval) and a mixture-of-experts / routing tree of small nets selected per call (pdg / energy / process), with soft/differentiable gating on continuous routing axes; a sampling-calorimeter (multi-material) dataset; and preferring **material + particle physical properties** over learned embeddings for conditioning. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).