Offset event_id per file to avoid cross-file collisions

Each input parquet file is one Geant4 job (scripts/steps_to_parquet.py), and a job's event_id numbering always restarts from 0 — so loading multiple files together (a directory or .manifest) let same-numbered events from different files collapse into one during the event index scan and train/val split, corrupting both. Every per-file event_id now gets offset by file index * EVENT_ID_FILE_STRIDE (giant/data/loader.py), threaded through the setup-cache event index, the streaming dataset, and predict/rollout seeding. Bumps the setup-cache format version so stale sidecars computed pre-fix are invalidated. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Merge pull request 'Cache giant train's setup stage in a sidecar file' (#21 ) from perf/setup-stage into master
2026-07-30 11:36:48 +02:00 · 2026-07-30 11:02:11 +02:00 · 2026-07-30 10:57:48 +02:00 · 2026-07-30 10:53:31 +02:00 · 2026-07-30 10:29:19 +02:00 · 2026-07-30 10:28:37 +02:00
31 changed files with 2997 additions and 1313 deletions
@@ -10,20 +10,17 @@ 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 --comment "..."   # scaffold a config.toml + run dir for a new training run
 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)
 giant train path/to/steps.parquet --router --router-type energy  # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap)
-giant train-submit path/to/steps.parquet --config run/config.toml --accounting-group cms   # train as a remote-GPU HTCondor job (TOpAS/NEMO2)
 giant predict path/to/steps.parquet --checkpoint ckpt/best.pt   # per-step predictions
 giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl  # full showers
-giant rollout-submit path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl --accounting-group cms  # rollout as a remote-GPU HTCondor job
 giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs-reference analysis on HTCondor
 giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep/submit)
 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)
 ```

@@ -79,8 +76,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep

 **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.

-**GPU HTCondor submission** (`giant/condor.py`, `giant train-submit`/`giant rollout-submit`): submits a single `giant train`/`giant rollout` invocation as a remote-GPU HTCondor job (ETP's TOpAS V100/A100 or NEMO2 L40S workers — see the ETP HTCondor wiki's "GPU Jobs" section). GPU workers are remote-only, so submit descriptions always carry `+RemoteJob = True`/`RequestGPUs`, and reach `/ceph` via `requirements = TARGET.ProvidesEtpCeph =?= True` rather than the local-only `TARGET.ProvidesETPResources` `giant analyze submit`'s CPU jobs use — this means the `giant` checkout submitting these jobs (and its `uv sync --extra cuda` venv) needs to live under `/ceph`, not `/work`/`/home`. Each submission writes a self-contained `condor/` (train, nested in `--out`) or `condor_<tag>/` (rollout, sibling to the output parquet) directory: `run.sh` (the wrapper actually executed — for training this re-checks `out_dir/last.pt` on every invocation and adds `--resume`, so a preempted/retried job resumes rather than restarting), `job.sub`, and `submission.json` (a `CondorJobMeta`: accounting group, GPU request, docker image, the exact command, and the assigned cluster id once known) — enough on its own to `condor_history <cluster_id>` a run later. `train-submit` requires `--config` (a TOML) rather than mirroring `train`'s hyperparameter flags, since the config file is already the reproducible source of truth (`giant/config.py:save_config`/`merge_cli_overrides`) and passing the same one on every retry is what keeps a resumed run's architecture from drifting. `giant new-run` scaffolds that TOML: resolves a base `--config` (or built-in defaults) plus a handful of override flags into a fresh run dir via `gconfig.default_out_dir` (hyperparam-named, uuid-tagged so repeat/same-day runs never collide), and prints the matching `giant train`/`giant train-submit` invocations — `giant train` overwrites this scaffolded `config.toml` in place once it actually runs, filling in the real dataset-derived meta section.
-
 ## Roadmap

 **Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
@@ -2,19 +2,15 @@

 **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.

-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.
+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.

 Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.

 ## Architecture

-A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`):
+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.

- **`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, evaluated against a ~10× native-Geant4 latency budget.
-
-**Stage 1 — primary (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):**
+**Stage 1 — primary (9D, diffused):**

 | Index | Variable | Encoding |
 |-------|----------|----------|
@@ -23,20 +19,13 @@ A **two-stage model**, both stages checkpointed together, with a choice of gener
 | 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 (`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.
+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.

 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 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.
+**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.

-**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 (`RoutedDenoisingMLP`/`RoutedSecondaryDecoder`), soft-gated in training and hard-dispatched at eval. See `--router-type`/`--n-experts`/`--router-axis` on `giant train`/`giant new-run`.
+**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.)

 ## Roadmap

@@ -44,13 +33,11 @@ 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.

-**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.
-
-**In progress:** a WGAN-GP single-pass mode (`--mode wgan`) and mixture-of-experts routing (`--router`) are implemented and under evaluation as faster-eval alternatives to iterative flow/DDPM sampling, against a ~10× native-Geant4 latency budget. A multi-material sampling-calorimeter dataset is the target for both the routing/WGAN evaluation and the physical-property generalization comparison.
+**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.

 ## Data

