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, ruff, ty, bump-my-version, git-cliff, + all runtime extras)
uv sync --extra cpu --extra geometry  # add scikit-learn for the geometry oracle (giant rollout)
pytest               # run tests
giant new-run --hidden-dim 512 --lr 3e-4        # scaffold a config.toml + run dir ahead of training
giant train path/to/steps.parquet               # train (defaults: stage 1 flow, stage 2 wgan + autoregressive)
giant train path/to/steps.parquet --mode flow   # set both stages' generative objective at once
giant train path/to/steps.parquet --stage1-generator flow --stage2-generator wgan   # per-stage override
giant train path/to/steps.parquet --router --router-type energy  # MoE routing trunk (see Roadmap for status)
giant model summary --config config.toml         # build-only: parameter counts + which config keys actually bite
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
giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs-reference analysis on HTCondor
giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep/submit)
giant analyze metrics <train_run_dir>                      # training-progress plots from metrics.csv
dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
                                                  # bump-schema, status, update-manifest, create-manifest,
                                                  # make-root, build-geometry-oracle, warm-cache, hparam-scan
                                                  # (see giant/tools/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.

configs/ holds kept reference configs (baseline.toml, default.toml, the router/WGAN scan configs) — pass them with --config.

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.

Release tooling

Merges to master auto-bump the patch version, tag, and update CHANGELOG.md via the Gitea workflow in .gitea/workflows/ci.yml (bump-my-version + git-cliff). Don't hand-edit the version in pyproject.toml or write changelog entries by hand.

Compute environment

Work on this repo happens across three kinds of machine:

• Local dev machines (laptop + desktop, identical): repo at ~/Programming/giant, no access to /ceph — datasets, training results, and models aren't reachable here.
• Portal machines (portal1, deepthought, deepthought2, bms1, bms2, bms3): repo lives under /work, and /ceph holds ROOT/parquet files and trained models. These are shared with other users — stay strictly within /work/lbogner and /ceph/lbogner, and keep resource usage to roughly a quarter of CPU/RAM and a single GPU so as not to disturb other users' jobs.
• HTCondor worker nodes: never run or SSH onto these directly — the only sanctioned path is submitting jobs through condor (giant analyze submit, and the in-progress remote-GPU train/rollout submission on condor-gpu-train-rollout). /ceph is available there; /work is only sometimes mounted, depending on the node.

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 — including the variable-length list of secondary particles the step produces.

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, streaming variant included). Train/val split is by event_id (--seed-controlled) to avoid leaking correlated steps from the same shower. Loading a directory or .manifest of several parquet files offsets each file's event_ids by a per-file stride so ids stay globally unique. setup_cache.py persists the pre-epoch setup scan (vocab maps, event split, process maps, normalizer stats) as a sidecar so repeated runs over the same data path don't rescan (--cache-setup/--rebuild-setup-cache, precomputable with dwarf warm-cache --config ...).

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 7 physical-property 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 — the model predicts them. giant/cond_layout.py is the single source of truth for the cond_cont/cond_cat column layout shared by giant.data.transforms, giant.model.encoders, and giant.model.routers.

ConditionEncoder (giant/model/encoders.py) configures the particle and material identity axes independently (conditioning.particle / conditioning.material, each a ConditioningAxisConfig with type/emb_dim/n_layers), so they may mix freely. Three per-axis modes:

• "physical" (default): the axis's raw physical properties routed through a small MLP — computable for any PDG code / material name, which is what lets the surrogate generalize beyond the training menu. 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 real Geant4-11.4.1-derived values for every material the detector geometry actually produces; the sole exception is G4_LYSO (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays MaterialProperties(None, ...) and raises loudly (MaterialPropertiesNotFilledError) rather than silently defaulting.
• "embedding": a learned nn.Embedding per PDG code / material name, indexed by a dataset-scoped dense vocab. Memorizes the training menu; the generalization-comparison baseline, and the only mode compatible with stage2_model.particle_type.target = "embedding".
• "onehot": a fixed, unlearned vector over the top emb_dim - 1 codes by training-set count plus one "other" bin. Not a reparameterization of "embedding" — the vocabulary cap is the real difference.

