Bump version to 0.2.0

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

CLAUDE.md

@@ -10,6 +10,7 @@ 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 new-run --hidden-dim 512 --lr 3e-4        # scaffold a config.toml + run dir ahead of training
 giant train path/to/steps.parquet --mode flow   # train (flow matching)
 giant train path/to/steps.parquet --mode ddpm   # train (DDPM baseline)
 giant train path/to/steps.parquet --mode wgan   # train (WGAN-GP, single-pass eval; implemented, not yet tested)
@@ -20,7 +21,7 @@ giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs
 giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep/submit)
 dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
                                                   # bump-schema, status, update-manifest, create-manifest,
-                                                  # make-root, build-geometry-oracle, hparam-scan
+                                                  # make-root, build-geometry-oracle, warm-cache, hparam-scan
                                                   # (see scripts/dwarf.py)
 ```
 

README.md

@@ -2,15 +2,19 @@
 
 **G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
 
-Proof-of-concept surrogate model for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained conditional generative model.
+Conditional generative surrogate for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the variable-length list of secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained generative model. A trained checkpoint autoregressively rolls out full showers, stepping each primary and pushing secondaries as new tracks.
 
 Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
 
 ## Architecture
 
-A **two-stage conditional flow matching** model (Lipman et al. 2022): a small MLP learns a vector field mapping noise → step outcomes in ~10 ODE steps per sample. Falls back to DDPM for comparison.
+A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`):
 
-**Stage 1 — primary (9D, diffused):**
+- **`flow`** (default) — conditional flow matching (Lipman et al. 2022): an MLP learns a vector field mapping noise → step outcomes, sampled via ODE integration in ~10 steps.
+- **`ddpm`** — a standard denoising diffusion baseline for comparison (`giant/model/schedule.py:CosineSchedule`).
+- **`wgan`** — a single-pass Wasserstein-GAN-GP generator/critic (`giant/model/wgan.py`), trading iterative sampling for one forward pass; implemented, not yet validated against the flow-matching baseline.
+
+**Stage 1 — primary (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):**
 
 | Index | Variable | Encoding |
 |-------|----------|----------|
@@ -19,13 +23,20 @@ A **two-stage conditional flow matching** model (Lipman et al. 2022): a small ML
 | 3–5 | `post_dir` in local frame | unit vector |
 | 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
 
-The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss. Stage 1 also has a classifier head predicting the number of secondaries `n_sec ∈ {0..15}` from the conditioning alone.
+The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss (`energy_simplex_decode`). Stage 1 also has a classifier head (`predict_n_sec`) predicting the number of secondaries `n_sec ∈ {0..K_MAX}` (`K_MAX = 15`) from the conditioning alone, no diffusion noise involved.
 
 Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
 
-**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second flow net generates all `K_MAX = 15` secondary slots at once. Each slot carries a stick-breaking energy fraction, a local-frame direction, and a continuous particle-type embedding (snapped to the nearest PDG at inference), ordered by descending energy; slots beyond the predicted `n_sec` are masked. The secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the full chain conserves energy. Each secondary's momentum is reconstructed afterward from `(energy, direction, species)` rather than predicted.
+**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second net generates all `K_MAX` secondary slots at once — `(stick-breaking energy logit, local-frame direction, log-mass, charge)` per slot, 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, so the whole chain conserves energy. A secondary's mass/charge are regressed directly against its ground-truth PDG code's physical values (`giant.particles.particle_mass_charge`) and used as-is at inference — including for its own conditioning if it takes further steps in a rollout. No snapping to a known PDG code happens in the model path; `giant.particles.nearest_known_pdg` is a reporting-only lookup used to populate a nominal `pdg` label on output rows.
 
-**Conditioning:** PDG code (embedding), pre-step position, log(pre-energy), pre-step direction, material (embedding), layer ID. (`n_sec` / `e_sec` are outputs now, not inputs.)
+**Conditioning (`--conditioning`, per-checkpoint):** pre-step position, log(pre-energy), pre-step direction, layer ID, plus particle/material physical properties — mass/charge (`giant/particles.py`) and Z_eff/A_eff/density/X0/λ_int (`giant/materials.py`). Two mutually exclusive modes:
+
+- **`physical`** (default) — the physical-property columns are routed through small MLPs, computable for any PDG code / material, letting the surrogate generalize to species/materials outside the training menu.
+- **`embedding`** — the original design: a learned `nn.Embedding` per PDG code / material, kept as a generalization-comparison baseline (memorizes the training menu).
+
+`n_sec` and `e_sec` are model outputs, not conditioning inputs — a rollout is self-contained and never injects ground truth.
+
+**Mixture-of-experts routing (`--router`, opt-in):** `giant/model/network.py` also implements a pluggable `Router` contract (`ROUTER_REGISTRY`: `energy`, `pdg`, `process`, plus a `composed` router combining several axes) that splits `DenoisingMLP`/`SecondaryDecoder` into per-expert trunks, soft-gated in training and top-1 dispatched at eval. Implemented; first rollout benchmark needs a retrain with a load-balancing loss and better-seeded router centers (see Roadmap). See `--router-type`/`--n-experts`/`--router-axis` on `giant train`/`giant new-run`.
 
 ## Roadmap
 
@@ -33,11 +44,15 @@ Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pr
 
 **Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
 
-**Next directions:** faster-eval architectures against a ~10× native-Geant4 budget (Wasserstein-GAN, mixture-of-experts routing tree), a multi-material sampling-calorimeter dataset, and physical-property conditioning over learned embeddings.
+**Physical-property conditioning (implemented):** replaces learned PDG/material embeddings with physical-property MLPs (see above); Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. Not yet done: the held-out-material/species generalization comparison against the `embedding` baseline — the natural dataset for that is the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes).
+
+**Faster-eval architectures (implemented, validation in progress):** both target a ~10× native-Geant4 eval budget. WGAN-GP (`--mode wgan`) has no rollout-vs-reference analysis run against it yet. The MoE router (`--router`) had its first rollout benchmark diverge from Geant4 despite matching bulk deposited energy — the experts weren't specializing (near-uniform gating), traced to a missing load-balance loss and a center-init that didn't match the real energy distribution; both are now fixable via `lambda_balance > 0` and quantile-seeded router centers, but a re-run to confirm hasn't happened yet.
+
+A multi-material sampling-calorimeter dataset is a planned future direction, not yet built.
 
 ## Data
 
-Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `uv run dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle) to avoid leaking correlated steps from the same shower.
+Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle, and `--seed`-controlled) to avoid leaking correlated steps from the same shower.
 
 ## Project structure
 
@@ -49,21 +64,32 @@ giant/
 │   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
 │   │   └── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
 │   ├── model/
-│   │   ├── network.py      # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP, SecondaryDecoder
-│   │   └── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
+│   │   ├── network.py      # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic
+│   │   ├── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
+│   │   └── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
 │   ├── constants.py        # output/conditioning dims, K_MAX, secondary slot layout, schema keys
+│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; nearest-known-PDG lookup
+│   ├── materials.py        # material name → (Z_eff, A_eff, density, X0, λ_int)
 │   ├── config.py           # default hyperparameters, TOML config merging, device autodetect
-│   ├── pipeline.py         # builds datasets/normalizers and kicks off a training run
-│   ├── train.py            # two-stage training loop, checkpointing, graceful shutdown
-│   ├── sample.py           # DDPM / DDIM / flow matching samplers + secondary sampling
+│   ├── pipeline.py         # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
+│   ├── train.py            # two-stage training loop, checkpointing, graceful shutdown, W&B logging
+│   ├── sample.py           # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
 │   ├── geometry.py         # GeometryOracle: position → (material, layer_id, escaped) for rollout
 │   ├── rollout.py          # autoregressive shower rollout driver
 │   ├── validate.py         # step-level marginal + KL-divergence validation
-│   ├── analysis.py         # step- and shower-level diagnostics: marginals, correlations, rollout observables
-│   └── cli.py              # `giant train` / `predict` / `rollout` Typer app
-├── scripts/                            # dataset/tooling logic, unified under the `dwarf` CLI (`uv run dwarf --help`)
+│   ├── analysis/           # rollout-vs-reference analysis pipeline (see `giant analyze` below)
+│   │   ├── sources.py      #   canonical LazyFrames + secondary view
+│   │   ├── reduce.py       #   streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
+│   │   ├── grouping.py     #   fixed bin edges + energy/pdg/material group sets
+│   │   ├── context.py      #   resolves grouping into `shared.json` once per run
+│   │   ├── catalog.py      #   declarative PlotSpec registry
+│   │   ├── condor.py       #   prep / compute-one / submit-description plumbing
+│   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
+│   └── cli.py              # `giant train` / `new-run` / `predict` / `rollout` / `analyze` Typer app
+├── scripts/                            # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
 │   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
-│   │                                    #   update-manifest, create-manifest, make-root, hparam-scan
+│   │                                    #   update-manifest, create-manifest, make-root,
+│   │                                    #   build-geometry-oracle, warm-cache, hparam-scan
 │   ├── steps_to_parquet.py             # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
 │   ├── steps_to_parquet_parallel.py    # fan out conversion over several ROOT files — `dwarf convert --jobs N`
 │   ├── migrate_geant_steps.py          # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
@@ -71,7 +97,9 @@ giant/
 │   │                                    #   `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
 │   ├── create_root_files.py            # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
 │   ├── geometry_oracle.py              # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
-│   └── hparam_scan.py                  # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
+│   ├── warm_setup_cache.py             # precompute `giant train`'s setup-stage sidecar — `dwarf warm-cache`
+│   ├── hparam_scan.py                  # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
+│   └── profile_analysis_costs.py       # profiling helper for the `giant analyze` reduction pipeline
 └── tests/
 ```
 
@@ -80,27 +108,37 @@ giant/
 ```bash
 uv sync --extra cpu               # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
 uv sync --extra cpu --extra dev   # add dev tools (pytest, ruff, ty)
+uv sync --extra cpu --extra geometry  # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
 ```
 
-`cpu` and `cuda` are mutually exclusive extras selecting the torch build; plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
+`cpu` and `cuda` are mutually exclusive extras selecting the torch build (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
 
 ## Training, prediction, and rollout
 
 ```bash
-giant train path/to/steps.parquet --mode flow
+giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir for a new run
+giant train path/to/steps.parquet --mode flow               # train (flow matching; also --mode ddpm / wgan)
 giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
 
 # Full-shower rollout needs a geometry oracle (position → material/layer_id):
-uv sync --extra cpu --extra geometry
 dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
 giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
 ```
 
-`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
+`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant train --wandb` logs per-epoch metrics (the same ones written to `metrics.csv`) to Weights & Biases; requires `uv sync --extra wandb`. A repeat `giant train` against the same dataset (e.g. a hyperparameter sweep) reuses a cached setup-stage sidecar (vocab maps, event split, normalizer stats) unless `--no-cache-setup`/`--rebuild-setup-cache`; `dwarf warm-cache` precomputes it ahead of time. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
 
-## Validation
+## Validation and analysis
 
-`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`). For deeper diagnostics on a trained checkpoint — stratified marginals, correlation structure, physical-constraint violations, and shower-level rollout observables (longitudinal/transverse profiles, PDG energy shares) — see `giant.analysis`.
+`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`).
+
+For deeper rollout-vs-reference diagnostics — marginals stratified by energy/pdg/material, per-event totals, shower profiles, species share, leakage, and secondaries — `giant analyze` runs a streaming compute/render pipeline against a `giant rollout` YAML sidecar:
+
+```bash
+giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot (compute only)
+giant analyze render <run_dir> --gallery                   # local: styled PDFs + HTML gallery (needs LaTeX)
+```
+
+`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` imports plotstyle/LaTeX, so it always runs locally.
 
 ## Development
 

configs/router_energy_n10_embedding.toml

@@ -0,0 +1,23 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning      = "embedding"
+dropout           = 0.0
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true

configs/router_energy_n10_gumbel_centers_physical.toml

@@ -0,0 +1,22 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true
+gumbel            = true

configs/router_energy_n10_gumbel_centers_temp_physical.toml

@@ -0,0 +1,23 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true
+learn_temperature = true
+gumbel            = true

configs/router_energy_n10_gumbel_physical.toml

@@ -0,0 +1,22 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning = "physical"
+dropout      = 0.0
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = false
+gumbel            = true

configs/router_energy_n10_physical.toml

@@ -0,0 +1,23 @@
+[train]
+mode          = "flow"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+conditioning      = "physical"
+dropout           = 0.0
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+
+[model.router]
+enabled           = true
+type              = "energy"
+n_experts         = 10
+expert_hidden_dim = 128
+expert_n_blocks   = 4
+temperature       = 0.05
+lambda_balance    = 0.035
+learn_centers     = true

configs/wgan_h128_b4_physical.toml

@@ -0,0 +1,13 @@
+[train]
+mode          = "wgan"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+hidden_dim   = 128
+n_blocks     = 4
+dropout      = 0.0
+conditioning = "physical"

configs/wgan_h512_b6_physical.toml

@@ -0,0 +1,13 @@
+[train]
+mode          = "wgan"
+epochs        = 30
+lr            = 3e-4
+warmup_epochs = 3
+val_fraction  = 0.1
+num_workers   = 4
+
+[model]
+hidden_dim   = 512
+n_blocks     = 6
+dropout      = 0.0
+conditioning = "physical"

giant/analysis/condor.py

@@ -42,6 +42,8 @@ HTCondor file transfer of the multi-GB inputs.
 from __future__ import annotations
 
 import json
+import shutil
+import sys
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -194,11 +196,23 @@ def prep(
     resolved once here rather than re-passed (and risking disagreement) at every
     later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
     resolve the actual directory.
+
+    Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
+    this same ``run_dir``: partial files carry no record of what context
+    (``n_chunks``, bin edges, group sets) they were computed under, so
+    re-prepping with a different ``n_chunks``/``**ctx_kwargs`` (or after the
+    rollout/reference files changed) would otherwise let ``merge_one`` silently
+    merge stale partials against the new ``shared.json``.
     """
     y = load_rollout_yaml(rollout_yaml)
     run_path = derive_run_dir(y, run_dir, default_base=default_base)
     run_path.mkdir(parents=True, exist_ok=True)
 
+    for stale in ("reduced_partial", "reduced"):
+        stale_dir = run_path / stale
+        if stale_dir.exists():
+            shutil.rmtree(stale_dir)
+
     rollout, reference = y["output"], y["dataset"]
     ctx = build_context(rollout, reference, **ctx_kwargs)
     ctx.save(run_path / "shared.json")
@@ -343,7 +357,7 @@ class SubmitConfig:
 _WRAPPER = """#!/bin/bash
 set -euo pipefail
 cd {repo_dir}
-exec {repo_dir}/.venv/bin/giant analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
+exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
 """
 
 
@@ -392,6 +406,29 @@ def _job_walltimes(
     return jobs
 
 
+def _resolve_giant_executable(repo_dir: Path) -> Path:
+    """Path to the ``giant`` entry point to bake into the condor wrapper script.
+
+    Prefers the venv currently running this process (``sys.executable``'s
+    sibling ``giant``) so a submit from a non-default venv (e.g. ``--extra
+    cuda`` on a dev box) doesn't silently pick up a different one; falls back
+    to ``repo_dir/.venv/bin/giant`` for the case this is invoked from outside
+    any venv (e.g. a system Python).
+    """
+    active = Path(sys.executable).parent / "giant"
+    if active.exists():
+        return active
+    venv_giant = repo_dir / ".venv" / "bin" / "giant"
+    if not venv_giant.exists():
+        raise FileNotFoundError(
+            f"no `giant` executable found next to {sys.executable} or at "
+            f"{venv_giant} — condor jobs run it directly (no `uv` on the "
+            f"worker image), so run `uv sync --extra cpu` in {repo_dir} "
+            "before submitting."
+        )
+    return venv_giant
+
+
 def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     """Write the wrapper script, (plot, chunk) job list, and HTCondor submit
     description.
@@ -403,23 +440,33 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     ``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``).
     Returns the submit description path (``<run_dir>/analyze.sub``). Does not
     submit — call ``condor_submit`` on the returned file.
+
+    ``cfg.n_chunks`` and the run directory's own ``RunMeta.n_chunks`` (fixed by
+    ``prep``, and what ``RunMeta.rows_per_chunk`` was sized against) are two
+    independent values — checked equal up front so a mismatch is a clear error
+    here rather than an ``IndexError`` out of ``_job_walltimes``.
     """
-    venv_giant = cfg.repo_dir / ".venv" / "bin" / "giant"
-    if not venv_giant.exists():
-        raise FileNotFoundError(
-            f"{venv_giant} not found — condor jobs run it directly (no `uv` on "
-            f"the worker image), so run `uv sync --extra cpu` in {cfg.repo_dir} "
-            "before submitting."
-        )
+    giant_exe = _resolve_giant_executable(cfg.repo_dir)
 
     ids = ids or catalog_ids()
     run_dir = cfg.run_dir
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    if cfg.n_chunks != meta.n_chunks:
+        raise ValueError(
+            f"SubmitConfig.n_chunks={cfg.n_chunks} does not match the "
+            f"n_chunks this run directory was prepped with "
+            f"(RunMeta.n_chunks={meta.n_chunks} in {run_dir}/run_meta.json) — "
+            "re-run `prep` with the desired n_chunks, or fix cfg.n_chunks to "
+            "match it."
+        )
     (run_dir / "logs").mkdir(parents=True, exist_ok=True)
     (run_dir / "reduced").mkdir(parents=True, exist_ok=True)
     (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)
 
     wrapper = run_dir / "run_compute.sh"
-    wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, run_dir=run_dir))
+    wrapper.write_text(
+        _WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir)
+    )
     wrapper.chmod(0o755)
 
     jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)

giant/cli.py

@@ -1,5 +1,5 @@
 from collections import Counter
-from datetime import date, datetime, timezone
+from datetime import datetime, timezone
 from enum import Enum
 import math
 from pathlib import Path
@@ -26,6 +26,7 @@ from giant.constants import (
     ROLLOUT_COORD_VALUE,
 )
 from giant.data.loader import (
+    event_id_offset,
     find_parquet_files,
     iter_file_chunks,
     iter_cond_chunks,
@@ -87,8 +88,9 @@ def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int
     """
     router_cfg = model_cfg.get("router")
     if router_cfg and router_cfg.get("enabled"):
-        hidden_dim = model_cfg.get("expert_hidden_dim", 128)
-        n_blocks = model_cfg.get("expert_n_blocks", 3)
+        hidden_dim, n_blocks = gconfig.resolve_expert_dims(
+            router_cfg, model_cfg["hidden_dim"], model_cfg["n_blocks"]
+        )
         if training:
             n_blocks *= _router_total_experts(router_cfg)
         return hidden_dim, n_blocks
@@ -132,6 +134,30 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
     return out
 
 
+def _router_cli_overrides(
+    router: bool | None,
+    router_type: str | None,
+    n_experts: int | None,
+    router_axis: list[str] | None,
+) -> dict[str, object]:
+    """Build the `model.router` override dict from `--router`/`--router-type`/
+    `--n-experts`/`--router-axis` flags (empty if none were given). Shared by
+    `train` and `new-run` so both resolve router overrides identically.
+    """
+    cli_router: dict[str, object] = {
+        k: v
+        for k, v in {
+            "enabled": router,
+            "type": router_type,
+            "n_experts": n_experts,
+        }.items()
+        if v is not None
+    }
+    if router_axis:
+        cli_router.update(_parse_router_axis_flags(router_axis))
+    return cli_router
+
+
 _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")
 
 
@@ -185,9 +211,10 @@ class Mode(str, Enum):
     wgan = "wgan"
 
 
-class Conditioning(str, Enum):
-    physical = "physical"
-    embedding = "embedding"
+# Conditioning itself lives in giant.config (imported below as gconfig) —
+# shared with scripts/dwarf.py's Typer commands so the two CLIs can't
+# silently drift apart on the option's valid values.
+Conditioning = gconfig.Conditioning
 
 
 class Coord(str, Enum):
@@ -384,10 +411,33 @@ def train(
             "--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"
         ),
     ] = 65536,
+    cache_setup: Annotated[
+        bool,
+        typer.Option(
+            "--cache-setup/--no-cache-setup",
+            help="Cache the training setup stage's expensive per-file "
+            "precomputation (vocab maps, event split index, normalizer stats) "
+            "in a JSON sidecar next to the data, so a repeat `giant train` "
+            "against the same dataset (e.g. a hyperparameter sweep) can skip "
+            "re-deriving it",
+        ),
+    ] = True,
+    rebuild_setup_cache: Annotated[
+        bool,
+        typer.Option(
+            "--rebuild-setup-cache/--no-rebuild-setup-cache",
+            help="Ignore any existing setup cache sidecar and recompute every "
+            "section fresh for this run (still writes the refreshed sections "
+            "back to the sidecar for later runs; no effect if --no-cache-setup)",
+        ),
+    ] = False,
     out: Annotated[
         Optional[Path],
         typer.Option(
-            "--out", "-o", help="Checkpoint dir (default: auto from hyperparams)"
+            "--out",
+            "-o",
+            help="Checkpoint dir (default: timestamped dir from hyperparams, "
+            "or the --resume checkpoint's own dir when resuming)",
         ),
     ] = None,
     device: Annotated[
@@ -399,6 +449,30 @@ def train(
         Optional[Path],
         typer.Option("--resume", "-r", help="Checkpoint .pt to resume training from"),
     ] = None,
+    wandb: Annotated[
+        Optional[bool],
+        typer.Option(
+            "--wandb/--no-wandb",
+            help="Log per-epoch training metrics to Weights & Biases "
+            "(requires `uv sync --extra wandb`)",
+        ),
+    ] = None,
+    wandb_project: Annotated[
+        Optional[str],
+        typer.Option("--wandb-project", help="W&B project name (default: giant)"),
+    ] = None,
+    wandb_run_name: Annotated[
+        Optional[str],
+        typer.Option("--wandb-run-name", help="W&B run name (default: out_dir name)"),
+    ] = None,
+    wandb_log_every: Annotated[
+        Optional[int],
+        typer.Option(
+            "--wandb-log-every",
+            help="Log batch-level loss/grad_norm/lr to W&B every N optimizer "
+            "steps (default: 50); per-epoch metrics always log in full",
+        ),
+    ] = None,
 ) -> None:
     """Train the GIANT surrogate model."""
     batch_size_auto = False
@@ -436,6 +510,10 @@ def train(
             "n_critic": n_critic,
             "gp_weight": gp_weight,
             "critic_lr": critic_lr,
+            "wandb": wandb,
+            "wandb_project": wandb_project,
+            "wandb_run_name": wandb_run_name,
+            "wandb_log_every": wandb_log_every,
         }.items()
         if v is not None
     }
@@ -451,17 +529,7 @@ def train(
         }.items()
         if v is not None
     }