-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) 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 `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.

 ## Project structure

@@ -62,33 +49,21 @@ giant/
 │   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
 │   │   └── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
 │   ├── model/
-│   │   ├── 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
+│   │   ├── network.py      # SinusoidalEmbedding, ConditionEncoder, DenoisingMLP, SecondaryDecoder
+│   │   └── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
 │   ├── 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 / WGAN samplers + secondary sampling
+│   ├── sample.py           # DDPM / DDIM / flow matching 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
-│   ├── condor.py           # HTCondor submission for `train-submit`/`rollout-submit` (remote GPU workers)
-│   ├── 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` / `predict` / `rollout` / `new-run` / `*-submit` / `analyze` Typer app
-├── scripts/                            # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
+│   ├── 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`)
 │   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
-│   │                                    #   update-manifest, create-manifest, make-root, build-geometry-oracle,
-│   │                                    #   hparam-scan
+│   │                                    #   update-manifest, create-manifest, make-root, 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`
@@ -105,48 +80,27 @@ 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 (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
+`cpu` and `cuda` are mutually exclusive extras selecting the torch build; plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.

 ## Training, prediction, and rollout

 ```bash
-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 train path/to/steps.parquet --mode flow
 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`. `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.

-### Remote-GPU training and rollout (HTCondor)
+## Validation

-`giant train-submit`/`giant rollout-submit` run a `giant train`/`giant rollout` invocation as a remote-GPU HTCondor job (ETP's TOpAS V100/A100 or NEMO2 L40S workers). GPU workers are remote-only and reach `/ceph` (not `/work`/`/home`) — the checkout submitting these jobs, and its venv, must live under `/ceph`.
-
-```bash
-giant train-submit path/to/steps.parquet --config run/config.toml --accounting-group cms
-giant rollout-submit path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl --accounting-group cms
-```
-
-Each submission writes a self-contained `condor/`/`condor_<tag>/` directory (`run.sh`, `job.sub`, `submission.json`). Training jobs re-check `out_dir/last.pt` on every invocation and resume automatically if preempted. `train-submit` requires `--config` (rather than mirroring `train`'s hyperparameter flags) since the TOML is the reproducible source of truth, and passing the same one on every retry is what keeps a resumed run's architecture from drifting.
-
-## Validation and 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.
+`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`.

 ## Development

@@ -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
@@ -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
@@ -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"
@@ -4,7 +4,6 @@ from enum import Enum
 import math
 from pathlib import Path
 import re
-import shlex
 from typing import Optional
 import uuid as uuid_mod

@@ -27,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,
@@ -88,38 +88,15 @@ 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
    return model_cfg["hidden_dim"], model_cfg["n_blocks"]


-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
-
-
 def _coerce_scalar(value: str) -> object:
    """Best-effort str -> bool/int/float, else leave as str.

@@ -160,21 +137,15 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
 _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")


-def _resolve_prediction_output(
-    data: Path, out: Path | None, pred_uuid: str | None = None
-) -> tuple[Path, Path, str]:
+def _resolve_prediction_output(data: Path, out: Path | None) -> tuple[Path, Path, str]:
    """Return (out_path, resolved_dataset_path, pred_uuid).

    When *out* is None the output path is derived from *data*:
    - under /ceph/ → fixed central store with a UUID filename
    - elsewhere   → sibling of *data* with a UUID filename
-
-    *pred_uuid* can be pinned by the caller (e.g. `rollout-submit`, which
-    needs the same id for its condor run-dir name, the output filename, and
-    the YAML sidecar) — default is a freshly generated one, as before.
    """
    dataset_path = data.resolve()
-    pred_uuid = pred_uuid or str(uuid_mod.uuid4())
+    pred_uuid = str(uuid_mod.uuid4())
    if out is None:
        if str(dataset_path).startswith("/ceph/"):
            out = _CEPH_PREDICTIONS / f"{pred_uuid}.parquet"
@@ -415,10 +386,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[
@@ -430,6 +424,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
@@ -467,6 +485,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
    }
@@ -482,7 +504,17 @@ def train(
        }.items()
        if v is not None
    }
-    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
+    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))
    if cli_router:
        cli_model["router"] = cli_router
    cfg = gconfig.merge_cli_overrides(
@@ -505,7 +537,28 @@ def train(
            f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)"
        )

-    out_dir = out or gconfig.default_out_dir(cfg)
+    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.
+        base_name = gconfig.default_out_dir_name(cfg)
+        out_dir = Path("checkpoints") / base_name
+        suffix = 2
+        while out_dir.exists():
+            out_dir = Path("checkpoints") / f"{base_name}_{suffix}"
+            suffix += 1

    typer.echo(f"device:  {_device}")
    typer.echo(f"out_dir: {out_dir}")
@@ -518,153 +571,12 @@ 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 commands "
-            "(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
-    and `giant train-submit --config ...` requires. `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.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}")
-    typer.echo(
-        f"  giant train-submit {data_arg} --config {config_path} --out {run_dir} "
-        "--accounting-group <group>"
-    )
-
-
@app.command()
 def predict(
    data: Annotated[
@@ -1000,8 +912,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"]])
@@ -1050,8 +962,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)):
@@ -1153,15 +1065,6 @@ def rollout(
        Optional[int],
        typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
    ] = None,
-    prediction_id: Annotated[
-        Optional[str],
-        typer.Option(
-            "--prediction-id",
-            help="Pin the prediction uuid (output filename, YAML sidecar name) "
-            "instead of generating a fresh one — used by `rollout-submit` so "
-            "its condor run-dir shares an id with the run it submitted",
-        ),
-    ] = None,
 ) -> None:
    """Roll the surrogate forward into full showers (autoregressive)."""
    if seed is not None:
@@ -1215,9 +1118,7 @@ def rollout(
    seeds = _seed_from_data(files, n_events)
    typer.echo(f"seeded {len(seeds['event_id']):,} shower(s)")

-    out, dataset_path, pred_uuid = _resolve_prediction_output(
-        data, out, pred_uuid=prediction_id
-    )
+    out, dataset_path, pred_uuid = _resolve_prediction_output(data, out)
    out.parent.mkdir(parents=True, exist_ok=True)

    # Written incrementally as each batch of steps is produced, rather than
@@ -1483,352 +1384,5 @@ def analyze_submit(
    subprocess.run(["condor_submit", str(sub)], check=True)


-def _warn_if_not_ceph(path: Path) -> None:
-    """Remote GPU condor workers only reliably reach /ceph (see giant/condor.py)."""
-    if not str(path).startswith("/ceph/"):
-        typer.echo(
-            f"warning: {path} is not under /ceph/ — remote GPU condor workers "
-            "may not be able to reach it (see TARGET.ProvidesEtpCeph)",
-            err=True,
-        )
-
-
-@app.command("train-submit")
-def train_submit(
-    data: Annotated[
-        Path,
-        typer.Argument(
-            help="Parquet file or directory (must resolve under /ceph so a "
-            "remote GPU worker can reach it)"
-        ),
-    ],
-    config: Annotated[
-        Path,
-        typer.Option(
-            "--config",
-            "-c",
-            help="TOML config — the full, reproducible spec for this run. "
-            "Hyperparams aren't exposed as flags here; edit the TOML instead "
-            "(this is also what makes resume-after-preemption safe: the same "
-            "file is passed on every condor retry, so the architecture never "
-            "drifts from the checkpoint being resumed).",
-        ),
-    ],
-    accounting_group: Annotated[str, typer.Option("--accounting-group")],
-    out: Annotated[
-        Optional[Path],
-        typer.Option(
-            "--out", "-o", help="Checkpoint dir (default: auto from hyperparams)"
-        ),
-    ] = None,
-    request_gpus: Annotated[int, typer.Option("--request-gpus")] = 1,
-    gpu_type: Annotated[
-        Optional[str],
-        typer.Option("--gpu-type", help='Pin GPU model, e.g. "Tesla V100-PCIE-32GB"'),
-    ] = None,
-    gpu_memory_mb: Annotated[Optional[int], typer.Option("--gpu-memory-mb")] = None,
-    request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 16384,
-    request_walltime: Annotated[
-        int, typer.Option("--request-walltime", help="Seconds (default: 2 days)")
-    ] = 172800,
-    docker_image: Annotated[
-        str, typer.Option("--docker-image")
-    ] = "mschnepf/slc7-condocker",
-    repo_dir: Annotated[
-        Optional[Path],
-        typer.Option(
-            "--repo-dir",
-            help="The /ceph checkout `uv run` executes from (default: cwd)",
-        ),
-    ] = None,
-    dry_run: Annotated[
-        bool, typer.Option("--dry-run", help="Write files but don't condor_submit")
-    ] = False,
-) -> None:
-    """Submit `giant train` as a single remote-GPU HTCondor job.
-
-    GPU workers (TOpAS/NEMO2) are remote-only, so this always sets
-    `+RemoteJob = True` / `RequestGPUs` and reaches data via
-    `TARGET.ProvidesEtpCeph` rather than the local-only
-    `TARGET.ProvidesETPResources` `giant analyze submit` uses — see
-    giant/condor.py and the ETP HTCondor wiki's "GPU Jobs" section.
-    """
-    from giant.condor import (
-        CondorJobMeta,
-        GpuSubmitConfig,
-        parse_cluster_id,
-        write_gpu_submit,
-    )
-
-    data_path = data.resolve()
-    config_path = config.resolve()
-    _warn_if_not_ceph(data_path)
-
-    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config_path, {}, {})
-    out_dir = (out or gconfig.default_out_dir(cfg)).resolve()
-    out_dir.mkdir(parents=True, exist_ok=True)
-    _warn_if_not_ceph(out_dir)
-
-    repo_path = (repo_dir or Path.cwd()).resolve()
-    run_dir = out_dir / "condor"
-
-    q = shlex.quote
-    last_ckpt = out_dir / "last.pt"
-    command = (
-        f"uv run giant train {q(str(data_path))} --config {q(str(config_path))} "
-        f"--out {q(str(out_dir))} --device cuda"
-    )
-    wrapper_body = (
-        "#!/bin/bash\n"
-        "set -euo pipefail\n"
-        f"cd {q(str(repo_path))}\n"
-        'RESUME=""\n'
-        f'[ -f {q(str(last_ckpt))} ] && RESUME="--resume {last_ckpt}"\n'
-        f"exec {command} $RESUME\n"
-    )
-
-    submit_cfg = GpuSubmitConfig(
-        run_dir=run_dir,
-        accounting_group=accounting_group,
-        repo_dir=repo_path,
-        command=command,
-        docker_image=docker_image,
-        request_memory_mb=request_memory,
-        request_gpus=request_gpus,
-        gpu_type=gpu_type,
-        gpu_memory_mb=gpu_memory_mb,
-        request_walltime_s=request_walltime,
-    )
-    sub = write_gpu_submit(submit_cfg, wrapper_body)
-
-    meta = CondorJobMeta(
-        accounting_group=accounting_group,
-        request_gpus=request_gpus,
-        gpu_type=gpu_type,
-        gpu_memory_mb=gpu_memory_mb,
-        request_memory_mb=request_memory,
-        request_walltime_s=request_walltime,
-        docker_image=docker_image,
-        repo_dir=str(repo_path),
-        command=command,
-        submitted_at=datetime.now(timezone.utc).isoformat(),
-    )
-    meta_path = run_dir / "submission.json"
-    meta.save(meta_path)
-
-    typer.echo(f"out dir: {out_dir}")
-    typer.echo(f"wrote submit description: {sub}")
-    if dry_run:
-        typer.echo("dry-run: not submitting")
-        return
-
-    import subprocess
-
-    result = subprocess.run(
-        ["condor_submit", str(sub)], check=True, capture_output=True, text=True
-    )
-    typer.echo(result.stdout)
-    meta.cluster_id = parse_cluster_id(result.stdout)
-    meta.save(meta_path)
-    if meta.cluster_id is not None:
-        typer.echo(f"cluster id: {meta.cluster_id}")
-
-
-@app.command("rollout-submit")
-def rollout_submit(
-    data: Annotated[
-        Path,
-        typer.Argument(
-            help="Parquet file/dir to seed showers from (must resolve under "
-            "/ceph so a remote GPU worker can reach it)"
-        ),
-    ],
-    checkpoint: Annotated[
-        Path,
-        typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"),
-    ],
-    geometry: Annotated[
-        Path,
-        typer.Option(
-            "--geometry",
-            "-g",
-            help="Geometry oracle .pkl (dwarf build-geometry-oracle)",
-        ),
-    ],
-    accounting_group: Annotated[str, typer.Option("--accounting-group")],
-    energy_cutoff: Annotated[
-        float,
-        typer.Option(
-            "--energy-cutoff",
-            help="Stop a track when its energy drops below this [MeV]",
-        ),
-    ] = 0.1,
-    max_steps: Annotated[
-        int, typer.Option("--max-steps", help="Max steps per individual track")
-    ] = 1000,
-    steps: Annotated[
-        int,
-        typer.Option(
-            "--steps",
-            "-s",
-            help="Flow matching ODE steps per model call (ignored for a wgan checkpoint)",
-        ),
-    ] = 10,
-    weights: Annotated[
-        Weights,
-        typer.Option("--weights", help="raw or ema (see `giant rollout --help`)"),
-    ] = Weights.raw,
-    batch_size: Annotated[
-        int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")
-    ] = 4096,
-    max_tracks_per_event: Annotated[
-        Optional[int],
-        typer.Option(
-            "--max-tracks-per-event",
-            help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)",
-        ),
-    ] = None,
-    escape_threshold: Annotated[
-        Optional[float],
-        typer.Option(
-            "--escape-threshold",
-            help="Override the oracle's NN-distance escape threshold [mm]",
-        ),
-    ] = None,
-    n_events: Annotated[
-        Optional[int], typer.Option("--n-events", help="Cap number of seed events")
-    ] = None,
-    seed: Annotated[
-        Optional[int],
-        typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
-    ] = None,
-    out: Annotated[
-        Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")
-    ] = None,
-    request_gpus: Annotated[int, typer.Option("--request-gpus")] = 1,
-    gpu_type: Annotated[
-        Optional[str],
-        typer.Option("--gpu-type", help='Pin GPU model, e.g. "Tesla V100-PCIE-32GB"'),
-    ] = None,
-    gpu_memory_mb: Annotated[Optional[int], typer.Option("--gpu-memory-mb")] = None,
-    request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 16384,
-    request_walltime: Annotated[
-        int, typer.Option("--request-walltime", help="Seconds (default: 2 days)")
-    ] = 172800,
-    docker_image: Annotated[
-        str, typer.Option("--docker-image")
-    ] = "mschnepf/slc7-condocker",
-    repo_dir: Annotated[
-        Optional[Path],
-        typer.Option(
-            "--repo-dir",
-            help="The /ceph checkout `uv run` executes from (default: cwd)",
-        ),
-    ] = None,
-    dry_run: Annotated[
-        bool, typer.Option("--dry-run", help="Write files but don't condor_submit")
-    ] = False,
-) -> None:
-    """Submit `giant rollout` as a single remote-GPU HTCondor job.
-
-    No resume logic (unlike `train-submit`) — a retried rollout just starts
-    over, which is fine since it's deterministic given `--seed` and has no
-    epoch-loop state to preserve.
-    """
-    from giant.condor import (
-        CondorJobMeta,
-        GpuSubmitConfig,
-        parse_cluster_id,
-        write_gpu_submit,
-    )
-
-    data_path = data.resolve()
-    checkpoint_path = checkpoint.resolve()
-    geometry_path = geometry.resolve()
-    _warn_if_not_ceph(data_path)
-    _warn_if_not_ceph(checkpoint_path)
-
-    out_path, _dataset_path, pred_uuid = _resolve_prediction_output(data_path, out)
-    out_path.parent.mkdir(parents=True, exist_ok=True)
-    _warn_if_not_ceph(out_path)
-
-    repo_path = (repo_dir or Path.cwd()).resolve()
-    run_dir = out_path.parent / f"condor_{pred_uuid[:8]}"
-
-    q = shlex.quote
-    flags = [
-        f"--checkpoint {q(str(checkpoint_path))}",
-        f"--geometry {q(str(geometry_path))}",
-        f"--energy-cutoff {energy_cutoff}",
-        f"--max-steps {max_steps}",
-        f"--steps {steps}",
-        f"--weights {weights.value}",
-        f"--batch-size {batch_size}",
-        f"--out {q(str(out_path))}",
-        f"--prediction-id {pred_uuid}",
-        "--device cuda",
-    ]
-    if max_tracks_per_event is not None:
-        flags.append(f"--max-tracks-per-event {max_tracks_per_event}")
-    if escape_threshold is not None:
-        flags.append(f"--escape-threshold {escape_threshold}")
-    if n_events is not None:
-        flags.append(f"--n-events {n_events}")
-    if seed is not None:
-        flags.append(f"--seed {seed}")
-
-    command = f"uv run giant rollout {q(str(data_path))} " + " ".join(flags)
-    wrapper_body = (
-        f"#!/bin/bash\nset -euo pipefail\ncd {q(str(repo_path))}\nexec {command}\n"
-    )
-
-    submit_cfg = GpuSubmitConfig(
-        run_dir=run_dir,
-        accounting_group=accounting_group,
-        repo_dir=repo_path,
-        command=command,
-        docker_image=docker_image,
-        request_memory_mb=request_memory,
-        request_gpus=request_gpus,
-        gpu_type=gpu_type,
-        gpu_memory_mb=gpu_memory_mb,
-        request_walltime_s=request_walltime,
-    )
-    sub = write_gpu_submit(submit_cfg, wrapper_body)
-
-    meta = CondorJobMeta(
-        accounting_group=accounting_group,
-        request_gpus=request_gpus,
-        gpu_type=gpu_type,
-        gpu_memory_mb=gpu_memory_mb,
-        request_memory_mb=request_memory,
-        request_walltime_s=request_walltime,
-        docker_image=docker_image,
-        repo_dir=str(repo_path),
-        command=command,
-        submitted_at=datetime.now(timezone.utc).isoformat(),
-    )
-    meta_path = run_dir / "submission.json"
-    meta.save(meta_path)
-
-    typer.echo(f"prediction id: {pred_uuid}")
-    typer.echo(f"output: {out_path}")
-    typer.echo(f"wrote submit description: {sub}")
-    if dry_run:
-        typer.echo("dry-run: not submitting")
-        return
-
-    import subprocess
-
-    result = subprocess.run(
-        ["condor_submit", str(sub)], check=True, capture_output=True, text=True
-    )
-    typer.echo(result.stdout)
-    meta.cluster_id = parse_cluster_id(result.stdout)
-    meta.save(meta_path)
-    if meta.cluster_id is not None:
-        typer.echo(f"cluster id: {meta.cluster_id}")
-
-
 if __name__ == "__main__":
    app()