conditioning.share_stages decides whether the two stages get one shared encoder instance or two identically-configured independent ones.

Model (giant/model/, both stages checkpointed together). network.py is only a re-export shim now; the real code is split by concern:

• layers.py — ResBlock/AdaLNResBlock + BLOCK_REGISTRY (conditioning-injection mechanism is selectable), SinusoidalEmbedding, ContextAdapter, build_mlp_head.
• encoders.py — ConditionEncoder (above).
• trunks.py — TRUNK_REGISTRY/build_trunk: everything downstream of the fused conditioning vector, as a registrable expert body (resmlp default, plus a none variant). RoutedTrunk builds router.n_experts instances of whichever body is named, so mixing is orthogonal to which body is mixed.
• routers.py — Router base + ROUTER_REGISTRY: energy/pdg/process/composed/none. Soft-mixed at train time, top-1 dispatched at eval time (each row runs exactly one small expert), which is the actual inference-speed win. EnergyRouter/PdgRouter gate on a quantity known at inference; ProcessRouter runs its own small classifier (process isn't known upfront); ComposedRouter gates jointly over outer-product expert cells via repeated --router-axis "type:key=val,...". The --router*/--n-experts CLI flags target stage1_model.router only; stage 2's router is config-file-only (stage2_model.router). EnergyRouter accepts centers_init, which giant/pipeline.py auto-populates from real data quantiles via a reservoir sample collected during the normalizer-fitting pass.
• history.py — HISTORY_REGISTRY/build_history: markov (previous token only), attention (causal self-attention, KV-cached at inference via init_cache/step), none. Stage-2 autoregressive only.
• objectives.py — Objective base + registry for flow/ddpm/wgan: answers in one place whether a stage needs a time embedding, is adversarial, folds the secondary type slice into its trunk output, what its trunk input is, and which loss it trains against.
• models.py — the composed stage models: Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel, all on a shared StageModel base.
• builders.py — build_models/build_critics, assembling the above from a config dict.
• schedule.py (CosineSchedule for DDPM + conditional-flow-matching losses), wgan.py (gradient penalty / critic / generator losses, Gulrajani et al. 2017), summary.py (giant model summary), _legacy.py (v0.2 checkpoint migration).

Stage 1 — primary step. Trunk (routed or not) over the fused conditioning, plus a SinusoidalEmbedding of the flow/diffusion time for non-adversarial objectives, predicting the 9D vector field. An n_sec classifier head over {0..k_max} runs on the condition encoding alone; stage2_model.n_sec.owner decides whether it lives on stage 1 (v0.2 checkpoints) or stage 2 (default).

Stage 2 — secondaries. Conditioned on the pre-step state plus a projected stage-1 outcome (stage2_model.context_dim; stage1_context selects ground-truth vs sampled context, annealable via ctx_p_start/ctx_p_end). Two decoders (stage2_model.decoder):

• autoregressive (default): one secondary at a time in descending-energy order, each token conditioned on a HistoryEncoder summary of prior tokens, with teacher forcing (always/scheduled/never, tf_p_start/tf_p_end). n_sec.mode = "stop_token" lets the length be emitted by the sequence itself instead of the classifier head.
• one_shot: all k_max slots in one pass, masked past the predicted n_sec (the v0.2 behaviour).

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. Particle identity is set by stage2_model.particle_type.target: "onehot" (default — categorical over the top n_classes - 1 PDG codes by training count plus "other", with configurable other_policy and class_weighting), "physical" (continuous (log-mass, charge) regressed against giant.particles.particle_mass_charge), or "embedding" (nearest-row snap into the conditioning embedding table; requires conditioning.particle.type = "embedding").

Samplers (giant/sample.py): DDPM, DDIM, flow matching (ODE integration, ~10 steps), and single-pass WGAN, plus the stage-2 secondary sampling loop (one-shot and autoregressive).

Training (giant/training/): loop.py (epoch loop, graceful shutdown, best-checkpoint selection), trainers.py (StageSpec + per-stage flow/ddpm and WGAN-GP trainers, and the MetricSpec declarations that define metrics.csv's columns), stage2_inputs.py (ground-truth stage-2 targets + teacher-forcing inputs), metrics.py (MetricsCollector: metrics.csv, W&B logging, progress/summary), checkpoint.py, amp.py (train.precision = fp32|bf16 autocast), plots.py (giant analyze metrics). Per-stage init_from/freeze lets one stage be retrained against a fixed, known-good other stage while still producing a complete rollout-capable checkpoint.

Config (giant/config.py): frozen dataclasses are the single source of truth; DEFAULT_CONFIG is generated from GiantConfig().to_dict() rather than hand-maintained. Blocks: [conditioning], [stage1_model], [stage2_model], [train], [meta]. Unknown keys are rejected on merge (with a did-you-mean suggestion), and tests/test_config_consumed_keys.py audits that every key is actually read somewhere.

Validation (giant/validate.py): step-level marginal + KL-divergence comparisons during training (--validate-every).

Analysis (giant/analysis/, giant analyze CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive giant rollout runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (pre_*/post_*/edep/step_length/pdg/material/event_id), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: sources.py (canonical LazyFrames + RolloutSpec/Side — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is generation>0 & step_no==0 rollout tracks vs exploded sec_*_list reference columns), variables.py (the per-step value expressions shared by range sizing and the plot registry), reduce.py (the streaming primitives — a single hist1d group_by([group,bin]).len() pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), grouping.py/context.py (fixed bin edges + energy-quantile/pdg/material group sets resolved once by prep into shared.json over the union of the reference and every rollout, so every compute job is one pass with no range scan), reduced.py (Partial/Reduced — the compact self-describing JSON a compute job emits), catalog.py (the declarative PlotSpec registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles/containment, species/leakage, secondaries, distance summaries, router and type-embedding diagnostics; giant analyze list prints every id), runtime_estimate.py (per-(plot, chunk) walltime estimates for the submit description), and render.py (the only module importing ETPlot's plotstyle/LaTeX; dispatches on Reduced.kind, writes PDFs + metadata.yaml; each rollout gets a stable ps.get_color(i) slot by its position in series, the reference always draws in one fixed dashed-ink style). Bundle.rollouts is a name-keyed dict of Side, and every compute_partial/finalize builds a Reduced.payload["series"] dict keyed the same way, with payload["reference"] as the one distinguished non-rollout entry. The heatmap-shaped specs (marginal_distance_summary, sec_count_per_step_by_species — the latter also drawing the reference as its own panel) and the checkpoint-bound diagnostics (router_gating.py, type_embedding_distance.py) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout.

Input is one or more giant rollout YAML sidecars (condor.py:load_rollout_yamls): each YAML's output/dataset keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same dataset, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated --label CLI flag, else the YAML stem (N>1), else "rollout" (a single YAML). prep creates a run directory (<cwd>/analysis_runs/analysis_<id>/ by default, --run-dir to override) holding shared.json, run_meta.json (RunMeta.rollouts: list[{name,path,plot_meta}], insertion order = CLI order = every plot's series order), reduced_partial/, reduced/, plots/. Compute/merge/render split: giant analyze submit a.yaml [b.yaml ...] --chunks N runs prep (recording N in run_meta.json) then submits one HTCondor job per (plot, chunk) pair (compute-one --id --chunk --run-dir, polars/numpy only — no LaTeX on workers), each streaming over an event_id-disjoint slice (event_id % N == chunk) of the reference and every rollout and writing a small reduced_partial/<id>__<chunk>.json; every PlotSpec splits into a compute_partial/finalize pair so chunks can be summed/concatenated back per rollout (chunkable=False specs — the checkpoint-bound diagnostics, already bounded/subsampled — always run as a single chunk). The local giant analyze render <run_dir> first joins every plot's chunk partials into reduced/<id>.json (merge_all, a no-op join when N=1; merge-one does a single plot for debugging), then turns those into the styled PDF/gallery tree. giant analyze metrics <train_run_dir> is a separate, unrelated entry point: training-progress plots straight from a run's metrics.csv.

Shower rollout (giant/rollout.py, giant rollout CLI): autoregressively steps the two-stage model into a full shower, advancing tracks breadth-first (every sweep steps all active tracks once, in batch_size chunks, so many tracks share each forward pass). 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 one of the TERM_* reasons in constants.py (energy cutoff, max steps, escape, natural end, unknown pdg, max tracks); energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. giant/checkpoint_io.py is the shared checkpoint → ready-to-run-models path used by both predict and rollout.

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 (done): the two-stage model 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).

Physical-property conditioning (implemented, default): conditioning.particle.type / conditioning.material.type = physical | embedding | onehot. Not yet done: the actual held-out-material/species generalization comparison against the "embedding" baseline is unrun — the 34GB multi-material dataset (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.

v0.3.0 — Stage-2 autoregressive redesign (implemented, released; on master since 2026-08-13): motivated by the 2026-08-03 WGAN rollout benchmark, which failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated -14 muon antineutrinos). Stage 2 became autoregressive in descending-energy order with teacher forcing, and the particle-type representation went back to categorical (particle_type.target = "onehot"), reversing the 2026-07-17 continuous (log-mass, charge) target. The config break ([conditioning]/[stage1_model]/[stage2_model]/[train] replacing the flat train.mode + [model]) makes per-stage generators, stage-2-only training, and one-shot-vs-autoregressive comparison all expressible, and the network.py refactor into composable parts (encoder × trunk × objective) also makes routed WGAN work for the first time.

v0.2 configs and checkpoints are auto-migrated (config.migrate_config, model._legacy._migrate_legacy_model_config, both drawing on shared facts in giant/_migration.py). v0.2 checkpoint-loading support has no expiry decided yet: /ceph still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or tests/legacy/network_v02_snapshot.py (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.

Faster-eval architectures — both implemented, neither validated. Target is a ~10× native-Geant4 eval budget; no eval-latency number exists for any configuration yet, so that budget is unverified across the board.

• WGAN-GP (--stage2-generator wgan, now the stage-2 default): first rollout benchmark 2026-08-03 failed with secondary-species mode collapse — the failure v0.3.0 was designed to address. No post-v0.3.0 benchmark has been run. Writeup: /home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md.
• MoE routing trunk (--router): first rollout benchmark 2026-07-22 diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. Cause identified as a bad config, not a bad idea: lambda_balance=0.0 (no load-balancing loss) plus EnergyRouter's torch.linspace(-2, 2, n_experts) center init assuming a roughly uniform z-normalized energy distribution — so the ten experts overlapped across ~5 decades of energy instead of partitioning it, and eval-time top-1 dispatch chose among near-ties rather than real specialists. Both prerequisites are fixed in code (quantile-seeded centers_init from pipeline.py, lambda_balance exposed). Next step: retrain with lambda_balance > 0 and quantile-seeded centers (consider a from-scratch run rather than a short fine-tune), then re-check whether router_gating sharpens up. Writeup: /home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md.

A sampling-calorimeter (multi-material) dataset track is still open and unblocked, not yet started. See the knowledge base (/home/lars/knowledge-base/meta/roadmap.md).

Condor-submitted GPU training/rollout (in progress, condor-gpu-train-rollout branch, not yet merged): moves giant train/giant rollout off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2. Partway between "needs major features" and feature-complete — not ready to merge yet.