-    cli_router: dict[str, object] = {
-        k: v
-        for k, v in {
-            "enabled": router,
-            "type": router_type,
-            "n_experts": n_experts,
-        }.items()
-        if v is not None
-    }
-    if router_axis:
-        cli_router.update(_parse_router_axis_flags(router_axis))
+    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
     if cli_router:
         cli_model["router"] = cli_router
     cfg = gconfig.merge_cli_overrides(
@@ -484,16 +552,24 @@ def train(
             f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
         )
 
-    out_dir = out or Path(
-        f"checkpoints/{date.today().strftime('%Y%m%d')}"
-        f"_{t['mode']}"
-        f"_h{m['hidden_dim']}"
-        f"_b{m['n_blocks']}"
-        f"_e{m['emb_dim']}"
-        f"_c{m['conditioning']}"
-        f"_lr{t['lr']}"
-        f"_bs{t['batch_size']}"
-    )
+    if out is not None:
+        out_dir = out
+    elif resume is not None:
+        # Continue writing into the resumed checkpoint's own directory
+        # rather than recomputing a hyperparam-derived name — the latter
+        # would (a) collide with the original run's dir only by accident
+        # (same day, unchanged hyperparams) and now never collides at all
+        # since the fresh-run name below is timestamped to the second, and
+        # (b) silently start a fresh directory if a resumed run tweaks any
+        # hyperparam baked into the name (e.g. --lr for a fine-tune).
+        out_dir = resume.parent
+    else:
+        # Name only encodes what's non-default (see default_out_dir_name), so
+        # two runs with identical hyperparams in the same to-the-minute
+        # timestamp would otherwise collide on this name — which also
+        # doubles as the W&B run id (giant.train) — hence the suffix loop in
+        # resolve_default_out_dir.
+        out_dir = gconfig.resolve_default_out_dir(cfg)
 
     typer.echo(f"device:  {_device}")
     typer.echo(f"out_dir: {out_dir}")
@@ -506,10 +582,150 @@ def train(
         shuffle_buffer=shuffle_buffer,
         num_workers=t["num_workers"],
         resume=resume,
+        cache_setup=cache_setup,
+        rebuild_setup_cache=rebuild_setup_cache,
         echo=typer.echo,
     )
 
 
+@app.command("new-run")
+def new_run(
+    config: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--config",
+            "-c",
+            help="Base TOML to start from (default: built-in defaults)",
+        ),
+    ] = None,
+    mode: Annotated[Optional[Mode], typer.Option("--mode", "-m")] = None,
+    epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
+    batch_size: Annotated[Optional[int], typer.Option("--batch-size", "-b")] = None,
+    lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
+    hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
+    n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
+    emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
+    dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
+    conditioning: Annotated[
+        Optional[Conditioning], typer.Option("--conditioning")
+    ] = None,
+    router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
+    router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
+    n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None,
+    router_axis: Annotated[Optional[list[str]], typer.Option("--router-axis")] = None,
+    out: Annotated[
+        Optional[Path],
+        typer.Option("--out", "-o", help="Run dir (default: auto from hyperparams)"),
+    ] = None,
+    comment: Annotated[
+        Optional[str],
+        typer.Option(
+            "--comment", help="Free-text note recorded in config.toml's meta section"
+        ),
+    ] = None,
+    data: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--data",
+            help="Dataset path to fill in the printed next-step command "
+            "(not stored in the config)",
+        ),
+    ] = None,
+    force: Annotated[
+        bool,
+        typer.Option(
+            "--force",
+            help="Overwrite config.toml even if --out already has checkpoints",
+        ),
+    ] = False,
+    dry_run: Annotated[
+        bool,
+        typer.Option(
+            "--dry-run", help="Print the resolved config without writing anything"
+        ),
+    ] = False,
+) -> None:
+    """Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir.
+
+    This is the config-file-first counterpart to hand-editing a TOML: start
+    from a base --config (or built-in defaults), override a few hyperparams
+    inline, and this resolves+writes the full `config.toml` into a fresh (or
+    explicit --out) run dir — the same file `giant train --config ...` reads.
+    `giant train` itself overwrites this file in place once it actually runs
+    (with the full dataset-derived meta section), so this scaffold's meta
+    section is just a placeholder recording what was asked for and when.
+    """
+    cli_train = {
+        k: v
+        for k, v in {
+            "mode": mode.value if mode is not None else None,
+            "epochs": epochs,
+            "batch_size": batch_size,
+            "lr": lr,
+        }.items()
+        if v is not None
+    }
+    cli_model: dict[str, object] = {
+        k: v
+        for k, v in {
+            "hidden_dim": hidden_dim,
+            "n_blocks": n_blocks,
+            "emb_dim": emb_dim,
+            "dropout": dropout,
+            "conditioning": conditioning.value if conditioning is not None else None,
+        }.items()
+        if v is not None
+    }
+    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
+    if cli_router:
+        cli_model["router"] = cli_router
+
+    cfg = gconfig.merge_cli_overrides(
+        gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
+    )
+    run_dir = (out or gconfig.resolve_default_out_dir(cfg)).resolve()
+
+    if not force:
+        existing = [n for n in ("last.pt", "best.pt") if (run_dir / n).exists()]
+        if existing:
+            typer.echo(
+                f"error: {run_dir} already has {', '.join(existing)} — pass "
+                "--force to overwrite its config.toml anyway",
+                err=True,
+            )
+            raise typer.Exit(1)
+
+    typer.echo(f"run dir: {run_dir}")
+
+    if dry_run:
+        typer.echo("dry-run: not writing anything. Resolved config:")
+        for section in ("train", "model"):
+            typer.echo(f"[{section}]")
+            for k, v in cfg[section].items():
+                if k == "router":
+                    continue
+                typer.echo(f"  {k} = {v}")
+        return
+
+    meta = {
+        "git_hash": gconfig.git_hash(),
+        "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+        "created_by": "giant new-run",
+    }
+    if comment:
+        meta["comment"] = comment
+
+    run_dir.mkdir(parents=True, exist_ok=True)
+    gconfig.save_config(cfg, run_dir, meta)
+    config_path = run_dir / "config.toml"
+    typer.echo(f"wrote {config_path}")
+
+    data_arg = str(data) if data is not None else "<data.parquet>"
+    typer.echo("")
+    typer.echo("next:")
+    typer.echo(f"  giant train {data_arg} --config {config_path} --out {run_dir}")
+
+
 @app.command()
 def predict(
     data: Annotated[
@@ -845,8 +1061,8 @@ def predict(
     buffer: dict[str, np.ndarray] | None = None
 
     bar = tqdm(total=total_rows, desc="predict", unit="row", dynamic_ncols=True)
-    for path in files:
-        for chunk in chunk_iter(path):
+    for i, path in enumerate(files):
+        for chunk in chunk_iter(path, offset=event_id_offset(i)):
             N_in = len(chunk["event_id"])
 
             pdg_mask = np.array([int(p) in pdg_map for p in chunk["pdg"]])
@@ -895,8 +1111,8 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
     """
     best_E: dict[int, float] = {}
     best: dict[int, tuple] = {}
-    for path in files:
-        for chunk in iter_cond_chunks(path):
+    for file_idx, path in enumerate(files):
+        for chunk in iter_cond_chunks(path, offset=event_id_offset(file_idx)):
             ev = chunk["event_id"]
             pe = chunk["pre_E"]
             for i in range(len(ev)):

giant/config.py

@@ -1,13 +1,26 @@
+import hashlib
 import random
 import subprocess
 import sys
 import tomllib
 from datetime import datetime, timezone
+from enum import Enum
 from pathlib import Path
 
 import numpy as np
 import torch
 
+
+class Conditioning(str, Enum):
+    """`model.conditioning` choices — shared by `giant.cli` and `scripts.dwarf`'s
+    Typer commands so the two CLIs can't silently drift apart on the option's
+    valid values (see DEFAULT_CONFIG["model"]["conditioning"] for what each
+    value means)."""
+
+    physical = "physical"
+    embedding = "embedding"
+
+
 DEFAULT_CONFIG: dict = {
     "train": {
         "mode": "flow",
@@ -35,6 +48,19 @@ DEFAULT_CONFIG: dict = {
         "n_critic": 5,
         "gp_weight": 10.0,
         "critic_lr": 0.0,
+        # Weights & Biases per-epoch metric logging (opt-in; see giant.train).
+        # "" for wandb_run_name means "use the checkpoint out_dir name" — not
+        # None, since save_config's TOML writer has no null literal to
+        # round-trip.
+        "wandb": False,
+        "wandb_project": "giant",
+        "wandb_run_name": "",
+        # Batch-granularity metrics (loss/grad_norm/lr) are logged every N
+        # optimizer steps, not every batch — a single epoch can be tens of
+        # thousands of steps (see steps_per_epoch above), and logging every
+        # one of them would flood the run with points the UI has to downsample
+        # anyway. Per-epoch metrics (the metrics.csv row) always log in full.
+        "wandb_log_every": 50,
     },
     "model": {
         "hidden_dim": 256,
@@ -54,11 +80,50 @@ DEFAULT_CONFIG: dict = {
             "enabled": False,
             "type": "energy",  # selects the Router impl from ROUTER_REGISTRY
             "n_experts": 4,
-            "expert_hidden_dim": 128,
-            "expert_n_blocks": 3,
+            # 0 means "inherit model.hidden_dim/n_blocks" (see
+            # resolve_expert_dims below) — not a fixed 128/3, which silently
+            # ignored --hidden-dim/--n-blocks whenever routing was enabled.
+            # TOML has no null literal to round-trip (same pattern as
+            # critic_lr/wandb_run_name above), hence 0 rather than None.
+            "expert_hidden_dim": 0,
+            "expert_n_blocks": 0,
             "temperature": 0.5,  # energy/pdg-router kwarg
             "learn_centers": True,  # energy/pdg-router kwarg
+            # energy-router kwargs: mutually exclusive optional learnable
+            # gate-sharpness modes (see giant.model.network.EnergyRouter).
+            # learn_width generalizes the shared `temperature` to one
+            # learnable width per expert; learn_temperature instead makes
+            # the single shared `temperature` itself learnable. Both are
+            # bounded to [width_min_ratio, width_max_ratio] * temperature
+            # (sigmoid-parameterized, warm-started to reproduce `temperature`
+            # exactly at init) so gate sharpness can't run away to a
+            # collapse-inducing extreme during training.
+            "learn_width": False,
+            "learn_temperature": False,
+            "width_min_ratio": 0.1,
+            "width_max_ratio": 10.0,
             "lambda_balance": 0.0,  # optional load-balance aux loss weight
+            # optional entropy-regularization aux loss weight (generic
+            # Router.entropy_loss, penalizes uniform/collapsed gating) — a
+            # secondary guard against all experts' widths/temperature
+            # co-inflating together, which lambda_balance alone can't see
+            # since per-expert usage shares stay even throughout that
+            # failure mode. Off by default; bounding above is the primary
+            # defense. See giant.model.network.Router.entropy_loss.
+            "lambda_entropy": 0.0,
+            # Opt-in straight-through Gumbel-softmax train-time combine weights
+            # (see giant.model.network.Router.combine_weights): the training
+            # forward pass samples a hard one-hot combination — matching
+            # eval-time top-1 dispatch exactly — while the backward pass still
+            # flows a smooth gradient to every expert. Targets the train/eval
+            # mismatch identified as a likely contributor to experts
+            # overlapping instead of partitioning (see CLAUDE.md roadmap).
+            # gumbel_tau_start/_end are annealed linearly over training
+            # (giant.train._gumbel_tau); off by default, no effect unless
+            # gumbel = true.
+            "gumbel": False,
+            "gumbel_tau_start": 1.0,
+            "gumbel_tau_end": 0.1,
             "emb_dim": 8,  # process/pdg-router kwarg: own pdg(/mat) embedding width
             "hidden_dim": 64,  # process-router kwarg: its classifier's hidden width
             "lambda_proc": 0.0,  # process-router kwarg: supervised process-CE weight
@@ -250,6 +315,151 @@ def merge_cli_overrides(
     return cfg
 
 
+def resolve_expert_dims(
+    router_cfg: dict, hidden_dim: int, n_blocks: int
+) -> tuple[int, int]:
+    """Resolve a router's expert hidden_dim/n_blocks, inheriting from the
+    monolith's when left at the 0 ("unset") sentinel.
+
+    Used by both `giant.pipeline` (to build the checkpoint's `model_config`)
+    and `giant.cli`'s batch-size auto-estimate, so `--hidden-dim`/`--n-blocks`
+    size the experts the same way in both places unless
+    `router.expert_hidden_dim`/`expert_n_blocks` are explicitly overridden.
+    """
+    expert_hidden_dim = router_cfg.get("expert_hidden_dim") or hidden_dim
+    expert_n_blocks = router_cfg.get("expert_n_blocks") or n_blocks
+    return expert_hidden_dim, expert_n_blocks
+
+
+_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb"}
+
+
+# Priority-ordered candidate fields for default_out_dir_name: (label, getter,
+# formatter). `getter(train, model)` returns None when the field is at its
+# default (and so should be omitted); otherwise formatter(value) renders the
+# name token. The router is a single unit gated on `router.enabled` rather
+# than one candidate per router key, since its type/n_experts are meaningless
+# while disabled.
+def _mode_candidate(train, model):
+    return None if train["mode"] == DEFAULT_CONFIG["train"]["mode"] else train["mode"]
+
+
+def _router_candidate(train, model):
+    router = model["router"]
+    if router["enabled"] == DEFAULT_CONFIG["model"]["router"]["enabled"]:
+        return None
+    return f"r-{router['type']}{router['n_experts']}"
+
+
+def _router_flag_candidate(field, token_map):
+    """Candidate factory for a boolean `model.router` sub-field.
+
+    Gated on `router.enabled` like `_router_candidate` (a disabled router's
+    sub-fields are meaningless), then omitted unless `field` differs from
+    its DEFAULT_CONFIG value — same "only show non-default" rule as every
+    other candidate. `token_map` need only cover the non-default value(s),
+    since the default value always yields None.
+    """
+
+    def _candidate(train, model):
+        router = model["router"]
+        default_router = DEFAULT_CONFIG["model"]["router"]
+        if router["enabled"] == default_router["enabled"]:
+            return None
+        value = router[field]
+        if value == default_router[field]:
+            return None
+        return token_map[value]
+
+    return _candidate
+
+
+def _conditioning_candidate(train, model):
+    if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]:
+        return None
+    code = _CONDITIONING_CODE.get(model["conditioning"], model["conditioning"])
+    return f"c{code}"
+
+
+def _default_field_candidate(section_key, field, prefix):
+    def _candidate(train, model):
+        section = train if section_key == "train" else model
+        value = section[field]
+        if value == DEFAULT_CONFIG[section_key][field]:
+            return None
+        return f"{prefix}{value}"
+
+    return _candidate
+
+
+_OUT_DIR_NAME_CANDIDATES = [
+    ("mode", _mode_candidate),
+    ("router", _router_candidate),
+    ("gumbel", _router_flag_candidate("gumbel", {True: "gum"})),
+    ("learn_centers", _router_flag_candidate("learn_centers", {False: "nolc"})),
+    ("learn_width", _router_flag_candidate("learn_width", {True: "lw"})),
+    ("learn_temperature", _router_flag_candidate("learn_temperature", {True: "lt"})),
+    ("conditioning", _conditioning_candidate),
+    ("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")),
+    ("n_blocks", _default_field_candidate("model", "n_blocks", "b")),
+    ("emb_dim", _default_field_candidate("model", "emb_dim", "e")),
+    ("lr", _default_field_candidate("train", "lr", "lr")),
+    ("batch_size", _default_field_candidate("train", "batch_size", "bs")),
+    ("seed", _default_field_candidate("train", "seed", "seed")),
+    ("epochs", _default_field_candidate("train", "epochs", "ep")),
+]
+
+_OUT_DIR_NAME_MAX_FIELDS = 6
+
+
+def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
+    """Build a default checkpoint out_dir name from what's non-default in `cfg`.
+
+    Only fields that differ from DEFAULT_CONFIG are included, so a fully
+    default run's name is just its timestamp — see
+    `_OUT_DIR_NAME_CANDIDATES` for the fixed, priority-ordered field list.
+    Beyond `_OUT_DIR_NAME_MAX_FIELDS` non-default fields, the remainder
+    collapse into a short deterministic hash suffix rather than growing the
+    name unboundedly. This name doubles as the run's W&B id (see
+    giant.train), which is the reason a timestamp is always included.
+    """
+    now = now or datetime.now()
+    train, model = cfg["train"], cfg["model"]
+    tokens = []
+    overflow = []
+    for label, candidate in _OUT_DIR_NAME_CANDIDATES:
+        token = candidate(train, model)
+        if token is None:
+            continue
+        if len(tokens) < _OUT_DIR_NAME_MAX_FIELDS:
+            tokens.append(token)
+        else:
+            overflow.append(f"{label}={token}")
+
+    name = now.strftime("%Y%m%d_%H%M")
+    if tokens:
+        name += "_" + "_".join(tokens)
+    if overflow:
+        digest = hashlib.md5("|".join(sorted(overflow)).encode()).hexdigest()[:6]
+        name += f"_+{len(overflow)}more-{digest}"
+    return name
+
+
+def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path:
+    """Auto-derived out dir from cfg's hyperparams (see `default_out_dir_name`),
+    with a numeric suffix loop so two runs whose name collides (same
+    non-default hyperparams, same to-the-minute timestamp) don't clobber each
+    other's directory. Shared by `giant train` and `giant new-run`.
+    """
+    base_name = default_out_dir_name(cfg)
+    out_dir = base / base_name
+    suffix = 2
+    while out_dir.exists():
+        out_dir = base / f"{base_name}_{suffix}"
+        suffix += 1
+    return out_dir
+
+
 def seed_everything(seed: int) -> None:
     random.seed(seed)
     np.random.seed(seed)

giant/data/dataset.py

@@ -6,8 +6,8 @@ import numpy as np
 import torch
 from torch.utils.data import IterableDataset
 
-from giant.data.loader import iter_file_chunks
-from giant.data.transforms import Normalizer, build_features
+from giant.data.loader import event_id_offset, iter_file_chunks
+from giant.data.transforms import Normalizer, build_features, sorted_membership
 
 
 def make_event_split(
@@ -19,7 +19,10 @@ def make_event_split(
     rng = np.random.default_rng(seed)
     unique = np.unique(all_event_ids)
     rng.shuffle(unique)
-    n_val = max(1, int(len(unique) * val_fraction))
+    # max(1, ...) only applies when a validation split was actually
+    # requested — val_fraction=0.0 is an explicit "train on everything"
+    # request and must not be silently overridden into holding out 1 event.
+    n_val = max(1, int(len(unique) * val_fraction)) if val_fraction > 0 else 0
     val_set = set(unique[:n_val].tolist())
     train_set = set(unique[n_val:].tolist())
     return train_set, val_set
@@ -65,6 +68,7 @@ class StreamingStepsDataset(IterableDataset):
         sec_phys_normalizer: Normalizer | None = None,
     ) -> None:
         self.files = list(files)
+        self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
         self.split_events = split_events
         self._events_arr = np.array(sorted(split_events))
         self.pdg_map = pdg_map
@@ -97,8 +101,8 @@ class StreamingStepsDataset(IterableDataset):
         buf_n = 0
 
         for path in files:
-            for chunk in iter_file_chunks(path):
-                mask = np.isin(chunk["event_id"], self._events_arr)
+            for chunk in iter_file_chunks(path, offset=self._offsets[path]):
+                mask = sorted_membership(chunk["event_id"], self._events_arr)
                 if not mask.any():
                     continue
                 chunk = {k: v[mask] for k, v in chunk.items()}

giant/data/loader.py

@@ -12,6 +12,40 @@ import pyarrow.parquet as pq
 # dataset tree is moved or copied elsewhere intact.
 MANIFEST_SUFFIX = ".manifest"
 
+# Each input parquet file is a separate Geant4 job converted 1:1 from its own
+# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
+# always restarts from 0 — so when multiple files are loaded together (a
+# directory or .manifest), raw event_id values collide across files even
+# though they refer to unrelated events. Every per-file event_id column gets
+# offset by its file's index in the (deterministically ordered) files list
+# so ids stay globally unique across a multi-file load; the stride is far
+# larger than any realistic per-file event count.
+EVENT_ID_FILE_STRIDE = 1_000_000
+
+
+def event_id_offset(file_index: int) -> int:
+    return file_index * EVENT_ID_FILE_STRIDE
+
+
+def _offset_event_id(raw_ids: np.ndarray, offset: int) -> np.ndarray:
+    """Add this file's `event_id_offset`, after checking the raw ids fit in one stride block.
+
+    Without this check, a file whose own raw event_id numbering reaches
+    `EVENT_ID_FILE_STRIDE` (an unusually large job, or non-contiguous
+    numbering) would silently collide into the next file's offset block,
+    merging unrelated events across files — reintroducing exactly the
+    train/val event leakage this offset scheme exists to prevent.
+    """
+    raw_ids = np.asarray(raw_ids, dtype=np.int64)
+    if raw_ids.size and int(raw_ids.max()) >= EVENT_ID_FILE_STRIDE:
+        raise ValueError(
+            f"event_id {int(raw_ids.max())} >= EVENT_ID_FILE_STRIDE "
+            f"({EVENT_ID_FILE_STRIDE}) — this file has a larger event_id "
+            "than the per-file offset scheme can support without colliding "
+            "with the next file's id block."
+        )
+    return raw_ids + offset
+
 
 def _read_manifest(path: Path) -> list[Path]:
     files = []
@@ -78,13 +112,13 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
     return out
 
 
-def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
+def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     from giant.constants import K_MAX
 
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
-        "event_id": df["event_id"].to_numpy(),
+        "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
@@ -121,20 +155,23 @@ def _df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
     return d
 
 
-def load_steps(path: str | Path) -> dict[str, np.ndarray]:
-    return _df_to_dict(pd.read_parquet(path))
+def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]:
+    return _df_to_dict(pd.read_parquet(path), offset=offset)
 
 
-def load_event_ids(path: str | Path) -> np.ndarray:
+def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
     """Read only the event_id column — cheap scan for split assignment."""
-    return pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
+    ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
+    return _offset_event_id(ids, offset)
 
 
-def iter_file_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
+def iter_file_chunks(
+    path: str | Path, offset: int = 0
+) -> Iterator[dict[str, np.ndarray]]:
     """Yield one parquet row-group at a time so a large file never fully loads."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _df_to_dict(pf.read_row_group(i).to_pandas())
+        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset)
 
 
 _COND_COLS = [
@@ -154,9 +191,9 @@ _COND_COLS = [
 ]
 
 
-def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
+def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     return {
-        "event_id": df["event_id"].to_numpy(),
+        "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
@@ -168,11 +205,15 @@ def _cond_df_to_dict(df: pd.DataFrame) -> dict[str, np.ndarray]:
     }
 
 
-def iter_cond_chunks(path: str | Path) -> Iterator[dict[str, np.ndarray]]:
+def iter_cond_chunks(
+    path: str | Path, offset: int = 0
+) -> Iterator[dict[str, np.ndarray]]:
     """Yield conditioning-only row-groups (no post-step columns read from disk)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas())
+        yield _cond_df_to_dict(
+            pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
+        )
 
 
 def build_index_maps(

giant/data/setup_cache.py

@@ -0,0 +1,330 @@
+"""Sidecar cache for `giant train`'s setup stage (vocab maps, event-id split
+index, process maps, normalizer stats).
+
+The setup stage scans the full training dataset before a single epoch runs
+(see giant/pipeline.py:run_train_job); on multi-hundred-million-row datasets
+that scan is itself expensive, and it's pure waste to repeat when the same
+`data` path is reused across runs (hyperparameter sweeps via `dwarf
+hparam-scan`, repeated manual training attempts, ...). This module persists
+those scan outputs to a JSON file next to `data`, validated by a file
+fingerprint + fixed dimension constants + a manually-bumped format version
+before reuse — see `load`/`save`.
+"""
+
+from __future__ import annotations
+
+import fcntl
+import json
+import os
+from dataclasses import dataclass, field
+from pathlib import Path
+
+import numpy as np
+
+from giant import config
+from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
+from giant.data.loader import event_id_offset, load_event_ids
+from giant.data.transforms import Normalizer, sorted_membership
+
+# Bump manually on a change to the data-encoding semantics (e.g. a future
+# energy_simplex_encode bugfix) that doesn't also move one of _DIMS below —
+# a dims change already hard-invalidates on its own.
+# v2: event_id is now offset per-file (see loader.event_id_offset) to avoid
+# cross-file collisions, so a v1 sidecar's event_index/normalizers were
+# computed against collided ids and must not be reused.
+# v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
+# energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
+# v2 sidecar has no such grid to fall back on, so it must be recomputed.
+_CACHE_FORMAT_VERSION = 3
+
+_DIMS = {
+    "COND_DIM": COND_DIM,
+    "X_DIM": X_DIM,
+    "K_MAX": K_MAX,
+    "PARTICLE_PHYS_DIM": PARTICLE_PHYS_DIM,
+    "SEC_SLOT_DIM": SEC_SLOT_DIM,
+}
+
+# Resolution of the stored energy-quantile summary (see NormalizerEntry).
+# Only a handful of quantile *levels* (one per EnergyRouter expert) are ever
+# consumed (see pipeline.py), so a dense fixed grid of quantile values is
+# enough to reconstruct any level via interpolation (energy_quantile_at) —
+# at roughly 1/100th the storage of the raw 100k-value reservoir sample it
+# replaces, with negligible loss of resolution for that use.
+ENERGY_QUANTILE_LEVELS = 1001
+
+
+def energy_quantiles_from_sample(sample: np.ndarray) -> np.ndarray:
+    """Collapse a raw reservoir sample into the fixed grid stored on disk."""
+    if sample.size == 0:
+        return np.empty(0, dtype=np.float32)
+    levels = np.linspace(0.0, 1.0, ENERGY_QUANTILE_LEVELS)
+    return np.quantile(sample, levels).astype(np.float32)
+
+
+def energy_quantile_at(energy_quantiles: np.ndarray, levels: np.ndarray) -> np.ndarray:
+    """Interpolate quantile values at arbitrary probability `levels` from the
+    stored grid (e.g. `np.linspace(0, 1, n_experts)` for router centers)."""
+    grid_levels = np.linspace(0.0, 1.0, len(energy_quantiles))
+    return np.interp(levels, grid_levels, energy_quantiles).astype(np.float32)
+
+
+def sidecar_path(data: str | Path) -> Path:
+    """The cache sidecar for `data`, always a sibling of `data` itself.
+
+    A directory `data` gets a sidecar *next to* it (not inside), since the
+    directory may be a shared/read-only dataset mount, and other code globs
+    `*.parquet` directly inside it.
+    """
+    p = Path(data)
+    return p.parent / f"{p.name}.giant_train_cache.json"
+
+
+def fingerprint_files(files: list[Path]) -> list[list]:
+    """`[[resolved_path_str, size, mtime_ns], ...]`, in `files` order (not sorted).
+
+    Order must be preserved rather than normalized (e.g. sorted): file scan
+    order affects `build_process_map_from_files`'s tie-breaking (see
+    tests/test_loader.py), so the cached fingerprint has to reflect the same
+    order `find_parquet_files` produced.
+    """
+    out = []
+    for f in files:
+        resolved = Path(f).resolve()
+        st = resolved.stat()
+        out.append([str(resolved), st.st_size, st.st_mtime_ns])
+    return out
+
+
+def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str:
+    # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
+    # spurious cache misses between runs with the "same" val_fraction.
+    return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}"
+
+
+@dataclass
+class NormalizerEntry:
+    cond_norm: Normalizer
+    tgt_norm: Normalizer
+    sec_phys_norm: Normalizer
+    n_train_steps: int
+    energy_quantiles: np.ndarray
+    """Fixed ENERGY_QUANTILE_LEVELS-point quantile grid of the raw (pre-
+    normalization) pre-step energy column — see energy_quantiles_from_sample
+    / energy_quantile_at."""
+
+    def to_json(self) -> dict:
+        return {
+            "cond_norm": self.cond_norm.to_dict(),
+            "tgt_norm": self.tgt_norm.to_dict(),
+            "sec_phys_norm": self.sec_phys_norm.to_dict(),
+            "n_train_steps": self.n_train_steps,
+            "energy_quantiles": np.asarray(
+                self.energy_quantiles, dtype=np.float32
+            ).tolist(),
+        }
+
+    @classmethod
+    def from_json(cls, d: dict) -> "NormalizerEntry":
+        return cls(
+            cond_norm=Normalizer.from_dict(d["cond_norm"]),
+            tgt_norm=Normalizer.from_dict(d["tgt_norm"]),
+            sec_phys_norm=Normalizer.from_dict(d["sec_phys_norm"]),
+            n_train_steps=int(d["n_train_steps"]),
+            energy_quantiles=np.array(d["energy_quantiles"], dtype=np.float32),
+        )
+
+
+@dataclass
+class SetupCache:
+    fingerprint: list
+    git_hash: str = field(default_factory=config.git_hash)
+    vocab: tuple[dict[int, int], dict[str, int]] | None = None
+    event_index: tuple[np.ndarray, np.ndarray] | None = None
+    proc_maps: dict[int, dict[str, int]] = field(default_factory=dict)
+    normalizers: dict[str, NormalizerEntry] = field(default_factory=dict)
+
+    @classmethod
+    def empty(cls, files: list[Path]) -> "SetupCache":
+        return cls(fingerprint=fingerprint_files(files))
+
+    def to_json(self) -> dict:
+        d: dict = {
+            "format_version": _CACHE_FORMAT_VERSION,
+            "dims": dict(_DIMS),
+            "git_hash": self.git_hash,
+            "fingerprint": self.fingerprint,
+            "proc_maps": {str(k): v for k, v in self.proc_maps.items()},
+            "normalizers": {k: v.to_json() for k, v in self.normalizers.items()},
+        }
+        if self.vocab is not None:
+            pdg_map, mat_map = self.vocab
+            d["vocab"] = {
+                "pdg_map": {str(k): v for k, v in pdg_map.items()},
+                "mat_map": dict(mat_map),
+            }
+        if self.event_index is not None:
+            unique_ids, counts = self.event_index
+            d["event_index"] = {
+                "event_ids": np.asarray(unique_ids).tolist(),
+                "counts": np.asarray(counts).tolist(),
+            }
+        return d
+
+    @classmethod
+    def from_json(cls, d: dict) -> "SetupCache":
+        vocab = None
+        if "vocab" in d:
+            pdg_map = {int(k): v for k, v in d["vocab"]["pdg_map"].items()}
+            mat_map = dict(d["vocab"]["mat_map"])
+            vocab = (pdg_map, mat_map)
+        event_index = None
+        if "event_index" in d:
+            event_index = (
+                np.array(d["event_index"]["event_ids"], dtype=np.int64),
+                np.array(d["event_index"]["counts"], dtype=np.int64),
+            )
+        proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
+        normalizers = {
+            k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
+        }
+        return cls(
+            fingerprint=d["fingerprint"],
+            git_hash=d.get("git_hash", "unknown"),
+            vocab=vocab,
+            event_index=event_index,
+            proc_maps=proc_maps,
+            normalizers=normalizers,
+        )
+
+    def merge(self, other: "SetupCache") -> "SetupCache":
+        """Union of both caches; `other`'s populated fields win on a shared key.
+
+        Used by `save` to combine freshly-computed sections with whatever a
+        concurrent writer already persisted, so two runs against the same
+        dataset with different (e.g.) val_fraction don't clobber each
+        other's normalizer entries.
+        """
+        return SetupCache(
+            fingerprint=other.fingerprint,
+            git_hash=other.git_hash,
+            vocab=other.vocab if other.vocab is not None else self.vocab,
+            event_index=(
+                other.event_index if other.event_index is not None else self.event_index
+            ),
+            proc_maps={**self.proc_maps, **other.proc_maps},
+            normalizers={**self.normalizers, **other.normalizers},
+        )
+
+
+def load(
+    data: str | Path, files: list[Path], echo=lambda *a, **k: None
+) -> SetupCache | None:
+    """Load and validate the sidecar for `data`; `None` on any miss (never raises).
+
+    A missing file, corrupt JSON, format-version mismatch, dimension-constant
+    mismatch, or file-fingerprint mismatch are all clean misses. A git-hash
+    mismatch alone is a soft warning only (see
+    `config.warn_if_git_hash_mismatch`) — most commits to this repo don't
+    touch data-encoding semantics, so hard-invalidating on every one would
+    defeat the cache.
+    """
+    path = sidecar_path(data)
+    if not path.exists():
+        return None
+    try:
+        raw = json.loads(path.read_text())
+    except (json.JSONDecodeError, OSError) as exc:
+        echo(f"setup cache: {path} is corrupt ({exc}) — ignoring")
+        return None
+
+    try:
+        if raw.get("format_version") != _CACHE_FORMAT_VERSION:
+            echo("setup cache: format version changed — ignoring stale cache")
+            return None
+        if raw.get("dims") != _DIMS:
+            echo(
+                "setup cache: model dimension constants changed — ignoring stale cache"
+            )
+            return None
+        fp = fingerprint_files(files)
+        if raw.get("fingerprint") != fp:
+            echo("setup cache: input files changed — ignoring stale cache")
+            return None
+        cache = SetupCache.from_json(raw)
+    except (KeyError, TypeError, ValueError) as exc:
+        echo(f"setup cache: {path} is malformed ({exc}) — ignoring")
+        return None
+
+    config.warn_if_git_hash_mismatch({"meta": {"git_hash": cache.git_hash}}, path)
+    return cache
+
+
+def save(
+    data: str | Path,
+    files: list[Path],
+    sections: SetupCache,
+    echo=lambda *a, **k: None,
+) -> None:
+    """Merge `sections` into the on-disk sidecar and write it atomically.
+
+    Best-effort: any OSError (permission denied on a read-only mount, disk
+    full, ...) is caught, echoed as a warning, and swallowed — a failure to
+    cache must never fail training.
+
+    The load-merge-write is serialized with an exclusive flock on a sidecar
+    lockfile: `os.replace` alone only guarantees the *file* is never
+    corrupt, not that concurrent writers don't race. Without the lock, two
+    concurrent `giant train`/condor jobs against the same `data` path (this
+    repo's shared-portal/condor usage makes that a real scenario, not just
+    theoretical) could both `load()` the same base state, merge their own
+    `sections` in independently, and whichever `os.replace()` lands last
+    silently discards the other's freshly-computed section.
+    """
+    path = sidecar_path(data)
+    lock_path = path.parent / f".{path.name}.lock"
+    tmp = path.parent / f".{path.name}.tmp.{os.getpid()}"
+    try:
+        with open(lock_path, "a") as lock_file:
+            fcntl.flock(lock_file, fcntl.LOCK_EX)
+            try:
+                base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(
+                    files
+                )
+                merged = base.merge(sections)
+                payload = json.dumps(merged.to_json(), separators=(",", ":"))
+                tmp.write_text(payload)
+                os.replace(tmp, path)
+            finally:
+                fcntl.flock(lock_file, fcntl.LOCK_UN)
+    except OSError as exc:
+        echo(
+            f"setup cache: could not write {path} ({exc}) — continuing without caching"
+        )
+        try:
+            tmp.unlink(missing_ok=True)
+        except OSError:
+            pass
+
+
+def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.ndarray]:
+    """Unique event ids + per-event row (step) counts, across all `files`."""
+    if not files:
+        return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
+    all_ids = np.concatenate(
+        [load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
+    )
+    unique_ids, counts = np.unique(all_ids, return_counts=True)
+    return unique_ids, counts
+
+
+def n_train_steps_for_split(
+    unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
+) -> int:
+    """Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.
+
+    `train_events_arr` must be ascending and duplicate-free (as produced by
+    `np.array(sorted(train_events))` in giant/pipeline.py).
+    """
+    mask = sorted_membership(unique_ids, train_events_arr)
+    return int(counts[mask].sum())

giant/data/transforms.py

@@ -12,7 +12,17 @@ _SIMPLEX_FLOOR = 1e-5
 
 
 def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
-    return np.log(np.asarray(x, dtype=np.float32) + eps)
+    x = np.asarray(x, dtype=np.float32)
+    y = np.log(x + eps)
+    if not np.all(np.isfinite(y)):
+        bad = int(np.sum(~np.isfinite(y)))
+        raise ValueError(
+            f"log_transform: {bad} value(s) produced non-finite output (input "
+            f"< -eps={eps:g}, or already NaN/Inf); every quantity this is "
+            "applied to should be non-negative, so this indicates upstream "
+            "data corruption rather than expected float noise."
+        )
+    return y
 
 
 def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
@@ -110,8 +120,22 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
         [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
     )
     axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)  # (N,1)
-    # Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
-    # are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
+    # axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
+    # pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
+    # ẑ) vanishes at both. The "choice is irrelevant" claim below only holds
+    # at +ẑ, where sin_t~0 AND (1-cos_t)~0 so every axis-dependent Rodrigues
+    # term vanishes. At -ẑ, sin_t~0 but (1-cos_t)~2 — not negligible — so
+    # snapping to a fixed x̂ there is a genuine (if physically rare)
+    # modeling choice, not a no-op: it picks one representative out of an
+    # inherently ambiguous family of 180°-about-any-transverse-axis
+    # rotations (no single-valued frame convention can be continuous through
+    # this antipode — same obstruction as a sphere's tangent frame having no
+    # continuous choice at a pole). x̂ is still fine to use — it's a fixed,
+    # self-consistent convention that `local_frame_rotation`/
+    # `inv_local_frame_rotation` (same threshold) round-trip correctly
+    # through — but steps whose pre_dir falls in this tiny near-backscatter
+    # cone get a discontinuous "roll" relative to their non-degenerate
+    # neighbors, injecting a small amount of label noise there.
     safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
     return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
 
@@ -134,8 +158,20 @@ def _validate_unit_pre_dir(pre_dir: np.ndarray) -> np.ndarray:
     drift is corrected silently; a near-zero-norm row has no well-defined
     direction, so it's raised loudly instead of producing a meaningless
     rotation (previously it fell through to an arbitrary axis with no error).
+
+    NaN/Inf rows are also raised on explicitly: `norm < 1e-6` is False for a
+    NaN norm, so without this check a non-finite row would silently pass
+    through and poison everything downstream (e.g. the persisted normalizer
+    stats in `setup_cache`, if the row is swept into a Welford accumulator).
     """
     pre_dir = np.asarray(pre_dir, dtype=np.float32)
+    if not np.all(np.isfinite(pre_dir)):
+        bad = int(np.sum(~np.all(np.isfinite(pre_dir), axis=1)))
+        raise ValueError(
+            f"pre_dir has {bad} row(s) with non-finite (NaN/Inf) components; "
+            "local/inv_local_frame_rotation require a well-defined incoming "
+            "direction for every row."
+        )
     norm = np.linalg.norm(pre_dir, axis=1, keepdims=True)
     if np.any(norm < 1e-6):
         raise ValueError(
@@ -196,7 +232,7 @@ class Normalizer:
 
 
 class _WelfordAccumulator:
-    """Streaming mean/variance (Welford's online algorithm, batch update).
+    """Streaming mean/variance (Chan/Golub/LeVeque 1979 parallel algorithm).
 
     Use to fit a Normalizer over data that doesn't fit in memory:
         acc = _WelfordAccumulator(n_features)
@@ -211,13 +247,23 @@ class _WelfordAccumulator:
         self._M2 = np.zeros(n_features, dtype=np.float64)
 
     def update(self, X: np.ndarray) -> None:
+        # Computes the chunk's own local mean/M2 (two passes over X, no
+        # reference to the running mean) and merges it into the running
+        # totals with the O(F) Chan/Golub/LeVeque combination formula.
+        # Equivalent to the textbook single-pass streaming update (which
+        # instead re-derives two full (B, F) arrays from the running mean,
+        # before and after updating it) but ~40% cheaper here since it
+        # avoids one of those (B, F) passes and its temporary array.
         X = np.asarray(X, dtype=np.float64)
         B = X.shape[0]
+        mean_b = X.mean(0)
+        diff = X - mean_b
+        M2_b = np.einsum("ij,ij->j", diff, diff)
+
         new_n = self.n + B
-        delta = X - self._mean
-        self._mean += delta.sum(0) / new_n
-        delta2 = X - self._mean
-        self._M2 += (delta * delta2).sum(0)
+        delta = mean_b - self._mean
+        self._mean += delta * (B / new_n)
+        self._M2 += M2_b + delta * delta * (self.n * B / new_n)
         self.n = new_n
 
     def to_normalizer(self) -> "Normalizer":
@@ -276,6 +322,58 @@ class _ReservoirSampler:
         return self._reservoir.astype(np.float32)
 
 
+def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
+    """Boolean membership of `values` (any order) in `sorted_arr` (ascending, unique).
+
+    Equivalent to `np.isin(values, sorted_arr)`, but `np.isin`'s default path
+    sorts both inputs on every call — costly when `sorted_arr` is a large,
+    already-sorted array (e.g. all train-split event ids) reused across many
+    chunks. This does one `searchsorted` per call instead. `values` need not
+    be sorted; `sorted_arr` must be ascending and duplicate-free.
+    """
+    values = np.asarray(values)
+    if sorted_arr.size == 0:
+        return np.zeros(values.shape, dtype=bool)
+    idx = np.searchsorted(sorted_arr, values)
+    idx = np.clip(idx, 0, len(sorted_arr) - 1)
+    return sorted_arr[idx] == values
+
+
+def _vectorized_map_lookup(
+    values: np.ndarray, mapping: dict, strict: bool = True
+) -> np.ndarray:
+    """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.
+
+    Replaces a per-element Python dict lookup with one `searchsorted` call.
+    Raises `KeyError` if any value in `values` isn't a key of `mapping`,
+    matching the dict-comprehension it replaces (never silently misassigns)
+    — unless `strict=False`, in which case unmapped values get a dummy index
+    of 0 instead. Only pass `strict=False` where the caller has independently
+    verified the resulting index is never actually read (e.g.
+    `build_cond_features` under `conditioning="physical"`, where
+    `ConditionEncoder` ignores `cond_cat` entirely); it exists so a rollout
+    can be seeded with a species/material outside the training vocab without
+    a spurious `KeyError`, which is the entire point of physical-property
+    conditioning.
+    """
+    keys = np.asarray(list(mapping.keys()))
+    vals = np.asarray(list(mapping.values()), dtype=np.int64)
+    order = np.argsort(keys, kind="stable")
+    keys_sorted, vals_sorted = keys[order], vals[order]
+    values = np.asarray(values)
+    pos = np.searchsorted(keys_sorted, values)
+    pos = np.clip(pos, 0, len(keys_sorted) - 1)
+    found = keys_sorted[pos] == values
+    if not found.all():
+        if not strict:
+            out = np.zeros(values.shape, dtype=np.int64)
+            out[found] = vals_sorted[pos[found]]
+            return out
+        missing = np.unique(values[~found])
+        raise KeyError(f"value(s) not in mapping: {missing[:10].tolist()}")
+    return vals_sorted[pos]
+
+
 def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
     """World-frame unit vector pointing from pre_pos to post_pos.
 
@@ -338,6 +436,7 @@ def encode_secondaries(
     e_sec: np.ndarray,
     pre_dir: np.ndarray,
     sec_pdg_list: np.ndarray | None = None,
+    phys_only: bool = False,
 ) -> np.ndarray:
     """Encode per-secondary attributes into continuous per-slot targets.
 
@@ -359,38 +458,73 @@ def encode_secondaries(
     `sec_pdg_list` is optional so callers that only need the continuous
     stick/dir block (e.g. inference-time re-encoding) can omit it; omitting
     it zero-fills the last two columns, matching the padding-slot convention.
+
+    `phys_only=True` skips the stick-breaking and direction-rotation blocks
+    (zero-filling them instead) and computes only log_mass/charge — for
+    callers (normalizer fitting) that discard the other four columns anyway,
+    so computing them would be wasted work repeated over the whole dataset.
     """
     N, K = sec_E_list.shape
     e_sec = np.asarray(e_sec, dtype=np.float64)
 
-    stick_logits = np.zeros((N, K), dtype=np.float32)
-    for i in range(K):
-        if i == 0:
-            remaining = np.maximum(e_sec, _EPS)
-        else:
-            remaining = np.maximum(e_sec - sec_E_list[:, :i].sum(axis=1), _EPS)
-        f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
-        logit = np.log(f / (1.0 - f)).astype(np.float32)
-        # Last valid slot: give it the full remaining budget
-        is_last = sec_valid[:, i] & ~(
-            sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
-        )
-        logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
-        logit = np.where(
-            sec_valid[:, i], np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP), 0.0
-        )
-        stick_logits[:, i] = logit.astype(np.float32)
-
-    # Rotate each slot's direction into the local frame of the primary.
-    # pre_dir is broadcast across all K slots.
-    dir_local = np.zeros((N, K, 3), dtype=np.float32)
-    for i in range(K):
-        # Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
-        valid_mask = sec_valid[:, i]
-        if valid_mask.any():
-            dir_local[valid_mask, i] = local_frame_rotation(
-                pre_dir[valid_mask], sec_dir_list[valid_mask, i]
+    if phys_only:
+        stick_logits = np.zeros((N, K), dtype=np.float32)
+        dir_local = np.zeros((N, K, 3), dtype=np.float32)
+    else:
+        cumsum = np.cumsum(sec_E_list.astype(np.float64), axis=1)
+        stick_logits = np.zeros((N, K), dtype=np.float32)
+        # A valid slot whose cumulative secondary energy so far exceeds
+        # e_sec by more than float noise means sec_E_list sums to more than
+        # e_sec — a real upstream data mismatch, not something to paper
+        # over. Flagged once after the loop rather than let `remaining`'s
+        # np.maximum(..., _EPS) floor silently absorb it by saturating that
+        # slot's stick-breaking logit with no signal that anything was off.
+        _SHORTFALL_TOL = 1e-3
+        shortfall_flagged = np.zeros(N, dtype=bool)
+        for i in range(K):
+            if i == 0:
+                remaining_raw = e_sec
+            else:
+                remaining_raw = e_sec - cumsum[:, i - 1]
+            shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
+            remaining = np.maximum(remaining_raw, _EPS)
+            f = np.clip(
+                sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS
             )
+            logit = np.log(f / (1.0 - f)).astype(np.float32)
+            # Last valid slot: give it the full remaining budget
+            is_last = sec_valid[:, i] & ~(
+                sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
+            )
+            logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
+            logit = np.where(
+                sec_valid[:, i],
+                np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP),
+                0.0,
+            )
+            stick_logits[:, i] = logit.astype(np.float32)
+
+        if shortfall_flagged.any():
+            n = int(shortfall_flagged.sum())
+            warnings.warn(
+                f"encode_secondaries: {n}/{N} row(s) have sec_E_list summing "
+                "to more than e_sec (beyond float noise) — the overflowing "
+                "slot(s)' stick-breaking logit was saturated instead of "
+                "reflecting a real fraction; check upstream secondary "
+                "energy accounting for these rows.",
+                stacklevel=2,
+            )
+
+        # Rotate each slot's direction into the local frame of the primary.
+        # pre_dir is broadcast across all K slots.
+        dir_local = np.zeros((N, K, 3), dtype=np.float32)
+        for i in range(K):
+            # Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
+            valid_mask = sec_valid[:, i]
+            if valid_mask.any():
+                dir_local[valid_mask, i] = local_frame_rotation(
+                    pre_dir[valid_mask], sec_dir_list[valid_mask, i]
+                )
 
     if sec_pdg_list is not None:
         from giant.particles import particle_phys_array
@@ -573,8 +707,15 @@ def build_cond_features(
         [cond_cont, _physical_cond_columns(data, conditioning)]
     ).astype(np.float32)
 
-    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
-    mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
+    # In "physical" mode cond_cat is only a reporting/router convenience —
+    # ConditionEncoder never reads it (giant/model/network.py) — so a
+    # species/material outside the training vocab (the whole point of
+    # physical-property conditioning) gets a dummy index instead of raising.
+    # In "embedding" mode cond_cat IS the conditioning signal, so an unmapped
+    # value must still raise loudly rather than silently misassign.
+    strict = conditioning == "embedding"
+    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict)
+    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict)
     cond_cat = np.column_stack([pdg_idx, mat_idx])
 
     if cond_normalizer is not None:
@@ -627,6 +768,7 @@ def build_features(
     proc_map: dict[str, int] | None = None,
     require_secondaries: bool = False,
     conditioning: str = "embedding",
+    sec_phys_only: bool = False,
 ) -> tuple[
     np.ndarray,
     np.ndarray,
@@ -653,6 +795,11 @@ def build_features(
         per-secondary list columns are absent (a mis-converted file that would
         otherwise silently zero all Stage-2 targets). Training paths set this;
         Stage-1-only callers (e.g. `giant predict`) leave it False.
+
+    sec_phys_only: passed straight through to `encode_secondaries` — skips
+        the stick-breaking/direction-rotation blocks of `sec_cont` (zero-filled
+        instead) for callers (normalizer fitting) that only read
+        `sec_cont[:, :, 4:6]` and would otherwise discard that work.
     """
     from giant.constants import K_MAX
 
@@ -687,8 +834,8 @@ def build_features(
         [cond_cont, _physical_cond_columns(data, conditioning)]
     ).astype(np.float32)  # (N, COND_DIM=15)
 
-    pdg_idx = np.array([pdg_map[int(p)] for p in data["pdg"]], dtype=np.int64)
-    mat_idx = np.array([mat_map[str(m)] for m in data["material"]], dtype=np.int64)
+    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map)
+    mat_idx = _vectorized_map_lookup(data["material"], mat_map)
     cond_cat = np.column_stack([pdg_idx, mat_idx])  # (N, 2)
 
     n_sec_raw = data["n_sec"].astype(
@@ -715,6 +862,7 @@ def build_features(
             data["e_sec"],
             data["pre_dir"],
             sec_pdg_list=sec_pdg_list,
+            phys_only=sec_phys_only,
         )  # (N, K_MAX, 6)
     else:
         # Guard against silently training Stage 2 on zeroed targets: if any step
@@ -755,7 +903,7 @@ def build_features(
 
     process = data.get("process")
     if proc_map is not None and process is not None:
-        proc_idx = np.array([proc_map[str(p)] for p in process], dtype=np.int64)
+        proc_idx = _vectorized_map_lookup(process, proc_map)
     else:
         proc_idx = np.zeros(len(cond_cat), dtype=np.int64)
 

giant/model/network.py

@@ -537,11 +537,49 @@ class Router(nn.Module):
     def __init__(self, n_experts: int) -> None:
         super().__init__()
         self.n_experts = n_experts
+        # Opt-in straight-through Gumbel-softmax combine weights (see
+        # combine_weights below) — off by default, set from model.router.gumbel
+        # by _build_router_from_cfg. gumbel_tau is annealed per training step
+        # by giant.train (model.router.gumbel_tau_start/_end); neither is an
+        # nn.Parameter/buffer since neither is learned or needs checkpointing —
+        # the tau schedule is deterministic in global_step, so it recomputes
+        # correctly on resume.
+        self.gumbel = False
+        self.gumbel_tau = 1.0
 
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """(B, n_experts) soft weights, rows summing to 1."""
         raise NotImplementedError
 
+    def combine_weights(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+    ) -> torch.Tensor:
+        """(B, n_experts) train-time expert-combination weights.
+
+        Default (`gumbel=False`): identical to `gate()` — the original dense
+        soft-mixture combination. Opt-in straight-through Gumbel-softmax
+        (`gumbel=True`, train mode only): samples a Gumbel-perturbed
+        categorical draw from the same distribution `gate()` defines
+        (`log(gate())` is a valid unnormalized-logit input to
+        `F.gumbel_softmax` since softmax is shift-invariant, so no subclass
+        needs to expose separate pre-softmax logits), then hardens it to a
+        one-hot vector on the forward pass while keeping the soft sample's
+        gradient on the backward pass. This makes the training-time forward
+        combination match eval-time top-1 dispatch exactly (one expert's
+        output, unweighted) instead of the smooth blend `gate()` gives —
+        intended to close the train/eval mismatch identified as a likely
+        cause of experts overlapping instead of partitioning (see the
+        router_gating write-up referenced in CLAUDE.md's roadmap).
+        `gate()` itself is untouched and still backs `balance_loss`/
+        `entropy_loss`/`gate_stats`, so those diagnostics keep reading the
+        smooth distribution rather than a noisy sample.
+        """
+        probs = self.gate(cond_cont, cond_cat)
+        if not (self.gumbel and self.training):
+            return probs
+        log_probs = torch.log(probs.clamp_min(1e-8))
+        return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)
+
     def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         """(B,) hard expert index, used for eval-time grouped dispatch."""
         return self.gate(cond_cont, cond_cat).argmax(dim=-1)
@@ -566,6 +604,49 @@ class Router(nn.Module):
         """
         return torch.zeros((), device=cond_cont.device)
 
+    def entropy_loss(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+    ) -> torch.Tensor:
+        """Optional auxiliary loss rewarding sharper (lower-entropy) routing.
+
+        Reuses `gate_stats`'s `norm_entropy` (already in [0, 1], 1.0 =
+        uniform/collapsed) directly as the loss, so minimizing it pushes
+        every router's gate toward decisiveness. A generic base-class
+        default — works for any Router via gate_stats, no per-subclass
+        override needed. Off by default (see `lambda_entropy` in
+        giant.train): bounded width/temperature (EnergyRouter's
+        `learn_width`/`learn_temperature`) is the primary defense against
+        gate collapse; this is a secondary, use-with-caution lever, since
+        indiscriminately penalizing entropy can also suppress legitimate
+        soft ambiguity near a router's own decision boundary.
+        """
+        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
+        return norm_entropy
+
+    def gate_stats(
+        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Diagnostics for catching a router that fails to specialize.
+
+        Returns `(norm_entropy, importance)`:
+          - `norm_entropy`: scalar, the batch-mean of each row's gate entropy
+            divided by `log(n_experts)`, in [0, 1] and comparable across
+            routers with different `n_experts` (1.0 = uniform/collapsed
+            gating, 0.0 = fully hard routing).
+          - `importance`: (n_experts,) tensor, `gate(...).sum(dim=0)` — the
+            *unnormalized* per-expert weight mass for this batch. Callers
+            wanting a global utilization share across many batches must sum
+            this across batches first and normalize once at the end;
+            averaging per-batch shares instead would treat every batch as
+            equally important regardless of size and understate a
+            rarely-but-fully-used expert.
+        """
+        gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
+        row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
+        norm_entropy = row_entropy.mean() / math.log(self.n_experts)
+        importance = gate.sum(dim=0)  # (n_experts,)
+        return norm_entropy, importance
+
 
 ROUTER_REGISTRY: dict[str, type[Router]] = {}
 
@@ -596,6 +677,22 @@ def build_router(name: str, n_experts: int, **kwargs) -> Router:
     return cls(n_experts=n_experts, **filtered)
 
 
+def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
+    """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient
+    bound (unlike `clamp`, which zeroes gradient past the boundary) used for
+    EnergyRouter's `learn_width`/`learn_temperature` modes."""
+    return lo + (hi - lo) * torch.sigmoid(raw)
+
+
+def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
+    """Inverse of `_bounded_interp`, used once at construction to warm-start
+    `raw` so `_bounded_interp(raw, lo, hi) == value` — lets `learn_width`/
+    `learn_temperature` start out exactly reproducing the fixed-`temperature`
+    gate before any training moves them."""
+    p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
+    return math.log(p / (1 - p))
+
+
 @register_router("energy")
 class EnergyRouter(Router):
     """Soft turn-on gate over normalized pre-step log-energy.
@@ -609,6 +706,28 @@ class EnergyRouter(Router):
     `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, differentiable in e; as
     tau -> 0 this hardens to nearest-center (Voronoi) selection, which is
     exactly what `top1` uses at eval.
+
+    `temperature` is normally a single fixed scalar shared by every expert.
+    Two mutually exclusive optional modes generalize it:
+      - `learn_width`: each expert gets its own learnable width, so
+        `gate(e) = softmax_i(-(e - c_i)^2 / width_i)` — experts can learn
+        independently how much of the energy axis they cover.
+      - `learn_temperature`: the single shared `temperature` itself becomes
+        learnable (still one scalar for every expert).
+    Both parameterize their raw learnable value through a sigmoid bounded
+    into `[width_min_ratio, width_max_ratio] * temperature` (see
+    `_bounded_interp`), warm-started so the initial effective width/
+    temperature exactly equals `temperature` — enabling either mode is a
+    no-op at init. The bound is deliberately not raw `softplus`/`exp`
+    (unbounded above): an unbounded width lets one expert's width run away
+    to infinity, making its logit `-d2/width -> 0` almost everywhere so it
+    wins nearly every row regardless of true distance to its center — the
+    same "experts overlap instead of partitioning" failure this whole
+    router design is trying to avoid, just via a new mechanism. See
+    `Router.entropy_loss`/`giant.train`'s `lambda_entropy` for a secondary,
+    optional guard against all experts' widths co-inflating together
+    (which bounding caps but doesn't forbid, and which the load-balance
+    loss alone can't see since usage shares stay even throughout).
     """
 
     def __init__(
@@ -618,10 +737,31 @@ class EnergyRouter(Router):
         learn_centers: bool = True,
         energy_idx: int = 3,
         centers_init: Sequence[float] | None = None,
+        learn_width: bool = False,
+        learn_temperature: bool = False,
+        width_min_ratio: float = 0.1,
+        width_max_ratio: float = 10.0,
     ) -> None:
         super().__init__(n_experts)
+        if learn_width and learn_temperature:
+            raise ValueError("learn_width and learn_temperature are mutually exclusive")
         self.temperature = temperature
         self.energy_idx = energy_idx
+        self.learn_width = learn_width
+        self.learn_temperature = learn_temperature
+        if learn_width or learn_temperature:
+            if not (width_min_ratio < 1.0 < width_max_ratio):
+                raise ValueError(
+                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio "
+                    f"({width_max_ratio}) must bracket 1.0"
+                )
+            self._width_lo = width_min_ratio * temperature
+            self._width_hi = width_max_ratio * temperature
+            raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
+            if learn_width:
+                self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
+            else:
+                self.raw_temperature = nn.Parameter(torch.tensor(raw0))
         if centers_init is None:
             centers = torch.linspace(-2.0, 2.0, n_experts)
         else:
@@ -636,10 +776,20 @@ class EnergyRouter(Router):
         else:
             self.register_buffer("centers", centers)
 
+    def effective_width(self) -> torch.Tensor | float:
+        """Softmax denominator used by `gate()`: a fixed scalar `temperature`
+        (default), a per-expert `(n_experts,)` bounded width (`learn_width`),
+        or a single bounded learnable scalar (`learn_temperature`)."""
+        if self.learn_width:
+            return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
+        if self.learn_temperature:
+            return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
+        return self.temperature
+
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         e = cond_cont[:, self.energy_idx].unsqueeze(-1)  # (B, 1)
         d2 = (e - self.centers.unsqueeze(0)) ** 2  # (B, n_experts)
-        return torch.softmax(-d2 / self.temperature, dim=-1)
+        return torch.softmax(-d2 / self.effective_width(), dim=-1)
 
 
 @register_router("pdg")
@@ -858,13 +1008,17 @@ def _route_forward(
 ) -> torch.Tensor:
     """Shared dispatch for both Routed* trunks.
 
-    Train mode: full soft mixture `sum_i gate_i * expert_i(x)` — fully
-    differentiable, N-expert compute. Eval mode: grouped top-1 dispatch —
-    each row runs exactly one (small) expert, which is the actual source
-    of the per-call speedup this architecture is for.
+    Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
+    N-expert dense compute, fully differentiable. `weight` is
+    `router.combine_weights(...)`: the plain soft `gate()` by default, or (see
+    `Router.combine_weights`) a straight-through Gumbel-softmax one-hot sample
+    when `router.gumbel` is enabled — either way, no change to the compute
+    cost of this branch. Eval mode: grouped top-1 dispatch — each row runs
+    exactly one (small) expert, which is the actual source of the per-call
+    speedup this architecture is for.
     """
     if training:
-        weights = router.gate(cond_cont, cond_cat)  # (B, n_experts)
+        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts)
         out = torch.zeros_like(x)
         for i, expert in enumerate(experts):
             out = out + weights[:, i : i + 1] * expert(x, cond)
@@ -1074,15 +1228,62 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]:
     return [axes[i] for i in range(len(axes))]
 
 