@@ -1,129 +0,0 @@
-"""HTCondor GPU-job submission for `giant train` / `giant rollout`.
-
-Trains and rollouts run as single condor jobs (no plot-style fan-out) on the
-ETP cluster's remote GPU resources (TOpAS V100/A100, NEMO2 L40S — see the ETP
-HTCondor wiki's "GPU Jobs" section). GPU workers are remote-only, so these
-jobs always carry ``+RemoteJob = True`` and ``RequestGPUs``; the local-only
-``TARGET.ProvidesETPResources`` requirement ``giant.analysis.condor``'s CPU
-jobs use is replaced by ``TARGET.ProvidesEtpCeph =?= True``, the
-wiki-documented way for a remote job to reach ``/ceph`` without HTCondor
-file-transferring multi-GB checkpoints/datasets.
-
-`giant/cli.py`'s `train-submit`/`rollout-submit` commands build a
-`GpuSubmitConfig`, write the wrapper + submit description via
-`write_gpu_submit`, run `condor_submit`, and record what was submitted (and,
-once known, the assigned cluster id) in a `CondorJobMeta` JSON sidecar next
-to the job files — so the run directory alone is enough to understand what
-ran, on what resources, and how to look it up later
-(``condor_history <cluster_id>``).
-"""
-
-from __future__ import annotations
-
-import json
-import re
-from dataclasses import asdict, dataclass
-from pathlib import Path
-
-
-@dataclass
-class GpuSubmitConfig:
-    """Everything needed to write one single-job GPU submit description."""
-
-    run_dir: Path
-    accounting_group: str
-    repo_dir: Path
-    command: str
-    docker_image: str = "mschnepf/slc7-condocker"
-    request_memory_mb: int = 16384
-    request_gpus: int = 1
-    gpu_type: str | None = None  # -> TARGET.GPUs_DeviceName =?= "..."
-    gpu_memory_mb: int | None = None  # -> TARGET.GPUs_GlobalMemoryMb >= N
-    request_walltime_s: int = 172800  # 2 days
-
-
-def _gpu_requirements(cfg: GpuSubmitConfig) -> str:
-    """`TARGET.ProvidesEtpCeph` (remote /ceph access) ANDed with any GPU pin."""
-    clauses = ["TARGET.ProvidesEtpCeph =?= True"]
-    if cfg.gpu_type is not None:
-        clauses.append(f'TARGET.GPUs_DeviceName =?= "{cfg.gpu_type}"')
-    if cfg.gpu_memory_mb is not None:
-        clauses.append(f"TARGET.GPUs_GlobalMemoryMb >= {cfg.gpu_memory_mb}")
-    return " && ".join(clauses)
-
-
-def _gpu_submit_description(cfg: GpuSubmitConfig, wrapper: Path) -> str:
-    return (
-        "universe = docker\n"
-        f"docker_image = {cfg.docker_image}\n"
-        f"executable = {wrapper}\n"
-        "should_transfer_files = YES\n"
-        "when_to_transfer_output = ON_EXIT\n"
-        f"request_memory = {cfg.request_memory_mb}\n"
-        f"RequestGPUs = {cfg.request_gpus}\n"
-        f"+RequestWalltime = {cfg.request_walltime_s}\n"
-        f"accounting_group = {cfg.accounting_group}\n"
-        "+RemoteJob = True\n"
-        f"requirements = ({_gpu_requirements(cfg)})\n"
-        f"output = {cfg.run_dir}/logs/job.out\n"
-        f"error = {cfg.run_dir}/logs/job.err\n"
-        f"log = {cfg.run_dir}/logs/job.log\n"
-        "queue 1\n"
-    )
-
-
-def write_gpu_submit(cfg: GpuSubmitConfig, wrapper_body: str) -> Path:
-    """Write the wrapper script + submit description under ``cfg.run_dir``.
-
-    Returns the submit description path (``<run_dir>/job.sub``). Does not
-    call ``condor_submit`` — that's the caller's job (``giant/cli.py``), same
-    contract as ``giant.analysis.condor.write_submit``.
-    """
-    run_dir = cfg.run_dir
-    (run_dir / "logs").mkdir(parents=True, exist_ok=True)
-
-    wrapper = run_dir / "run.sh"
-    wrapper.write_text(wrapper_body)
-    wrapper.chmod(0o755)
-
-    sub = run_dir / "job.sub"
-    sub.write_text(_gpu_submit_description(cfg, wrapper))
-    return sub
-
-
-_CLUSTER_ID_RE = re.compile(r"submitted to cluster (\d+)")
-
-
-def parse_cluster_id(condor_submit_stdout: str) -> int | None:
-    """Extract the assigned cluster id from `condor_submit`'s stdout, if present."""
-    m = _CLUSTER_ID_RE.search(condor_submit_stdout)
-    return int(m.group(1)) if m else None
-
-
-@dataclass
-class CondorJobMeta:
-    """What was submitted — written to ``<run_dir>/submission.json``.
-
-    ``cluster_id`` starts unset and is filled in by the caller once
-    ``condor_submit``'s stdout has been parsed, so the run directory alone is
-    enough to ``condor_history <cluster_id>`` this job later.
-    """
-
-    accounting_group: str
-    request_gpus: int
-    gpu_type: str | None
-    gpu_memory_mb: int | None
-    request_memory_mb: int
-    request_walltime_s: int
-    docker_image: str
-    repo_dir: str
-    command: str
-    submitted_at: str
-    cluster_id: int | None = None
-
-    def save(self, path: str | Path) -> None:
-        Path(path).write_text(json.dumps(asdict(self), indent=2))
-
-    @classmethod
-    def load(cls, path: str | Path) -> "CondorJobMeta":
-        return cls(**json.loads(Path(path).read_text()))
@@ -1,9 +1,9 @@
+import hashlib
 import random
 import subprocess
 import sys
 import tomllib
