lars
72bd65ff9f
step_length already encodes |post_pos - pre_pos| by definition, so a raw post_pos target would duplicate that magnitude and could drift inconsistent with step_length during sampling. Instead add travel_dir, a unit vector (local frame) giving only the direction of pre_pos->post_pos; post_pos is reconstructed at inference as pre_pos + step_length * travel_dir, keeping the two self-consistent. Target grows from 6D to 9D. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2.6 KiB
CLAUDE.md
This file provides guidance to Claude Code (claude.ai/code) when working with code in this repository.
Commands
uv sync # install dependencies (including dev extras: uv sync --extra dev)
pytest # run tests
python scripts/train.py --data path/to/steps.parquet --mode flow # train (flow matching)
python scripts/train.py --data path/to/steps.parquet --mode ddpm # train (DDPM baseline)
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 validation is planned.
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.