-def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router:
+# Router types that read cond_cat's pdg index through their own
+# nn.Embedding(pdg_vocab, ...), regardless of the trunk's `conditioning`
+# mode — see _check_router_conditioning_compat.
+_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
+
+
+def _check_router_conditioning_compat(
+    router_types: list[str], conditioning: str
+) -> None:
+    """Reject a router axis that reintroduces a training-vocab PDG lookup
+    under `conditioning="physical"`.
+
+    `PdgRouter`/`ProcessRouter` always build their own dataset-scoped
+    `nn.Embedding(pdg_vocab, ...)` (network.py's PdgRouter/ProcessRouter),
+    independent of `ConditionEncoder`'s `conditioning` mode. Pairing either
+    with `conditioning="physical"` would silently reintroduce a
+    training-menu-scoped lookup at the routing layer — defeating the entire
+    point of physical-property conditioning, which is to generalize to a
+    species/material outside that menu (see giant/rollout.py's
+    `build_cond_features(strict=...)` gate for the same concern on the
+    trunk side). Raised loudly at model-build time rather than left to
+    surface as a confusing rollout/generalization-benchmark result.
+    """
+    bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
+    if bad and conditioning == "physical":
+        raise ValueError(
+            f"router type(s) {bad} always use a training-vocab PDG embedding, "
+            "which is incompatible with conditioning='physical' (whose whole "
+            "point is generalizing beyond that vocab) — pick a different "
+            "router type (e.g. 'energy') or use conditioning='embedding'."
+        )
+
+
+def _build_router_from_cfg(
+    router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding"
+) -> Router:
     """Resolve one `model.router` config into a `Router`, single-axis or composed.
 
     `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys
     (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair.
+
+    `gumbel` is set as a post-construction attribute here rather than a
+    per-subclass constructor kwarg, same reasoning as `lambda_balance`/
+    `lambda_proc`/`lambda_entropy` living in `router_cfg` without being a
+    `Router` subclass constructor param: it's a training-time toggle shared by
+    every router type, not a per-type hyperparameter (`build_router`'s
+    kwarg-filtering would otherwise just silently drop it).
     """
     shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
     if router_cfg["type"] == "composed":
-        return build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab)
+        axes = _parse_composed_axes(router_cfg)
+        _check_router_conditioning_compat([a["type"] for a in axes], conditioning)
+        router = build_composed_router(axes, **shared_vocab)
+        router.gumbel = bool(router_cfg.get("gumbel", False))
+        return router
+    _check_router_conditioning_compat([router_cfg["type"]], conditioning)
     router_kwargs = {
         k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")
     }
@@ -1092,7 +1293,9 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) ->
     # vocab, same as the trunk's ConditionEncoder.
     router_kwargs.setdefault("pdg_vocab", pdg_vocab)
     router_kwargs.setdefault("mat_vocab", mat_vocab)
-    return build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router.gumbel = bool(router_cfg.get("gumbel", False))
+    return router
 
 
 def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
@@ -1130,19 +1333,26 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
         shared = dict(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
-            expert_hidden_dim=model_config.get("expert_hidden_dim", 128),
-            expert_n_blocks=model_config.get("expert_n_blocks", 3),
+            expert_hidden_dim=model_config.get("expert_hidden_dim")
+            or model_config.get("hidden_dim", 128),
+            expert_n_blocks=model_config.get("expert_n_blocks")
+            or model_config.get("n_blocks", 3),
             emb_dim=model_config.get("emb_dim", EMB_DIM),
             dropout=model_config.get("dropout", 0.1),
             conditioning=model_config.get("conditioning", "embedding"),
         )
+        conditioning = shared["conditioning"]
         stage1 = RoutedDenoisingMLP(
-            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
+            router=_build_router_from_cfg(
+                router_cfg, pdg_vocab, mat_vocab, conditioning
+            ),
             k_max=model_config.get("k_max", K_MAX),
             **shared,
         )
         sec_decoder = RoutedSecondaryDecoder(
-            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab),
+            router=_build_router_from_cfg(
+                router_cfg, pdg_vocab, mat_vocab, conditioning
+            ),
             **shared,
         )
         return stage1, sec_decoder

giant/model/wgan.py