-import uuid
-from datetime import date, datetime, timezone
+from datetime import datetime, timezone
 from pathlib import Path

 import numpy as np
@@ -36,6 +36,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,
@@ -55,8 +68,13 @@ 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
            "lambda_balance": 0.0,  # optional load-balance aux loss weight
@@ -251,6 +269,109 @@ 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 _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),
+    ("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 seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
@@ -295,26 +416,6 @@ def save_config(cfg: dict, out_dir: Path, meta: dict) -> None:
    (out_dir / "config.toml").write_text("\n".join(lines))


-def default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path:
-    """Auto-derived checkpoint dir: hyperparams for readability + a short
-    uuid tag so same-day/same-hyperparam runs (including a repeat condor
-    submission) never collide on an existing directory.
-    """
-    t, m = cfg["train"], cfg["model"]
-    tag = uuid.uuid4().hex[:8]
-    return base / (
-        f"{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']}"
-        f"_{tag}"
-    )
-
-
 def build_run_meta(
    data: Path,
    seed: int,
@@ -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(
@@ -65,6 +65,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 +98,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()}
@@ -12,6 +12,20 @@ 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 _read_manifest(path: Path) -> list[Path]:
    files = []
@@ -78,13 +92,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": df["event_id"].to_numpy().astype(np.int64) + 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 +135,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 ids.astype(np.int64) + 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 +171,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": df["event_id"].to_numpy().astype(np.int64) + 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 +185,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(
@@ -0,0 +1,285 @@
+"""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 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.
+_CACHE_FORMAT_VERSION = 2
+
+_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,
+}
+
+
+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_reservoir_sample: np.ndarray
+
+    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_reservoir_sample": np.asarray(
+                self.energy_reservoir_sample, 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_reservoir_sample=np.array(
+                d["energy_reservoir_sample"], 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.
+    """
+    path = sidecar_path(data)
+    tmp = path.parent / f".{path.name}.tmp.{os.getpid()}"
+    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)
+    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())
@@ -196,7 +196,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 +211,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 +286,44 @@ 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) -> 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).
+    """
+    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():
+        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 +386,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 +408,52 @@ 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)
+        for i in range(K):
+            if i == 0:
+                remaining = np.maximum(e_sec, _EPS)
+            else:
+                remaining = np.maximum(e_sec - cumsum[:, i - 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 sec_pdg_list is not None:
        from giant.particles import particle_phys_array
@@ -573,8 +636,8 @@ 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)
+    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])

    if cond_normalizer is not None:
@@ -627,6 +690,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 +717,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 +756,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 +784,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 +825,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)

@@ -566,6 +566,30 @@ class Router(nn.Module):
        """
        return torch.zeros((), device=cond_cont.device)

