giant/CLAUDE.md

CLAUDE.md

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

Commands

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

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