@@ -17,7 +17,11 @@ def gradient_penalty(
     is for Stage 2's variable-length slot vector: both the interpolate and the
     critic's gradient are zeroed on padded dims first, so the norm target of 1
     is only ever asked of genuine content, not the padding convention shared
-    by both `real` and `fake`.
+    by both `real` and `fake`. Rows fully masked out (e.g. `n_sec == 0`, so
+    every slot is padding) have no real content to constrain the gradient
+    norm to 1 — `x_hat`/`grad` are forced to all-zero for such a row, which
+    would otherwise contribute a constant `(||0|| - 1)^2 == 1` bias to the
+    mean regardless of critic behavior — so they're excluded from the mean.
     """
     eps = torch.rand(real.size(0), 1, device=real.device)
     x_hat = eps * real + (1 - eps) * fake
@@ -28,7 +32,11 @@ def gradient_penalty(
     grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0]
     if mask is not None:
         grad = grad * mask
-    return ((grad.norm(2, dim=1) - 1) ** 2).mean()
+    penalty = (grad.norm(2, dim=1) - 1) ** 2
+    if mask is not None:
+        valid = (mask.sum(dim=1) > 0).float()
+        return (penalty * valid).sum() / valid.sum().clamp_min(1.0)
+    return penalty.mean()
 
 
 def critic_loss(

giant/pipeline.py

@@ -1,3 +1,5 @@
+import os
+from dataclasses import dataclass
 from pathlib import Path
 
 import numpy as np
@@ -13,19 +15,231 @@ from giant.constants import (
     SEC_SLOT_DIM,
     X_DIM,
 )
+from giant.data import setup_cache
 from giant.data.loader import (
+    event_id_offset,
     find_parquet_files,
-    load_event_ids,
     iter_file_chunks,
     build_index_maps_from_files,
     build_process_map_from_files,
 )
-from giant.data.transforms import build_features, _WelfordAccumulator, _ReservoirSampler
+from giant.data.transforms import (
+    Normalizer,
+    build_features,
+    _WelfordAccumulator,
+    _ReservoirSampler,
+    sorted_membership,
+)
 from giant.data.dataset import make_event_split, StreamingStepsDataset
 from giant.model.network import build_models, build_critics
 from giant.train import train as run_training
 
 
+@dataclass
+class SetupStageResult:
+    """Everything `run_train_job`'s pre-epoch setup stage derives from `data`.
+
+    Also returned standalone by `run_setup_stage` for callers (e.g. `dwarf
+    warm-cache`) that only want to populate/refresh the setup cache sidecar
+    without actually training a model.
+    """
+
+    files: list[Path]
+    pdg_map: dict[int, int]
+    mat_map: dict[str, int]
+    proc_map: dict[str, int] | None
+    cond_norm: Normalizer
+    tgt_norm: Normalizer
+    sec_phys_norm: Normalizer
+    train_events: set
+    val_events: set
+    n_train_steps: int
+
+
+def run_setup_stage(
+    data: str | Path,
+    val_fraction: float,
+    seed: int,
+    conditioning: str,
+    router_cfg: dict,
+    cache_setup: bool = True,
+    rebuild_setup_cache: bool = False,
+    echo=print,
+) -> SetupStageResult:
+    """Scan `data` for everything training needs before the epoch loop: the
+    train/val event split, pdg/material vocab maps, an optional process map
+    (`router_cfg.type == "process"`), and the Stage-1/Stage-2 normalizers.
+
+    Reads from and writes to the `giant.data.setup_cache` sidecar when
+    `cache_setup` is set (`rebuild_setup_cache` ignores — but still
+    refreshes — any existing sidecar content). `router_cfg` may be mutated
+    in place: an active `EnergyRouter` (`router_cfg["type"] == "energy"`)
+    gets its `centers_init` seeded from real data quantiles here.
+    """
+    files = find_parquet_files(data)
+    echo(f"found {len(files)} parquet file(s)")
+
+    cache: setup_cache.SetupCache | None = None
+    if cache_setup:
+        if rebuild_setup_cache:
+            echo("setup cache: --rebuild-setup-cache given, recomputing all sections")
+            loaded = None
+        else:
+            loaded = setup_cache.load(data, files, echo=echo)
+        cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)
+
+    if cache is not None and cache.event_index is not None:
+        unique_ids, counts = cache.event_index
+        echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
+    else:
+        echo("scanning event IDs …")
+        unique_ids, counts = setup_cache.compute_event_index_from_files(files)
+        if cache is not None:
+            cache.event_index = (unique_ids, counts)
+
+    train_events, val_events = make_event_split(
+        unique_ids, val_fraction=val_fraction, seed=seed
+    )
+    events_arr = np.array(sorted(train_events))
+    n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
+    echo(
+        f"  {int(counts.sum()):,} steps  |  "
+        f"{len(train_events)} train events  |  "
+        f"{len(val_events)} val events"
+    )
+
+    if cache is not None and cache.vocab is not None:
+        pdg_map, mat_map = cache.vocab
+        echo(
+            f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
+            f"{len(mat_map)} materials)"
+        )
+    else:
+        echo("building vocabulary maps …")
+        pdg_map, mat_map = build_index_maps_from_files(files)
+        echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
+        if cache is not None:
+            cache.vocab = (pdg_map, mat_map)
+
+    proc_map: dict[str, int] | None = None
+    if router_cfg.get("enabled") and router_cfg.get("type") == "process":
+        n_experts = router_cfg["n_experts"]
+        cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
+        if cached_proc_map is not None:
+            proc_map = cached_proc_map
+            echo(
+                f"process vocabulary: cache hit ({len(proc_map)} labels, "
+                f"{n_experts} experts)"
+            )
+        else:
+            echo("building process vocabulary …")
+            proc_map = build_process_map_from_files(files, n_experts=n_experts)
+            echo(f"  {len(proc_map)} process labels mapped to {n_experts} experts")
+            if cache is not None:
+                cache.proc_maps[n_experts] = proc_map
+
+    energy_router_active = (
+        router_cfg.get("enabled") and router_cfg.get("type") == "energy"
+    )
+    energy_idx = router_cfg.get("energy_idx", 3)
+    norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning)
+    entry = cache.normalizers.get(norm_key) if cache is not None else None
+
+    if entry is not None:
+        echo(f"normalizer: cache hit (key={norm_key!r})")
+        cond_norm = entry.cond_norm
+        tgt_norm = entry.tgt_norm
+        sec_phys_norm = entry.sec_phys_norm
+        energy_quantiles = entry.energy_quantiles
+    else:
+        echo("fitting normalizer (streaming) …")
+        cond_acc = _WelfordAccumulator(COND_DIM)
+        tgt_acc = _WelfordAccumulator(X_DIM)
+        sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
+        # EnergyRouter's default center spread (linspace over [-2, 2]) assumes
+        # the z-normalized energy column is roughly uniform, which real energy
+        # spectra rarely are — collect a reservoir sample here (reusing this
+        # same pass, not a second scan), then collapse it to a fixed quantile
+        # grid (setup_cache.energy_quantiles_from_sample) so centers can
+        # instead be seeded from actual data quantiles below. Collected
+        # whenever the setup cache is being populated, not only when *this*
+        # run's router is energy-typed, so a later run enabling
+        # --router-type energy against this same (val_fraction, seed,
+        # conditioning) key never needs to rescan just to seed centers.
+        collect_energy_sample = energy_router_active or cache is not None
+        energy_sampler = (
+            _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
+        )
+        for i, path in enumerate(files):
+            for chunk in iter_file_chunks(path, offset=event_id_offset(i)):
+                mask = sorted_membership(chunk["event_id"], events_arr)
+                if not mask.any():
+                    continue
+                chunk_tr = {k: v[mask] for k, v in chunk.items()}
+                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
+                    chunk_tr,
+                    pdg_map,
+                    mat_map,
+                    proc_map=proc_map,
+                    require_secondaries=True,
+                    conditioning=conditioning,
+                    sec_phys_only=True,
+                )
+                cond_acc.update(cond_cont)
+                tgt_acc.update(target_s1)
+                if energy_sampler is not None:
+                    energy_sampler.update(cond_cont[:, energy_idx])
+                sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
+                sec_phys = sec_cont[:, :, 4:6][sec_valid]
+                if len(sec_phys) > 0:
+                    sec_phys_acc.update(sec_phys)
+        cond_norm = cond_acc.to_normalizer()
+        tgt_norm = tgt_acc.to_normalizer()
+        sec_phys_norm = sec_phys_acc.to_normalizer()
+        energy_quantiles = (
+            setup_cache.energy_quantiles_from_sample(energy_sampler.sample)
+            if energy_sampler is not None
+            else np.empty(0, dtype=np.float32)
+        )
+        if cache is not None:
+            cache.normalizers[norm_key] = setup_cache.NormalizerEntry(
+                cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles
+            )
+
+    if energy_router_active and energy_quantiles.size > 0:
+        assert cond_norm.mean is not None and cond_norm.std is not None
+        levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
+        raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
+        centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
+            energy_idx
+        ]
+        router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
+        echo(
+            f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
+        )
+    elif energy_router_active:
+        echo(
+            "  warning: no energy samples collected — EnergyRouter falls back to "
+            "default centers"
+        )
+
+    if cache is not None:
+        setup_cache.save(data, files, cache, echo=echo)
+
+    return SetupStageResult(
+        files=files,
+        pdg_map=pdg_map,
+        mat_map=mat_map,
+        proc_map=proc_map,
+        cond_norm=cond_norm,
+        tgt_norm=tgt_norm,
+        sec_phys_norm=sec_phys_norm,
+        train_events=train_events,
+        val_events=val_events,
+        n_train_steps=n_train_steps,
+    )
+
+
 def run_train_job(
     data: Path,
     cfg: dict,
@@ -34,6 +248,8 @@ def run_train_job(
     shuffle_buffer: int,
     num_workers: int,
     resume: Path | None = None,
+    cache_setup: bool = True,
+    rebuild_setup_cache: bool = False,
     echo=print,
 ) -> None:
     t, m = cfg["train"], cfg["model"]
@@ -41,26 +257,19 @@ def run_train_job(
 
     out_dir = Path(out_dir)
 
-    files = find_parquet_files(data)
-    echo(f"found {len(files)} parquet file(s)")
-
-    echo("scanning event IDs …")
-    all_event_ids = np.concatenate([load_event_ids(f) for f in files])
-    train_events, val_events = make_event_split(
-        all_event_ids, val_fraction=t["val_fraction"]
-    )
-    events_arr = np.array(sorted(train_events))
-    n_train_steps = int(np.isin(all_event_ids, events_arr).sum())
-    total_train_batches = n_train_steps // t["batch_size"]
-    echo(
-        f"  {len(all_event_ids):,} steps  |  "
-        f"{len(train_events)} train events (~{n_train_steps:,} steps, ~{total_train_batches:,} batches)  |  "
-        f"{len(val_events)} val events"
-    )
-
-    echo("building vocabulary maps …")
-    pdg_map, mat_map = build_index_maps_from_files(files)
-    echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
+    # Soft warning (never blocks) — CLAUDE.md's Compute environment section
+    # asks that shared portal machines (portal1/deepthought{,2}/bms{1..3})
+    # stay within ~1/4 of CPU/RAM so as not to disturb other users' jobs;
+    # DataLoader's num_workers has no awareness of that on its own.
+    cpu_count = os.cpu_count() or 1
+    quota = max(1, cpu_count // 4)
+    if num_workers > quota:
+        echo(
+            f"warning: --num-workers={num_workers} exceeds ~1/4 of this "
+            f"machine's {cpu_count} CPU(s) ({quota}) — portal machines are "
+            "shared with other users (see CLAUDE.md's Compute environment "
+            "section)"
+        )
 
     router_cfg = m["router"]
     if t["mode"] == "wgan" and router_cfg.get("enabled"):
@@ -68,70 +277,33 @@ def run_train_job(
             "--mode wgan does not support --router (no routed WGAN generator/"
             "critic exists) — disable one or the other"
         )
-    proc_map: dict[str, int] | None = None
-    if router_cfg.get("enabled") and router_cfg.get("type") == "process":
-        echo("building process vocabulary …")
-        proc_map = build_process_map_from_files(
-            files, n_experts=router_cfg["n_experts"]
-        )
-        echo(
-            f"  {len(proc_map)} process labels mapped to {router_cfg['n_experts']} experts"
-        )
 
-    echo("fitting normalizer (streaming) …")
     conditioning = m["conditioning"]
-    cond_acc = _WelfordAccumulator(COND_DIM)
-    tgt_acc = _WelfordAccumulator(X_DIM)
-    sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
-    # EnergyRouter's default center spread (linspace over [-2, 2]) assumes
-    # the z-normalized energy column is roughly uniform, which real energy
-    # spectra rarely are — collect a reservoir sample here (reusing this
-    # same pass, not a second scan) so centers can instead be seeded from
-    # actual data quantiles below.
-    energy_router_active = (
-        router_cfg.get("enabled") and router_cfg.get("type") == "energy"
+    setup = run_setup_stage(
+        data,
+        val_fraction=t["val_fraction"],
+        seed=t["seed"],
+        conditioning=conditioning,
+        router_cfg=router_cfg,
+        cache_setup=cache_setup,
+        rebuild_setup_cache=rebuild_setup_cache,
+        echo=echo,
     )
-    energy_idx = router_cfg.get("energy_idx", 3)
-    energy_sampler = (
-        _ReservoirSampler(capacity=100_000) if energy_router_active else None
+    files = setup.files
+    pdg_map, mat_map, proc_map = setup.pdg_map, setup.mat_map, setup.proc_map
+    cond_norm, tgt_norm, sec_phys_norm = (
+        setup.cond_norm,
+        setup.tgt_norm,
+        setup.sec_phys_norm,
+    )
+    train_events, val_events, n_train_steps = (
+        setup.train_events,
+        setup.val_events,
+        setup.n_train_steps,
     )
-    for path in files:
-        for chunk in iter_file_chunks(path):
-            mask = np.isin(chunk["event_id"], events_arr)
-            if not mask.any():
-                continue
-            chunk_tr = {k: v[mask] for k, v in chunk.items()}
-            cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
-                chunk_tr,
-                pdg_map,
-                mat_map,
-                proc_map=proc_map,
-                require_secondaries=True,
-                conditioning=conditioning,
-            )
-            cond_acc.update(cond_cont)
-            tgt_acc.update(target_s1)
-            if energy_sampler is not None:
-                energy_sampler.update(cond_cont[:, energy_idx])
-            sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
-            sec_phys = sec_cont[:, :, 4:6][sec_valid]
-            if len(sec_phys) > 0:
-                sec_phys_acc.update(sec_phys)
-    cond_norm = cond_acc.to_normalizer()
-    tgt_norm = tgt_acc.to_normalizer()
-    sec_phys_norm = sec_phys_acc.to_normalizer()
 
-    if energy_sampler is not None and energy_sampler.n_seen > 0:
-        assert cond_norm.mean is not None and cond_norm.std is not None
-        normalized_sample = (
-            energy_sampler.sample - cond_norm.mean[energy_idx]
-        ) / cond_norm.std[energy_idx]
-        quantiles = np.linspace(0.0, 1.0, router_cfg["n_experts"])
-        centers_init = np.quantile(normalized_sample, quantiles).astype(np.float32)
-        router_cfg["centers_init"] = centers_init.tolist()
-        echo(
-            f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
-        )
+    total_train_batches = n_train_steps // t["batch_size"]
+    echo(f"  ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches")
 
     train_ds = StreamingStepsDataset(
         files=files,
@@ -176,6 +348,25 @@ def run_train_job(
     )
 
     emb_dim = m.get("emb_dim", EMB_DIM)
+    expert_hidden_dim, expert_n_blocks = config.resolve_expert_dims(
+        router_cfg, m["hidden_dim"], m["n_blocks"]
+    )
+    if router_cfg.get("enabled") and (expert_hidden_dim, expert_n_blocks) != (
+        m["hidden_dim"],
+        m["n_blocks"],
+    ):
+        # Only reachable via an explicit router.expert_hidden_dim/n_blocks
+        # override (the 0/"unset" sentinel always resolves to m["hidden_dim"]/
+        # ["n_blocks"] — see resolve_expert_dims), so this is never a false
+        # positive from inheritance, only a deliberate narrow/wide-experts
+        # config the checkpoint dir name (_h{hidden_dim}_b{n_blocks}) won't
+        # reflect.
+        echo(
+            f"  warning: experts are {expert_hidden_dim}x{expert_n_blocks}, "
+            f"different from model.hidden_dim/n_blocks ({m['hidden_dim']}x"
+            f"{m['n_blocks']}) — the checkpoint dir name reflects the latter, "
+            "not the experts actually being trained"
+        )
 
     model_config = {
         "pdg_vocab": len(pdg_map),
@@ -188,8 +379,8 @@ def run_train_job(
         "sec_slot_dim": SEC_SLOT_DIM,
         "conditioning": conditioning,
         "router": dict(router_cfg),
-        "expert_hidden_dim": router_cfg["expert_hidden_dim"],
-        "expert_n_blocks": router_cfg["expert_n_blocks"],
+        "expert_hidden_dim": expert_hidden_dim,
+        "expert_n_blocks": expert_n_blocks,
         # Read by `predict`/`rollout` (which never receive their own --mode
         # flag) to auto-detect which sampler a checkpoint needs.
         "mode": t["mode"],
@@ -240,6 +431,9 @@ def run_train_job(
         lambda_s2=t.get("lambda_s2", 1.0),
         lambda_balance=router_cfg.get("lambda_balance", 0.0),
         lambda_proc=router_cfg.get("lambda_proc", 0.0),
+        lambda_entropy=router_cfg.get("lambda_entropy", 0.0),
+        gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0),
+        gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1),
         normalizer_dict={
             "cond": cond_norm.to_dict(),
             "target": tgt_norm.to_dict(),
@@ -259,4 +453,8 @@ def run_train_job(
         n_critic=t.get("n_critic", 5),
         gp_weight=t.get("gp_weight", 10.0),
         critic_lr=t.get("critic_lr") or None,
+        use_wandb=t.get("wandb", False),
+        wandb_project=t.get("wandb_project", "giant"),
+        wandb_run_name=t.get("wandb_run_name", ""),
+        wandb_log_every=t.get("wandb_log_every", 50),
     )

giant/rollout.py

@@ -407,7 +407,16 @@ def _step_chunk(
     tr["_material"] = material
     tr["_layer_id"] = layer_id
 
-    known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool)
+    if conditioning == "physical":
+        # Under physical-property conditioning, mass/charge (already resolved
+        # on every track — see the cond_dict comment below) drive the model,
+        # not a training-vocab PDG embedding — build_cond_features passes
+        # strict=False for exactly this mode, so an out-of-vocab species no
+        # longer raises. Terminating on it here would defeat the entire
+        # point of physical conditioning: generalizing to a held-out species.
+        known_pdg = np.ones(n, dtype=bool)
+    else:
+        known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool)
 
     # --- Pre-step termination gates (in priority order; each track picks one) ---
     stop = np.zeros(n, dtype=bool)

giant/train.py

@@ -32,6 +32,8 @@ _METRICS_FIELDS = [
     "train_loss_s2",
     "train_loss_balance",
     "train_loss_proc",
+    "train_loss_entropy",
+    "train_nsec_acc",
     "d_loss",
     "g_loss",
     "wasserstein_estimate",
@@ -42,9 +44,25 @@ _METRICS_FIELDS = [
     "val_loss_s2",
     "val_loss_balance",
     "val_loss_proc",
+    "val_loss_entropy",
+    "val_nsec_acc",
     "val_marginal_kl",
+    "router_s1_entropy",
+    "router_s1_util_min",
+    "router_s1_util_max",
+    "router_s1_util_std",
+    "router_s2_entropy",
+    "router_s2_util_min",
+    "router_s2_util_max",
+    "router_s2_util_std",
     "lr",
+    "critic_lr",
     "grad_norm",
+    "grad_norm_d",
+    "grad_norm_g",
+    "gpu_mem_mb",
+    "samples_per_sec",
+    "is_best",
     "epoch_time_s",
 ]
 
@@ -96,6 +114,85 @@ def _update_ema(
         ema_p.mul_(decay).add_(p, alpha=1 - decay)
 
 
+def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -> float:
+    """Linear anneal of the straight-through Gumbel-softmax temperature.
+
+    Deterministic in `step`/`total_steps` alone (no extra state), so it
+    recomputes correctly on `--resume` from a checkpoint's saved `global_step`
+    without needing to persist anything new (see
+    giant.model.network.Router.combine_weights).
+    """
+    progress = min(step / max(total_steps, 1), 1.0)
+    return tau_start + (tau_end - tau_start) * progress
+
+
+def _wandb_run_config(
+    *,
+    mode: str,
+    epochs: int,
+    lr: float,
+    warmup_epochs: int,
+    weight_decay: float,
+    ema_decay: float,
+    lambda_nsec: float,
+    lambda_s2: float,
+    lambda_balance: float,
+    lambda_proc: float,
+    lambda_entropy: float,
+    gumbel_tau_start: float,
+    gumbel_tau_end: float,
+    n_critic: int,
+    gp_weight: float,
+    model_config: dict | None,
+    stage1_params: int,
+    sec_decoder_params: int,
+    critic_params: int,
+    sec_critic_params: int,
+    total_params: int,
+) -> dict:
+    """Build the dict logged as a wandb run's `config`.
+
+    Router-only knobs (`lambda_balance`/`lambda_proc`/`lambda_entropy`/
+    `gumbel_tau_start`/`gumbel_tau_end`) and WGAN-only knobs (`n_critic`/
+    `gp_weight`) are omitted unless actually active, so a run's wandb config
+    doesn't imply hyperparameters from an inactive code path (a disabled
+    router's fine-tuning knobs, or GAN critic settings for a flow/DDPM run).
+    The full `model_config` (including its `router` sub-dict, whatever the
+    router type/state) is always included, so no information is lost — this
+    only trims the flattened top-level convenience duplicates.
+    """
+    router_enabled = bool((model_config or {}).get("router", {}).get("enabled", False))
+    cfg = {
+        "mode": mode,
+        "epochs": epochs,
+        "lr": lr,
+        "warmup_epochs": warmup_epochs,
+        "weight_decay": weight_decay,
+        "ema_decay": ema_decay,
+        "lambda_nsec": lambda_nsec,
+        "lambda_s2": lambda_s2,
+        "model": model_config or {},
+        "stage1_params": stage1_params,
+        "sec_decoder_params": sec_decoder_params,
+        "critic_params": critic_params,
+        "sec_critic_params": sec_critic_params,
+        "total_params": total_params,
+    }
+    if router_enabled:
+        cfg.update(
+            {
+                "lambda_balance": lambda_balance,
+                "lambda_proc": lambda_proc,
+                "lambda_entropy": lambda_entropy,
+                "gumbel_tau_start": gumbel_tau_start,
+                "gumbel_tau_end": gumbel_tau_end,
+            }
+        )
+    if mode == "wgan":
+        cfg.update({"n_critic": n_critic, "gp_weight": gp_weight})
+    return cfg
+
+
 def _compute_losses(
     stage1_model: torch.nn.Module,
     sec_decoder: torch.nn.Module,
@@ -107,10 +204,18 @@ def _compute_losses(
     lambda_s2: float,
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
+    lambda_entropy: float = 0.0,
 ) -> tuple[
-    torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
 ]:
-    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc) for one batch."""
+    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, L_entropy, nsec_acc) for one batch."""
     cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch
     cond_cont = cond_cont.to(device)
     cond_cat = cond_cat.to(device)
@@ -129,6 +234,7 @@ def _compute_losses(
     # n_sec classification loss
     n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat)
     l_nsec = F.cross_entropy(n_sec_logits, n_sec)
+    nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean()
 
     # Stage-2 secondary flow loss
     # Use a noiseless Stage-1 target as context (detach to avoid back-prop
@@ -162,16 +268,25 @@ def _compute_losses(
         l_proc = stage1_model.router.classify_loss(
             cond_cont, cond_cat, proc_idx
         ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx)
+        # Optional entropy-regularization aux loss (see Router.entropy_loss):
+        # penalizes uniform/collapsed gating, a secondary guard against
+        # gate-sharpness collapse that lambda_balance alone can't see.
+        l_entropy = stage1_model.router.entropy_loss(
+            cond_cont, cond_cat
+        ) + sec_decoder.router.entropy_loss(cond_cont, cond_cat)
     else:
         l_balance = torch.zeros((), device=device)
         l_proc = torch.zeros((), device=device)
+        l_entropy = torch.zeros((), device=device)
 
     total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
     if lambda_balance > 0:
         total = total + lambda_balance * l_balance
     if lambda_proc > 0:
         total = total + lambda_proc * l_proc
-    return total, l_s1, l_nsec, l_s2, l_balance, l_proc
+    if lambda_entropy > 0:
+        total = total + lambda_entropy * l_entropy
+    return total, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc
 
 
 def _wgan_train_step(
@@ -263,7 +378,9 @@ def _wgan_train_step(
 
     # --- Generator (+ n_sec) step ---
     did_g_step = step_count % n_critic == 0
-    l_nsec = F.cross_entropy(generator.predict_n_sec(cond_cont, cond_cat), n_sec)
+    n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat)
+    l_nsec = F.cross_entropy(n_sec_logits, n_sec)
+    nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean()
     optimizer_g.zero_grad()
     if did_g_step:
         g1 = generator_loss(critic_fn1, fake1)
@@ -283,8 +400,11 @@ def _wgan_train_step(
         "wasserstein_estimate": wasserstein_estimate,
         "gp_loss": (gp1 + lambda_s2 * gp2).detach(),
         "l_nsec": l_nsec.detach(),
+        "nsec_acc": nsec_acc.detach(),
         "did_g_step": did_g_step,
         "grad_norm": grad_norm_d.item() + grad_norm_g.item(),
+        "grad_norm_d": grad_norm_d.item(),
+        "grad_norm_g": grad_norm_g.item(),
     }
 
 
@@ -305,6 +425,9 @@ def train(
     lambda_s2: float = 1.0,
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
+    lambda_entropy: float = 0.0,
+    gumbel_tau_start: float = 1.0,
+    gumbel_tau_end: float = 0.1,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -320,10 +443,67 @@ def train(
     n_critic: int = 5,
     gp_weight: float = 10.0,
     critic_lr: float | None = None,
+    use_wandb: bool = False,
+    wandb_project: str = "giant",
+    wandb_run_name: str = "",
+    wandb_log_every: int = 50,
 ) -> None:
     out_dir = Path(out_dir)
     out_dir.mkdir(parents=True, exist_ok=True)
 
+    stage1_params = sum(p.numel() for p in stage1_model.parameters())
+    sec_decoder_params = sum(p.numel() for p in sec_decoder.parameters())
+    critic_params = (
+        sum(p.numel() for p in critic.parameters()) if critic is not None else 0
+    )
+    sec_critic_params = (
+        sum(p.numel() for p in sec_critic.parameters()) if sec_critic is not None else 0
+    )
+    total_params = (
+        stage1_params + sec_decoder_params + critic_params + sec_critic_params
+    )
+
+    wandb_run = None
+    if use_wandb:
+        try:
+            import wandb
+        except ImportError as exc:
+            raise RuntimeError(
+                "train.wandb = true (--wandb) requires the 'wandb' package — "
+                "install it via `uv sync --extra wandb`"
+            ) from exc
+        # `id` is derived from out_dir so resuming a run (--resume) reattaches
+        # to the same wandb run instead of starting a new one.
+        wandb_run = wandb.init(
+            project=wandb_project,
+            name=wandb_run_name or out_dir.name,
+            id=out_dir.name,
+            resume="allow",
+            config=_wandb_run_config(
+                mode=mode,
+                epochs=epochs,
+                lr=lr,
+                warmup_epochs=warmup_epochs,
+                weight_decay=weight_decay,
+                ema_decay=ema_decay,
+                lambda_nsec=lambda_nsec,
+                lambda_s2=lambda_s2,
+                lambda_balance=lambda_balance,
+                lambda_proc=lambda_proc,
+                lambda_entropy=lambda_entropy,
+                gumbel_tau_start=gumbel_tau_start,
+                gumbel_tau_end=gumbel_tau_end,
+                n_critic=n_critic,
+                gp_weight=gp_weight,
+                model_config=model_config,
+                stage1_params=stage1_params,
+                sec_decoder_params=sec_decoder_params,
+                critic_params=critic_params,
+                sec_critic_params=sec_critic_params,
+                total_params=total_params,
+            ),
+        )
+
     stage1_model = stage1_model.to(device)
     sec_decoder = sec_decoder.to(device)
     if mode == "wgan":
@@ -333,6 +513,13 @@ def train(
         critic = critic.to(device)
         sec_critic = sec_critic.to(device)
 
+    # MoE routing trunk (RoutedDenoisingMLP/RoutedSecondaryDecoder) is
+    # optional and orthogonal to `mode` — both stages carry a `.router`
+    # when enabled. Each router is an independent instance (their
+    # `n_experts` need not match), used both for the batch-level gate
+    # entropy snapshot below and the val-level gate stats further down.
+    has_router = hasattr(stage1_model, "router") and hasattr(sec_decoder, "router")
+
     # Flow-matching/diffusion models sample noticeably better from an EMA of
     # the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed
     # sinusoidal-embedding freqs, or non-learned router centers) never change
@@ -377,7 +564,9 @@ def train(
     # would never finish and cosine decay would barely move.
     steps_per_epoch = max(total_train_batches, 1)
     if mode == "wgan":
-        steps_per_epoch = max(total_train_batches // (n_critic + 1), 1)
+        # Generator steps fire every n_critic-th batch (did_g_step =
+        # step_count % n_critic == 0 in _wgan_train_step), not n_critic + 1.
+        steps_per_epoch = max(total_train_batches // n_critic, 1)
     warmup_steps = warmup_epochs * steps_per_epoch
     total_steps = max(epochs * steps_per_epoch, 1)
 
@@ -392,8 +581,44 @@ def train(
 
     ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None
 
+    def _build_checkpoint(epoch: int, global_step: int, best_val_loss: float) -> dict:
+        ckpt: dict = {
+            "model": stage1_model.state_dict(),
+            "sec_decoder": sec_decoder.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "lr_sched": lr_sched.state_dict(),
+            "epoch": epoch,
+            "best_val_loss": best_val_loss,
+            "global_step": global_step,
+        }
+        if mode == "wgan":
+            assert (
+                critic is not None
+                and sec_critic is not None
+                and optimizer_d is not None
+            )
+            ckpt["critic"] = critic.state_dict()
+            ckpt["sec_critic"] = sec_critic.state_dict()
+            ckpt["optimizer_d"] = optimizer_d.state_dict()
+        if ema_decay > 0:
+            assert ema_stage1_model is not None and ema_sec_decoder is not None
+            ckpt["model_ema"] = ema_stage1_model.state_dict()
+            ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict()
+        if normalizer_dict is not None:
+            ckpt["normalizer"] = normalizer_dict
+        if pdg_map is not None:
+            ckpt["pdg_map"] = pdg_map
+        if mat_map is not None:
+            ckpt["mat_map"] = mat_map
+        if proc_map is not None:
+            ckpt["proc_map"] = proc_map
+        if model_config is not None:
+            ckpt["model_config"] = model_config
+        return ckpt
+
     start_epoch = 1
     best_val_loss = float("inf")
+    resumed_global_step = 0
     if resume_path is not None:
         ckpt = torch.load(resume_path, map_location=device, weights_only=False)
         stage1_model.load_state_dict(ckpt["model"])
@@ -413,10 +638,20 @@ def train(
             critic.load_state_dict(ckpt["critic"])
             sec_critic.load_state_dict(ckpt["sec_critic"])
             optimizer_d.load_state_dict(ckpt["optimizer_d"])
+            # Mirrors the `lr` fixup below for the generator optimizer:
+            # optimizer_d.load_state_dict() above restores the checkpoint's
+            # own critic LR, which would otherwise silently override an
+            # explicit `--critic-lr` passed on this resume. optimizer_d has
+            # no LR scheduler (unlike `optimizer`/`lr_sched`), so this is a
+            # flat set rather than a schedule-relative one.
+            resumed_critic_lr = critic_lr if critic_lr is not None else lr
+            for group in optimizer_d.param_groups:
+                group["lr"] = resumed_critic_lr
         optimizer.load_state_dict(ckpt["optimizer"])
         lr_sched.load_state_dict(ckpt["lr_sched"])
         start_epoch = ckpt.get("epoch", 0) + 1
         best_val_loss = ckpt.get("best_val_loss", float("inf"))
+        resumed_global_step = ckpt.get("global_step", 0)
 
         # optimizer/lr_sched.load_state_dict() above restore the checkpoint's
         # own base LR, which would otherwise silently override an explicit
@@ -446,10 +681,16 @@ def train(
 
     epoch_w = len(str(epochs))
     last_completed_epoch = start_epoch - 1
-    global_step = 0
+    # Restored from the checkpoint on --resume so wandb_run.log(..., step=...)
+    # keeps advancing monotonically instead of restarting at 0 mid-run (a
+    # reattached wandb run — see wandb.init(id=..., resume="allow") below —
+    # would otherwise silently drop every post-resume point).
+    global_step = resumed_global_step
     with _GracefulShutdown() as shutdown:
         for epoch in range(start_epoch, epochs + 1):
             epoch_start = time.monotonic()
+            if device.type == "cuda":
+                torch.cuda.reset_peak_memory_stats(device)
             stage1_model.train()
             sec_decoder.train()
             if mode == "wgan":
@@ -462,10 +703,14 @@ def train(
             train_s2_sum = 0.0
             train_balance_sum = 0.0
             train_proc_sum = 0.0
+            train_entropy_sum = 0.0
             train_d_sum = 0.0
             train_g_sum = 0.0
             train_wasserstein_sum = 0.0
             train_gp_sum = 0.0
+            train_nsec_acc_sum = 0.0
+            train_grad_norm_d_sum = 0.0
+            train_grad_norm_g_sum = 0.0
             train_n = 0
             train_batches = 0
             grad_norm_sum = 0.0
@@ -480,6 +725,13 @@ def train(
                 dynamic_ncols=True,
             )
             for batch in bar:
+                if has_router:
+                    gumbel_tau = _gumbel_tau(
+                        global_step, total_steps, gumbel_tau_start, gumbel_tau_end
+                    )
+                    stage1_model.router.gumbel_tau = gumbel_tau
+                    sec_decoder.router.gumbel_tau = gumbel_tau
+
                 if mode == "wgan":
                     assert (
                         critic is not None
@@ -503,7 +755,6 @@ def train(
                         lambda_nsec,
                         lambda_s2,
                     )
-                    global_step += 1
                     if stats["did_g_step"]:
                         lr_sched.step()
                         if ema_decay > 0:
@@ -523,18 +774,24 @@ def train(
                     train_g_sum += stats["g_loss"].item() * B
                     train_wasserstein_sum += stats["wasserstein_estimate"].item() * B
                     train_gp_sum += stats["gp_loss"].item() * B
+                    train_nsec_acc_sum += stats["nsec_acc"].item() * B
+                    train_grad_norm_d_sum += stats["grad_norm_d"] * B
+                    train_grad_norm_g_sum += stats["grad_norm_g"] * B
                 else:
-                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
-                        stage1_model,
-                        sec_decoder,
-                        batch,
-                        mode,
-                        ddpm_schedule,
-                        device,
-                        lambda_nsec,
-                        lambda_s2,
-                        lambda_balance,
-                        lambda_proc,
+                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc = (
+                        _compute_losses(
+                            stage1_model,
+                            sec_decoder,
+                            batch,
+                            mode,
+                            ddpm_schedule,
+                            device,
+                            lambda_nsec,
+                            lambda_s2,
+                            lambda_balance,
+                            lambda_proc,
+                            lambda_entropy,
+                        )
                     )
                     optimizer.zero_grad()
                     loss.backward()
@@ -557,6 +814,8 @@ def train(
                     train_s2_sum += l_s2.item() * B
                     train_balance_sum += l_balance.item() * B
                     train_proc_sum += l_proc.item() * B
+                    train_entropy_sum += l_entropy.item() * B
+                    train_nsec_acc_sum += nsec_acc.item() * B
 
                 train_n += B
                 train_batches += 1
@@ -573,16 +832,80 @@ def train(
                     f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False
                 )
 
+                global_step += 1
+                if (
+                    wandb_run is not None
+                    and wandb_log_every > 0
+                    and global_step % wandb_log_every == 0
+                ):
+                    log_payload = {
+                        "batch/epoch": epoch,
+                        "batch/loss": batch_loss,
+                        "batch/loss_ema": ema_loss,
+                        "batch/grad_norm": batch_grad_norm,
+                        "batch/lr": optimizer.param_groups[0]["lr"],
+                        "batch/critic_lr": (
+                            optimizer_d.param_groups[0]["lr"]
+                            if optimizer_d is not None
+                            else 0.0
+                        ),
+                    }
+                    if has_router:
+                        # Cheap re-use of the batch already in hand — no
+                        # extra data loading, just a small forward through
+                        # each router's own gate function. Only entropy is
+                        # logged at this granularity (not per-expert
+                        # utilization): a single batch's importance sum is
+                        # too noisy as a "global share" estimate, whereas
+                        # the val-loop aggregate (below) sums over the
+                        # whole val set for that. Batch-level entropy alone
+                        # is still enough to see a router collapsing in
+                        # real time, mid-epoch, rather than only at the
+                        # next validation pass.
+                        with torch.no_grad():
+                            cond_cont_b = batch[0].to(device)
+                            cond_cat_b = batch[1].to(device)
+                            s1_entropy, _ = stage1_model.router.gate_stats(
+                                cond_cont_b, cond_cat_b
+                            )
+                            s2_entropy, _ = sec_decoder.router.gate_stats(
+                                cond_cont_b, cond_cat_b
+                            )
+                        log_payload["batch/router_s1_entropy"] = s1_entropy.item()
+                        log_payload["batch/router_s2_entropy"] = s2_entropy.item()
+                        log_payload["batch/gumbel_tau"] = gumbel_tau
+                    wandb_run.log(log_payload, step=global_step)
+
                 if shutdown.requested:
                     break
             bar.close()
 
             if shutdown.requested:
+                # Epoch was interrupted mid-loop, so there's no val_loss to
+                # weigh a "best" checkpoint against — save the in-progress
+                # weights as last.pt only, under the last *fully completed*
+                # epoch number so --resume restarts this epoch from scratch
+                # rather than skipping it (weights/optimizer state are still
+                # kept, so those partial-epoch batches aren't wasted work).
+                ckpt = _build_checkpoint(epoch - 1, global_step, best_val_loss)
+                torch.save(ckpt, out_dir / "last.pt")
+                last_completed_epoch = epoch - 1
+                print(
+                    f"saved in-progress weights from partway through epoch "
+                    f"{epoch} to {out_dir / 'last.pt'} "
+                    f"(resume will restart epoch {epoch})"
+                )
                 break
 
             train_loss = train_loss_sum / max(train_n, 1)
             train_grad_norm = grad_norm_sum / max(train_batches, 1)
+            train_nsec_acc = train_nsec_acc_sum / max(train_n, 1)
+            train_grad_norm_d = train_grad_norm_d_sum / max(train_n, 1)
+            train_grad_norm_g = train_grad_norm_g_sum / max(train_n, 1)
             current_lr = optimizer.param_groups[0]["lr"]
+            critic_lr_value = (
+                optimizer_d.param_groups[0]["lr"] if optimizer_d is not None else 0.0
+            )
 
             stage1_model.eval()
             sec_decoder.eval()
@@ -618,8 +941,12 @@ def train(
                 val_loss = val_marginal_kl
                 val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = (
                     val_proc_sum
-                ) = 0.0
+                ) = val_entropy_sum = val_nsec_acc_sum = 0.0
                 val_n = 1
+                val_nsec_acc = 0.0
+                router_s1_entropy = router_s2_entropy = 0.0
+                router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0
+                router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0
             else:
                 val_loss_sum = 0.0
                 val_s1_sum = 0.0
@@ -627,12 +954,34 @@ def train(
                 val_s2_sum = 0.0
                 val_balance_sum = 0.0
                 val_proc_sum = 0.0
+                val_entropy_sum = 0.0
+                val_nsec_acc_sum = 0.0
                 val_n = 0
+                if has_router:
+                    n_experts_s1 = stage1_model.router.n_experts
+                    n_experts_s2 = sec_decoder.router.n_experts
+                    val_router_s1_entropy_sum = 0.0
+                    val_router_s2_entropy_sum = 0.0
+                    val_router_s1_importance_sum = torch.zeros(
+                        n_experts_s1, device=device
+                    )
+                    val_router_s2_importance_sum = torch.zeros(
+                        n_experts_s2, device=device
+                    )
                 with torch.no_grad():
                     for val_batch_idx, batch in enumerate(val_loader):
                         if max_val_batches > 0 and val_batch_idx >= max_val_batches:
                             break
-                        loss, l_s1, l_nsec, l_s2, l_balance, l_proc = _compute_losses(
+                        (
+                            loss,
+                            l_s1,
+                            l_nsec,
+                            l_s2,
+                            l_balance,
+                            l_proc,
+                            l_entropy,
+                            nsec_acc,
+                        ) = _compute_losses(
                             stage1_model,
                             sec_decoder,
                             batch,
@@ -643,6 +992,7 @@ def train(
                             lambda_s2,
                             lambda_balance,
                             lambda_proc,
+                            lambda_entropy,
                         )
                         B = batch[0].size(0)
                         val_loss_sum += loss.item() * B
@@ -651,8 +1001,48 @@ def train(
                         val_s2_sum += l_s2.item() * B
                         val_balance_sum += l_balance.item() * B
                         val_proc_sum += l_proc.item() * B
+                        val_entropy_sum += l_entropy.item() * B
+                        val_nsec_acc_sum += nsec_acc.item() * B
+                        if has_router:
+                            cond_cont = batch[0].to(device)
+                            cond_cat = batch[1].to(device)
+                            s1_entropy, s1_importance = stage1_model.router.gate_stats(
+                                cond_cont, cond_cat
+                            )
+                            s2_entropy, s2_importance = sec_decoder.router.gate_stats(
+                                cond_cont, cond_cat
+                            )
+                            val_router_s1_entropy_sum += s1_entropy.item() * B
+                            val_router_s2_entropy_sum += s2_entropy.item() * B
+                            val_router_s1_importance_sum += s1_importance
+                            val_router_s2_importance_sum += s2_importance
                         val_n += B
                 val_loss = val_loss_sum / max(val_n, 1)
+                val_nsec_acc = val_nsec_acc_sum / max(val_n, 1)
+
+                if has_router:
+                    router_s1_entropy = val_router_s1_entropy_sum / max(val_n, 1)
+                    router_s2_entropy = val_router_s2_entropy_sum / max(val_n, 1)
+                    s1_util = val_router_s1_importance_sum / (
+                        val_router_s1_importance_sum.sum().clamp_min(1e-8)
+                    )
+                    s2_util = val_router_s2_importance_sum / (
+                        val_router_s2_importance_sum.sum().clamp_min(1e-8)
+                    )
+                    router_s1_util_min = s1_util.min().item()
+                    router_s1_util_max = s1_util.max().item()
+                    router_s1_util_std = (
+                        s1_util.std().item() if n_experts_s1 > 1 else 0.0
+                    )
+                    router_s2_util_min = s2_util.min().item()
+                    router_s2_util_max = s2_util.max().item()
+                    router_s2_util_std = (
+                        s2_util.std().item() if n_experts_s2 > 1 else 0.0
+                    )
+                else:
+                    router_s1_entropy = router_s2_entropy = 0.0
+                    router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0
+                    router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0
 
                 if validate_every > 0 and epoch % validate_every == 0:
                     print(f"[epoch {epoch}] marginal validation:")
@@ -668,6 +1058,11 @@ def train(
                     val_marginal_kl = float(np.mean(marginal_result["kl_divergence"]))
 
             epoch_time = time.monotonic() - epoch_start
+            gpu_mem_mb = (
+                torch.cuda.max_memory_allocated(device) / (1024 * 1024)
+                if device.type == "cuda"
+                else 0.0
+            )
 
             is_best = val_loss < best_val_loss
             marker = "  [best]" if is_best else ""
@@ -679,70 +1074,64 @@ def train(
                 f"  s2={train_s2_sum / max(train_n, 1):.3f}"
                 f"  bal={train_balance_sum / max(train_n, 1):.3f}"
                 f"  proc={train_proc_sum / max(train_n, 1):.3f}"
+                f"  entropy={train_entropy_sum / max(train_n, 1):.3f}"
                 f"  d={train_d_sum / max(train_n, 1):.3f}"
                 f"  g={train_g_sum / max(train_n, 1):.3f})"
                 f"  val {val_loss:.4f}"
                 f"  lr {current_lr:.2e}  gnorm {train_grad_norm:.3f}"
                 f"  {epoch_time:.1f}s{marker}"
             )
-            metrics_writer.writerow(
-                {
-                    "epoch": epoch,
-                    "train_loss": train_loss,
-                    "train_loss_s1": train_s1_sum / max(train_n, 1),
-                    "train_loss_nsec": train_nsec_sum / max(train_n, 1),
-                    "train_loss_s2": train_s2_sum / max(train_n, 1),
-                    "train_loss_balance": train_balance_sum / max(train_n, 1),
-                    "train_loss_proc": train_proc_sum / max(train_n, 1),
-                    "d_loss": train_d_sum / max(train_n, 1),
-                    "g_loss": train_g_sum / max(train_n, 1),
-                    "wasserstein_estimate": train_wasserstein_sum / max(train_n, 1),
-                    "gp_loss": train_gp_sum / max(train_n, 1),
-                    "val_loss": val_loss,
-                    "val_loss_s1": val_s1_sum / max(val_n, 1),
-                    "val_loss_nsec": val_nsec_sum / max(val_n, 1),
-                    "val_loss_s2": val_s2_sum / max(val_n, 1),
-                    "val_loss_balance": val_balance_sum / max(val_n, 1),
-                    "val_loss_proc": val_proc_sum / max(val_n, 1),
-                    "val_marginal_kl": val_marginal_kl,
-                    "lr": current_lr,
-                    "grad_norm": train_grad_norm,
-                    "epoch_time_s": epoch_time,
-                }
-            )
-            metrics_file.flush()
-
-            ckpt: dict = {
-                "model": stage1_model.state_dict(),
-                "sec_decoder": sec_decoder.state_dict(),
-                "optimizer": optimizer.state_dict(),
-                "lr_sched": lr_sched.state_dict(),
+            metrics_row = {
                 "epoch": epoch,
-                "best_val_loss": best_val_loss,
+                "train_loss": train_loss,
+                "train_loss_s1": train_s1_sum / max(train_n, 1),
+                "train_loss_nsec": train_nsec_sum / max(train_n, 1),
+                "train_loss_s2": train_s2_sum / max(train_n, 1),
+                "train_loss_balance": train_balance_sum / max(train_n, 1),
+                "train_loss_proc": train_proc_sum / max(train_n, 1),
+                "train_loss_entropy": train_entropy_sum / max(train_n, 1),
+                "train_nsec_acc": train_nsec_acc,
+                "d_loss": train_d_sum / max(train_n, 1),
+                "g_loss": train_g_sum / max(train_n, 1),
+                "wasserstein_estimate": train_wasserstein_sum / max(train_n, 1),
+                "gp_loss": train_gp_sum / max(train_n, 1),
+                "val_loss": val_loss,
+                "val_loss_s1": val_s1_sum / max(val_n, 1),
+                "val_loss_nsec": val_nsec_sum / max(val_n, 1),
+                "val_loss_s2": val_s2_sum / max(val_n, 1),
+                "val_loss_balance": val_balance_sum / max(val_n, 1),
+                "val_loss_proc": val_proc_sum / max(val_n, 1),
+                "val_loss_entropy": val_entropy_sum / max(val_n, 1),
+                "val_nsec_acc": val_nsec_acc,
+                "val_marginal_kl": val_marginal_kl,
+                "router_s1_entropy": router_s1_entropy,
+                "router_s1_util_min": router_s1_util_min,
+                "router_s1_util_max": router_s1_util_max,
+                "router_s1_util_std": router_s1_util_std,
+                "router_s2_entropy": router_s2_entropy,
+                "router_s2_util_min": router_s2_util_min,
+                "router_s2_util_max": router_s2_util_max,
+                "router_s2_util_std": router_s2_util_std,
+                "lr": current_lr,
+                "critic_lr": critic_lr_value,
+                "grad_norm": train_grad_norm,
+                "grad_norm_d": train_grad_norm_d,
+                "grad_norm_g": train_grad_norm_g,
+                "gpu_mem_mb": gpu_mem_mb,
+                "samples_per_sec": train_n / max(epoch_time, 1e-8),
+                "is_best": int(is_best),
+                "epoch_time_s": epoch_time,
             }
-            if mode == "wgan":
-                assert (
-                    critic is not None
-                    and sec_critic is not None
-                    and optimizer_d is not None
-                )
-                ckpt["critic"] = critic.state_dict()
-                ckpt["sec_critic"] = sec_critic.state_dict()
-                ckpt["optimizer_d"] = optimizer_d.state_dict()
-            if ema_decay > 0:
-                assert ema_stage1_model is not None and ema_sec_decoder is not None
-                ckpt["model_ema"] = ema_stage1_model.state_dict()
-                ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict()
-            if normalizer_dict is not None:
-                ckpt["normalizer"] = normalizer_dict
-            if pdg_map is not None:
-                ckpt["pdg_map"] = pdg_map
-            if mat_map is not None:
-                ckpt["mat_map"] = mat_map
-            if proc_map is not None:
-                ckpt["proc_map"] = proc_map
-            if model_config is not None:
-                ckpt["model_config"] = model_config
+            metrics_writer.writerow(metrics_row)
+            metrics_file.flush()
+            if wandb_run is not None:
+                # Shares the same monotonic step axis as the per-batch
+                # `batch/*` logs above (global_step) rather than `epoch`,
+                # since a wandb run's `step` argument across `log()` calls
+                # must never decrease.
+                wandb_run.log(metrics_row, step=global_step)
+
+            ckpt = _build_checkpoint(epoch, global_step, best_val_loss)
 
             if val_loss < best_val_loss:
                 best_val_loss = val_loss
@@ -756,6 +1145,8 @@ def train(
                 break
 
     metrics_file.close()
+    if wandb_run is not None:
+        wandb_run.finish()
 
     if shutdown.requested:
         print(

giant/validate.py

@@ -215,6 +215,8 @@ def validate_marginals(
     print("-" * 68)
     for j, name in enumerate(_SEC_PHYS_NAMES):
         r, g = phys_real[:, j], phys_gen[:, j]
+        if len(r) == 0 or len(g) == 0:
+            continue
         print(
             f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
             f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.1.0"
+version = "0.2.0"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -25,11 +25,14 @@ dev = [
     "pytest>=8,<10",
     "ruff>=0.15,<1",
     "ty>=0.0.50,<0.1",
-    "giant[convert,analysis,geometry]",
+    "giant[convert,analysis,geometry,wandb]",
 ]
 geometry = [
     "scikit-learn>=1.4,<2",
 ]
+wandb = [
+    "wandb>=0.16,<1",
+]
 convert = [
     "uproot>=5.3,<6",
     "awkward>=2.6,<3",

scripts/bump_dataset_version.py

@@ -85,7 +85,12 @@ def _git_user_name() -> str | None:
         out = subprocess.run(
             ["git", "config", "user.name"], capture_output=True, text=True, timeout=2
         )
-    except OSError:
+    except (OSError, subprocess.SubprocessError):
+        # OSError (e.g. git not on PATH) and subprocess.SubprocessError
+        # (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
+        # is not an OSError, so catching only OSError (as before) let a
+        # slow/loaded NFS-backed portal machine crash this instead of
+        # degrading to by=None as intended.
         return None
     name = out.stdout.strip()
     return name or None
@@ -730,6 +735,7 @@ def run_create_manifest(
     pool: str | None = None,
     type_: str | None = None,
     root: str = "/ceph/lbogner/geant_steps",
+    force: bool = False,
 ) -> None:
     if (output is None) == (pool is None):
         raise SystemExit("error: exactly one of --output or --pool is required")
@@ -745,6 +751,12 @@ def run_create_manifest(
     parquet_files = [Path(f) for f in files]
     lines, missing, resolved = plan_create_manifest(output_path, parquet_files)
     overlaps = check_holdout_overlap(output_path, resolved)
+    # Unlike missing/overlaps this is a hard stop even without --execute
+    # reaching the write, since create_manifest has no in-place "update" mode
+    # (unlike update_manifest) — a second run against the same output_path
+    # (e.g. holdout.manifest, the file check_holdout_overlap exists to
+    # protect) would otherwise silently clobber it with no diff/backup.
+    already_exists = output_path.exists() and not force
 
     print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
     print(f"manifest: {output_path.resolve()}")
@@ -761,7 +773,10 @@ def run_create_manifest(
         for name, f in overlaps:
             print(f"  {f}  (also in {name})")
 
-    if (missing or overlaps) and execute:
+    if already_exists:
+        print(f"\n{output_path} already exists — pass --force to overwrite it.")
+
+    if (missing or overlaps or already_exists) and execute:
         raise SystemExit("error: refusing to write manifest (see above)")
 
     if not execute:

scripts/dwarf.py

@@ -5,6 +5,7 @@ simulation-fanout tools into one Typer app so there's a single command name
 (and `--help`) to remember instead of five differently-hyphenated ones.
 """
 
+import os
 from enum import Enum
 from pathlib import Path
 from typing import Optional
@@ -12,6 +13,7 @@ from typing import Optional
 import typer
 from typing_extensions import Annotated
 