+    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]] = {}

@@ -1130,8 +1154,10 @@ 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"),
@@ -1,3 +1,4 @@
+from dataclasses import dataclass
 from pathlib import Path

 import numpy as np
@@ -13,19 +14,230 @@ 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_sample = entry.energy_reservoir_sample
+    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) 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_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_sample
+            )
+
+    if energy_router_active and energy_sample.size > 0:
+        assert cond_norm.mean is not None and cond_norm.std is not None
+        normalized_sample = (
+            energy_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']}"
+        )
+    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 +246,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,97 +255,39 @@ 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")
-
    router_cfg = m["router"]
    if t["mode"] == "wgan" and router_cfg.get("enabled"):
        raise ValueError(
            "--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 +332,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 +363,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"],
@@ -259,4 +434,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),
    )
@@ -32,6 +32,7 @@ _METRICS_FIELDS = [
    "train_loss_s2",
    "train_loss_balance",
    "train_loss_proc",
+    "train_nsec_acc",
    "d_loss",
    "g_loss",
    "wasserstein_estimate",
@@ -42,9 +43,24 @@ _METRICS_FIELDS = [
    "val_loss_s2",
    "val_loss_balance",
    "val_loss_proc",
+    "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",
 ]

@@ -108,9 +124,15 @@ def _compute_losses(
    lambda_balance: float = 0.0,
    lambda_proc: 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,
 ]:
-    """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, 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 +151,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
@@ -171,7 +194,7 @@ def _compute_losses(
        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
+    return total, l_s1, l_nsec, l_s2, l_balance, l_proc, nsec_acc


 def _wgan_train_step(
@@ -263,7 +286,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 +308,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(),
    }


@@ -320,10 +348,64 @@ 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={
+                "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,
+                "n_critic": n_critic,
+                "gp_weight": gp_weight,
+                "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,
+            },
+        )
+
    stage1_model = stage1_model.to(device)
    sec_decoder = sec_decoder.to(device)
    if mode == "wgan":
@@ -333,6 +415,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
@@ -394,6 +483,7 @@ def train(

    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"])
@@ -417,6 +507,7 @@ def train(
        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 +537,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":
@@ -466,6 +563,9 @@ def train(
            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
@@ -503,7 +603,6 @@ def train(
                        lambda_nsec,
                        lambda_s2,
                    )
-                    global_step += 1
                    if stats["did_g_step"]:
                        lr_sched.step()
                        if ema_decay > 0:
@@ -523,18 +622,23 @@ 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, nsec_acc = (
+                        _compute_losses(
+                            stage1_model,
+                            sec_decoder,
+                            batch,
+                            mode,
+                            ddpm_schedule,
+                            device,
+                            lambda_nsec,
+                            lambda_s2,
+                            lambda_balance,
+                            lambda_proc,
+                        )
                    )
                    optimizer.zero_grad()
                    loss.backward()
@@ -557,6 +661,7 @@ def train(
                    train_s2_sum += l_s2.item() * B
                    train_balance_sum += l_balance.item() * B
                    train_proc_sum += l_proc.item() * B
+                    train_nsec_acc_sum += nsec_acc.item() * B

                train_n += B
                train_batches += 1
@@ -573,6 +678,49 @@ 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()
+                    wandb_run.log(log_payload, step=global_step)
+
                if shutdown.requested:
                    break
            bar.close()
@@ -582,7 +730,13 @@ def train(

            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 +772,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_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,22 +785,36 @@ def train(
                val_s2_sum = 0.0
                val_balance_sum = 0.0
                val_proc_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(
-                            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, nsec_acc = (
+                            _compute_losses(
+                                stage1_model,
+                                sec_decoder,
+                                batch,
+                                mode,
+                                ddpm_schedule,
+                                device,
+                                lambda_nsec,
+                                lambda_s2,
+                                lambda_balance,
+                                lambda_proc,
+                            )
                        )
                        B = batch[0].size(0)
                        val_loss_sum += loss.item() * B
@@ -651,8 +823,47 @@ def train(
                        val_s2_sum += l_s2.item() * B
                        val_balance_sum += l_balance.item() * B
                        val_proc_sum += l_proc.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 +879,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 ""
@@ -685,32 +901,53 @@ def train(
                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_row = {
+                "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),
+                "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_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,
+            }
+            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: dict = {
                "model": stage1_model.state_dict(),
@@ -719,6 +956,7 @@ def train(
                "lr_sched": lr_sched.state_dict(),
                "epoch": epoch,
                "best_val_loss": best_val_loss,
+                "global_step": global_step,
            }
            if mode == "wgan":
                assert (
@@ -756,6 +994,8 @@ def train(
                break

    metrics_file.close()
+    if wandb_run is not None:
+        wandb_run.finish()

    if shutdown.requested:
        print(
@@ -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}"
@@ -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",
@@ -25,6 +25,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)

@@ -471,6 +472,80 @@ def build_geometry_oracle(
    )


+class Conditioning(str, Enum):
+    physical = "physical"
+    embedding = "embedding"
+
+
+@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,
@@ -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 reservoir sample 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.")
@@ -1,118 +0,0 @@
-"""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-submit" 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()
@@ -58,17 +58,6 @@ def test_each_call_produces_a_distinct_uuid(tmp_path):
    assert uuid1 != uuid2


-def test_pinned_pred_uuid_is_used_as_is(tmp_path):
-    # `rollout-submit` pins the uuid at submit time so its condor run-dir,
-    # the output filename, and the eventual YAML sidecar all agree.
-    data = tmp_path / "data.parquet"
-    pinned = "abcd1234-abcd-4abc-9abc-abcdabcdabcd"
-    out, _, pred_uuid = _resolve_prediction_output(data, None, pred_uuid=pinned)
-
-    assert pred_uuid == pinned
-    assert out.name == f"{pinned}.parquet"
-
-
 def test_dataset_path_is_resolved(tmp_path):
    data = tmp_path / "data.parquet"
    _, dataset_path, _ = _resolve_prediction_output(data, None)
@@ -1,216 +0,0 @@
-"""Tests for GPU-job HTCondor submission (`giant train-submit` / `giant rollout-submit`)."""
-
-from __future__ import annotations
-
-from pathlib import Path
-
-from typer.testing import CliRunner
-
-from giant import config as gconfig
-from giant.cli import app
-from giant.condor import (
-    CondorJobMeta,
-    GpuSubmitConfig,
-    parse_cluster_id,
-    write_gpu_submit,
-)
-
-runner = CliRunner()
-
-_MINIMAL_TOML = """\
-[train]
-mode = "flow"
-epochs = 1
-
-[model]
-hidden_dim = 8
-n_blocks = 2
-"""
-
-
-# ---------------------------------------------------------------------------
-# default_out_dir
-# ---------------------------------------------------------------------------
-
-
-def test_default_out_dir_encodes_hyperparams():
-    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {})
-    out_dir = gconfig.default_out_dir(cfg)
-    name = out_dir.name
-    assert f"_h{cfg['model']['hidden_dim']}_" in name
-    assert f"_b{cfg['model']['n_blocks']}_" in name
-    assert f"_c{cfg['model']['conditioning']}_" in name
-
-
-def test_default_out_dir_avoids_collisions():
-    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}, {})
-    a = gconfig.default_out_dir(cfg)
-    b = gconfig.default_out_dir(cfg)
-    assert a != b
-    # same hyperparam-derived prefix, differing only in the uuid tag
-    assert a.name.rsplit("_", 1)[0] == b.name.rsplit("_", 1)[0]
-
-
-# ---------------------------------------------------------------------------
-# giant/condor.py: GpuSubmitConfig / write_gpu_submit
-# ---------------------------------------------------------------------------
-
-
-def test_write_gpu_submit_description(tmp_path: Path):
-    cfg = GpuSubmitConfig(
-        run_dir=tmp_path / "condor",
-        accounting_group="cms",
-        repo_dir=tmp_path,
-        command="uv run giant train data.parquet --config c.toml --out out --device cuda",
-    )
-    sub = write_gpu_submit(cfg, "#!/bin/bash\necho hi\n")
-    txt = sub.read_text()
-
-    assert "universe = docker" in txt
-    assert "docker_image = mschnepf/slc7-condocker" in txt
-    assert "+RemoteJob = True" in txt
-    assert "RequestGPUs = 1" in txt
-    assert "accounting_group = cms" in txt
-    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in txt
-    assert "ProvidesETPResources" not in txt
-    assert "queue 1" in txt
-
-    wrapper = cfg.run_dir / "run.sh"
-    assert wrapper.exists()
-    assert wrapper.stat().st_mode & 0o111
-    assert wrapper.read_text() == "#!/bin/bash\necho hi\n"
-
-
-def test_gpu_requirements_include_type_and_memory_pins(tmp_path: Path):
-    cfg = GpuSubmitConfig(
-        run_dir=tmp_path / "condor",
-        accounting_group="cms",
-        repo_dir=tmp_path,
-        command="uv run giant train ...",
-        request_gpus=2,
-        gpu_type="Tesla V100-PCIE-32GB",
-        gpu_memory_mb=16000,
-    )
-    txt = write_gpu_submit(cfg, "#!/bin/bash\n").read_text()
-
-    assert "RequestGPUs = 2" in txt
-    assert 'TARGET.GPUs_DeviceName =?= "Tesla V100-PCIE-32GB"' in txt
-    assert "TARGET.GPUs_GlobalMemoryMb >= 16000" in txt
-
-
-# ---------------------------------------------------------------------------
-# CondorJobMeta / parse_cluster_id
-# ---------------------------------------------------------------------------
-
-
-def test_condor_job_meta_round_trip(tmp_path: Path):
-    meta = CondorJobMeta(
-        accounting_group="cms",
-        request_gpus=1,
-        gpu_type=None,
-        gpu_memory_mb=None,
-        request_memory_mb=16384,
-        request_walltime_s=172800,
-        docker_image="mschnepf/slc7-condocker",
-        repo_dir="/ceph/lbogner/giant",
-        command="uv run giant train ...",
-        submitted_at="2026-07-24T00:00:00+00:00",
-    )
-    path = tmp_path / "submission.json"
-    meta.save(path)
-    loaded = CondorJobMeta.load(path)
-    assert loaded == meta
-    assert loaded.cluster_id is None
-
-    loaded.cluster_id = 123456
-    loaded.save(path)
-    assert CondorJobMeta.load(path).cluster_id == 123456
-
-
-def test_parse_cluster_id():
-    assert parse_cluster_id("1 job(s) submitted to cluster 123456.\n") == 123456
-    assert parse_cluster_id("ERROR: something went wrong\n") is None
-
-
-# ---------------------------------------------------------------------------
-# CLI: giant train-submit / giant rollout-submit (--dry-run only, no cluster contact)
-# ---------------------------------------------------------------------------
-
-
-def test_train_submit_dry_run_writes_condor_dir(tmp_path: Path):
-    config = tmp_path / "config.toml"
-    config.write_text(_MINIMAL_TOML)
-    data = tmp_path / "train.parquet"
-    out_dir = tmp_path / "ckpt"
-
-    result = runner.invoke(
-        app,
-        [
-            "train-submit",
-            str(data),
-            "--config",
-            str(config),
-            "--accounting-group",
-            "cms",
-            "--out",
-            str(out_dir),
-            "--dry-run",
-        ],
-    )
-    assert result.exit_code == 0, result.output
-
-    run_dir = out_dir / "condor"
-    sub_txt = (run_dir / "job.sub").read_text()
-    assert "+RemoteJob = True" in sub_txt
-    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt
-
-    wrapper = (run_dir / "run.sh").read_text()
-    assert "--device cuda" in wrapper
-    assert "last.pt" in wrapper
-    assert "RESUME" in wrapper
-
-    meta = CondorJobMeta.load(run_dir / "submission.json")
-    assert meta.accounting_group == "cms"
-    assert meta.cluster_id is None
-
-
-def test_rollout_submit_dry_run_writes_condor_dir(tmp_path: Path):
-    data = tmp_path / "seed.parquet"
-    checkpoint = tmp_path / "best.pt"
-    geometry = tmp_path / "oracle.pkl"
-    out = tmp_path / "rollout.parquet"
-
-    result = runner.invoke(
-        app,
-        [
-            "rollout-submit",
-            str(data),
-            "--checkpoint",
-            str(checkpoint),
-            "--geometry",
-            str(geometry),
-            "--accounting-group",
-            "cms",
-            "--out",
-            str(out),
-            "--dry-run",
-        ],
-    )
-    assert result.exit_code == 0, result.output
-
-    condor_dirs = list(tmp_path.glob("condor_*"))
-    assert len(condor_dirs) == 1
-    run_dir = condor_dirs[0]
-
-    sub_txt = (run_dir / "job.sub").read_text()
-    assert "+RemoteJob = True" in sub_txt
-    assert "requirements = (TARGET.ProvidesEtpCeph =?= True)" in sub_txt
-
-    wrapper = (run_dir / "run.sh").read_text()
-    assert "--device cuda" in wrapper
-    assert "--prediction-id" in wrapper
-    assert "last.pt" not in wrapper
-    assert "RESUME" not in wrapper
-
-    meta = CondorJobMeta.load(run_dir / "submission.json")
-    assert meta.accounting_group == "cms"
@@ -1,3 +1,5 @@
+from datetime import datetime
+
 from giant import config as gconfig


@@ -115,3 +117,130 @@ 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_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_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
@@ -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():
@@ -31,3 +36,94 @@ def test_make_event_split_reproducible():
    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())
@@ -1,6 +1,8 @@
 from typer.testing import CliRunner

+from giant.data import setup_cache
 from scripts.dwarf import app
+from test_pipeline import _make_synthetic_steps

 runner = CliRunner()

@@ -66,3 +68,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
@@ -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,267 @@ 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_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])
@@ -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
@@ -0,0 +1,196 @@
+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] = []
+    run_train_job(
+        data=data,
+        cfg=cfg or _tiny_cfg(),
+        out_dir=out_dir,
+        device=torch.device("cpu"),
+        shuffle_buffer=64,
+        num_workers=0,
+        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_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"]
@@ -0,0 +1,259 @@
+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_reservoir_sample, [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
+
+
+# ── 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])
@@ -12,7 +12,10 @@ from giant.data.transforms import (
    log_transform,
    Normalizer,
    reconstruct_post_pos,
+    sorted_membership,
    travel_direction,
+    _vectorized_map_lookup,
+    _WelfordAccumulator,
 )


@@ -440,3 +443,120 @@ 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)
+
+
+# ── _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)
@@ -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"
@@ -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"