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 (15D continuous, COND_DIM): pre-step position, log(pre-energy), pre-step direction, layer ID (COND_DIM_BASE=8) — plus, since particle/material physical-property conditioning (model.conditioning, see below), 7 more columns: particle log(mass)/charge (PARTICLE_PHYS_DIM=2, giant/particles.py) and material Z_eff/A_eff/log(density)/log(X0)/log(λ_int) (MATERIAL_PHYS_DIM=5, giant/materials.py). n_sec and e_sec are not conditioning inputs (that was Phase 1 / the energy-conservation PoC); the model predicts them.

ConditionEncoder/SecondaryConditionEncoder (giant/model/network.py) support two mutually exclusive conditioning modes, selected per-checkpoint (model_config["conditioning"], defaulting to "embedding" for old checkpoints without the key, "physical" for new giant train runs — see --conditioning):

• "embedding" (original Phase 2 design): a learned nn.Embedding per PDG code / material name, indexed by a dataset-scoped dense vocab (pdg_map/mat_map). Memorizes the training menu.
• "physical" (default): the 7 physical-property columns above are each routed through a small MLP (particle_mlp/material_mlp) to the same emb_dim width the embedding tables would have produced — a drop-in replacement computable for any PDG code / material name, not just ones seen in training, which is what lets the surrogate generalize to a held-out material or species. giant/particles.py decodes nuclear/ion PDG codes (the 10LZZZAAAI scheme) via the scikit-HEP particle package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. giant/materials.py ships as an intentionally-unfilled stub (MaterialProperties(None, ...) per material) that raises loudly (MaterialPropertiesNotFilledError) rather than silently defaulting — a physicist must populate real values before "physical" mode can train.

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, log-mass, charge) = SEC_SLOT_DIM=6, 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. A secondary's mass/charge are regressed directly against a fixed physics-derived target (its ground-truth PDG code's giant.particles.particle_mass_charge) — not a learned/moving embedding target, so nothing needs detaching. No snapping at inference: the predicted (mass, charge) are used as-is as the secondary's physical identity, including for its own future conditioning if it goes on to take further steps in a rollout. A separate, reporting-only nearest-known-PDG lookup (giant.particles.nearest_known_pdg) is used purely to populate a nominal pdg label for output rows / "embedding"-mode fallback conditioning — it never feeds back into the model.

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

Physical-property conditioning (implemented): model.conditioning = "physical" | "embedding" (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. "embedding" stays available as the generalization-comparison baseline. Not yet done: giant/materials.py's table needs real physicist-supplied values before "physical" mode can train (currently unfilled, fails loudly if used); once filled, the actual held-out-material/species generalization comparison against the "embedding" baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.

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. See the knowledge base (/home/lars/knowledge-base/meta/roadmap.md).