+from giant.config import Conditioning
 from scripts.bump_dataset_version import (
     run_bump_gen,
     run_bump_schema,
@@ -25,6 +27,7 @@ from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
 from scripts.migrate_geant_steps import run_migration
 from scripts.steps_to_parquet import convert_steps_to_parquet
 from scripts.steps_to_parquet_parallel import run_parallel_job
+from scripts.warm_setup_cache import run_warm_setup_cache
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -36,6 +39,25 @@ def _main() -> None:
     """dwarf — dataset/tooling CLI (ROOT<->parquet conversion, dataset versioning, sim fanout)."""
 
 
+def _warn_if_exceeds_shared_quota(n: int, flag: str) -> None:
+    """Soft warning (never blocks) when a worker/job count looks likely to
+    grab more than this repo's documented shared-portal-machine etiquette
+    (CLAUDE.md's Compute environment: stay within ~1/4 of CPU/RAM and a
+    single GPU, since portal1/deepthought{,2}/bms{1..3} are shared with
+    other users). Not a hard cap — a legitimate big machine or a
+    deliberately aggressive run is still the caller's call.
+    """
+    cpu_count = os.cpu_count() or 1
+    quota = max(1, cpu_count // 4)
+    if n > quota:
+        typer.echo(
+            f"warning: {flag}={n} exceeds ~1/4 of this machine's "
+            f"{cpu_count} CPU(s) ({quota}) — portal machines are shared "
+            "with other users (see CLAUDE.md's Compute environment section)",
+            err=True,
+        )
+
+
 class Compression(str, Enum):
     snappy = "snappy"
     lz4 = "lz4"
@@ -105,6 +127,7 @@ def convert(
     if jobs < 1:
         typer.echo("error: --jobs must be >= 1", err=True)
         raise typer.Exit(1)
+    _warn_if_exceeds_shared_quota(jobs, "--jobs")
 
     compression_value = (
         "uncompressed" if compression is Compression.none else compression.value
@@ -326,6 +349,10 @@ def create_manifest(
     execute: Annotated[
         bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
     ] = False,
+    force: Annotated[
+        bool,
+        typer.Option("--force", help="Overwrite the manifest if it already exists"),
+    ] = False,
 ) -> None:
     """Create a new manifest from a list of parquet files."""
     run_create_manifest(
@@ -335,6 +362,7 @@ def create_manifest(
         pool=pool,
         type_=type_.value if type_ is not None else None,
         root=str(root),
+        force=force,
     )
 
 
@@ -390,6 +418,7 @@ def make_root(
     ] = False,
 ) -> None:
     """Generate new ROOT shards via a minicalosim executable."""
+    _warn_if_exceeds_shared_quota(jobs, "--jobs")
     run_make_root(
         executable=executable,
         detector=detector,
@@ -471,6 +500,75 @@ def build_geometry_oracle(
     )
 
 
+@app.command("warm-cache")
+def warm_cache(
+    data: Annotated[
+        Path,
+        typer.Argument(
+            help="Parquet file, directory, or .manifest — same as `giant train`'s"
+        ),
+    ],
+    val_fraction: Annotated[
+        float,
+        typer.Option(
+            "--val-fraction",
+            "-f",
+            help="Must match the `giant train` run(s) to warm for",
+        ),
+    ] = 0.1,
+    seed: Annotated[
+        int,
+        typer.Option(
+            "--seed", "-s", help="Must match the `giant train` run(s) to warm for"
+        ),
+    ] = 0,
+    conditioning: Annotated[
+        Conditioning,
+        typer.Option(
+            "--conditioning", help="Must match the `giant train` run(s) to warm for"
+        ),
+    ] = Conditioning.physical,
+    router: Annotated[
+        bool,
+        typer.Option(
+            "--router/--no-router",
+            help="Warm the process vocabulary too (only takes effect with "
+            "--router-type process)",
+        ),
+    ] = False,
+    router_type: Annotated[
+        str, typer.Option("--router-type", help="Router implementation name")
+    ] = "energy",
+    n_experts: Annotated[
+        int, typer.Option("--n-experts", help="Number of routed experts")
+    ] = 4,
+    rebuild: Annotated[
+        bool,
+        typer.Option(
+            "--rebuild", help="Ignore any existing sidecar and recompute every section"
+        ),
+    ] = False,
+) -> None:
+    """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
+
+    Warms the vocab maps, event-id split index, and the normalizer entry for
+    the given --val-fraction/--seed/--conditioning, so a later `giant train`
+    run (or a `dwarf hparam-scan` sweep, which shares one such entry across
+    every run) skips straight to training. See giant/data/setup_cache.py.
+    """
+    run_warm_setup_cache(
+        data=str(data),
+        val_fraction=val_fraction,
+        seed=seed,
+        conditioning=conditioning.value,
+        router_enabled=router,
+        router_type=router_type,
+        n_experts=n_experts,
+        rebuild=rebuild,
+        echo=typer.echo,
+    )
+
+
 @app.command("hparam-scan")
 def hparam_scan(
     data: Annotated[str, typer.Option("--data")] = DATA_DEFAULT,

scripts/warm_setup_cache.py

@@ -0,0 +1,52 @@
+"""dwarf warm-cache — precompute `giant train`'s setup-stage sidecar ahead of time.
+
+Thin wrapper around `giant.pipeline.run_setup_stage` so a dataset's vocab
+maps, event-id split index, and normalizer stats can be warmed once — e.g.
+right after `dwarf convert`, or before kicking off a `dwarf hparam-scan`
+sweep — without needing to also start training. See giant/data/setup_cache.py
+for the sidecar itself.
+"""
+
+from pathlib import Path
+
+from giant.pipeline import run_setup_stage
+
+
+def run_warm_setup_cache(
+    data: str,
+    val_fraction: float = 0.1,
+    seed: int = 0,
+    conditioning: str = "physical",
+    router_enabled: bool = False,
+    router_type: str = "energy",
+    n_experts: int = 4,
+    rebuild: bool = False,
+    echo=print,
+) -> None:
+    """Populate (or refresh) the setup cache sidecar for `data`.
+
+    `val_fraction`/`seed`/`conditioning` select the normalizer cache entry
+    (`giant.data.setup_cache.normalizer_key`) — pass the same values a later
+    `giant train` invocation will use so it hits this warmed entry.
+    `router_enabled`/`router_type`/`n_experts` only matter for
+    `router_type == "process"` (warms that `n_experts`'s process map); the
+    energy-router quantile summary is always collected regardless, so a
+    later `--router-type energy` run never needs to rescan just to seed
+    centers.
+    """
+    router_cfg = {
+        "enabled": router_enabled,
+        "type": router_type,
+        "n_experts": n_experts,
+    }
+    run_setup_stage(
+        Path(data),
+        val_fraction=val_fraction,
+        seed=seed,
+        conditioning=conditioning,
+        router_cfg=router_cfg,
+        cache_setup=True,
+        rebuild_setup_cache=rebuild,
+        echo=echo,
+    )
+    echo("setup cache warmed.")

tests/test_bump_dataset_version.py

@@ -1,4 +1,6 @@
 import os
+import subprocess
+
 from scripts import bump_dataset_version
 
 plan_bump_gen = bump_dataset_version.plan_bump_gen
@@ -11,6 +13,14 @@ apply_create_manifest = bump_dataset_version.apply_create_manifest
 check_holdout_overlap = bump_dataset_version.check_holdout_overlap
 
 
+def test_git_user_name_returns_none_on_timeout(monkeypatch):
+    def _raise_timeout(*args, **kwargs):
+        raise subprocess.TimeoutExpired(cmd=["git"], timeout=2)
+
+    monkeypatch.setattr(subprocess, "run", _raise_timeout)
+    assert bump_dataset_version._git_user_name() is None
+
+
 def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
     dirs, log_line = plan_bump_gen(
         tmp_path, "steps", "first generation", None, "2026-01-01"
@@ -401,6 +411,36 @@ def test_create_manifest_creates_parent_dirs(tmp_path):
     assert output.exists()
 
 
+def test_run_create_manifest_refuses_to_overwrite_existing_output(tmp_path):
+    pq = tmp_path / "a.parquet"
+    pq.touch()
+    output = tmp_path / "pools" / "pbwo4" / "holdout.manifest"
+    output.parent.mkdir(parents=True)
+    output.write_text("original contents\n")
+
+    try:
+        bump_dataset_version.run_create_manifest(
+            [str(pq)], execute=True, output=str(output)
+        )
+        assert False, "expected SystemExit"
+    except SystemExit:
+        pass
+    assert output.read_text() == "original contents\n"
+
+
+def test_run_create_manifest_force_overwrites_existing_output(tmp_path):
+    pq = tmp_path / "a.parquet"
+    pq.touch()
+    output = tmp_path / "pools" / "pbwo4" / "holdout.manifest"
+    output.parent.mkdir(parents=True)
+    output.write_text("original contents\n")
+
+    bump_dataset_version.run_create_manifest(
+        [str(pq)], execute=True, output=str(output), force=True
+    )
+    assert output.read_text() != "original contents\n"
+
+
 # ---------------------------------------------------------------------------
 # check_holdout_overlap
 # ---------------------------------------------------------------------------

tests/test_cli_new_run.py

@@ -0,0 +1,118 @@
+"""Tests for `giant new-run` (config.toml + run-dir scaffolding)."""
+
+from __future__ import annotations
+
+import tomllib
+from pathlib import Path
+
+from typer.testing import CliRunner
+
+from giant.cli import app
+
+runner = CliRunner()
+
+
+def test_writes_config_with_overrides_applied(tmp_path: Path):
+    out_dir = tmp_path / "run1"
+    result = runner.invoke(
+        app,
+        [
+            "new-run",
+            "--out",
+            str(out_dir),
+            "--mode",
+            "ddpm",
+            "--hidden-dim",
+            "128",
+            "--n-blocks",
+            "4",
+            "--lr",
+            "0.0005",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    config_path = out_dir / "config.toml"
+    assert config_path.exists()
+    with open(config_path, "rb") as f:
+        cfg = tomllib.load(f)
+
+    assert cfg["train"]["mode"] == "ddpm"
+    assert cfg["train"]["lr"] == 0.0005
+    assert cfg["model"]["hidden_dim"] == 128
+    assert cfg["model"]["n_blocks"] == 4
+    # untouched defaults still present
+    assert cfg["train"]["epochs"] == 100
+    assert "router" in cfg["model"]
+
+    assert str(out_dir) in result.output
+    assert "<data.parquet>" in result.output
+    assert "giant train" in result.output
+
+
+def test_comment_and_provenance_recorded_in_meta(tmp_path: Path):
+    out_dir = tmp_path / "run2"
+    result = runner.invoke(
+        app,
+        ["new-run", "--out", str(out_dir), "--comment", "quick test"],
+    )
+    assert result.exit_code == 0, result.output
+
+    with open(out_dir / "config.toml", "rb") as f:
+        cfg = tomllib.load(f)
+
+    assert cfg["meta"]["comment"] == "quick test"
+    assert cfg["meta"]["created_by"] == "giant new-run"
+    assert "created_at" in cfg["meta"]
+    assert "git_hash" in cfg["meta"]
+
+
+def test_data_flag_fills_printed_next_step_commands(tmp_path: Path):
+    out_dir = tmp_path / "run3"
+    result = runner.invoke(
+        app,
+        ["new-run", "--out", str(out_dir), "--data", "/ceph/lbogner/train.parquet"],
+    )
+    assert result.exit_code == 0, result.output
+    assert "/ceph/lbogner/train.parquet" in result.output
+    assert "<data.parquet>" not in result.output
+
+
+def test_dry_run_writes_nothing(tmp_path: Path):
+    out_dir = tmp_path / "run4"
+    result = runner.invoke(
+        app,
+        ["new-run", "--out", str(out_dir), "--hidden-dim", "512", "--dry-run"],
+    )
+    assert result.exit_code == 0, result.output
+    assert "dry-run" in result.output
+    assert "hidden_dim = 512" in result.output
+    assert not out_dir.exists()
+
+
+def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
+    out_dir = tmp_path / "run5"
+    out_dir.mkdir()
+    (out_dir / "last.pt").touch()
+
+    result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm"])
+    assert result.exit_code != 0
+    assert "already has last.pt" in result.output
+    assert not (out_dir / "config.toml").exists()
+
+    result = runner.invoke(
+        app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"]
+    )
+    assert result.exit_code == 0, result.output
+    assert (out_dir / "config.toml").exists()
+
+
+def test_default_out_dir_used_when_out_omitted(tmp_path: Path, monkeypatch):
+    monkeypatch.chdir(tmp_path)
+    result = runner.invoke(app, ["new-run", "--hidden-dim", "64"])
+    assert result.exit_code == 0, result.output
+
+    checkpoints_dir = tmp_path / "checkpoints"
+    run_dirs = list(checkpoints_dir.iterdir()) if checkpoints_dir.exists() else []
+    assert len(run_dirs) == 1
+    assert (run_dirs[0] / "config.toml").exists()

tests/test_condor.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import sys
 from pathlib import Path
 
 import pyarrow.parquet as pq
@@ -121,6 +122,26 @@ def test_prep_splits_rows_per_chunk(tmp_path: Path):
     assert sum(meta.rows_per_chunk) == meta.total_rows == 8
 
 
+def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Path):
+    """Re-prepping with a different n_chunks must not leave old chunk
+    partials on disk for merge_one to silently merge against the new
+    context (they'd be keyed/sized for the old n_chunks)."""
+    yaml_path = _write_inputs(tmp_path)
+    run_dir = _prep(yaml_path, chunks=2)
+    compute_one("marginal_edep", run_dir, chunk_index=0)
+    compute_one("marginal_edep", run_dir, chunk_index=1)
+    stale = run_dir / "reduced_partial" / "marginal_edep__0.json"
+    assert stale.exists()
+    (run_dir / "reduced").mkdir(exist_ok=True)
+    (run_dir / "reduced" / "marginal_edep.json").write_text("{}")
+
+    _prep(yaml_path, run_dir, chunks=1)
+
+    assert not stale.exists()
+    assert not (run_dir / "reduced" / "marginal_edep.json").exists()
+    assert (run_dir / "shared.json").exists()  # prep's own fresh output untouched
+
+
 def test_compute_one_from_run_dir(tmp_path: Path):
     run_dir = _prep(_write_inputs(tmp_path))
     out = compute_one("marginal_edep", run_dir)
@@ -203,9 +224,16 @@ def test_write_submit_description(tmp_path: Path):
     assert "--chunk" in body and "--run-dir" in body
 
 
-def test_write_submit_requires_synced_venv(tmp_path: Path):
+def test_write_submit_requires_synced_venv(
+    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
+):
     run_dir = _prep(_write_inputs(tmp_path))
     cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
+    # No `giant` next to the (fake) active interpreter, so this falls through
+    # to repo_dir/.venv/bin/giant, which _write_inputs/_prep also didn't create.
+    monkeypatch.setattr(
+        sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python")
+    )
     with pytest.raises(FileNotFoundError, match="uv sync"):
         write_submit(cfg)
 
@@ -237,6 +265,20 @@ def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
     assert counts["router_gating"] == 1  # chunkable=False, ignores n_chunks
 
 
+def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
+    """cfg.n_chunks must match the n_chunks the run_dir was actually prepped
+    with — RunMeta.rows_per_chunk is sized to the prepped value, so a
+    mismatch would otherwise surface as a confusing IndexError deep inside
+    _job_walltimes instead of a clear error here."""
+    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    _fake_venv(tmp_path)
+    cfg = SubmitConfig(
+        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
+    )
+    with pytest.raises(ValueError, match="n_chunks"):
+        write_submit(cfg)
+
+
 def test_estimate_runtime_s_scales_with_rows_and_margin():
     from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
     from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S

tests/test_config.py

@@ -1,3 +1,5 @@
+from datetime import datetime
+
 from giant import config as gconfig
 
 
@@ -115,3 +117,198 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
 
     gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
     assert capsys.readouterr().err == ""
+
+
+def test_resolve_expert_dims_default_config_inherits_hidden_dim_and_n_blocks():
+    # The unset sentinel (expert_hidden_dim/n_blocks == 0 in DEFAULT_CONFIG)
+    # is exactly the bug fixed by resolve_expert_dims: it must not silently
+    # fall back to some other hardcoded default, only to the monolith's own
+    # hidden_dim/n_blocks, so --hidden-dim/--n-blocks reach the experts too.
+    router_cfg = dict(gconfig.DEFAULT_CONFIG["model"]["router"])
+    assert router_cfg["expert_hidden_dim"] == 0
+    assert router_cfg["expert_n_blocks"] == 0
+
+    hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
+    assert (hidden_dim, n_blocks) == (512, 6)
+
+
+def test_resolve_expert_dims_missing_keys_also_inherit():
+    hidden_dim, n_blocks = gconfig.resolve_expert_dims({}, 512, 6)
+    assert (hidden_dim, n_blocks) == (512, 6)
+
+
+def test_default_config_gumbel_router_defaults_off():
+    # Straight-through Gumbel-softmax combine weights (giant.model.network.
+    # Router.combine_weights) must be opt-in — existing routed configs and
+    # checkpoints should be unaffected unless gumbel is explicitly enabled.
+    router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"]
+    assert router_cfg["gumbel"] is False
+    assert router_cfg["gumbel_tau_start"] == 1.0
+    assert router_cfg["gumbel_tau_end"] == 0.1
+
+
+def test_resolve_expert_dims_explicit_override_wins():
+    router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3}
+    hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
+    assert (hidden_dim, n_blocks) == (128, 3)
+
+
+def test_resolve_expert_dims_partial_override_mixes_explicit_and_inherited():
+    router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 0}
+    hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
+    assert (hidden_dim, n_blocks) == (128, 6)
+
+
+def _default_cfg(**overrides):
+    train_overrides = {
+        k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["train"]
+    }
+    model_overrides = {
+        k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["model"]
+    }
+    router_overrides = overrides.get("router")
+    if router_overrides:
+        model_overrides["router"] = router_overrides
+    return gconfig.merge_cli_overrides(
+        gconfig.DEFAULT_CONFIG, None, train_overrides, model_overrides
+    )
+
+
+_NOW = datetime(2026, 7, 29, 14, 30)
+
+
+def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
+    cfg = _default_cfg()
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
+
+
+def test_default_out_dir_name_single_non_default_field():
+    cfg = _default_cfg(hidden_dim=512)
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_h512"
+
+
+def test_default_out_dir_name_conditioning_embedding_shown_abbreviated():
+    cfg = _default_cfg(conditioning="embedding")
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb"
+
+
+def test_default_out_dir_name_conditioning_default_omitted():
+    cfg = _default_cfg(conditioning="physical")
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
+
+
+def test_default_out_dir_name_router_enabled_shown_as_unit():
+    cfg = _default_cfg(router={"enabled": True, "type": "energy", "n_experts": 8})
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8"
+
+
+def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault():
+    cfg = _default_cfg(router={"enabled": False, "type": "pdg", "n_experts": 8})
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
+
+
+def test_default_out_dir_name_router_gumbel_shown_when_enabled():
+    cfg = _default_cfg(
+        router={"enabled": True, "type": "energy", "n_experts": 8, "gumbel": True}
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_gum"
+
+
+def test_default_out_dir_name_router_gumbel_omitted_when_router_disabled():
+    cfg = _default_cfg(
+        router={"enabled": False, "type": "energy", "n_experts": 8, "gumbel": True}
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
+
+
+def test_default_out_dir_name_router_learn_centers_shown_only_when_disabled():
+    cfg_default = _default_cfg(
+        router={"enabled": True, "type": "energy", "n_experts": 8}
+    )
+    assert (
+        gconfig.default_out_dir_name(cfg_default, now=_NOW) == "20260729_1430_r-energy8"
+    )
+
+    cfg_off = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_centers": False,
+        }
+    )
+    assert (
+        gconfig.default_out_dir_name(cfg_off, now=_NOW)
+        == "20260729_1430_r-energy8_nolc"
+    )
+
+
+def test_default_out_dir_name_router_learn_width_and_temperature_shown():
+    cfg = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_width": True,
+        }
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_lw"
+
+    cfg2 = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_temperature": True,
+        }
+    )
+    assert gconfig.default_out_dir_name(cfg2, now=_NOW) == "20260729_1430_r-energy8_lt"
+
+
+def test_default_out_dir_name_mode_shown_bare_no_prefix():
+    cfg = _default_cfg(mode="wgan")
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan"
+
+
+def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
+    cfg = _default_cfg(
+        mode="wgan",
+        router={"enabled": True, "type": "energy", "n_experts": 8},
+        conditioning="embedding",
+        hidden_dim=512,
+        n_blocks=8,
+        emb_dim=32,
+        lr=1e-3,
+        batch_size=2048,
+        seed=3,
+        epochs=200,
+    )
+    name = gconfig.default_out_dir_name(cfg, now=_NOW)
+    # First 6 by priority: mode, router, conditioning, hidden_dim, n_blocks, emb_dim.
+    assert name.startswith("20260729_1430_wgan_r-energy8_cemb_h512_b8_e32_+4more-")
+    digest = name.split("-")[-1]
+    assert len(digest) == 6
+
+
+def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
+    base = dict(
+        mode="wgan",
+        router={"enabled": True, "type": "energy", "n_experts": 8},
+        conditioning="embedding",
+        hidden_dim=512,
+        n_blocks=8,
+        emb_dim=32,
+        lr=1e-3,
+        batch_size=2048,
+        seed=3,
+        epochs=200,
+    )
+    name_a = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW)
+    name_b = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW)
+    assert name_a == name_b  # stable across calls with the same overflow set
+
+    changed = dict(base, epochs=999)
+    name_c = gconfig.default_out_dir_name(_default_cfg(**changed), now=_NOW)
+    assert (
+        name_c != name_a
+    )  # differs when an overflowed value changes  # n_blocks inherited, hidden_dim not

tests/test_dataset.py

@@ -1,5 +1,10 @@
 import numpy as np
-from giant.data.dataset import make_event_split
+import pandas as pd
+
+from giant.constants import COND_DIM, X_DIM
+from giant.data import setup_cache
+from giant.data.dataset import StreamingStepsDataset, make_event_split
+from giant.data.transforms import Normalizer
 
 
 def test_make_event_split_sizes():
@@ -25,9 +30,110 @@ def test_make_event_split_no_empty_sets():
     assert len(val_set) > 0
 
 
+def test_make_event_split_val_fraction_zero_holds_out_nothing():
+    """val_fraction=0.0 is an explicit "train on everything" request and
+    must not be silently overridden into holding out 1 event."""
+    rng = np.random.default_rng(3)
+    event_ids = rng.integers(0, 50, size=1000)
+    train_set, val_set = make_event_split(event_ids, val_fraction=0.0)
+    assert val_set == set()
+    assert train_set == set(np.unique(event_ids).tolist())
+
+
 def test_make_event_split_reproducible():
     event_ids = np.arange(100)
     a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
     b_tr, b_val = make_event_split(event_ids, val_fraction=0.1, seed=42)
     assert a_tr == b_tr
     assert a_val == b_val
+
+
+# ── StreamingStepsDataset: cross-file event_id offsetting ──────────────────
+
+
+def _steps_df(event_ids, n_per_event=3, pre_E=100.0):
+    """A schema-complete but minimal steps DataFrame — no secondaries, so
+    `require_secondaries=True` never needs the per-secondary list columns."""
+    rows = []
+    for eid in event_ids:
+        for s in range(n_per_event):
+            rows.append(
+                {
+                    "event_id": eid,
+                    "pdg": 11,
+                    "pre_x": 0.0,
+                    "pre_y": 0.0,
+                    "pre_z": 0.0,
+                    "pre_E": pre_E,
+                    "pre_dx": 0.0,
+                    "pre_dy": 0.0,
+                    "pre_dz": 1.0,
+                    "material": "G4_AIR",
+                    "layer_id": s,
+                    "child_track_ids": [],
+                    "e_sec": 0.0,
+                    "step_length": 1.0,
+                    "post_E": pre_E * 0.9,
+                    "edep": pre_E * 0.1,
+                    "post_dx": 0.0,
+                    "post_dy": 0.0,
+                    "post_dz": 1.0,
+                    "post_x": 0.0,
+                    "post_y": 0.0,
+                    "post_z": 1.0,
+                }
+            )
+    return pd.DataFrame(rows)
+
+
+def _dummy_normalizer(width):
+    norm = Normalizer()
+    norm.mean = np.zeros(width, dtype=np.float32)
+    norm.std = np.ones(width, dtype=np.float32)
+    return norm
+
+
+def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
+    """Two files that each restart event_id from 0 (one Geant4 job per file,
+    see scripts/steps_to_parquet.py) must not have their same-numbered events
+    collapsed together: every row from every file must show up in exactly one
+    of train/val, and the number of distinct events must be the sum across
+    files, not the union of raw ids."""
+    n_events, n_per_event = 5, 3
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    _steps_df(range(n_events), n_per_event=n_per_event).to_parquet(path_a)
+    _steps_df(range(n_events), n_per_event=n_per_event).to_parquet(path_b)
+    files = [path_a, path_b]
+
+    unique_ids, counts = setup_cache.compute_event_index_from_files(files)
+    assert len(unique_ids) == 2 * n_events
+
+    train_events, val_events = make_event_split(unique_ids, val_fraction=0.4, seed=0)
+    assert train_events.isdisjoint(val_events)
+
+    pdg_map, mat_map = {11: 0}, {"G4_AIR": 0}
+    cond_norm = _dummy_normalizer(COND_DIM)
+    tgt_norm = _dummy_normalizer(X_DIM)
+
+    def _count_rows(split_events):
+        ds = StreamingStepsDataset(
+            files=files,
+            split_events=split_events,
+            pdg_map=pdg_map,
+            mat_map=mat_map,
+            cond_normalizer=cond_norm,
+            target_normalizer=tgt_norm,
+            batch_size=4,
+            shuffle=False,
+            conditioning="embedding",
+        )
+        return sum(len(batch[0]) for batch in ds)
+
+    n_train = _count_rows(train_events)
+    n_val = _count_rows(val_events)
+
+    total_rows = 2 * n_events * n_per_event
+    assert n_train + n_val == total_rows
+    assert n_train == int(counts[np.isin(unique_ids, list(train_events))].sum())
+    assert n_val == int(counts[np.isin(unique_ids, list(val_events))].sum())

tests/test_dwarf.py

@@ -1,10 +1,35 @@
 from typer.testing import CliRunner
 
+from giant import cli as giant_cli
+from giant.config import Conditioning
+from giant.data import setup_cache
+from scripts import dwarf
 from scripts.dwarf import app
+from test_pipeline import _make_synthetic_steps
 
 runner = CliRunner()
 
 
+def test_conditioning_enum_shared_across_both_clis():
+    """giant.cli and scripts.dwarf must use the one giant.config.Conditioning
+    enum, not independently redefined copies that could silently drift apart
+    on valid --conditioning values."""
+    assert dwarf.Conditioning is Conditioning
+    assert giant_cli.Conditioning is Conditioning
+
+
+def test_warn_if_exceeds_shared_quota_warns_over_quarter_cpu(monkeypatch, capsys):
+    monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8)  # quota = 2
+    dwarf._warn_if_exceeds_shared_quota(3, "--jobs")
+    assert "warning: --jobs=3 exceeds" in capsys.readouterr().err
+
+
+def test_warn_if_exceeds_shared_quota_silent_within_quota(monkeypatch, capsys):
+    monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8)  # quota = 2
+    dwarf._warn_if_exceeds_shared_quota(2, "--jobs")
+    assert capsys.readouterr().err == ""
+
+
 def test_convert_rejects_jobs_below_one(tmp_path):
     root_file = tmp_path / "shard.root"
     root_file.touch()
@@ -66,3 +91,81 @@ def test_status_reports_missing_root(tmp_path):
     result = runner.invoke(app, ["status", "--root", str(missing)])
     assert result.exit_code != 0
     assert "is not a directory" in result.output
