larsandClaude Opus 4.8
26a176aeaa
Closes the loop from single-step prediction into full showers: - giant/geometry.py + `dwarf build-geometry-oracle`: learn position -> (material, layer_id) from data (KNN/SVM) to supply the conditioning the surrogate does not predict; flag detector escape by NN distance. - giant/rollout.py: breadth-first batched frontier that steps all active tracks, spawns secondaries as new tracks, and terminates 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 exactly. - `giant rollout` CLI: seed from real events (argmax pre_E), load checkpoint, write a world-frame steps parquet + YAML sidecar. - giant/analysis.py: compute_rollout_observables + plot_rollout_* for single-sided longitudinal/transverse/total-energy shower profiles; analysis/export_rollout_observables.py driver. - scikit-learn added as an optional `geometry` extra (lazy-imported). - Tests: tests/test_geometry.py, tests/test_rollout.py. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
55 lines
4.4 KiB
Markdown
55 lines
4.4 KiB
Markdown
# 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)
|
|
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
|
|
|
|
```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.
|
|
|
|
**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 `ResBlock`s, 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.
|