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)
dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
                                                  # bump-schema, status, update-manifest, create-manifest,
                                                  # make-root, 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.

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.

Output space (9D): step_length (log), ΔE (log), edep (log), 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). post_pos itself 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: PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID, number of secondaries (Phase 1 only — see Roadmap below).

Model (giant/model/): DenoisingMLP built from ResBlocks, with a SinusoidalEmbedding for the diffusion/flow time variable and a ConditionEncoder that fuses all conditioning inputs. 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. Shower-level (rollout) observables live in giant/analysis.py (compute_rollout_observables + plot_rollout_*), fed by giant rollout output.

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 (current): number of secondaries is a conditioning input — model predicts only 9D post-step kinematics (including derived post_pos).

Phase 2 (target): model must jointly predict the number of secondaries and all their properties (energy, direction, species), requiring an extended output space and likely a set-based or autoregressive generation scheme for the variable-length secondary list.