+
+
+def test_warm_cache_writes_sidecar(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+
+    result = runner.invoke(app, ["warm-cache", str(data)])
+
+    assert result.exit_code == 0, result.output
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    assert loaded.vocab is not None
+    assert loaded.event_index is not None
+    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
+
+
+def test_warm_cache_second_run_hits_cache(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    runner.invoke(app, ["warm-cache", str(data)])
+
+    result = runner.invoke(app, ["warm-cache", str(data)])
+
+    assert result.exit_code == 0, result.output
+    assert "event index: cache hit" in result.output
+    assert "vocabulary maps: cache hit" in result.output
+    assert "normalizer: cache hit" in result.output
+
+
+def test_warm_cache_router_process_warms_proc_map(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+
+    result = runner.invoke(
+        app,
+        [
+            "warm-cache",
+            str(data),
+            "--router",
+            "--router-type",
+            "process",
+            "--n-experts",
+            "3",
+        ],
+    )
+
+    assert result.exit_code == 0, result.output
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    assert 3 in loaded.proc_maps
+
+
+def test_warm_cache_rebuild_ignores_existing(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    files = [data]
+    stale = setup_cache.SetupCache.empty(files)
+    stale.vocab = ({999999: 0}, {"G4_AIR": 0})  # deliberately wrong
+    setup_cache.save(data, files, stale)
+
+    result = runner.invoke(app, ["warm-cache", str(data), "--rebuild"])
+
+    assert result.exit_code == 0, result.output
+    loaded = setup_cache.load(data, files)
+    assert loaded is not None
+    assert loaded.vocab is not None
+    assert set(loaded.vocab[0].keys()) == {11, 22}
+
+
+def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
+    data = _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+    runner.invoke(app, ["warm-cache", str(data), "--val-fraction", "0.1"])
+
+    result = runner.invoke(app, ["warm-cache", str(data), "--val-fraction", "0.3"])
+
+    assert result.exit_code == 0, result.output
+    assert "vocabulary maps: cache hit" in result.output
+    assert "fitting normalizer (streaming)" in result.output
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
+    assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers

tests/test_loader.py

@@ -1,7 +1,19 @@
+import numpy as np
 import pandas as pd
 import pytest
 
-from giant.data.loader import build_process_map_from_files, find_parquet_files
+from giant.data.loader import (
+    EVENT_ID_FILE_STRIDE,
+    build_index_maps,
+    build_index_maps_from_files,
+    build_process_map_from_files,
+    event_id_offset,
+    find_parquet_files,
+    iter_cond_chunks,
+    iter_file_chunks,
+    load_event_ids,
+    load_steps,
+)
 
 
 def _touch(path):
@@ -90,3 +102,276 @@ def test_build_process_map_from_files_spans_multiple_files(tmp_path):
     assert proc_map["phot"] == 0
     assert proc_map["eIoni"] == 1
     assert proc_map["compt"] == 2
+
+
+def test_build_process_map_from_files_tie_breaking_pins_first_seen_order(tmp_path):
+    """When two processes end up with equal total counts, ranking falls back
+    to whichever was accumulated first (`sorted(..., reverse=True)` is stable,
+    and `counts` is built in file/row-scan order) — this is implementation-
+    defined, not a documented contract, so pin it explicitly: a future
+    rewrite (e.g. a polars-based single-scan) that ties differently would
+    silently reshuffle which processes get their own expert slot across a
+    retrain, and this test is what should catch that."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path)
+
+    proc_map = build_process_map_from_files([path], n_experts=3)
+
+    assert proc_map == {"compt": 0, "phot": 1}
+
+
+def test_build_process_map_from_files_tie_breaking_favors_first_scanned_file(
+    tmp_path,
+):
+    """Same total-count tie as above, but split across two files with equal
+    per-file counts — the file listed first wins the tie."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"process": ["zzz", "zzz"]}).to_parquet(path_a)
+    pd.DataFrame({"process": ["aaa", "aaa"]}).to_parquet(path_b)
+
+    forward = build_process_map_from_files([path_a, path_b], n_experts=3)
+    backward = build_process_map_from_files([path_b, path_a], n_experts=3)
+
+    assert forward == {"zzz": 0, "aaa": 1}
+    assert backward == {"aaa": 0, "zzz": 1}
+
+
+def test_build_process_map_from_files_fewer_processes_than_experts(tmp_path):
+    """When there are fewer distinct processes than expert slots, every
+    process gets its own index and the shared "other" bucket goes unused."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"process": ["eIoni", "phot"]}).to_parquet(path)
+
+    proc_map = build_process_map_from_files([path], n_experts=5)
+
+    assert proc_map == {"eIoni": 0, "phot": 1}
+    assert 4 not in proc_map.values()  # the "other" slot (n_experts - 1) is unused
+
+
+def test_build_process_map_from_files_n_experts_one_buckets_everything(tmp_path):
+    """n_experts=1 leaves no room for a "most frequent" slot — every process
+    (however frequent) is bucketed into the single shared index 0."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"process": ["eIoni"] * 10 + ["phot"] * 1}).to_parquet(path)
+
+    proc_map = build_process_map_from_files([path], n_experts=1)
+
+    assert proc_map == {"eIoni": 0, "phot": 0}
+
+
+def test_build_process_map_from_files_three_files_partial_overlap(tmp_path):
+    """Counts for a process appearing in only some of several files must sum
+    correctly, not just match the two-file case already covered above."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    path_c = tmp_path / "c.parquet"
+    pd.DataFrame({"process": ["eIoni"] * 2}).to_parquet(path_a)
+    pd.DataFrame({"process": ["phot"] * 3}).to_parquet(path_b)
+    pd.DataFrame({"process": ["eIoni"] * 2 + ["compt"] * 1}).to_parquet(path_c)
+
+    # eIoni: 2+2=4 > phot: 3 > compt: 1
+    proc_map = build_process_map_from_files([path_a, path_b, path_c], n_experts=3)
+
+    assert proc_map["eIoni"] == 0
+    assert proc_map["phot"] == 1
+    assert proc_map["compt"] == 2
+
+
+# ── build_index_maps (in-memory) ────────────────────────────────────────────
+
+
+def test_build_index_maps_sorts_numerically_not_lexicographically():
+    """10-digit nuclear/ion PDG codes must sort numerically — a lexicographic
+    sort would place "1000060120" before "22" since '1' < '2'."""
+    data = {
+        "pdg": np.array([22, 1000060120, 11], dtype=np.int64),
+        "material": np.array(["G4_AIR", "PbWO4", "G4_Fe"], dtype=object),
+    }
+    pdg_map, mat_map = build_index_maps(data)
+    assert list(pdg_map.keys()) == [11, 22, 1000060120]
+    assert mat_map == {"G4_AIR": 0, "G4_Fe": 1, "PbWO4": 2}
+
+
+def test_build_index_maps_dedups_repeated_values():
+    data = {
+        "pdg": np.array([11, 11, 22, 22, 22], dtype=np.int64),
+        "material": np.array(["PbWO4"] * 5, dtype=object),
+    }
+    pdg_map, mat_map = build_index_maps(data)
+    assert pdg_map == {11: 0, 22: 1}
+    assert mat_map == {"PbWO4": 0}
+
+
+def test_build_index_maps_handles_negative_pdg_codes():
+    """Antiparticle codes (negative) must sort numerically, not by magnitude."""
+    data = {
+        "pdg": np.array([-13, 11, -11, 13], dtype=np.int64),
+        "material": np.array(["X"] * 4, dtype=object),
+    }
+    pdg_map, _ = build_index_maps(data)
+    assert list(pdg_map.keys()) == [-13, -11, 11, 13]
+
+
+def test_build_index_maps_indices_are_dense_and_bijective():
+    data = {
+        "pdg": np.array([5, 1, 9, 1, 5], dtype=np.int64),
+        "material": np.array(["a", "b", "c", "a", "b"], dtype=object),
+    }
+    pdg_map, mat_map = build_index_maps(data)
+    assert sorted(pdg_map.values()) == list(range(len(pdg_map)))
+    assert sorted(mat_map.values()) == list(range(len(mat_map)))
+
+
+# ── build_index_maps_from_files ─────────────────────────────────────────────
+
+
+def test_build_index_maps_from_files_unions_and_dedups_across_files(tmp_path):
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"pdg": [11, 22], "material": ["G4_AIR", "PbWO4"]}).to_parquet(path_a)
+    pd.DataFrame({"pdg": [22, 2112], "material": ["PbWO4", "G4_Fe"]}).to_parquet(path_b)
+
+    pdg_map, mat_map = build_index_maps_from_files([path_a, path_b])
+
+    assert pdg_map == {11: 0, 22: 1, 2112: 2}
+    assert mat_map == {"G4_AIR": 0, "G4_Fe": 1, "PbWO4": 2}
+
+
+def test_build_index_maps_from_files_ordering_independent_of_file_order(tmp_path):
+    """Index assignment comes from the globally sorted union, not file-scan
+    order — swapping which file is scanned first must not change the map,
+    since the map is baked into a trained checkpoint's vocabulary."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"pdg": [22], "material": ["PbWO4"]}).to_parquet(path_a)
+    pd.DataFrame({"pdg": [11], "material": ["G4_AIR"]}).to_parquet(path_b)
+
+    forward = build_index_maps_from_files([path_a, path_b])
+    backward = build_index_maps_from_files([path_b, path_a])
+
+    assert forward == backward
+    assert forward == ({11: 0, 22: 1}, {"G4_AIR": 0, "PbWO4": 1})
+
+
+def test_build_index_maps_from_files_numeric_sort_for_nuclear_codes(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(
+        path
+    )
+
+    pdg_map, _ = build_index_maps_from_files([path])
+    assert list(pdg_map.keys()) == [11, 22, 1000060120]
+
+
+def test_build_index_maps_from_files_single_file(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": [11, 11, 22], "material": ["PbWO4"] * 3}).to_parquet(path)
+    pdg_map, mat_map = build_index_maps_from_files([path])
+    assert pdg_map == {11: 0, 22: 1}
+    assert mat_map == {"PbWO4": 0}
+
+
+def test_build_index_maps_from_files_matches_build_index_maps(tmp_path):
+    """Sanity-pin: the file-scanning and in-memory variants must agree on the
+    same data, since a future single-pass rewrite (pyarrow/polars) may
+    replace one but not the other."""
+    rng = np.random.default_rng(0)
+    pdg = rng.choice([11, -11, 22, 2112, 1000060120], size=200)
+    material = rng.choice(["G4_AIR", "PbWO4", "G4_Fe"], size=200)
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": pdg, "material": material}).to_parquet(path)
+
+    from_files = build_index_maps_from_files([path])
+    from_memory = build_index_maps({"pdg": pdg, "material": material})
+    assert from_files == from_memory
+
+
+# ── event_id_offset / per-file event_id offsetting ─────────────────────────
+
+
+def _steps_df(event_ids):
+    """Minimal schema-complete steps rows (no secondaries) for _df_to_dict."""
+    return pd.DataFrame(
+        [
+            {
+                "event_id": eid,
+                "pdg": 11,
+                "pre_x": 0.0,
+                "pre_y": 0.0,
+                "pre_z": 0.0,
+                "pre_E": 100.0,
+                "pre_dx": 0.0,
+                "pre_dy": 0.0,
+                "pre_dz": 1.0,
+                "material": "G4_AIR",
+                "layer_id": 0,
+                "child_track_ids": [],
+                "e_sec": 0.0,
+                "step_length": 1.0,
+                "post_E": 90.0,
+                "edep": 10.0,
+                "post_dx": 0.0,
+                "post_dy": 0.0,
+                "post_dz": 1.0,
+                "post_x": 0.0,
+                "post_y": 0.0,
+                "post_z": 1.0,
+            }
+            for eid in event_ids
+        ]
+    )
+
+
+def test_event_id_offset_scales_by_file_index():
+    assert event_id_offset(0) == 0
+    assert event_id_offset(1) == EVENT_ID_FILE_STRIDE
+    assert event_id_offset(3) == 3 * EVENT_ID_FILE_STRIDE
+
+
+def test_load_event_ids_default_offset_is_zero(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [5, 6, 7]}).to_parquet(path)
+    np.testing.assert_array_equal(load_event_ids(path), [5, 6, 7])
+
+
+def test_load_event_ids_applies_offset(tmp_path):
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path)
+    offset = event_id_offset(1)
+    np.testing.assert_array_equal(
+        load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2]
+    )
+
+
+def test_load_event_ids_raises_when_event_id_reaches_stride(tmp_path):
+    """A raw event_id >= EVENT_ID_FILE_STRIDE would collide into the next
+    file's offset block if silently allowed through — must raise instead."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [0, 1, EVENT_ID_FILE_STRIDE]}).to_parquet(path)
+    with pytest.raises(ValueError, match="EVENT_ID_FILE_STRIDE"):
+        load_event_ids(path)
+
+
+def test_load_steps_applies_offset_to_event_id(tmp_path):
+    path = tmp_path / "a.parquet"
+    _steps_df([0, 1]).to_parquet(path)
+    offset = event_id_offset(2)
+    d = load_steps(path, offset=offset)
+    np.testing.assert_array_equal(d["event_id"], [offset, offset + 1])
+
+
+def test_iter_file_chunks_applies_offset(tmp_path):
+    path = tmp_path / "a.parquet"
+    _steps_df([0, 1, 2]).to_parquet(path)
+    offset = event_id_offset(1)
+    ids = np.concatenate([c["event_id"] for c in iter_file_chunks(path, offset=offset)])
+    np.testing.assert_array_equal(sorted(ids), [offset, offset + 1, offset + 2])
+
+
+def test_iter_cond_chunks_applies_offset(tmp_path):
+    path = tmp_path / "a.parquet"
+    _steps_df([0, 1]).to_parquet(path)
+    offset = event_id_offset(5)
+    ids = np.concatenate([c["event_id"] for c in iter_cond_chunks(path, offset=offset)])
+    np.testing.assert_array_equal(sorted(ids), [offset, offset + 1])

tests/test_phase2.py

@@ -231,6 +231,53 @@ def test_encode_secondaries_energy_conservation():
     assert np.isfinite(sec_cont).all()
 
 
+def test_encode_secondaries_stick_logits_match_naive_reference():
+    """Cumsum-based remaining-budget computation must match a naive
+    per-row, per-slot Python reference (no cumsum) within float tolerance."""
+    from giant.data.transforms import encode_secondaries, _EPS, _STICK_LOGIT_CLIP
+
+    rng = np.random.default_rng(11)
+    N = 25
+    n_sec = rng.integers(1, K_MAX + 1, size=N)
+    e_sec = rng.uniform(0.1, 10.0, size=N).astype(np.float32)
+
+    sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
+    sec_dir_list = np.zeros((N, K_MAX, 3), dtype=np.float32)
+    sec_dir_list[:, :, 2] = 1.0
+    sec_valid = np.zeros((N, K_MAX), dtype=bool)
+    for i in range(N):
+        k = n_sec[i]
+        energies = rng.dirichlet(np.ones(k)) * e_sec[i]
+        energies = np.sort(energies)[::-1]
+        sec_E_list[i, :k] = energies.astype(np.float32)
+        sec_valid[i, :k] = True
+
+    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
+
+    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
+    stick_logits = sec_cont[:, :, 0]
+
+    # Naive reference: recompute the remaining budget from scratch each slot,
+    # exactly what the pre-cumsum implementation did.
+    expected = np.zeros((N, K_MAX), dtype=np.float64)
+    for row in range(N):
+        for i in range(K_MAX):
+            if not sec_valid[row, i]:
+                continue
+            remaining = max(float(e_sec[row]) - float(sec_E_list[row, :i].sum()), _EPS)
+            f = min(max(float(sec_E_list[row, i]) / remaining, _EPS), 1.0 - _EPS)
+            logit = np.log(f / (1.0 - f))
+            is_last = not (i + 1 < K_MAX and sec_valid[row, i + 1])
+            if is_last:
+                logit = _STICK_LOGIT_CLIP
+            logit = np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP)
+            expected[row, i] = logit
+
+    np.testing.assert_allclose(
+        stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4
+    )
+
+
 def test_encode_secondaries_direction_encoding():
     """Local-frame secondary directions should be unit vectors for valid slots."""
     from giant.data.transforms import encode_secondaries
@@ -272,6 +319,57 @@ def test_encode_secondaries_physical_columns_without_pdg_list():
     np.testing.assert_allclose(sec_cont[:, :, 4:6], 0.0)
 
 
+def test_encode_secondaries_phys_only_matches_full_and_zero_fills_rest():
+    """phys_only=True must reproduce the mass/charge columns exactly and
+    zero-fill the stick-logit/direction columns it skips computing."""
+    from giant.data.transforms import encode_secondaries
+
+    rng = np.random.default_rng(3)
+    N = 30
+    n_sec = rng.integers(1, 5, size=N)
+    e_sec = rng.uniform(0.1, 10.0, size=N).astype(np.float32)
+
+    sec_E_list = np.zeros((N, K_MAX), dtype=np.float32)
+    sec_dir_list = rng.standard_normal((N, K_MAX, 3)).astype(np.float32)
+    norms = np.linalg.norm(sec_dir_list, axis=-1, keepdims=True)
+    sec_dir_list /= np.where(norms > 0, norms, 1.0)
+    sec_pdg_list = np.zeros((N, K_MAX), dtype=np.int64)
+    sec_valid = np.zeros((N, K_MAX), dtype=bool)
+    for i in range(N):
+        k = n_sec[i]
+        energies = rng.dirichlet(np.ones(k)) * e_sec[i]
+        energies = np.sort(energies)[::-1]
+        sec_E_list[i, :k] = energies.astype(np.float32)
+        sec_valid[i, :k] = True
+        sec_pdg_list[i, :k] = 11  # electron — resolvable by giant.particles
+
+    pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
+    pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
+
+    full = encode_secondaries(
+        sec_E_list,
+        sec_dir_list,
+        sec_valid,
+        e_sec,
+        pre_dir,
+        sec_pdg_list=sec_pdg_list,
+        phys_only=False,
+    )
+    phys_only = encode_secondaries(
+        sec_E_list,
+        sec_dir_list,
+        sec_valid,
+        e_sec,
+        pre_dir,
+        sec_pdg_list=sec_pdg_list,
+        phys_only=True,
+    )
+
+    np.testing.assert_array_equal(phys_only[:, :, 4:6], full[:, :, 4:6])
+    np.testing.assert_array_equal(phys_only[:, :, 0], np.zeros((N, K_MAX)))
+    np.testing.assert_array_equal(phys_only[:, :, 1:4], np.zeros((N, K_MAX, 3)))
+
+
 def test_encode_secondaries_physical_columns_match_ground_truth_pdg():
     """log_mass/charge for a valid slot match giant.particles for that PDG."""
     from giant.data.transforms import encode_secondaries, log_transform

tests/test_pipeline.py

@@ -0,0 +1,212 @@
+import copy
+
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+from giant import config as gconfig
+from giant.data import setup_cache
+from giant.pipeline import run_train_job
+
+
+def _unit(v):
+    v = np.asarray(v, dtype=np.float64)
+    n = np.linalg.norm(v)
+    return v / n if n > 1e-9 else np.array([0.0, 0.0, 1.0])
+
+
+def _make_synthetic_steps(path, n_events=20, seed=0):
+    """A tiny but schema-complete synthetic steps parquet for run_train_job.
+
+    pdg/material/process are assigned deterministically by row index (not
+    random) so tests that assert on the resulting vocab/proc maps aren't
+    flaky; only continuous quantities (positions/energies/directions) are
+    drawn from `rng`.
+    """
+    rng = np.random.default_rng(seed)
+    materials = ["G4_AIR", "G4_Fe"]
+    pdgs = [11, 22]
+    processes = ["eIoni", "phot", "compt"]
+    rows = []
+    row_idx = 0
+    for event_id in range(n_events):
+        n_steps = int(rng.integers(2, 4))
+        for s in range(n_steps):
+            pre_E = float(rng.uniform(50.0, 500.0))
+            n_sec = int(rng.integers(0, 3))
+            frac_dep = float(rng.uniform(0.05, 0.3))
+            frac_sec = float(rng.uniform(0.05, 0.2)) if n_sec > 0 else 0.0
+            frac_post = 1.0 - frac_dep - frac_sec
+            edep = pre_E * frac_dep
+            e_sec = pre_E * frac_sec
+            post_E = pre_E * frac_post
+            pre_pos = rng.uniform(-10, 10, size=3)
+            step_length = float(rng.uniform(0.1, 5.0))
+            pre_dir = np.array([0.0, 0.0, 1.0])
+            post_dir = _unit(rng.normal(size=3))
+            post_pos = pre_pos + step_length * pre_dir
+            sec_energies = (
+                list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
+            )
+            sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)]
+            sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)]
+            rows.append(
+                {
+                    "event_id": event_id,
+                    "pdg": pdgs[row_idx % 2],
+                    "pre_x": pre_pos[0],
+                    "pre_y": pre_pos[1],
+                    "pre_z": pre_pos[2],
+                    "pre_E": pre_E,
+                    "pre_dx": pre_dir[0],
+                    "pre_dy": pre_dir[1],
+                    "pre_dz": pre_dir[2],
+                    "material": materials[row_idx % 2],
+                    "layer_id": s,
+                    "child_track_ids": list(range(n_sec)),
+                    "e_sec": e_sec,
+                    "process": processes[row_idx % 3],
+                    "step_length": step_length,
+                    "post_E": post_E,
+                    "edep": edep,
+                    "post_dx": post_dir[0],
+                    "post_dy": post_dir[1],
+                    "post_dz": post_dir[2],
+                    "post_x": post_pos[0],
+                    "post_y": post_pos[1],
+                    "post_z": post_pos[2],
+                    "sec_E_list": sec_energies,
+                    "sec_pdg_list": sec_pdgs,
+                    "sec_dx_list": [d[0] for d in sec_dirs],
+                    "sec_dy_list": [d[1] for d in sec_dirs],
+                    "sec_dz_list": [d[2] for d in sec_dirs],
+                }
+            )
+            row_idx += 1
+    pd.DataFrame(rows).to_parquet(path)
+    return path
+
+
+def _tiny_cfg(**train_overrides):
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["train"].update(
+        {
+            "epochs": 1,
+            "batch_size": 8,
+            "val_fraction": 0.2,
+            "seed": 0,
+            "warmup_epochs": 0,
+            "validate_every": 0,
+            "max_val_batches": 1,
+        }
+    )
+    cfg["train"].update(train_overrides)
+    cfg["model"].update({"hidden_dim": 8, "n_blocks": 1, "emb_dim": 4, "dropout": 0.0})
+    return cfg
+
+
+def _run(data, out_dir, cfg=None, **kwargs):
+    echoed: list[str] = []
+    kwargs.setdefault("num_workers", 0)
+    run_train_job(
+        data=data,
+        cfg=cfg or _tiny_cfg(),
+        out_dir=out_dir,
+        device=torch.device("cpu"),
+        shuffle_buffer=64,
+        echo=echoed.append,
+        **kwargs,
+    )
+    return echoed
+
+
+@pytest.fixture
+def data(tmp_path):
+    return _make_synthetic_steps(tmp_path / "data.parquet", n_events=20)
+
+
+def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
+    echo1 = _run(data, tmp_path / "out1")
+    assert any("fitting normalizer (streaming)" in m for m in echo1)
+
+    def _forbidden(*a, **k):
+        raise AssertionError("should be served from cache, not recomputed")
+
+    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+    monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden)
+
+    echo2 = _run(data, tmp_path / "out2")
+    joined = "\n".join(echo2)
+    assert "event index: cache hit" in joined
+    assert "vocabulary maps: cache hit" in joined
+    assert "normalizer: cache hit" in joined
+
+
+def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
+    tmp_path, data, monkeypatch
+):
+    monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8)  # quota = 2
+    echo = _run(data, tmp_path / "out", num_workers=3)
+    assert any("num-workers=3" in m and "exceeds" in m for m in echo)
+
+
+def test_run_train_job_no_warning_when_num_workers_within_shared_quota(
+    tmp_path, data, monkeypatch
+):
+    monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8)  # quota = 2
+    echo = _run(data, tmp_path / "out", num_workers=2)
+    assert not any("exceeds" in m for m in echo)
+
+
+def test_run_train_job_no_cache_setup_never_writes_sidecar(tmp_path, data):
+    _run(data, tmp_path / "out", cache_setup=False)
+    assert not setup_cache.sidecar_path(data).exists()
+
+
+def test_run_train_job_rebuild_setup_cache_ignores_existing(tmp_path, data):
+    files = [data]
+    stale = setup_cache.SetupCache.empty(files)
+    stale.vocab = ({999999: 0}, {"G4_AIR": 0})  # deliberately wrong
+    setup_cache.save(data, files, stale)
+
+    _run(data, tmp_path / "out", rebuild_setup_cache=True)
+
+    loaded = setup_cache.load(data, files)
+    assert loaded is not None
+    assert loaded.vocab is not None
+    assert set(loaded.vocab[0].keys()) == {11, 22}
+    assert set(loaded.vocab[1].keys()) == {"G4_AIR", "G4_Fe"}
+
+
+def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypatch):
+    _run(data, tmp_path / "out1", cfg=_tiny_cfg(val_fraction=0.1))
+
+    def _forbidden(*a, **k):
+        raise AssertionError("vocab should be served from cache")
+
+    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+
+    echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3))
+    joined = "\n".join(echo2)
+    assert "vocabulary maps: cache hit" in joined
+    assert "fitting normalizer (streaming)" in joined
+
+
+def test_run_train_job_matches_uncached_output(tmp_path, data):
+    _run(data, tmp_path / "uncached", cache_setup=False)
+    _run(data, tmp_path / "cached1", cache_setup=True)
+    _run(data, tmp_path / "cached2", cache_setup=True)  # second is a cache hit
+
+    uncached = torch.load(tmp_path / "uncached" / "last.pt", weights_only=False)
+    cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False)
+
+    for key in ("cond", "target", "sec_phys"):
+        np.testing.assert_allclose(
+            uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]
+        )
+        np.testing.assert_allclose(
+            uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]
+        )
+    assert uncached["pdg_map"] == cached["pdg_map"]
+    assert uncached["mat_map"] == cached["mat_map"]

tests/test_rollout.py

@@ -119,6 +119,21 @@ def test_rollout_physical_conditioning_end_to_end(fake_material_props):
     assert set(np.unique(rec["pdg"]).tolist()) <= set(PDG_MAP.keys())
 
 
+def test_rollout_physical_conditioning_generalizes_to_out_of_vocab_pdg(
+    fake_material_props,
+):
+    """A real, giant.particles-resolvable species outside the training PDG
+    vocab (muon, 13) must run through physical-property conditioning rather
+    than terminate via TERM_UNKNOWN_PDG — that generalization is the entire
+    point of "physical" mode (see build_cond_features(strict=...))."""
+    seeds = _seeds(6)
+    seeds["pdg"] = np.full(6, 13, dtype=np.int64)
+    assert 13 not in PDG_MAP
+    rec = _run(seeds=seeds, conditioning="physical")
+    assert len(rec["event_id"]) > 0
+    assert TERM_UNKNOWN_PDG not in set(rec["termination_reason"].tolist())
+
+
 def test_seed_frontier_track_ids():
     seeds = _seeds(3)
     fr, counts = make_seed_frontier(**seeds)

tests/test_router.py

@@ -1,5 +1,6 @@
 """Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""
 
+import pytest
 import torch
 
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
@@ -141,6 +142,255 @@ def test_build_router_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown router type")
 
 
+# ── EnergyRouter learn_width / learn_temperature ────────────────────────────
+
+
+def test_energy_router_learn_width_matches_fixed_temperature_at_init():
+    """Enabling learn_width should be a no-op at init — the warm-started
+    per-expert width must reproduce the fixed-temperature gate exactly."""
+    centers_init = [-1.0, 0.0, 0.5, 1.5]
+    fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
+    learned = EnergyRouter(
+        n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True
+    )
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        learned.gate(cond_cont, cond_cat),
+        fixed.gate(cond_cont, cond_cat),
+        atol=1e-5,
+        rtol=0,
+    )
+
+
+def test_energy_router_learn_temperature_matches_fixed_temperature_at_init():
+    centers_init = [-1.0, 0.0, 0.5, 1.5]
+    fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
+    learned = EnergyRouter(
+        n_experts=4,
+        temperature=0.3,
+        centers_init=centers_init,
+        learn_temperature=True,
+    )
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        learned.gate(cond_cont, cond_cat),
+        fixed.gate(cond_cont, cond_cat),
+        atol=1e-5,
+        rtol=0,
+    )
+
+
+def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
+    try:
+        EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True)
+    except ValueError:
+        return
+    raise AssertionError(
+        "expected ValueError for learn_width and learn_temperature both set"
+    )
+
+
+def test_energy_router_width_ratio_bounds_must_bracket_one_raises():
+    try:
+        EnergyRouter(
+            n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0
+        )
+    except ValueError:
+        return
+    raise AssertionError(
+        "expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0"
+    )
+
+
+def test_energy_router_effective_width_stays_within_bounds():
+    router = EnergyRouter(
+        n_experts=4,
+        temperature=0.5,
+        learn_width=True,
+        width_min_ratio=0.1,
+        width_max_ratio=10.0,
+    )
+    lo, hi = 0.1 * 0.5, 10.0 * 0.5
+    with torch.no_grad():
+        router.raw_width.fill_(1e6)
+    width = router.effective_width()
+    assert isinstance(width, torch.Tensor)
+    assert torch.all(width <= hi + 1e-4)
+    with torch.no_grad():
+        router.raw_width.fill_(-1e6)
+    width = router.effective_width()
+    assert isinstance(width, torch.Tensor)
+    assert torch.all(width >= lo - 1e-4)
+
+
+def test_energy_router_learn_width_gate_still_partition_of_unity():
+    router = EnergyRouter(n_experts=4, learn_width=True)
+    with torch.no_grad():
+        router.raw_width.copy_(torch.randn(4) * 3)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    torch.testing.assert_close(g.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
+
+
+def test_energy_router_learn_width_hardens_when_pushed_to_floor():
+    """Pushing every expert's width toward the (tiny) floor should harden the
+    gate to a one-hot at the nearest center, generalizing the fixed-
+    temperature->0 hardening test to the per-expert path."""
+    router = EnergyRouter(
+        n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0
+    )
+    with torch.no_grad():
+        router.raw_width.fill_(-1e6)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    e = cond_cont[:, router.energy_idx].unsqueeze(-1)
+    d2 = (e - router.centers.unsqueeze(0)) ** 2
+    onehot = torch.nn.functional.one_hot(d2.argmin(dim=-1), num_classes=4).float()
+    torch.testing.assert_close(g, onehot, atol=1e-3, rtol=0)
+
+
+def test_energy_router_own_width_controls_own_coverage_independent_of_others():
+    """Widening one expert's width should monotonically grow only that
+    expert's own gate share, without needing to touch any other expert's
+    width — the "each expert learns its own coverage independently" property
+    this feature is meant to add."""
+    router = EnergyRouter(
+        n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]
+    )
+    cond_cont, cond_cat = _cond(4)
+    cond_cont[:, 3] = 3.0  # fixed energy, unequal distance to each center
+
+    shares = []
+    with torch.no_grad():
+        for raw in torch.linspace(-8.0, 8.0, 9):
+            router.raw_width[0] = raw
+            shares.append(router.gate(cond_cont, cond_cat)[0, 0].item())
+
+    assert all(a <= b + 1e-6 for a, b in zip(shares, shares[1:]))
+
+
+def test_build_router_threads_learn_width_kwargs_through():
+    router = build_router(
+        "energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0
+    )
+    assert isinstance(router, EnergyRouter)
+    assert router.learn_width is True
+    assert isinstance(router.raw_width, torch.nn.Parameter)
+    assert router.raw_width.shape == (4,)
+
+
+def test_router_entropy_loss_is_nonnegative_bounded_scalar():
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    loss = router.entropy_loss(cond_cont, cond_cat)
+    assert loss.shape == ()
+    assert 0.0 <= loss.item() <= 1.0
+
+
+# ── Router.combine_weights (straight-through Gumbel-softmax) ───────────────
+
+
+def test_combine_weights_defaults_to_gate():
+    """gumbel=False (the default) must be a pure pass-through to gate()."""
+    router = EnergyRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        router.combine_weights(cond_cont, cond_cat),
+        router.gate(cond_cont, cond_cat),
+    )
+
+
+def test_combine_weights_gumbel_train_mode_is_hard_one_hot():
+    router = EnergyRouter(n_experts=4)
+    router.gumbel = True
+    router.gumbel_tau = 0.5
+    router.train()
+    cond_cont, cond_cat = _cond(16)
+    weights = router.combine_weights(cond_cont, cond_cat)
+    assert weights.shape == (16, 4)
+    torch.testing.assert_close(weights.sum(dim=-1), torch.ones(16), atol=1e-5, rtol=0)
+    assert torch.all((weights.max(dim=-1).values - 1.0).abs() < 1e-5)
+
+
+def test_combine_weights_gumbel_eval_mode_falls_back_to_gate():
+    """No Gumbel noise at eval — combine_weights must match gate() exactly,
+    same as the gumbel=False path, once the router is in eval mode."""
+    router = EnergyRouter(n_experts=4)
+    router.gumbel = True
+    router.eval()
+    cond_cont, cond_cat = _cond(16)
+    torch.testing.assert_close(
+        router.combine_weights(cond_cont, cond_cat),
+        router.gate(cond_cont, cond_cat),
+    )
+
+
+def test_combine_weights_gumbel_straight_through_gradient_reaches_centers():
+    router = EnergyRouter(n_experts=4, learn_centers=True)
+    router.gumbel = True
+    router.gumbel_tau = 0.5
+    router.train()
+    cond_cont, cond_cat = _cond(16)
+    weights = router.combine_weights(cond_cont, cond_cat)
+    weights.sum().backward()
+    assert router.centers.grad is not None
+    assert torch.any(router.centers.grad != 0.0)
+
+
+def test_build_router_from_cfg_sets_gumbel_from_config():
+    from giant.model.network import _build_router_from_cfg
+
+    router = _build_router_from_cfg(
+        {"enabled": True, "type": "energy", "n_experts": 4, "gumbel": True},
+        pdg_vocab=3,
+        mat_vocab=2,
+    )
+    assert router.gumbel is True
+
+    router_off = _build_router_from_cfg(
+        {"enabled": True, "type": "energy", "n_experts": 4},
+        pdg_vocab=3,
+        mat_vocab=2,
+    )
+    assert router_off.gumbel is False
+
+
+def test_build_router_from_cfg_sets_gumbel_for_composed_router():
+    from giant.model.network import _build_router_from_cfg
+
+    router = _build_router_from_cfg(
+        {
+            "enabled": True,
+            "type": "composed",
+            "gumbel": True,
+            "axis0_type": "energy",
+            "axis0_n_experts": 4,
+            "axis1_type": "pdg",
+            "axis1_n_experts": 3,
+        },
+        pdg_vocab=5,
+        mat_vocab=2,
+    )
+    assert isinstance(router, ComposedRouter)
+    assert router.gumbel is True
+
+
+def test_routed_denoising_mlp_forward_runs_with_gumbel_enabled():
+    """End-to-end forward through _route_forward's train branch with
+    straight-through Gumbel-softmax combine weights enabled."""
+    B = 8
+    model = _routed_stage1(n_experts=3)
+    model.router.gumbel = True
+    model.router.gumbel_tau = 0.5
+    model.train()
+    x_t = torch.randn(B, X_DIM)
+    t = torch.rand(B)
+    cond_cont, cond_cat = _cond(B)
+    out = model(x_t, t, cond_cont, cond_cat)
+    assert out.shape == (B, X_DIM)
+    assert torch.isfinite(out).all()
+
+
 # ── PdgRouter ────────────────────────────────────────────────────────────────
 
 
@@ -235,6 +485,49 @@ def test_build_models_routed_with_pdg_router():
     assert stage1.router.pdg_emb.num_embeddings == 4
 
 
+def test_build_models_rejects_pdg_router_with_physical_conditioning():
+    """conditioning="physical" is meant to generalize beyond the training PDG
+    vocab; PdgRouter always uses a training-vocab nn.Embedding regardless of
+    conditioning, so the combination must raise rather than silently building
+    a model that can't actually generalize the way it claims to."""
+    model_config = dict(
+        pdg_vocab=4,
+        mat_vocab=2,
+        emb_dim=16,
+        dropout=0.1,
+        k_max=K_MAX,
+        expert_hidden_dim=16,
+        expert_n_blocks=2,
+        conditioning="physical",
+        router={"enabled": True, "type": "pdg", "n_experts": 3},
+    )
+    with pytest.raises(ValueError, match="physical"):
+        build_models(model_config)
+
+
+def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditioning():
+    model_config = dict(
+        pdg_vocab=4,
+        mat_vocab=2,
+        emb_dim=16,
+        dropout=0.1,
+        k_max=K_MAX,
+        expert_hidden_dim=16,
+        expert_n_blocks=2,
+        conditioning="physical",
+        router={
+            "enabled": True,
+            "type": "composed",
+            "axis0_type": "energy",
+            "axis0_n_experts": 2,
+            "axis1_type": "pdg",
+            "axis1_n_experts": 3,
+        },
+    )
+    with pytest.raises(ValueError, match="physical"):
+        build_models(model_config)
+
+
 # ── ProcessRouter ────────────────────────────────────────────────────────────
 
 

tests/test_setup_cache.py

@@ -0,0 +1,327 @@
+import json
+import os
+import stat
+
+import numpy as np
+import pandas as pd
+import pytest
+
+from giant.data import setup_cache
+from giant.data.setup_cache import NormalizerEntry, SetupCache
+from giant.data.transforms import Normalizer
+
+
+def _touch_parquet(path, n=1):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    pd.DataFrame(
+        {"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}
+    ).to_parquet(path)
+    return path
+
+
+def _normalizer(width=3):
+    norm = Normalizer()
+    norm.mean = np.zeros(width, dtype=np.float32)
+    norm.std = np.ones(width, dtype=np.float32)
+    return norm
+
+
+def _entry(n_train_steps=100, sample=None):
+    sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) if sample is None else sample
+    return NormalizerEntry(
+        _normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample
+    )
+
+
+# ── sidecar_path ─────────────────────────────────────────────────────────
+
+
+def test_sidecar_path_single_file(tmp_path):
+    f = tmp_path / "shard.parquet"
+    assert (
+        setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
+    )
+
+
+def test_sidecar_path_directory(tmp_path):
+    d = tmp_path / "pbwo4"
+    assert setup_cache.sidecar_path(d) == tmp_path / "pbwo4.giant_train_cache.json"
+
+
+def test_sidecar_path_manifest(tmp_path):
+    m = tmp_path / "pools" / "full.manifest"
+    assert (
+        setup_cache.sidecar_path(m)
+        == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
+    )
+
+
+# ── fingerprint_files ────────────────────────────────────────────────────
+
+
+def test_fingerprint_files_order_preserving(tmp_path):
+    a = _touch_parquet(tmp_path / "a.parquet")
+    b = _touch_parquet(tmp_path / "b.parquet")
+
+    forward = setup_cache.fingerprint_files([a, b])
+    backward = setup_cache.fingerprint_files([b, a])
+
+    assert forward[0][0] == str(a.resolve())
+    assert forward[1][0] == str(b.resolve())
+    assert backward[0][0] == str(b.resolve())
+    assert forward != backward
+
+
+# ── save / load round trip ──────────────────────────────────────────────
+
+
+def test_save_load_round_trip(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+
+    cache = SetupCache.empty(files)
+    cache.vocab = ({11: 0, 22: 1}, {"G4_AIR": 0})
+    cache.event_index = (np.array([1, 2, 3]), np.array([10, 20, 30]))
+    cache.proc_maps[4] = {"eIoni": 0, "phot": 1}
+    cache.normalizers["valfrac=0.1_seed=0_cond=physical"] = _entry()
+
+    setup_cache.save(data, files, cache)
+    loaded = setup_cache.load(data, files)
+
+    assert loaded is not None
+    assert loaded.vocab == ({11: 0, 22: 1}, {"G4_AIR": 0})
+    assert loaded.event_index is not None
+    np.testing.assert_array_equal(loaded.event_index[0], [1, 2, 3])
+    np.testing.assert_array_equal(loaded.event_index[1], [10, 20, 30])
+    assert loaded.proc_maps == {4: {"eIoni": 0, "phot": 1}}
+    entry = loaded.normalizers["valfrac=0.1_seed=0_cond=physical"]
+    assert entry.cond_norm.mean is not None
+    np.testing.assert_allclose(entry.cond_norm.mean, np.zeros(3, dtype=np.float32))
+    assert entry.n_train_steps == 100
+    np.testing.assert_allclose(entry.energy_quantiles, [1.0, 2.0, 3.0])
+
+
+def test_load_missing_sidecar_returns_none(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    assert setup_cache.load(data, [data]) is None
+
+
+def test_load_corrupt_json_returns_none(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    setup_cache.sidecar_path(data).write_text("not valid json {{{")
+    assert setup_cache.load(data, [data]) is None
+
+
+def test_load_invalidates_on_dims_mismatch(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    path = setup_cache.sidecar_path(data)
+    raw = json.loads(path.read_text())
+    raw["dims"]["K_MAX"] = raw["dims"]["K_MAX"] + 1
+    path.write_text(json.dumps(raw))
+
+    assert setup_cache.load(data, files) is None
+
+
+def test_load_invalidates_on_format_version_mismatch(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    path = setup_cache.sidecar_path(data)
+    raw = json.loads(path.read_text())
+    raw["format_version"] = raw["format_version"] + 1
+    path.write_text(json.dumps(raw))
+
+    assert setup_cache.load(data, files) is None
+
+
+def test_load_invalidates_on_file_content_change(tmp_path):
+    data = tmp_path / "shard.parquet"
+    _touch_parquet(data, n=1)
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    _touch_parquet(data, n=50)  # different size -> fingerprint changes
+
+    assert setup_cache.load(data, files) is None
+
+
+def test_load_soft_warns_on_git_hash_mismatch_but_still_hits(tmp_path, capsys):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    cache = SetupCache.empty(files)
+    cache.git_hash = "not-a-real-git-hash"
+    setup_cache.save(data, files, cache)
+
+    loaded = setup_cache.load(data, files)
+
+    assert loaded is not None
+    err = capsys.readouterr().err
+    assert "not-a-real-git-hash" in err
+
+
+# ── save: atomicity / robustness ────────────────────────────────────────
+
+
+def test_save_is_atomic_no_stray_tmp_file(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    leftovers = [p for p in tmp_path.iterdir() if ".tmp." in p.name]
+    assert leftovers == []
+
+
+def test_save_degrades_gracefully_on_permission_error(tmp_path):
+    if os.geteuid() == 0:
+        pytest.skip("root bypasses directory permission bits")
+    data_dir = tmp_path / "ro"
+    data_dir.mkdir()
+    data = _touch_parquet(data_dir / "shard.parquet")
+    files = [data]
+
+    warnings = []
+    mode = data_dir.stat().st_mode
+    data_dir.chmod(stat.S_IREAD | stat.S_IEXEC)
+    try:
+        setup_cache.save(data, files, SetupCache.empty(files), echo=warnings.append)
+    finally:
+        data_dir.chmod(mode)
+
+    assert any("could not write" in w for w in warnings)
+    assert not setup_cache.sidecar_path(data).exists()
+
+
+def test_save_merges_non_colliding_normalizer_keys(tmp_path):
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+
+    cache1 = SetupCache.empty(files)
+    cache1.normalizers["k1"] = _entry(n_train_steps=1)
+    setup_cache.save(data, files, cache1)
+
+    cache2 = SetupCache.empty(files)
+    cache2.normalizers["k2"] = _entry(n_train_steps=2)
+    setup_cache.save(data, files, cache2)
+
+    loaded = setup_cache.load(data, files)
+    assert loaded is not None
+    assert set(loaded.normalizers.keys()) == {"k1", "k2"}
+    assert loaded.normalizers["k1"].n_train_steps == 1
+    assert loaded.normalizers["k2"].n_train_steps == 2
+
+
+def test_save_is_serialized_against_concurrent_writers(tmp_path):
+    """Without the flock in setup_cache.save(), two concurrent writers can
+    both load() the same base state and merge their own section in
+    independently, so whichever os.replace() lands last silently drops the
+    other's key — a lost-update race, not a corrupt file. Each of these
+    threads writes a distinct normalizer key many times over; if the
+    load-merge-write critical section isn't actually serialized, at least
+    one thread's key is likely to go missing from the final merged cache."""
+    import threading
+
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    n_writers, n_rounds = 6, 15
+
+    def _writer(idx: int) -> None:
+        for r in range(n_rounds):
+            cache = SetupCache.empty(files)
+            cache.normalizers[f"k{idx}"] = _entry(n_train_steps=r)
+            setup_cache.save(data, files, cache)
+
+    threads = [threading.Thread(target=_writer, args=(i,)) for i in range(n_writers)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    loaded = setup_cache.load(data, files)
+    assert loaded is not None
+    assert set(loaded.normalizers.keys()) == {f"k{i}" for i in range(n_writers)}
+    for i in range(n_writers):
+        assert loaded.normalizers[f"k{i}"].n_train_steps == n_rounds - 1
+
+
+# ── energy_quantiles_from_sample / energy_quantile_at ───────────────────
+
+
+def test_energy_quantiles_from_sample_empty():
+    result = setup_cache.energy_quantiles_from_sample(np.empty(0, dtype=np.float32))
+    assert result.size == 0
+
+
+def test_energy_quantiles_from_sample_has_fixed_grid_size():
+    sample = np.random.default_rng(0).normal(size=5000).astype(np.float32)
+    result = setup_cache.energy_quantiles_from_sample(sample)
+    assert result.shape == (setup_cache.ENERGY_QUANTILE_LEVELS,)
+    assert result[0] == pytest.approx(sample.min(), abs=1e-3)
+    assert result[-1] == pytest.approx(sample.max(), abs=1e-3)
+
+
+def test_energy_quantile_at_matches_direct_quantile_on_stored_grid():
+    sample = np.random.default_rng(1).exponential(size=20_000).astype(np.float32)
+    grid = setup_cache.energy_quantiles_from_sample(sample)
+
+    levels = np.linspace(0.0, 1.0, 5)
+    got = setup_cache.energy_quantile_at(grid, levels)
+    expected = np.quantile(sample, levels)
+
+    np.testing.assert_allclose(got, expected, rtol=0.05)
+
+
+def test_energy_quantile_at_median_of_two_points():
+    grid = np.array([0.0, 10.0], dtype=np.float32)
+    result = setup_cache.energy_quantile_at(grid, np.array([0.0, 0.5, 1.0]))
+    np.testing.assert_allclose(result, [0.0, 5.0, 10.0])
+
+
+# ── n_train_steps_for_split ──────────────────────────────────────────────
+
+
+def test_n_train_steps_for_split_matches_full_scan():
+    unique_ids = np.array([1, 2, 3, 4, 5])
+    counts = np.array([10, 20, 30, 40, 50])
+    train_events_arr = np.array([2, 4, 5])
+
+    result = setup_cache.n_train_steps_for_split(unique_ids, counts, train_events_arr)
+
+    assert result == 20 + 40 + 50
+
+
+# ── compute_event_index_from_files: cross-file event_id offsetting ─────────
+
+
+def test_compute_event_index_from_files_offsets_colliding_ids(tmp_path):
+    """Two files that each restart event_id from 0 (one Geant4 job per file)
+    must not have their same-numbered events collapsed into one by
+    np.unique — each file's ids get shifted by a distinct offset first (see
+    giant.data.loader.event_id_offset)."""
+    path_a = tmp_path / "a.parquet"
+    path_b = tmp_path / "b.parquet"
+    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path_a)
+    pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path_b)
+
+    unique_ids, counts = setup_cache.compute_event_index_from_files([path_a, path_b])
+
+    assert len(unique_ids) == 6
+    assert int(counts.sum()) == 6
+    assert np.all(counts == 1)
+
+
+def test_compute_event_index_from_files_single_file_unaffected(tmp_path):
+    """A single file's ids are offset by 0 (event_id_offset(0) == 0), so a
+    single-file load's unique ids/counts are unchanged by the offsetting."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [5, 5, 7]}).to_parquet(path)
+
+    unique_ids, counts = setup_cache.compute_event_index_from_files([path])
+
+    np.testing.assert_array_equal(unique_ids, [5, 7])
+    np.testing.assert_array_equal(counts, [2, 1])

tests/test_train.py

@@ -0,0 +1,97 @@
+"""Tests for giant/train.py helpers."""
+
+from giant.train import _gumbel_tau, _wandb_run_config
+
+
+def test_gumbel_tau_at_step_zero_is_start():
+    assert _gumbel_tau(0, 1000, 1.0, 0.1) == 1.0
+
+
+def test_gumbel_tau_at_total_steps_is_end():
+    assert abs(_gumbel_tau(1000, 1000, 1.0, 0.1) - 0.1) < 1e-9
+
+
+def test_gumbel_tau_interpolates_linearly_midway():
+    assert abs(_gumbel_tau(500, 1000, 1.0, 0.1) - 0.55) < 1e-9
+
+
+def test_gumbel_tau_clamps_beyond_total_steps():
+    assert _gumbel_tau(5000, 1000, 1.0, 0.1) == _gumbel_tau(1000, 1000, 1.0, 0.1)
+
+
+def test_gumbel_tau_handles_zero_total_steps():
+    # total_steps=0 is guarded to 1 internally: step=0 gives zero progress
+    # (still tau_start), any step>=1 immediately clamps to full progress.
+    assert _gumbel_tau(0, 0, 1.0, 0.1) == 1.0
+    assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9
+
+
+def _base_wandb_kwargs(**overrides):
+    kwargs = dict(
+        mode="flow",
+        epochs=30,
+        lr=3e-4,
+        warmup_epochs=3,
+        weight_decay=0.01,
+        ema_decay=0.9999,
+        lambda_nsec=0.1,
+        lambda_s2=1.0,
+        lambda_balance=0.035,
+        lambda_proc=0.0,
+        lambda_entropy=0.0,
+        gumbel_tau_start=1.0,
+        gumbel_tau_end=0.1,
+        n_critic=5,
+        gp_weight=10.0,
+        model_config={"router": {"enabled": False}},
+        stage1_params=100,
+        sec_decoder_params=50,
+        critic_params=0,
+        sec_critic_params=0,
+        total_params=150,
+    )
+    kwargs.update(overrides)
+    return kwargs
+
+
+def test_wandb_run_config_omits_router_knobs_when_router_disabled():
+    cfg = _wandb_run_config(**_base_wandb_kwargs())
+    for key in (
+        "lambda_balance",
+        "lambda_proc",
+        "lambda_entropy",
+        "gumbel_tau_start",
+        "gumbel_tau_end",
+    ):
+        assert key not in cfg
+    # still present, nested, regardless of router state
+    assert cfg["model"] == {"router": {"enabled": False}}
+
+
+def test_wandb_run_config_includes_router_knobs_when_router_enabled():
+    cfg = _wandb_run_config(
+        **_base_wandb_kwargs(model_config={"router": {"enabled": True}})
+    )
+    assert cfg["lambda_balance"] == 0.035
+    assert cfg["lambda_proc"] == 0.0
+    assert cfg["lambda_entropy"] == 0.0
+    assert cfg["gumbel_tau_start"] == 1.0
+    assert cfg["gumbel_tau_end"] == 0.1
+
+
+def test_wandb_run_config_omits_wgan_knobs_when_mode_is_not_wgan():
+    cfg = _wandb_run_config(**_base_wandb_kwargs(mode="flow"))
+    assert "n_critic" not in cfg
+    assert "gp_weight" not in cfg
+
+
+def test_wandb_run_config_includes_wgan_knobs_when_mode_is_wgan():
+    cfg = _wandb_run_config(**_base_wandb_kwargs(mode="wgan"))
+    assert cfg["n_critic"] == 5
+    assert cfg["gp_weight"] == 10.0
+
+
+def test_wandb_run_config_handles_missing_model_config():
+    cfg = _wandb_run_config(**_base_wandb_kwargs(model_config=None))
+    assert cfg["model"] == {}
+    assert "lambda_balance" not in cfg

tests/test_transforms.py

@@ -1,9 +1,12 @@
+import warnings
+
 import numpy as np
 import pytest
 from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX
 from giant.data.transforms import (
     build_cond_features,
     build_features,
+    encode_secondaries,
     energy_simplex_decode,
     energy_simplex_encode,
     inv_local_frame_rotation,
@@ -12,7 +15,10 @@ from giant.data.transforms import (
     log_transform,
     Normalizer,
     reconstruct_post_pos,
+    sorted_membership,
     travel_direction,
+    _vectorized_map_lookup,
+    _WelfordAccumulator,
 )
 
 
@@ -21,6 +27,50 @@ def test_log_transform_invertible():
     np.testing.assert_allclose(inv_log_transform(log_transform(x)), x, rtol=1e-5)
 
 
+def test_log_transform_raises_on_input_below_negative_eps():
+    """A meaningfully negative input (upstream data corruption, not float
+    noise near 0) must raise instead of silently returning NaN."""
+    x = np.array([1.0, -5.0], dtype=np.float32)
+    with np.errstate(invalid="ignore"), pytest.raises(ValueError, match="non-finite"):
+        log_transform(x)
+
+
+def test_log_transform_raises_on_nan_input():
+    x = np.array([1.0, np.nan], dtype=np.float32)
+    with pytest.raises(ValueError, match="non-finite"):
+        log_transform(x)
+
+
+def test_encode_secondaries_warns_when_sec_energies_exceed_e_sec():
+    """sec_E_list summing to more than e_sec (before the last slot is even
+    reached) is a real upstream data mismatch — must warn instead of
+    silently saturating the overflowing slot's stick-breaking logit via the
+    _EPS floor. (A single slot alone exceeding what's left of the budget is
+    the normal, expected "last slot takes the remainder" case and must NOT
+    warn — the mismatch here is the *cumulative* sum through an earlier
+    slot already exceeding e_sec.)"""
+    sec_E_list = np.array([[5.0, 4.0, 1.0]], dtype=np.float32)  # sums to 10
+    sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 3, 1)).astype(np.float32)
+    sec_valid = np.array([[True, True, True]])
+    e_sec = np.array([6.0], dtype=np.float32)  # cumsum already 9 by slot 2
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    with pytest.warns(UserWarning, match="sec_E_list summing to more than e_sec"):
+        encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
+
+
+def test_encode_secondaries_no_warning_when_energies_are_consistent():
+    sec_E_list = np.array([[3.0, 2.0]], dtype=np.float32)  # sums to 5
+    sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 2, 1)).astype(np.float32)
+    sec_valid = np.array([[True, True]])
+    e_sec = np.array([6.0], dtype=np.float32)  # >= 5, no shortfall
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
+
+
 def test_local_frame_rotation_noop_when_aligned():
     N = 8
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
@@ -65,10 +115,43 @@ def test_local_frame_rotation_rejects_near_zero_pre_dir():
     post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
     with pytest.raises(ValueError, match="near-zero norm"):
         local_frame_rotation(pre_dir, post_dir)
-    with pytest.raises(ValueError, match="near-zero norm"):
+
+
+def test_local_frame_rotation_rejects_nan_pre_dir():
+    """A NaN pre_dir must raise loudly — `norm < 1e-6` is False for NaN, so
+    without an explicit isfinite check this would silently poison the
+    rotation (and any normalizer stats it feeds) instead of erroring."""
+    pre_dir = np.array([[np.nan, 0.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32)
+    post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
+    with pytest.raises(ValueError, match="non-finite"):
+        local_frame_rotation(pre_dir, post_dir)
+    with pytest.raises(ValueError, match="non-finite"):
         inv_local_frame_rotation(pre_dir, post_dir)
 
 
+def test_local_frame_rotation_antipodal_pre_dir_uses_x_axis_convention():
+    """pre_dir ~ -ẑ (near-exact backscatter) is a second axis_norm~0
+    degeneracy besides pre_dir ~ +ẑ; unlike the forward case, the Rodrigues
+    axis-dependent terms are NOT negligible there ((1-cos_t)~2), so the x̂
+    fallback is a real (if arbitrary and physically rare) convention choice
+    rather than a no-op. Pin it explicitly — angle-preservation and the
+    round-trip property must still hold even though the "roll" is degenerate.
+    """
+    pre_dir = np.array([[0.0, 0.0, -1.0]], dtype=np.float32)
+    post_dir = np.array([[0.3, 0.4, 0.5]], dtype=np.float32)
+    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
+
+    rotated = local_frame_rotation(pre_dir, post_dir)
+
+    cos_before = (pre_dir * post_dir).sum(axis=1)
+    cos_after = rotated[:, 2]
+    np.testing.assert_allclose(cos_after, cos_before, atol=1e-5)
+    np.testing.assert_allclose(np.linalg.norm(rotated, axis=1), 1.0, atol=1e-5)
+
+    recovered = inv_local_frame_rotation(pre_dir, rotated)
+    np.testing.assert_allclose(recovered, post_dir, atol=1e-5)
+
+
 def test_local_frame_rotation_normalizes_non_unit_pre_dir():
     """A pre_dir with float32-drift norm (not exactly 1) must still produce the
     same result as its exactly-normalized counterpart, not a skewed frame."""
@@ -440,3 +523,165 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
             cond_normalizer=legacy_norm,
             conditioning="physical",
         )
+
+
+# ── sorted_membership / _vectorized_map_lookup ──────────────────────────────
+
+
+def test_sorted_membership_matches_np_isin():
+    rng = np.random.default_rng(0)
+    sorted_arr = np.unique(rng.integers(0, 10_000, size=500))
+    values = rng.integers(-100, 10_100, size=2_000)  # some in, some out of range
+    # values deliberately not sorted
+    rng.shuffle(values)
+
+    result = sorted_membership(values, sorted_arr)
+    expected = np.isin(values, sorted_arr)
+    np.testing.assert_array_equal(result, expected)
+
+
+def test_sorted_membership_empty_sorted_arr():
+    values = np.array([1, 2, 3])
+    sorted_arr = np.array([], dtype=np.int64)
+    result = sorted_membership(values, sorted_arr)
+    np.testing.assert_array_equal(result, np.zeros(3, dtype=bool))
+
+
+def test_vectorized_map_lookup_matches_dict_comprehension_int_keys():
+    rng = np.random.default_rng(1)
+    keys = np.unique(rng.integers(-1000, 1000, size=200))
+    mapping = {int(k): i for i, k in enumerate(keys)}
+    values = rng.choice(keys, size=500)
+
+    result = _vectorized_map_lookup(values, mapping)
+    expected = np.array([mapping[int(v)] for v in values], dtype=np.int64)
+    np.testing.assert_array_equal(result, expected)
+
+
+def test_vectorized_map_lookup_matches_dict_comprehension_str_keys():
+    mapping = {"PbWO4": 0, "G4_AIR": 1, "G4_Fe": 2}
+    values = np.array(["G4_Fe", "PbWO4", "G4_AIR", "PbWO4"], dtype=object)
+
+    result = _vectorized_map_lookup(values, mapping)
+    expected = np.array([mapping[str(v)] for v in values], dtype=np.int64)
+    np.testing.assert_array_equal(result, expected)
+
+
+def test_vectorized_map_lookup_raises_keyerror_on_missing_value():
+    mapping = {1: 0, 2: 1}
+    values = np.array([1, 2, 3])
+    with pytest.raises(KeyError):
+        _vectorized_map_lookup(values, mapping)
+
+
+def test_vectorized_map_lookup_strict_false_dummy_indexes_unmapped_values():
+    """strict=False must leave found values untouched and only dummy-index
+    (0) the unmapped ones — never raise, and never disturb a value that IS
+    in the mapping (e.g. one that happens to map to a nonzero index)."""
+    mapping = {1: 5, 2: 7}
+    values = np.array([1, 99, 2, 100])
+    result = _vectorized_map_lookup(values, mapping, strict=False)
+    np.testing.assert_array_equal(result, [5, 0, 7, 0])
+
+
+def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_material():
+    """conditioning="physical" must not KeyError on a pdg/material outside
+    the training-dataset vocab (mat_map/pdg_map) — that's the entire point
+    of the mode (see giant.rollout's known_pdg gate for the paired fix).
+    "embedding" mode must still raise, since cond_cat IS the conditioning
+    signal there. Note this is specifically about the dataset-scoped
+    vocab index, not giant.materials' physical-properties table — a
+    material must still be a real, known Geant4 material (e.g. "G4_Pb",
+    just not one *this* mat_map happened to include) for "physical" mode
+    to derive its Z_eff/A_eff/density/X0/λ_int; a genuinely unknown
+    material name correctly still raises via giant.materials, same as the
+    documented G4_LYSO precedent — that's a separate, intentional guard."""
+    pdg_map = {11: 0, 22: 1}
+    mat_map = {"G4_AIR": 0}
+    data = {
+        "pre_pos": np.zeros((1, 3), dtype=np.float32),
+        "pre_E": np.array([10.0], dtype=np.float32),
+        "pre_dir": np.array([[0.0, 0.0, 1.0]], dtype=np.float32),
+        "layer_id": np.array([0], dtype=np.int32),
+        "pdg": np.array([13], dtype=np.int64),  # not in pdg_map
+        "material": np.array(["G4_Pb"], dtype=object),  # not in mat_map
+        "mass": np.array([105.7], dtype=np.float32),
+        "charge": np.array([-1.0], dtype=np.float32),
+    }
+
+    cond_cont, cond_cat = build_cond_features(
+        data, pdg_map, mat_map, conditioning="physical"
+    )
+    assert cond_cont.shape[-1] == COND_DIM
+    np.testing.assert_array_equal(cond_cat, [[0, 0]])  # dummy indices, no raise
+
+    with pytest.raises(KeyError):
+        build_cond_features(data, pdg_map, mat_map, conditioning="embedding")
+
+
+# ── _WelfordAccumulator ──────────────────────────────────────────────────────
+
+
+def test_welford_accumulator_matches_direct_mean_std_over_many_chunks():
+    rng = np.random.default_rng(5)
+    F = 4
+    chunks = [rng.standard_normal((rng.integers(1, 50), F)) * 10 + 3 for _ in range(20)]
+    full = np.concatenate(chunks, axis=0)
+
+    acc = _WelfordAccumulator(F)
+    for chunk in chunks:
+        acc.update(chunk)
+    norm = acc.to_normalizer()
+
+    assert norm.mean is not None and norm.std is not None
+    np.testing.assert_allclose(norm.mean, full.mean(axis=0), rtol=1e-5, atol=1e-5)
+    np.testing.assert_allclose(norm.std, full.std(axis=0), rtol=1e-5, atol=1e-5)
+    assert acc.n == full.shape[0]
+
+
+def test_welford_accumulator_single_chunk():
+    rng = np.random.default_rng(6)
+    X = rng.standard_normal((100, 3)) * 5 - 2
+
+    acc = _WelfordAccumulator(3)
+    acc.update(X)
+    norm = acc.to_normalizer()
+
+    assert norm.mean is not None and norm.std is not None
+    np.testing.assert_allclose(norm.mean, X.mean(axis=0), rtol=1e-5)
+    np.testing.assert_allclose(norm.std, X.std(axis=0), rtol=1e-5)
+
+
+def test_welford_accumulator_matches_naive_running_mean_reference():
+    """The chunk-local-mean + Chan-merge formula must agree with the naive
+    textbook streaming update (subtract the *running* mean before and after
+    updating it) that it replaces, within float64 rounding tolerance."""
+    rng = np.random.default_rng(7)
+    F = 3
+    chunks = [rng.standard_normal((rng.integers(1, 40), F)) for _ in range(15)]
+
+    def naive_update(mean, M2, n, X):
+        X = np.asarray(X, dtype=np.float64)
+        B = X.shape[0]
+        new_n = n + B
+        delta = X - mean
+        mean = mean + delta.sum(0) / new_n
+        delta2 = X - mean
+        M2 = M2 + (delta * delta2).sum(0)
+        return mean, M2, new_n
+
+    naive_mean = np.zeros(F)
+    naive_M2 = np.zeros(F)
+    naive_n = 0
+    for chunk in chunks:
+        naive_mean, naive_M2, naive_n = naive_update(
+            naive_mean, naive_M2, naive_n, chunk
+        )
+
+    acc = _WelfordAccumulator(F)
+    for chunk in chunks:
+        acc.update(chunk)
+
+    assert acc.n == naive_n
+    np.testing.assert_allclose(acc._mean, naive_mean, rtol=1e-9, atol=1e-9)
+    np.testing.assert_allclose(acc._M2, naive_M2, rtol=1e-9, atol=1e-9)

uv.lock

@@ -36,6 +36,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/1e/d3/26bf1008eb3d2daa8ef4cacc7f3bfdc11818d111f7e2d0201bc6e3b49d45/annotated_doc-0.0.4-py3-none-any.whl", hash = "sha256:571ac1dc6991c450b25a9c2d84a3705e2ae7a53467b5d111c24fa8baabbed320", size = 5303, upload-time = "2025-11-10T22:07:40.673Z" },
 ]
 
+[[package]]
+name = "annotated-types"
+version = "0.8.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" },
+]
+
 [[package]]
 name = "appnope"
 version = "0.1.4"
@@ -125,6 +134,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/44/a1/70ebfffd6c6edc6034a547838ee46287c65ed89f710592ddc39c76b4a5a8/awkward_cpp-53-cp314-cp314t-win_arm64.whl", hash = "sha256:1be0c1d87d9f4fdf94b767a061df849f1bb21579d302b2996fb101527fc80a97", size = 551257, upload-time = "2026-06-08T12:31:56.319Z" },
 ]
 
+[[package]]
+name = "certifi"
+version = "2026.7.22"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" },
+]
+
 [[package]]
 name = "cffi"
 version = "2.0.0"
@@ -182,6 +200,79 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ae/3a/dbeec9d1ee0844c679f6bb5d6ad4e9f198b1224f4e7a32825f47f6192b0c/cffi-2.0.0-cp314-cp314t-win_arm64.whl", hash = "sha256:0a1527a803f0a659de1af2e1fd700213caba79377e27e4693648c2923da066f9", size = 184195, upload-time = "2025-09-08T23:23:43.004Z" },
 ]
 
+[[package]]
+name = "charset-normalizer"
+version = "3.4.9"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/bd/2a/23f34ec9d04624958e137efdc394888716353190e75f25dd22c7a2c7a8aa/charset_normalizer-3.4.9.tar.gz", hash = "sha256:673611bbd43f0810bec0b0f028ddeaaa501190339cac411f347ac76917c3ae7b", size = 152439, upload-time = "2026-07-07T14:34:58.454Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/70/4a/ecbd131485c07fcdfad54e28946d513e3da22ef3b4bd854dcafae54ec739/charset_normalizer-3.4.9-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:45b0cc4e3556cd875e09102988d1ab8356c998b596c9fced84547c8138b487a0", size = 319300, upload-time = "2026-07-07T14:33:15.666Z" },
+    { url = "https://files.pythonhosted.org/packages/ec/96/5d9364e3342d69f3a045e1777bc47c85c383e6e9466d561b33fdb419d1f9/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9b2aff1c7b3884512b9512c3eaadd9bab39fb45042ffaaa1dd08ff2b9f8109d9", size = 215802, upload-time = "2026-07-07T14:33:17.031Z" },
+    { url = "https://files.pythonhosted.org/packages/4b/4c/5361f9aa7f2cb58d94f2ab831b3d493f69efb1d239654b4744e3c09527cb/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:9104ed0bd76a429d46f9ec0dbc9b08ad1d2dcdf2b00a5a0daa1c145329b35b44", size = 237171, upload-time = "2026-07-07T14:33:18.576Z" },
+    { url = "https://files.pythonhosted.org/packages/50/78/ce342ca4ff30b2eb49fe6d9578df85974f90c67d294113e94efdd9664cbd/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:7b86a2b16095d250c6f58b3d9b2eee6f4147754344f3dab0922f7c9bf7d226c9", size = 233075, upload-time = "2026-07-07T14:33:20.084Z" },
+    { url = "https://files.pythonhosted.org/packages/01/c4/4fa4c8b3097a11f3c5f09a35b72ed6855fb1d332469504962ab7bafcc702/charset_normalizer-3.4.9-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5e226f6218febc71f6c1fc2fafb91c226f75bdc1d8fb12d66823716e891608fd", size = 224256, upload-time = "2026-07-07T14:33:21.747Z" },
+    { url = "https://files.pythonhosted.org/packages/87/3a/ad914516df7e358a81aae018caa5e0470ba827fa6d763b1d2e87d920a5f6/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_armv7l.whl", hash = "sha256:90c44bc373b7687f6948b693cceaea1348ae0975d7474746559494468e3c1d84", size = 208784, upload-time = "2026-07-07T14:33:23.313Z" },
+    { url = "https://files.pythonhosted.org/packages/d7/74/3c12f9755717dfe5c5c87da63f35d765fa0c00382ec26bf23f7fae34f2ba/charset_normalizer-3.4.9-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:9cdef90ae47919cae358d8ab15797a800ed41da7aba5d72419fb510729e2ed4b", size = 219928, upload-time = "2026-07-07T14:33:24.814Z" },
+    { url = "https://files.pythonhosted.org/packages/33/9a/895095b83e7907abd6d3d99aad3a38ad0d9686cc186cb0c94c24320fe63e/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:60f44ade2cf573dad7a277e6f8ca9a51a21dda572b13bd7d8539bb3cd5dbedde", size = 218489, upload-time = "2026-07-07T14:33:26.42Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/34/ef5c05f412f42520d7709b7d3784d19640839eb7366ded1755511585429f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_armv7l.whl", hash = "sha256:a1786910334ed46ab1dd73222f2cd1e05c2c3bb39f6dddb4f8b36fc382058a39", size = 210267, upload-time = "2026-07-07T14:33:27.952Z" },
+    { url = "https://files.pythonhosted.org/packages/83/dc/9b29fa4412b318bf3bfea985c35d67eb55e04b59a7c3f2237168b0e0be6f/charset_normalizer-3.4.9-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:03d07803992c6c7bbc976327f34b18b6160327fc81cb82c9d504720ac0be3b62", size = 226030, upload-time = "2026-07-07T14:33:29.397Z" },
+    { url = "https://files.pythonhosted.org/packages/0e/42/6dbc00b8cd16011691203e33570fa42ed5746599a2e878112d16eab403a3/charset_normalizer-3.4.9-cp312-cp312-win32.whl", hash = "sha256:78841cccf1af7b40f6f716338d50c0902dbe88d9f800b3c973b7a9a0a693a642", size = 151185, upload-time = "2026-07-07T14:33:30.781Z" },
+    { url = "https://files.pythonhosted.org/packages/80/cc/f920afd1a23c58ccd53c1d36085a71893a4737ff5e66e0371efab6809850/charset_normalizer-3.4.9-cp312-cp312-win_amd64.whl", hash = "sha256:4b3dac63058cc36820b0dd072f89898604e2d39686fe05321729d00d8ac185a0", size = 162557, upload-time = "2026-07-07T14:33:32.176Z" },
+    { url = "https://files.pythonhosted.org/packages/f0/e6/0386d43a261ff4e4b30c5857af7df877254b46bec7b9d1b74b6bf969a90b/charset_normalizer-3.4.9-cp312-cp312-win_arm64.whl", hash = "sha256:78fa18e436a1a0e58dbd7e02fc4473f3f32cceb12df9dfca542d075961c307d2", size = 152665, upload-time = "2026-07-07T14:33:33.711Z" },
+    { url = "https://files.pythonhosted.org/packages/b2/06/97ec2aeae780b31d742b6352218b43841a6871e2564578ca522dce4a45c3/charset_normalizer-3.4.9-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:440eede837960000d74978f0eba527be106b5b9aee0daf779d395276ed0b0614", size = 317688, upload-time = "2026-07-07T14:33:35.408Z" },
+    { url = "https://files.pythonhosted.org/packages/d0/39/8ff066c672434225f8d25f8b739f992af250944392173dcc88362681c9bf/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:21e764fd1e70b6a3e205a0e46f3051701f98a8cb3fad66eeb80e48bb502f8698", size = 214982, upload-time = "2026-07-07T14:33:36.996Z" },
+    { url = "https://files.pythonhosted.org/packages/92/8f/3a47a3667c83c2df9483d91644c6c107de3bf8874aa1793da9d3012eb986/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:e4fd89cc178bced6ad29cb3e6dd4aa63fa5017c3524dbd0b25998fb64a87cc8b", size = 236460, upload-time = "2026-07-07T14:33:38.536Z" },
+    { url = "https://files.pythonhosted.org/packages/f1/60/b22cdbee7e4013dab8b0d7647fc6181120fbbbc8f7025c226d15bd5a47fc/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:bd47ba7fc3ca94896759ea0109775132d3e7ab921fbf54038e1bab2e46c313c9", size = 232003, upload-time = "2026-07-07T14:33:40.059Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/f8/72eb13dcabe7257035cea8aefd922caad2f110d252bf9f67c4c2ca763aee/charset_normalizer-3.4.9-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:84fd18bcc17526fc2b3c1af7d2b9217d32c9c04448c16ec693b9b4f1985c3d33", size = 223149, upload-time = "2026-07-07T14:33:41.631Z" },
+    { url = "https://files.pythonhosted.org/packages/b0/3e/faee8f9de92b14ee1198e9163252bb15efee7301b31256a3b6d9ebfdd0dd/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_armv7l.whl", hash = "sha256:5b10cd92fc5c498b35a8635df6d5a100207f88b63a4dc1de7ef9a548e1e2cd63", size = 207901, upload-time = "2026-07-07T14:33:43.209Z" },
+    { url = "https://files.pythonhosted.org/packages/3a/25/45f30093ae27dd7b92a793b61882a38685f993700113ca36e0c9c14965e1/charset_normalizer-3.4.9-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a4fbdde9dd4a9ce5fd52c2b3a347bb50cc89483ef783f1cb00d408c13f7a96c0", size = 219176, upload-time = "2026-07-07T14:33:44.725Z" },
+    { url = "https://files.pythonhosted.org/packages/48/18/c8f397329c35e32f6a837e488986f4ae03bd2abebc453b48714991630c2f/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:416c229f77e5ea25b3dfd4b582f8d73d7e43c22320302b9ab128a2d3a0b38efe", size = 217356, upload-time = "2026-07-07T14:33:46.192Z" },
+    { url = "https://files.pythonhosted.org/packages/86/7e/5ce0bba863470fd1902d5e5843968951bddf38abe4742fc97116ef4598b3/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_armv7l.whl", hash = "sha256:75286256590a6320cf106a0d28970d3560aad9ee09aa7b34fb40524792436d35", size = 209614, upload-time = "2026-07-07T14:33:47.705Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/ef/2473d3c4d869155be4af1191111d59c4d5c4e0173026f7e85b176e23bf65/charset_normalizer-3.4.9-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:69b157c5d3292bcd443faca052f3096f637f1e074b98212a933c074ae23dc3b8", size = 224991, upload-time = "2026-07-07T14:33:49.238Z" },
+    { url = "https://files.pythonhosted.org/packages/d0/a3/53ddae3db108a088156aa8ddfafd411ebbc1340f48c5573f697b27f69a39/charset_normalizer-3.4.9-cp313-cp313-win32.whl", hash = "sha256:51307f5c71007673a2bf8232ad973483d281e74cb99c8c5a990af1eefa6277d9", size = 150622, upload-time = "2026-07-07T14:33:50.711Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/ef/6953a77c7cf2c2ff9998e6f575ab3e380119f100223381565a4f94c1f836/charset_normalizer-3.4.9-cp313-cp313-win_amd64.whl", hash = "sha256:fe2c7201c642b7c308f1675355ad7ff7b66acfe3541625efe5a3ad38f29d6115", size = 161947, upload-time = "2026-07-07T14:33:52.197Z" },
+    { url = "https://files.pythonhosted.org/packages/6e/fb/d560d1d1555debbfe7849d9cac6145c1b537709d79576bf22557ed803b82/charset_normalizer-3.4.9-cp313-cp313-win_arm64.whl", hash = "sha256:611057cc5d5c0afc743ba8be6bd828c17e0aaa8643f9d0a9b9bb7dea80eb8012", size = 152594, upload-time = "2026-07-07T14:33:53.486Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/8d/496817fa0944239ecae662dd57ea765cfeaec6a735f9f025d4b7b72e7143/charset_normalizer-3.4.9-cp314-cp314-macosx_10_15_universal2.whl", hash = "sha256:0327fcd59a935777d83410750c50600ee9571af2846f71ce40f25b13da1ef380", size = 317253, upload-time = "2026-07-07T14:33:54.994Z" },
+    { url = "https://files.pythonhosted.org/packages/2b/f9/ef4a69ea338ad3c0deceea0f5f7d2380ae8b52132b06d652cb0d2cd86706/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:8a79d9f4d8001473a30c163556b3c3bfebec837495a412dde78b51672f6134f9", size = 215898, upload-time = "2026-07-07T14:33:56.334Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/e7/5ddfd76fc061eb52de219658a4aa431cbacadf0a0219c8854f00da50d289/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:33bdcc2a32c0a0e861f60841a512c8acc658c87c2ac59d89e3a46dacf7d866e4", size = 236718, upload-time = "2026-07-07T14:33:57.9Z" },
+    { url = "https://files.pythonhosted.org/packages/49/ba/768fa3f36048d81c477a0ce61f813bc1454d80917ccfe550abd9f44f5e24/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:f840ed6d8ecba8255df8c42b87fadeda98ddfc6eeec05e2dc66e26d46dd6f58a", size = 232519, upload-time = "2026-07-07T14:33:59.811Z" },
+    { url = "https://files.pythonhosted.org/packages/f4/c4/b3e049d2aa3766180c78507110543d9d50894cc97f57de543f1be521dcdc/charset_normalizer-3.4.9-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c25fe15c70c59eb7c5ce8c06a1f3fa1da0ecc5ea1e7a5922c40fd2fa9b0d5046", size = 223143, upload-time = "2026-07-07T14:34:01.517Z" },
+    { url = "https://files.pythonhosted.org/packages/19/79/55c32d06d76ae4feafe053f061f3e3ab70bcf19f4007797ce8c3efda7830/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_armv7l.whl", hash = "sha256:f7fb7d750cfa0a070d2c24e831fd3481019a60dd317ea2b39acbcebc08b6ed81", size = 206742, upload-time = "2026-07-07T14:34:03.04Z" },
+    { url = "https://files.pythonhosted.org/packages/10/e0/47c079dd82d217c807479cd59ffd30af56307ea31c108b75758970459ad3/charset_normalizer-3.4.9-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:4d1c96a7a18b9690a4d46df09e3e3382406ae3213727cd1019ebade1c4a81917", size = 219191, upload-time = "2026-07-07T14:34:04.657Z" },
+    { url = "https://files.pythonhosted.org/packages/42/ab/b9bc2e77d6b44a7e46ef62ec5cac1c9a6ba7b9135a5d560f002696ec9995/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:a4cfde78a9f2880208d16a93b795726a3017d5977e08d1e162a7a31322479c41", size = 218328, upload-time = "2026-07-07T14:34:06.115Z" },
+    { url = "https://files.pythonhosted.org/packages/f1/78/c9c71d599f5aa2d42bcdd35cbbd46d7f535351a57e40ff7d8e5a7e219401/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_armv7l.whl", hash = "sha256:d4d6fcde76f94f5cb9e43e9e9a61f16dacefd228cbbf6f1a09bd9b219a92f1a1", size = 207406, upload-time = "2026-07-07T14:34:07.554Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/39/c914445c321a845097ce4f6ac7de9a18228a77b766272125a1ce00d851eb/charset_normalizer-3.4.9-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:898f0e9068ca27d37f8e83a5b962821df851532e6c4a7d615c1c033f9da6eedf", size = 225157, upload-time = "2026-07-07T14:34:09.061Z" },
+    { url = "https://files.pythonhosted.org/packages/9b/f2/c0d4b8508565a36bc5c624e88ed297f5b0b1095011034d7f5b83a69908b5/charset_normalizer-3.4.9-cp314-cp314-win32.whl", hash = "sha256:c1c948747b03be832dceed96ca815cef7360de9aa19d37c730f8e3f6101aca48", size = 151095, upload-time = "2026-07-07T14:34:10.901Z" },
+    { url = "https://files.pythonhosted.org/packages/49/fd/a1d26144398c67486422a72bf5812cda22cb4ccfcd95a290fb41ceb4b8e2/charset_normalizer-3.4.9-cp314-cp314-win_amd64.whl", hash = "sha256:16b65ea0f2465b6fb52aa22de5eca612aa964ddfec00a912e26f4656cbef890b", size = 162796, upload-time = "2026-07-07T14:34:12.47Z" },
+    { url = "https://files.pythonhosted.org/packages/20/95/d75e82f8ce9fd323ebf059c16c9aadefb22a1ecde13b7840b35835e4886c/charset_normalizer-3.4.9-cp314-cp314-win_arm64.whl", hash = "sha256:40a126142a56b2dfc0aacbad1de8310cbf60da7656db0e6b16eebd48e3e93519", size = 153334, upload-time = "2026-07-07T14:34:14.044Z" },
+    { url = "https://files.pythonhosted.org/packages/00/5e/17398df3a139985ba9d11ed072531986f408c8fca952835ef1ab1820c02b/charset_normalizer-3.4.9-cp314-cp314t-macosx_10_15_universal2.whl", hash = "sha256:609b3ba8fcc0fb5ab7af00719d0fb6ad0cb518e48e7712d12fd68f1327951198", size = 338848, upload-time = "2026-07-07T14:34:15.688Z" },
+    { url = "https://files.pythonhosted.org/packages/cd/91/7253a32e86b7e1d1239b1b36ba6dd0f021a21107ab33054b53119cc083b9/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:51447e9aa2684679af07ca5021c3db526e0284347ebf4ffcec1154c3350cfe32", size = 223022, upload-time = "2026-07-07T14:34:17.248Z" },
+    { url = "https://files.pythonhosted.org/packages/cb/32/2e64bd2be10e89c61e57ebe6a93fd98ae88eb7ebe414b5121f22c96c69eb/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:cc1b0fff8ead343dae06305f954eb8468ba0ec1a97881f42489d198e4ce3c632", size = 241590, upload-time = "2026-07-07T14:34:18.813Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/ef/d96ec496cfea0c21db43b0ad03891308b02388d054cc902cf0e5a1ad6a88/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fa36ec09ef71d158186bc79e359ff5fdd6e7996fe8ab638f00d6b93139ba4fcf", size = 239584, upload-time = "2026-07-07T14:34:20.52Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/ce/9af95f7876194bd7a14e3dfe4a4de2e0bff02666a3910d72beafd06cc297/charset_normalizer-3.4.9-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:df115d4d83168fdf2cae48ef1ff6d1cb4c466364e30861b37121de0f3bf1b990", size = 230224, upload-time = "2026-07-07T14:34:22.189Z" },
+    { url = "https://files.pythonhosted.org/packages/52/94/af74dde74a3996bd959c350709bfe50e297823d70a8c1cbd54b838880863/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_armv7l.whl", hash = "sha256:f86c6358749bd4fda175388691e3ba8c46e24c5347d0afd20f9b7edfc9faf07d", size = 212667, upload-time = "2026-07-07T14:34:23.857Z" },
+    { url = "https://files.pythonhosted.org/packages/ee/f0/f1c4fe746c395922961b5916ed1d7d6e7d4c84851d19ed43cc89980ec953/charset_normalizer-3.4.9-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:32286a2c8d167e897177b673176c1e3e00d4057caf5d2b64eef9a3666b03018e", size = 227179, upload-time = "2026-07-07T14:34:25.586Z" },
+    { url = "https://files.pythonhosted.org/packages/e4/56/6c745619ac397e8871e2bcd3cea1eec86b877488f33888b3aef5c3ed506e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:83aed2c10721ddd90f68140685391b50811a880af20654c59af6b6c66c40513c", size = 225372, upload-time = "2026-07-07T14:34:27.212Z" },
+    { url = "https://files.pythonhosted.org/packages/78/ad/98aae8630ac71f16711968e38a5acfecce41b778bf2f0312851020f565a8/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_armv7l.whl", hash = "sha256:cd6c3d4b783c556fa00bf540854e42f135e2f256abd29669fcd0da0f2dec79c2", size = 215222, upload-time = "2026-07-07T14:34:28.774Z" },
+    { url = "https://files.pythonhosted.org/packages/f7/40/9593d54209765207a7f11073c06494c1721e4ca4a0a426c597679bf7f91e/charset_normalizer-3.4.9-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:ee2f2a527e3c1a6e6411eb4209642e138b544a2d72fe5d0d76daf77b24063534", size = 231958, upload-time = "2026-07-07T14:34:30.345Z" },
+    { url = "https://files.pythonhosted.org/packages/b1/27/693ee5e8a18191eb38647360c51cd505013e2bd3b366aa43fd5344c21e3c/charset_normalizer-3.4.9-cp314-cp314t-win32.whl", hash = "sha256:0d861473f743244d349b50f850d10eb87aeb22bbdcc8e64f79273c94af5a8226", size = 155580, upload-time = "2026-07-07T14:34:31.884Z" },
+    { url = "https://files.pythonhosted.org/packages/80/3f/bd97d3d9c613013d07cb7733d299385b41df37f0471310f5a73dc359f0b8/charset_normalizer-3.4.9-cp314-cp314t-win_amd64.whl", hash = "sha256:9b8e0f3107e2200b76f6054de99016eac3ee6762713587b36baaa7e4bd2ae177", size = 167620, upload-time = "2026-07-07T14:34:33.438Z" },
+    { url = "https://files.pythonhosted.org/packages/3d/c6/eee9dca4439b1061f76373f06ea855678cc4a64c1c3c90b50e479edbb8eb/charset_normalizer-3.4.9-cp314-cp314t-win_arm64.whl", hash = "sha256:19ac87f93086ce37b86e098888555c4b4bc48102279bae3350098c0ed664b501", size = 158037, upload-time = "2026-07-07T14:34:35.018Z" },
+    { url = "https://files.pythonhosted.org/packages/98/2b/f97f1c193fb855c345d678f5077d6926034db0722df74c8f057020e05a25/charset_normalizer-3.4.9-py3-none-any.whl", hash = "sha256:68e5f26a1ad57ded6d1cfb85331d1c1a195314756471d97758c48498bb4dcdf5", size = 64538, upload-time = "2026-07-07T14:34:56.993Z" },
+]
+
+[[package]]
+name = "click"
+version = "8.4.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/76/d4/81420972a676e8ffea40450d8c8c92943e7218a78fe9b64359836cc9876b/click-8.4.2.tar.gz", hash = "sha256:9a6cea6e60b17ebe0a44c5cc636d94f09bd66142c1cd7d8b4cd731c4917a15f6", size = 338000, upload-time = "2026-06-24T17:45:15.148Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fb/e2/79c688af8b210d232694e31e59da9f6ec747bae31c3f5946e4e9b98860d5/click-8.4.2-py3-none-any.whl", hash = "sha256:e6f9f66136c816745b9d65817da91d61d957fb16e02e4dcd0552553c5a197b76", size = 119243, upload-time = "2026-06-24T17:45:13.73Z" },
+]
+
 [[package]]
 name = "colorama"
 version = "0.4.6"
@@ -443,7 +534,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.1.0"
+version = "0.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },
@@ -485,15 +576,19 @@ dev = [
     { name = "scikit-learn" },
     { name = "ty" },
     { name = "uproot" },
+    { name = "wandb" },
 ]
 geometry = [
     { name = "scikit-learn" },
 ]
+wandb = [
+    { name = "wandb" },
+]
 
 [package.metadata]
 requires-dist = [
     { name = "awkward", marker = "extra == 'convert'", specifier = ">=2.6,<3" },
-    { name = "giant", extras = ["convert", "analysis", "geometry"], marker = "extra == 'dev'" },
+    { name = "giant", extras = ["convert", "analysis", "geometry", "wandb"], marker = "extra == 'dev'" },
     { name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
     { name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
     { name = "numpy", specifier = ">=1.26,<3" },
@@ -513,8 +608,9 @@ requires-dist = [
     { name = "ty", marker = "extra == 'dev'", specifier = ">=0.0.50,<0.1" },
     { name = "typer", specifier = ">=0.12,<1" },
     { name = "uproot", marker = "extra == 'convert'", specifier = ">=5.3,<6" },
+    { name = "wandb", marker = "extra == 'wandb'", specifier = ">=0.16,<1" },
 ]
-provides-extras = ["cpu", "cuda", "dev", "geometry", "convert", "analysis"]
+provides-extras = ["cpu", "cuda", "dev", "geometry", "wandb", "convert", "analysis"]
 
 [[package]]
 name = "hepunits"
@@ -525,6 +621,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/85/10/7f9c58d1ec6a0b7f7783fe552f3593f39cda30c2e1d7a9d148ae711e748d/hepunits-2.4.6-py3-none-any.whl", hash = "sha256:089c52c3b84ef67a159b5e9ee9bdd50e1a442e3fd0c101303cc409c1e9011c4d", size = 17090, upload-time = "2026-06-16T09:23:35.35Z" },
 ]
 
+[[package]]
+name = "idna"
+version = "3.18"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/cd/63/9496c57188a2ee585e0f1db071d75089a11e98aa86eb99d9d7618fc1edce/idna-3.18.tar.gz", hash = "sha256:ffb385a7e039654cef1ab9ef32c6fafe283c0c0467bba1d9029738ce4a14a848", size = 196711, upload-time = "2026-06-02T14:34:07.794Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1e/5e/d4e9f1a599fb8e573b7b87160658329fbf28d19eac2718f51fc3def3aa5a/idna-3.18-py3-none-any.whl", hash = "sha256:7f952cbe720b688055e3f87de14f5c3e5fdaa8bc3928985c4077ca689de849a2", size = 65455, upload-time = "2026-06-02T14:34:06.319Z" },
+]
+
 [[package]]
 name = "iniconfig"
 version = "2.3.0"
@@ -1371,6 +1476,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" },
 ]
 
+[[package]]
+name = "protobuf"
+version = "7.35.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/10/03/8aeeb7458d22546bf64b5250ca1daeb5ff757d900e8e4a7476c6f0db843e/protobuf-7.35.1-cp310-abi3-macosx_10_9_universal2.whl", hash = "sha256:24f857477359a85c0c235261b8ba905fd51b2562f4a64ca1df5473f29850cbf6", size = 433226, upload-time = "2026-06-11T21:55:31.719Z" },
+    { url = "https://files.pythonhosted.org/packages/37/4b/dfb89eb0e652a1ff073c39a59fb5e3a83cfe9b57a2c83fa6d78270101767/protobuf-7.35.1-cp310-abi3-manylinux2014_aarch64.whl", hash = "sha256:11d6b0ec246892d85215b0a13ca6e0233cf5284b68f0ac02646427f4ff88a799", size = 328847, upload-time = "2026-06-11T21:55:34.035Z" },
+    { url = "https://files.pythonhosted.org/packages/0f/58/dc12f2cd484951524af6e3382c785869b9b3fb5e52ee95ae23add53ee8f9/protobuf-7.35.1-cp310-abi3-manylinux2014_s390x.whl", hash = "sha256:b73f9489a4b8b1c9cb1f8ed951c736392592edb24b9d6819f36d2e10b171d5b4", size = 344030, upload-time = "2026-06-11T21:55:34.941Z" },
+    { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" },
+    { url = "https://files.pythonhosted.org/packages/d8/bc/6d6c7ba8709c85f8f2c390b2b118d6fb08a783676a572271851bf45a7d22/protobuf-7.35.1-cp310-abi3-win32.whl", hash = "sha256:353652e4efd0bca5b5fc2656abf8307ef351f0cf938c9eba09f0e09c20a25c30", size = 428945, upload-time = "2026-06-11T21:55:37.034Z" },
+    { url = "https://files.pythonhosted.org/packages/0a/19/8d0cb6f20a1ef7b18f1c8986ad5783f22f84cce39c6ce9a6e645ea55192e/protobuf-7.35.1-cp310-abi3-win_amd64.whl", hash = "sha256:230a75ddfc2de4806e56696ce9640c1cdfdb6543b7cfce98d42a4c0a0e7bdb87", size = 439996, upload-time = "2026-06-11T21:55:38.123Z" },
+    { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" },
+]
+
 [[package]]
 name = "psutil"
 version = "7.2.2"
@@ -1469,6 +1589,96 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/0c/c3/44f3fbbfa403ea2a7c779186dc20772604442dde72947e7d01069cbe98e3/pycparser-3.0-py3-none-any.whl", hash = "sha256:b727414169a36b7d524c1c3e31839a521725078d7b2ff038656844266160a992", size = 48172, upload-time = "2026-01-21T14:26:50.693Z" },
 ]
 
+[[package]]
+name = "pydantic"
+version = "2.13.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "annotated-types" },
+    { name = "pydantic-core" },
+    { name = "typing-extensions" },
+    { name = "typing-inspection" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/18/a5/b60d21ac674192f8ab0ba4e9fd860690f9b4a6e51ca5df118733b487d8d6/pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6", size = 844775, upload-time = "2026-05-06T13:43:05.343Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fd/7b/122376b1fd3c62c1ed9dc80c931ace4844b3c55407b6fb2d199377c9736f/pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba", size = 472262, upload-time = "2026-05-06T13:43:02.641Z" },
+]
+
+[[package]]
+name = "pydantic-core"
+version = "2.46.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/9d/56/921726b776ace8d8f5db44c4ef961006580d91dc52b803c489fafd1aa249/pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1", size = 471464, upload-time = "2026-05-06T13:37:06.98Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ce/8c/af022f0af448d7747c5154288d46b5f2bc5f17366eaa0e23e9aa04d59f3b/pydantic_core-2.46.4-cp312-cp312-macosx_10_12_x86_64.whl", hash = "sha256:3245406455a5d98187ec35530fd772b1d799b26667980872c8d4614991e2c4a2", size = 2106158, upload-time = "2026-05-06T13:38:57.215Z" },
+    { url = "https://files.pythonhosted.org/packages/19/95/6195171e385007300f0f5574592e467c568becce2d937a0b6804f218bc49/pydantic_core-2.46.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:962ccbab7b642487b1d8b7df90ef677e03134cf1fd8880bf698649b22a69371f", size = 1951724, upload-time = "2026-05-06T13:37:02.697Z" },
+    { url = "https://files.pythonhosted.org/packages/8e/bc/f47d1ff9cbb1620e1b5b697eef06010035735f07820180e74178226b27b3/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8233f2947cf85404441fd7e0085f53b10c93e0ee78611099b5c7237e36aacbf7", size = 1975742, upload-time = "2026-05-06T13:37:09.448Z" },
+    { url = "https://files.pythonhosted.org/packages/5b/11/9b9a5b0306345664a2da6410877af6e8082481b5884b3ddd78d47c6013ce/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:3a233125ac121aa3ffba9a2b59edfc4a985a76092dc8279586ab4b71390875e7", size = 2052418, upload-time = "2026-05-06T13:37:38.234Z" },
+    { url = "https://files.pythonhosted.org/packages/f1/b7/a65fec226f5d78fc39f4a13c4cc0c768c22b113438f60c14adc9d2865038/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5b712b53160b79a5850310b912a5ef8e57e56947c8ad690c227f5c9d7e561712", size = 2232274, upload-time = "2026-05-06T13:38:27.753Z" },
+    { url = "https://files.pythonhosted.org/packages/68/f0/92039db98b907ef49269a8271f67db9cb78ae2fc68062ef7e4e77adb5f61/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:9401557acd873c3a7f3eb9383edef8ac4968f9510e340f4808d427e75667e7b4", size = 2309940, upload-time = "2026-05-06T13:38:05.353Z" },
+    { url = "https://files.pythonhosted.org/packages/5f/97/2aab507d3d00ca626e8e57c1eac6a79e4e5fbcc63eb99733ff55d1717f65/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce", size = 2094516, upload-time = "2026-05-06T13:39:10.577Z" },
+    { url = "https://files.pythonhosted.org/packages/22/37/a8aca44d40d737dde2bc05b3c6c07dff0de07ce6f82e9f3167aeaf4d5dea/pydantic_core-2.46.4-cp312-cp312-manylinux_2_31_riscv64.whl", hash = "sha256:56cb4851bcaf3d117eddcef4fe66afd750a50274b0da8e22be256d10e5611987", size = 2136854, upload-time = "2026-05-06T13:40:22.59Z" },
+    { url = "https://files.pythonhosted.org/packages/24/99/fcef1b79238c06a8cbec70819ac722ba76e02bc8ada9b0fd66eba40da01b/pydantic_core-2.46.4-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:c68fcd102d71ea85c5b2dfac3f4f8476eff42a9e078fd5faefff6d145063536b", size = 2180306, upload-time = "2026-05-06T13:40:10.666Z" },
+    { url = "https://files.pythonhosted.org/packages/ae/6c/fc44000918855b42779d007ae63b0532794739027b2f417321cddbc44f6a/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:b2f69dec1725e79a012d920df1707de5caf7ed5e08f3be4435e25803efc47458", size = 2190044, upload-time = "2026-05-06T13:40:43.231Z" },
+    { url = "https://files.pythonhosted.org/packages/6b/65/d9cadc9f1920d7a127ad2edba16c1db7916e59719285cd6c94600b0080ba/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_armv7l.whl", hash = "sha256:8d0820e8192167f80d88d64038e609c31452eeca865b4e1d9950a27a4609b00b", size = 2329133, upload-time = "2026-05-06T13:39:57.365Z" },
+    { url = "https://files.pythonhosted.org/packages/d0/cf/c873d91679f3a30bcf5e7ac280ce5573483e72295307685120d0d5ad3416/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c", size = 2374464, upload-time = "2026-05-06T13:38:06.976Z" },
+    { url = "https://files.pythonhosted.org/packages/47/bd/6f2fc8188f31bf10590f1e98e7b306336161fac930a8c514cd7bd828c7dc/pydantic_core-2.46.4-cp312-cp312-win32.whl", hash = "sha256:9aa768456404a8bf48a4406685ac2bec8e72b62c69313734fa3b73cf33b3a894", size = 1974823, upload-time = "2026-05-06T13:40:47.985Z" },
+    { url = "https://files.pythonhosted.org/packages/40/8c/985c1d41ea1107c2534abd9870e4ed5c8e7669b5c308297835c001e7a1c4/pydantic_core-2.46.4-cp312-cp312-win_amd64.whl", hash = "sha256:e9c26f834c65f5752f3f06cb08cb86a913ceb7274d0db6e267808a708b46bc89", size = 2072919, upload-time = "2026-05-06T13:39:21.153Z" },
+    { url = "https://files.pythonhosted.org/packages/c4/ba/f463d006e0c47373ca7ec5e1a261c59dc01ef4d62b2657af925fb0deee3a/pydantic_core-2.46.4-cp312-cp312-win_arm64.whl", hash = "sha256:4fc73cb559bdb54b1134a706a2802a4cddd27a0633f5abb7e53056268751ac6a", size = 2027604, upload-time = "2026-05-06T13:39:03.753Z" },
+    { url = "https://files.pythonhosted.org/packages/51/a2/5d30b469c5267a17b39dec53208222f76a8d351dfac4af661888c5aee77d/pydantic_core-2.46.4-cp313-cp313-macosx_10_12_x86_64.whl", hash = "sha256:5d5902252db0d3cedf8d4a1bc68f70eeb430f7e4c7104c8c476753519b423008", size = 2106306, upload-time = "2026-05-06T13:37:48.029Z" },
+    { url = "https://files.pythonhosted.org/packages/c1/81/4fa520eaffa8bd7d1525e644cd6d39e7d60b1592bc5b516693c7340b50f1/pydantic_core-2.46.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:c94f0688e7b8d0a67abf40e57a7eaaecd17cc9586706a31b76c031f63df052b4", size = 1951906, upload-time = "2026-05-06T13:37:17.012Z" },
+    { url = "https://files.pythonhosted.org/packages/03/d5/fd02da45b659668b05923b17ba3a0100a0a3d5541e3bd8fcc4ecb711309e/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f027324c56cd5406ca49c124b0db10e56c69064fec039acc571c29020cc87c76", size = 1976802, upload-time = "2026-05-06T13:37:35.113Z" },
+    { url = "https://files.pythonhosted.org/packages/21/f2/95727e1368be3d3ed485eaab7adbd7dda408f33f7a36e8b48e0144002b91/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:e739fee756ba1010f8bcccb534252e85a35fe45ae92c295a06059ce58b74ccd3", size = 2052446, upload-time = "2026-05-06T13:37:12.313Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/86/5d99feea3f77c7234b8718075b23db11532773c1a0dbd9b9490215dc2eeb/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:9d56801be94b86a9da183e5f3766e6310752b99ff647e38b09a9500d88e46e76", size = 2232757, upload-time = "2026-05-06T13:39:01.149Z" },
+    { url = "https://files.pythonhosted.org/packages/d2/3a/508ac615935ef7588cf6d9e9b91309fdc2da751af865e02a9098de88258c/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2412e734dcb48da14d4e4006b82b46b74f2518b8a26ee7e58c6844a6cd6d03c4", size = 2309275, upload-time = "2026-05-06T13:37:41.406Z" },
+    { url = "https://files.pythonhosted.org/packages/07/f8/41db9de19d7987d6b04715a02b3b40aea467000275d9d758ffaa31af7d50/pydantic_core-2.46.4-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9551187363ffc0de2a00b2e47c25aeaeb1020b69b668762966df15fc5659dd5a", size = 2094467, upload-time = "2026-05-06T13:39:18.847Z" },
+    { url = "https://files.pythonhosted.org/packages/2c/e2/f35033184cb11d0052daf4416e8e10a502ea2ac006fc4f459aee872727d1/pydantic_core-2.46.4-cp313-cp313-manylinux_2_31_riscv64.whl", hash = "sha256:0186750b482eefa11d7f435892b09c5c606193ef3375bcf94aa00ae6bfb66262", size = 2134417, upload-time = "2026-05-06T13:40:17.944Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/7b/6ceeb1cc90e193862f444ebe373d8fdf613f0a82572dde03fb10734c6c71/pydantic_core-2.46.4-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:5855698a4856556d86e8e6cd8434bc3ac0314ee8e12089ae0e143f64c6256e4e", size = 2179782, upload-time = "2026-05-06T13:40:32.618Z" },
+    { url = "https://files.pythonhosted.org/packages/5a/f2/c8d7773ede6af08036423a00ae0ceffce266c3c52a096c435d68c896083f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_aarch64.whl", hash = "sha256:cbaf13819775b7f769bf4a1f066cb6df7a28d4480081a589828ef190226881cd", size = 2188782, upload-time = "2026-05-06T13:36:51.018Z" },
+    { url = "https://files.pythonhosted.org/packages/59/31/0c864784e31f09f05cdd87606f08923b9c9e7f6e51dd27f20f62f975ce9f/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_armv7l.whl", hash = "sha256:633147d34cf4550417f12e2b1a0383973bdf5cdfde212cb09e9a581cf10820be", size = 2328334, upload-time = "2026-05-06T13:40:37.764Z" },
+    { url = "https://files.pythonhosted.org/packages/c2/eb/4f6c8a41efa30baa755590f4141abf3a8c370fab610915733e74134a7270/pydantic_core-2.46.4-cp313-cp313-musllinux_1_1_x86_64.whl", hash = "sha256:82cf5301172168103724d49a1444d3378cb20cdee30b116a1bd6031236298a5d", size = 2372986, upload-time = "2026-05-06T13:39:34.152Z" },
+    { url = "https://files.pythonhosted.org/packages/5b/24/b375a480d53113860c299764bfe9f349a3dc9108b3adc0d7f0d786492ebf/pydantic_core-2.46.4-cp313-cp313-win32.whl", hash = "sha256:9fa8ae11da9e2b3126c6426f147e0fba88d96d65921799bb30c6abd1cb2c97fb", size = 1973693, upload-time = "2026-05-06T13:37:55.072Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/e8/cff247591966f2d22ec8c003cd7587e27b7ba7b81ab2fb888e3ab75dc285/pydantic_core-2.46.4-cp313-cp313-win_amd64.whl", hash = "sha256:6b3ace8194b0e5204818c92802dcdca7fc6d88aabbb799d7c795540d9cd6d292", size = 2071819, upload-time = "2026-05-06T13:38:49.139Z" },
+    { url = "https://files.pythonhosted.org/packages/c6/1a/f4aee670d5670e9e148e0c82c7db98d780be566c6e6a97ee8035528ca0b3/pydantic_core-2.46.4-cp313-cp313-win_arm64.whl", hash = "sha256:184c081504d17f1c1066e430e117142b2c77d9448a97f7b65c6ac9fd9aee238d", size = 2027411, upload-time = "2026-05-06T13:40:45.796Z" },
+    { url = "https://files.pythonhosted.org/packages/8d/74/228a26ddad29c6672b805d9fd78e8d251cd04004fa7eed0e622096cd0250/pydantic_core-2.46.4-cp314-cp314-macosx_10_12_x86_64.whl", hash = "sha256:428e04521a40150c85216fc8b85e8d39fece235a9cf5e383761238c7fa9b96fb", size = 2102079, upload-time = "2026-05-06T13:38:41.019Z" },
+    { url = "https://files.pythonhosted.org/packages/ad/1f/8970b150a4b4365623ae00fc88603491f763c627311ae8031e3111356d6e/pydantic_core-2.46.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:23ace664830ee0bfe014a0c7bc248b1f7f25ed7ad103852c317624a1083af462", size = 1952179, upload-time = "2026-05-06T13:36:59.812Z" },
+    { url = "https://files.pythonhosted.org/packages/95/30/5211a831ae054928054b2f79731661087a2bc5c01e825c672b3a4a8f1b3e/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ce5c1d2a8b27468f433ca974829c44060b8097eedc39933e3c206a90ee49c4a9", size = 1978926, upload-time = "2026-05-06T13:37:39.933Z" },
+    { url = "https://files.pythonhosted.org/packages/57/e9/689668733b1eb67adeef047db3c2e8788fcf65a7fd9c9e2b46b7744fe245/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:7283d57845ecf5a163403eb0702dfc220cc4fbdd18919cb5ccea4f95ee1cdab4", size = 2046785, upload-time = "2026-05-06T13:38:01.995Z" },
+    { url = "https://files.pythonhosted.org/packages/60/d9/6715260422ff50a2109878fd24d948a6c3446bb2664f34ee78cd972b3acd/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:8daafc69c93ee8a0204506a3b6b30f586ef54028f52aeeeb5c4cfc5184fd5914", size = 2228733, upload-time = "2026-05-06T13:40:50.371Z" },
+    { url = "https://files.pythonhosted.org/packages/18/ae/fdb2f64316afca925640f8e70bb1a564b0ec2721c1389e25b8eb4bf9a299/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:cd2213145bcc2ba85884d0ac63d222fece9209678f77b9b4d76f054c561adb28", size = 2307534, upload-time = "2026-05-06T13:37:21.531Z" },
+    { url = "https://files.pythonhosted.org/packages/89/1d/8eff589b45bb8190a9d12c49cfad0f176a5cbd1534908a6b5125e2886239/pydantic_core-2.46.4-cp314-cp314-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7a5f930472650a82629163023e630d160863fce524c616f4e5186e5de9d9a49b", size = 2099732, upload-time = "2026-05-06T13:39:31.942Z" },
+    { url = "https://files.pythonhosted.org/packages/06/d5/ee5a3366637fee41dee51a1fc91562dcf12ddbc68fda34e6b253da2324bb/pydantic_core-2.46.4-cp314-cp314-manylinux_2_31_riscv64.whl", hash = "sha256:c1b3f518abeca3aa13c712fd202306e145abf59a18b094a6bafb2d2bbf59192c", size = 2129627, upload-time = "2026-05-06T13:37:25.033Z" },
+    { url = "https://files.pythonhosted.org/packages/94/33/2414be571d2c6a6c4d08be21f9292b6d3fdb08949a97b6dfe985017821db/pydantic_core-2.46.4-cp314-cp314-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:1a7dd0b3ee80d90150e3495a3a13ac34dbcbfd4f012996a6a1d8900e91b5c0fb", size = 2179141, upload-time = "2026-05-06T13:37:14.046Z" },
+    { url = "https://files.pythonhosted.org/packages/7b/79/7daa95be995be0eecc4cf75064cb33f9bbbfe3fe0158caf2f0d4a996a5c7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_aarch64.whl", hash = "sha256:3fb702cd90b0446a3a1c5e470bfa0dd23c0233b676a9099ddcc964fa6ca13898", size = 2184325, upload-time = "2026-05-06T13:36:53.615Z" },
+    { url = "https://files.pythonhosted.org/packages/9f/cb/d0a382f5c0de8a222dc61c65348e0ce831b1f68e0a018450d31c2cace3a5/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_armv7l.whl", hash = "sha256:b8458003118a712e66286df6a707db01c52c0f52f7db8e4a38f0da1d3b94fc4e", size = 2323990, upload-time = "2026-05-06T13:40:29.971Z" },
+    { url = "https://files.pythonhosted.org/packages/05/db/d9ba624cc4a5aced1598e88c04fdbd8310c8a69b9d38b9a3d39ce3a61ed7/pydantic_core-2.46.4-cp314-cp314-musllinux_1_1_x86_64.whl", hash = "sha256:372429a130e469c9cd698925ce5fc50940b7a1336b0d82038e63d5bbc4edc519", size = 2369978, upload-time = "2026-05-06T13:37:23.027Z" },
+    { url = "https://files.pythonhosted.org/packages/f2/20/d15df15ba918c423461905802bfd2981c3af0bfa0e40d05e13edbfa48bc3/pydantic_core-2.46.4-cp314-cp314-win32.whl", hash = "sha256:85bb3611ff1802f3ee7fdd7dbff26b56f343fb432d57a4728fdd49b6ef35e2f4", size = 1966354, upload-time = "2026-05-06T13:38:03.499Z" },
+    { url = "https://files.pythonhosted.org/packages/fc/b6/6b8de4c0a7d7ab3004c439c80c5c1e0a3e8d78bbae19379b01960383d9e5/pydantic_core-2.46.4-cp314-cp314-win_amd64.whl", hash = "sha256:811ff8e9c313ab425368bcbb36e5c4ebd7108c2bbf4e4089cfbb0b01eff63fac", size = 2072238, upload-time = "2026-05-06T13:39:40.807Z" },
+    { url = "https://files.pythonhosted.org/packages/32/36/51eb763beec1f4cf59b1db243a7dcc39cbb41230f050a09b9d69faaf0a48/pydantic_core-2.46.4-cp314-cp314-win_arm64.whl", hash = "sha256:bfec22eab3c8cc2ceec0248aec886624116dc079afa027ecc8ad4a7e62010f8a", size = 2018251, upload-time = "2026-05-06T13:37:26.72Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/91/855af51d625b23aa987116a19e231d2aaef9c4a415273ddc189b79a45fee/pydantic_core-2.46.4-cp314-cp314t-macosx_10_12_x86_64.whl", hash = "sha256:af8244b2bef6aaad6d92cda81372de7f8c8d36c9f0c3ea36e827c60e7d9467a0", size = 2099593, upload-time = "2026-05-06T13:39:47.682Z" },
+    { url = "https://files.pythonhosted.org/packages/fb/1b/8784a54c65edb5f49f0a14d6977cf1b209bba85a4c77445b255c2de58ab3/pydantic_core-2.46.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:5a4330cdbc57162e4b3aa303f588ba752257694c9c9be3e7ebb11b4aca659b5d", size = 1935226, upload-time = "2026-05-06T13:40:40.428Z" },
+    { url = "https://files.pythonhosted.org/packages/e8/e7/1955d28d1afc56dd4b3ad7cc0cf39df1b9852964cf16e5d13912756d6d6b/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:29c61fc04a3d840155ff08e475a04809278972fe6aef51e2720554e96367e34b", size = 1974605, upload-time = "2026-05-06T13:37:32.029Z" },
+    { url = "https://files.pythonhosted.org/packages/93/e2/3fedbf0ba7a22850e6e9fd78117f1c0f10f950182344d8a6c535d468fdd8/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_armv7l.manylinux2014_armv7l.whl", hash = "sha256:c50f2528cf200c5eed56faf3f4e22fcd5f38c157a8b78576e6ba3168ec35f000", size = 2030777, upload-time = "2026-05-06T13:38:55.239Z" },
+    { url = "https://files.pythonhosted.org/packages/f8/61/46be275fcaaba0b4f5b9669dd852267ce1ff616592dccf7a7845588df091/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:0cbe8b01f948de4286c74cdd6c667aceb38f5c1e26f0693b3983d9d74887c65e", size = 2236641, upload-time = "2026-05-06T13:37:08.096Z" },
+    { url = "https://files.pythonhosted.org/packages/60/db/12e93e46a8bac9988be3c016860f83293daea8c716c029c9ace279036f2f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:617d7e2ca7dcb8c5cf6bcb8c59b8832c94b36196bbf1cbd1bfb56ed341905edd", size = 2286404, upload-time = "2026-05-06T13:40:20.221Z" },
+    { url = "https://files.pythonhosted.org/packages/e2/4a/4d8b19008f38d31c53b8219cfedc2e3d5de5fe99d90076b7e767de29274f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7027560ee92211647d0d34e3f7cd6f50da56399d26a9c8ad0da286d3869a53f3", size = 2109219, upload-time = "2026-05-06T13:38:12.153Z" },
+    { url = "https://files.pythonhosted.org/packages/88/70/3cbc40978fefb7bb09c6708d40d4ad1a5d70fd7213c3d17f971de868ec1f/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_31_riscv64.whl", hash = "sha256:f99626688942fb746e545232e7726926f3be91b5975f8b55327665fafda991c7", size = 2110594, upload-time = "2026-05-06T13:40:02.971Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/20/b8d36736216e29491125531685b2f9e61aa5b4b2599893f8268551da3338/pydantic_core-2.46.4-cp314-cp314t-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:fc3e9034a63de20e15e8ade85358bc6efc614008cab72898b4b4952bea0509ff", size = 2159542, upload-time = "2026-05-06T13:39:27.506Z" },
+    { url = "https://files.pythonhosted.org/packages/1d/a2/367df868eb584dacf6bf82a389272406d7178e301c4ac82545ab98bc2dd9/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_aarch64.whl", hash = "sha256:97e7cf2be5c77b7d1a9713a05605d49460d02c6078d38d8bef3cbe323c548424", size = 2168146, upload-time = "2026-05-06T13:38:31.93Z" },
+    { url = "https://files.pythonhosted.org/packages/c1/b8/4460f77f7e201893f649a29ab355dddd3beee8a97bcb1a320db414f9a06e/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_armv7l.whl", hash = "sha256:3bf92c5d0e00fefaab325a4d27828fe6b6e2a21848686b5b60d2d9eeb09d76c6", size = 2306309, upload-time = "2026-05-06T13:37:44.717Z" },
+    { url = "https://files.pythonhosted.org/packages/64/c4/be2639293acd87dc8ddbcec41a73cee9b2ebf996fe6d892a1a74e88ad3f7/pydantic_core-2.46.4-cp314-cp314t-musllinux_1_1_x86_64.whl", hash = "sha256:3ecbc122d18468d06ca279dc26a8c2e2d5acb10943bb35e36ae92096dc3b5565", size = 2369736, upload-time = "2026-05-06T13:37:05.645Z" },
+    { url = "https://files.pythonhosted.org/packages/30/a6/9f9f380dbb301f67023bf8f707aaa75daadf84f7152d95c410fd7e81d994/pydantic_core-2.46.4-cp314-cp314t-win32.whl", hash = "sha256:e846ae7835bf0703ae43f534ab79a867146dadd59dc9ca5c8b53d5c8f7c9ef02", size = 1955575, upload-time = "2026-05-06T13:38:51.116Z" },
+    { url = "https://files.pythonhosted.org/packages/40/1f/f1eb9eb350e795d1af8586289746f5c5677d16043040d63710e22abc43c9/pydantic_core-2.46.4-cp314-cp314t-win_amd64.whl", hash = "sha256:2108ba5c1c1eca18030634489dc544844144ee36357f2f9f780b93e7ddbb44b5", size = 2051624, upload-time = "2026-05-06T13:38:21.672Z" },
+    { url = "https://files.pythonhosted.org/packages/f6/d2/42dd53d0a85c27606f316d3aa5d2869c4e8470a5ed6dec30e4a1abe19192/pydantic_core-2.46.4-cp314-cp314t-win_arm64.whl", hash = "sha256:4fcbe087dbc2068af7eda3aa87634eba216dbda64d1ae73c8684b621d33f6596", size = 2017325, upload-time = "2026-05-06T13:40:52.723Z" },
+    { url = "https://files.pythonhosted.org/packages/9d/1d/8987ad40f65ae1432753072f214fb5c74fe47ffbd0698bb9cbbb585664f8/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_10_12_x86_64.whl", hash = "sha256:1d8ba486450b14f3b1d63bc521d410ec7565e52f887b9fb671791886436a42f7", size = 2095527, upload-time = "2026-05-06T13:39:52.283Z" },
+    { url = "https://files.pythonhosted.org/packages/64/d3/84c282a7eee1d3ac4c0377546ef5a1ea436ce26840d9ac3b7ed54a377507/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-macosx_11_0_arm64.whl", hash = "sha256:3009f12e4e90b7f88b4f9adb1b0c4a3d58fe7820f3238c190047209d148026df", size = 1936024, upload-time = "2026-05-06T13:40:15.671Z" },
+    { url = "https://files.pythonhosted.org/packages/d7/ca/eac61596cdeb4d7e174d3dc0bd8a6238f14f75f97a24e7b7db4c7e7340a0/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ad785e92e6dc634c21555edc8bd6b64957ab844541bcb96a1366c202951ae526", size = 1990696, upload-time = "2026-05-06T13:38:34.717Z" },
+    { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" },
+]
+
 [[package]]
 name = "pygments"
 version = "2.20.0"
@@ -1604,6 +1814,21 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/81/d6/4bfbb40c9a0b42fc53c7cf442f6385db70b40f74a783130c5d0a5aa62228/pyzmq-27.1.0-cp314-cp314t-win_arm64.whl", hash = "sha256:dc5dbf68a7857b59473f7df42650c621d7e8923fb03fa74a526890f4d33cc4d7", size = 575170, upload-time = "2025-09-08T23:09:01.418Z" },
 ]
 
+[[package]]
+name = "requests"
+version = "2.34.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "certifi" },
+    { name = "charset-normalizer" },
+    { name = "idna" },
+    { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" },
+]
+
 [[package]]
 name = "rich"
 version = "15.0.0"
@@ -1732,6 +1957,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/d5/19/969dc072906c84dd0a3b05dcf57ea750936087d7873549e408b35cfc3f97/scipy-1.18.0-cp314-cp314t-win_arm64.whl", hash = "sha256:368e0a705903c466aa5f08eefb39e6b1b6b2d659e7352a31fd9e2438365be0f8", size = 25279661, upload-time = "2026-06-19T15:01:40.817Z" },
 ]
 
+[[package]]
+name = "sentry-sdk"
+version = "2.66.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "certifi" },
+    { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/7f/6f/d59cad0889d15fde85254cf58e701484de3f3f0406003b3197746910b19b/sentry_sdk-2.66.1.tar.gz", hash = "sha256:f882fb08710c5f8bfc603aafa3e901b384009a19cc3f76a572b863392ee81cdc", size = 940543, upload-time = "2026-07-22T12:26:54.553Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/89/d3/726bd88f0eece09ddf431bea4c9191c18e7a8d070b854eb0014d447712ee/sentry_sdk-2.66.1-py3-none-any.whl", hash = "sha256:86002793161d9a95ef04bdd8d442e9bfece5d989b755f05d6360215094a7aff6", size = 505555, upload-time = "2026-07-22T12:26:52.71Z" },
+]
+
 [[package]]
 name = "shellingham"
 version = "1.5.4"
@@ -1966,6 +2204,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/18/67/36e9267722cc04a6b9f15c7f3441c2363321a3ea07da7ae0c0707beb2a9c/typing_extensions-4.15.0-py3-none-any.whl", hash = "sha256:f0fa19c6845758ab08074a0cfa8b7aecb71c999ca73d62883bc25cc018c4e548", size = 44614, upload-time = "2025-08-25T13:49:24.86Z" },
 ]
 
+[[package]]
+name = "typing-inspection"
+version = "0.4.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/55/e3/70399cb7dd41c10ac53367ae42139cf4b1ca5f36bb3dc6c9d33acdb43655/typing_inspection-0.4.2.tar.gz", hash = "sha256:ba561c48a67c5958007083d386c3295464928b01faa735ab8547c5692e87f464", size = 75949, upload-time = "2025-10-01T02:14:41.687Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" },
+]
+
 [[package]]
 name = "tzdata"
 version = "2026.2"
@@ -1992,6 +2242,43 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/74/03/d426348a5f13514182c1d1afab2285ec25a94bacc8d2f8d2cc627496754a/uproot-5.7.4-py3-none-any.whl", hash = "sha256:497b7db1592f62edf05404884ec235f6cb804a50382a62c8df5f885d138c3695", size = 397455, upload-time = "2026-04-30T09:11:47.994Z" },
 ]
 
+[[package]]
+name = "urllib3"
+version = "2.7.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" },
+]
+
+[[package]]
+name = "wandb"
+version = "0.28.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "click" },
+    { name = "packaging" },
+    { name = "platformdirs" },
+    { name = "protobuf" },
+    { name = "pydantic" },
+    { name = "pyyaml" },
+    { name = "requests" },
+    { name = "sentry-sdk" },
+    { name = "typing-extensions" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/92/fb/8d3f96a8b143060d6fa145462d0785981373e04694e4152555ccb5d23939/wandb-0.28.1.tar.gz", hash = "sha256:870ccb1a01238b0ac07c6fd96a0810a1f79090aba04ea29f4ee012ac8327705d", size = 40578119, upload-time = "2026-07-16T18:47:05.413Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/06/21/8df50164d07623cfcefec19bbf9327d9be84b637a827cea1f0c06db005fd/wandb-0.28.1-py3-none-macosx_12_0_arm64.whl", hash = "sha256:da909a76e65c64c0d93acc485d2a19f66e336f1e3f725f1c98a070883e084943", size = 24277925, upload-time = "2026-07-16T18:46:42.383Z" },
+    { url = "https://files.pythonhosted.org/packages/8f/18/6c3da7e6cb215ad363324db8dc4d83b93626f5e339822b05b1c38a6097fd/wandb-0.28.1-py3-none-macosx_12_0_x86_64.whl", hash = "sha256:3da3db219c54bfd1082c00e9061c8ea894ba43e42733b5af00bb10c09d7158fe", size = 25480852, upload-time = "2026-07-16T18:46:45.102Z" },
+    { url = "https://files.pythonhosted.org/packages/e2/1a/d15bcfb4417fa69edcaa33db8ea012db733da1057e193b047e3f69fdd671/wandb-0.28.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:ae9ae6fb29e2e2b1d097ed8b75c0c0240c778c2a8cad1d996dee870a1e401c2c", size = 24832138, upload-time = "2026-07-16T18:46:47.433Z" },
+    { url = "https://files.pythonhosted.org/packages/b3/da/49924c7df2952dfd82c86c3779c339c0c3d6f6439387c03d97d0470c3658/wandb-0.28.1-py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:8cfb898b6a6c884d9c9294b02764e88bce65049f027a124d6bee53fe722469b6", size = 26486533, upload-time = "2026-07-16T18:46:49.839Z" },
+    { url = "https://files.pythonhosted.org/packages/11/c0/06b23518e29690784f1b3081e39c7679ca076cb0af094cb9b4bb309150f5/wandb-0.28.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:cf2b1533945395e4fdbe6182b272bb0ca8a02c10b3086a395e2d57686ae3ed0d", size = 25022635, upload-time = "2026-07-16T18:46:52.376Z" },
+    { url = "https://files.pythonhosted.org/packages/23/30/6de2f7995a8a6eecbd03d24c79a139a734c0168f5520cf4c7ccb43c1dbbc/wandb-0.28.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:7233061080507a4b4098bed1ccb381ce6f890c60397cd4153d060285bcb267bd", size = 27008895, upload-time = "2026-07-16T18:46:55.025Z" },
+    { url = "https://files.pythonhosted.org/packages/b2/83/49deab9447687625371435ca21b6da82f223f1c7d014d77386b9cb91833c/wandb-0.28.1-py3-none-win32.whl", hash = "sha256:4bc461cda3ce23a19d8df5e42981a664d95fa3231efb10fd1e85d9d4824c7d29", size = 24418398, upload-time = "2026-07-16T18:46:57.43Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/6f/ed6616b11ea15b8ceabedcaa567286c1c9ec65fa50563230a90bfb627cc5/wandb-0.28.1-py3-none-win_amd64.whl", hash = "sha256:d98a10370162b1e970850237114c56e9c4c58f3cb701e4b8cb38f36f6749fd52", size = 24418404, upload-time = "2026-07-16T18:47:00.427Z" },
+    { url = "https://files.pythonhosted.org/packages/07/78/75b6827a6665337a715c5347c5edbd84eca660f7a0f48d8d6d24d1f66bee/wandb-0.28.1-py3-none-win_arm64.whl", hash = "sha256:4aa07f13dd3bcac2c0524c8d0f49f76e83ab5c1054fd09f3b1a436cfcde146a6", size = 22299006, upload-time = "2026-07-16T18:47:02.71Z" },
+]
+
 [[package]]
 name = "wcwidth"
 version = "0.8.1"