Bump version to 0.2.0

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

CLAUDE.md

@@ -10,6 +10,7 @@ uv sync --extra cuda               # install dependencies with CUDA 11.8 torch
 uv sync --extra cpu --extra dev    # add dev extras (pytest, etc.)
 uv sync --extra cpu --extra geometry  # add scikit-learn for the geometry oracle (giant rollout)
 pytest               # run tests
+giant new-run --hidden-dim 512 --lr 3e-4        # scaffold a config.toml + run dir ahead of training
 giant train path/to/steps.parquet --mode flow   # train (flow matching)
 giant train path/to/steps.parquet --mode ddpm   # train (DDPM baseline)
 giant train path/to/steps.parquet --mode wgan   # train (WGAN-GP, single-pass eval; implemented, not yet tested)

README.md

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

configs/router_energy_n10_gumbel_centers_physical.toml

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

configs/router_energy_n10_gumbel_centers_temp_physical.toml

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

configs/router_energy_n10_gumbel_physical.toml

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

configs/wgan_h512_b6_physical.toml

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

giant/analysis/condor.py

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

giant/cli.py

@@ -134,6 +134,30 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
     return out
 
 
+def _router_cli_overrides(
+    router: bool | None,
+    router_type: str | None,
+    n_experts: int | None,
+    router_axis: list[str] | None,
+) -> dict[str, object]:
+    """Build the `model.router` override dict from `--router`/`--router-type`/
+    `--n-experts`/`--router-axis` flags (empty if none were given). Shared by
+    `train` and `new-run` so both resolve router overrides identically.
+    """
+    cli_router: dict[str, object] = {
+        k: v
+        for k, v in {
+            "enabled": router,
+            "type": router_type,
+            "n_experts": n_experts,
+        }.items()
+        if v is not None
+    }
+    if router_axis:
+        cli_router.update(_parse_router_axis_flags(router_axis))
+    return cli_router
+
+
 _CEPH_PREDICTIONS = Path("/ceph/lbogner/geant_steps/predictions")
 
 
@@ -187,9 +211,10 @@ class Mode(str, Enum):
     wgan = "wgan"
 
 
-class Conditioning(str, Enum):
-    physical = "physical"
-    embedding = "embedding"
+# Conditioning itself lives in giant.config (imported below as gconfig) —
+# shared with scripts/dwarf.py's Typer commands so the two CLIs can't
+# silently drift apart on the option's valid values.
+Conditioning = gconfig.Conditioning
 
 
 class Coord(str, Enum):
@@ -504,17 +529,7 @@ def train(
         }.items()
         if v is not None
     }
-    cli_router: dict[str, object] = {
-        k: v
-        for k, v in {
-            "enabled": router,
-            "type": router_type,
-            "n_experts": n_experts,
-        }.items()
-        if v is not None
-    }
-    if router_axis:
-        cli_router.update(_parse_router_axis_flags(router_axis))
+    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
     if cli_router:
         cli_model["router"] = cli_router
     cfg = gconfig.merge_cli_overrides(
@@ -552,13 +567,9 @@ def train(
         # 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
+        # doubles as the W&B run id (giant.train) — hence the suffix loop in
+        # resolve_default_out_dir.
+        out_dir = gconfig.resolve_default_out_dir(cfg)
 
     typer.echo(f"device:  {_device}")
     typer.echo(f"out_dir: {out_dir}")
@@ -577,6 +588,144 @@ def train(
     )
 
 
+@app.command("new-run")
+def new_run(
+    config: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--config",
+            "-c",
+            help="Base TOML to start from (default: built-in defaults)",
+        ),
+    ] = None,
+    mode: Annotated[Optional[Mode], typer.Option("--mode", "-m")] = None,
+    epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
+    batch_size: Annotated[Optional[int], typer.Option("--batch-size", "-b")] = None,
+    lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
+    hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
+    n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
+    emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
+    dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
+    conditioning: Annotated[
+        Optional[Conditioning], typer.Option("--conditioning")
+    ] = None,
+    router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
+    router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
+    n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None,
+    router_axis: Annotated[Optional[list[str]], typer.Option("--router-axis")] = None,
+    out: Annotated[
+        Optional[Path],
+        typer.Option("--out", "-o", help="Run dir (default: auto from hyperparams)"),
+    ] = None,
+    comment: Annotated[
+        Optional[str],
+        typer.Option(
+            "--comment", help="Free-text note recorded in config.toml's meta section"
+        ),
+    ] = None,
+    data: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--data",
+            help="Dataset path to fill in the printed next-step command "
+            "(not stored in the config)",
+        ),
+    ] = None,
+    force: Annotated[
+        bool,
+        typer.Option(
+            "--force",
+            help="Overwrite config.toml even if --out already has checkpoints",
+        ),
+    ] = False,
+    dry_run: Annotated[
+        bool,
+        typer.Option(
+            "--dry-run", help="Print the resolved config without writing anything"
+        ),
+    ] = False,
+) -> None:
+    """Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir.
+
+    This is the config-file-first counterpart to hand-editing a TOML: start
+    from a base --config (or built-in defaults), override a few hyperparams
+    inline, and this resolves+writes the full `config.toml` into a fresh (or
+    explicit --out) run dir — the same file `giant train --config ...` reads.
+    `giant train` itself overwrites this file in place once it actually runs
+    (with the full dataset-derived meta section), so this scaffold's meta
+    section is just a placeholder recording what was asked for and when.
+    """
+    cli_train = {
+        k: v
+        for k, v in {
+            "mode": mode.value if mode is not None else None,
+            "epochs": epochs,
+            "batch_size": batch_size,
+            "lr": lr,
+        }.items()
+        if v is not None
+    }
+    cli_model: dict[str, object] = {
+        k: v
+        for k, v in {
+            "hidden_dim": hidden_dim,
+            "n_blocks": n_blocks,
+            "emb_dim": emb_dim,
+            "dropout": dropout,
+            "conditioning": conditioning.value if conditioning is not None else None,
+        }.items()
+        if v is not None
+    }
+    cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis)
+    if cli_router:
+        cli_model["router"] = cli_router
+
+    cfg = gconfig.merge_cli_overrides(
+        gconfig.DEFAULT_CONFIG, config, cli_train, cli_model
+    )
+    run_dir = (out or gconfig.resolve_default_out_dir(cfg)).resolve()
+
+    if not force:
+        existing = [n for n in ("last.pt", "best.pt") if (run_dir / n).exists()]
+        if existing:
+            typer.echo(
+                f"error: {run_dir} already has {', '.join(existing)} — pass "
+                "--force to overwrite its config.toml anyway",
+                err=True,
+            )
+            raise typer.Exit(1)
+
+    typer.echo(f"run dir: {run_dir}")
+
+    if dry_run:
+        typer.echo("dry-run: not writing anything. Resolved config:")
+        for section in ("train", "model"):
+            typer.echo(f"[{section}]")
+            for k, v in cfg[section].items():
+                if k == "router":
+                    continue
+                typer.echo(f"  {k} = {v}")
+        return
+
+    meta = {
+        "git_hash": gconfig.git_hash(),
+        "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+        "created_by": "giant new-run",
+    }
+    if comment:
+        meta["comment"] = comment
+
+    run_dir.mkdir(parents=True, exist_ok=True)
+    gconfig.save_config(cfg, run_dir, meta)
+    config_path = run_dir / "config.toml"
+    typer.echo(f"wrote {config_path}")
+
+    data_arg = str(data) if data is not None else "<data.parquet>"
+    typer.echo("")
+    typer.echo("next:")
+    typer.echo(f"  giant train {data_arg} --config {config_path} --out {run_dir}")
+
+
 @app.command()
 def predict(
     data: Annotated[

giant/config.py

@@ -4,11 +4,23 @@ import subprocess
 import sys
 import tomllib
 from datetime import datetime, timezone
+from enum import Enum
 from pathlib import Path
 
 import numpy as np
 import torch
 
+
+class Conditioning(str, Enum):
+    """`model.conditioning` choices — shared by `giant.cli` and `scripts.dwarf`'s
+    Typer commands so the two CLIs can't silently drift apart on the option's
+    valid values (see DEFAULT_CONFIG["model"]["conditioning"] for what each
+    value means)."""
+
+    physical = "physical"
+    embedding = "embedding"
+
+
 DEFAULT_CONFIG: dict = {
     "train": {
         "mode": "flow",
@@ -77,7 +89,41 @@ DEFAULT_CONFIG: dict = {
             "expert_n_blocks": 0,
             "temperature": 0.5,  # energy/pdg-router kwarg
             "learn_centers": True,  # energy/pdg-router kwarg
+            # energy-router kwargs: mutually exclusive optional learnable
+            # gate-sharpness modes (see giant.model.network.EnergyRouter).
+            # learn_width generalizes the shared `temperature` to one
+            # learnable width per expert; learn_temperature instead makes
+            # the single shared `temperature` itself learnable. Both are
+            # bounded to [width_min_ratio, width_max_ratio] * temperature
+            # (sigmoid-parameterized, warm-started to reproduce `temperature`
+            # exactly at init) so gate sharpness can't run away to a
+            # collapse-inducing extreme during training.
+            "learn_width": False,
+            "learn_temperature": False,
+            "width_min_ratio": 0.1,
+            "width_max_ratio": 10.0,
             "lambda_balance": 0.0,  # optional load-balance aux loss weight
+            # optional entropy-regularization aux loss weight (generic
+            # Router.entropy_loss, penalizes uniform/collapsed gating) — a
+            # secondary guard against all experts' widths/temperature
+            # co-inflating together, which lambda_balance alone can't see
+            # since per-expert usage shares stay even throughout that
+            # failure mode. Off by default; bounding above is the primary
+            # defense. See giant.model.network.Router.entropy_loss.
+            "lambda_entropy": 0.0,
+            # Opt-in straight-through Gumbel-softmax train-time combine weights
+            # (see giant.model.network.Router.combine_weights): the training
+            # forward pass samples a hard one-hot combination — matching
+            # eval-time top-1 dispatch exactly — while the backward pass still
+            # flows a smooth gradient to every expert. Targets the train/eval
+            # mismatch identified as a likely contributor to experts
+            # overlapping instead of partitioning (see CLAUDE.md roadmap).
+            # gumbel_tau_start/_end are annealed linearly over training
+            # (giant.train._gumbel_tau); off by default, no effect unless
+            # gumbel = true.
+            "gumbel": False,
+            "gumbel_tau_start": 1.0,
+            "gumbel_tau_end": 0.1,
             "emb_dim": 8,  # process/pdg-router kwarg: own pdg(/mat) embedding width
             "hidden_dim": 64,  # process-router kwarg: its classifier's hidden width
             "lambda_proc": 0.0,  # process-router kwarg: supervised process-CE weight
@@ -305,6 +351,29 @@ def _router_candidate(train, model):
     return f"r-{router['type']}{router['n_experts']}"
 
 
+def _router_flag_candidate(field, token_map):
+    """Candidate factory for a boolean `model.router` sub-field.
+
+    Gated on `router.enabled` like `_router_candidate` (a disabled router's
+    sub-fields are meaningless), then omitted unless `field` differs from
+    its DEFAULT_CONFIG value — same "only show non-default" rule as every
+    other candidate. `token_map` need only cover the non-default value(s),
+    since the default value always yields None.
+    """
+
+    def _candidate(train, model):
+        router = model["router"]
+        default_router = DEFAULT_CONFIG["model"]["router"]
+        if router["enabled"] == default_router["enabled"]:
+            return None
+        value = router[field]
+        if value == default_router[field]:
+            return None
+        return token_map[value]
+
+    return _candidate
+
+
 def _conditioning_candidate(train, model):
     if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]:
         return None
@@ -326,6 +395,10 @@ def _default_field_candidate(section_key, field, prefix):
 _OUT_DIR_NAME_CANDIDATES = [
     ("mode", _mode_candidate),
     ("router", _router_candidate),
+    ("gumbel", _router_flag_candidate("gumbel", {True: "gum"})),
+    ("learn_centers", _router_flag_candidate("learn_centers", {False: "nolc"})),
+    ("learn_width", _router_flag_candidate("learn_width", {True: "lw"})),
+    ("learn_temperature", _router_flag_candidate("learn_temperature", {True: "lt"})),
     ("conditioning", _conditioning_candidate),
     ("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")),
     ("n_blocks", _default_field_candidate("model", "n_blocks", "b")),
@@ -372,6 +445,21 @@ def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
     return name
 
 
+def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path:
+    """Auto-derived out dir from cfg's hyperparams (see `default_out_dir_name`),
+    with a numeric suffix loop so two runs whose name collides (same
+    non-default hyperparams, same to-the-minute timestamp) don't clobber each
+    other's directory. Shared by `giant train` and `giant new-run`.
+    """
+    base_name = default_out_dir_name(cfg)
+    out_dir = base / base_name
+    suffix = 2
+    while out_dir.exists():
+        out_dir = base / f"{base_name}_{suffix}"
+        suffix += 1
+    return out_dir
+
+
 def seed_everything(seed: int) -> None:
     random.seed(seed)
     np.random.seed(seed)

giant/data/dataset.py

@@ -19,7 +19,10 @@ def make_event_split(
     rng = np.random.default_rng(seed)
     unique = np.unique(all_event_ids)
     rng.shuffle(unique)
-    n_val = max(1, int(len(unique) * val_fraction))
+    # max(1, ...) only applies when a validation split was actually
+    # requested — val_fraction=0.0 is an explicit "train on everything"
+    # request and must not be silently overridden into holding out 1 event.
+    n_val = max(1, int(len(unique) * val_fraction)) if val_fraction > 0 else 0
     val_set = set(unique[:n_val].tolist())
     train_set = set(unique[n_val:].tolist())
     return train_set, val_set

giant/data/loader.py

@@ -27,6 +27,26 @@ def event_id_offset(file_index: int) -> int:
     return file_index * EVENT_ID_FILE_STRIDE
 
 
+def _offset_event_id(raw_ids: np.ndarray, offset: int) -> np.ndarray:
+    """Add this file's `event_id_offset`, after checking the raw ids fit in one stride block.
+
+    Without this check, a file whose own raw event_id numbering reaches
+    `EVENT_ID_FILE_STRIDE` (an unusually large job, or non-contiguous
+    numbering) would silently collide into the next file's offset block,
+    merging unrelated events across files — reintroducing exactly the
+    train/val event leakage this offset scheme exists to prevent.
+    """
+    raw_ids = np.asarray(raw_ids, dtype=np.int64)
+    if raw_ids.size and int(raw_ids.max()) >= EVENT_ID_FILE_STRIDE:
+        raise ValueError(
+            f"event_id {int(raw_ids.max())} >= EVENT_ID_FILE_STRIDE "
+            f"({EVENT_ID_FILE_STRIDE}) — this file has a larger event_id "
+            "than the per-file offset scheme can support without colliding "
+            "with the next file's id block."
+        )
+    return raw_ids + offset
+
+
 def _read_manifest(path: Path) -> list[Path]:
     files = []
     for line in path.read_text().splitlines():
@@ -98,7 +118,7 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
-        "event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
+        "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
@@ -142,7 +162,7 @@ def load_steps(path: str | Path, offset: int = 0) -> dict[str, 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."""
     ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
-    return ids.astype(np.int64) + offset
+    return _offset_event_id(ids, offset)
 
 
 def iter_file_chunks(
@@ -173,7 +193,7 @@ _COND_COLS = [
 
 def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     return {
-        "event_id": df["event_id"].to_numpy().astype(np.int64) + offset,
+        "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
         "pdg": df["pdg"].to_numpy(dtype=np.int32),
         "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
         "pre_E": df["pre_E"].to_numpy(dtype=np.float32),

giant/data/setup_cache.py

@@ -13,6 +13,7 @@ before reuse — see `load`/`save`.
 
 from __future__ import annotations
 
+import fcntl
 import json
 import os
 from dataclasses import dataclass, field
@@ -31,7 +32,10 @@ from giant.data.transforms import Normalizer, sorted_membership
 # 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
+# v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
+# energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
+# v2 sidecar has no such grid to fall back on, so it must be recomputed.
+_CACHE_FORMAT_VERSION = 3
 
 _DIMS = {
     "COND_DIM": COND_DIM,
@@ -41,6 +45,29 @@ _DIMS = {
     "SEC_SLOT_DIM": SEC_SLOT_DIM,
 }
 
+# Resolution of the stored energy-quantile summary (see NormalizerEntry).
+# Only a handful of quantile *levels* (one per EnergyRouter expert) are ever
+# consumed (see pipeline.py), so a dense fixed grid of quantile values is
+# enough to reconstruct any level via interpolation (energy_quantile_at) —
+# at roughly 1/100th the storage of the raw 100k-value reservoir sample it
+# replaces, with negligible loss of resolution for that use.
+ENERGY_QUANTILE_LEVELS = 1001
+
+
+def energy_quantiles_from_sample(sample: np.ndarray) -> np.ndarray:
+    """Collapse a raw reservoir sample into the fixed grid stored on disk."""
+    if sample.size == 0:
+        return np.empty(0, dtype=np.float32)
+    levels = np.linspace(0.0, 1.0, ENERGY_QUANTILE_LEVELS)
+    return np.quantile(sample, levels).astype(np.float32)
+
+
+def energy_quantile_at(energy_quantiles: np.ndarray, levels: np.ndarray) -> np.ndarray:
+    """Interpolate quantile values at arbitrary probability `levels` from the
+    stored grid (e.g. `np.linspace(0, 1, n_experts)` for router centers)."""
+    grid_levels = np.linspace(0.0, 1.0, len(energy_quantiles))
+    return np.interp(levels, grid_levels, energy_quantiles).astype(np.float32)
+
 
 def sidecar_path(data: str | Path) -> Path:
     """The cache sidecar for `data`, always a sibling of `data` itself.
@@ -81,7 +108,10 @@ class NormalizerEntry:
     tgt_norm: Normalizer
     sec_phys_norm: Normalizer
     n_train_steps: int
-    energy_reservoir_sample: np.ndarray
+    energy_quantiles: np.ndarray
+    """Fixed ENERGY_QUANTILE_LEVELS-point quantile grid of the raw (pre-
+    normalization) pre-step energy column — see energy_quantiles_from_sample
+    / energy_quantile_at."""
 
     def to_json(self) -> dict:
         return {
@@ -89,8 +119,8 @@ class NormalizerEntry:
             "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
+            "energy_quantiles": np.asarray(
+                self.energy_quantiles, dtype=np.float32
             ).tolist(),
         }
 
@@ -101,9 +131,7 @@ class NormalizerEntry:
             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
-            ),
+            energy_quantiles=np.array(d["energy_quantiles"], dtype=np.float32),
         )
 
 
@@ -243,15 +271,32 @@ def save(
     Best-effort: any OSError (permission denied on a read-only mount, disk
     full, ...) is caught, echoed as a warning, and swallowed — a failure to
     cache must never fail training.
+
+    The load-merge-write is serialized with an exclusive flock on a sidecar
+    lockfile: `os.replace` alone only guarantees the *file* is never
+    corrupt, not that concurrent writers don't race. Without the lock, two
+    concurrent `giant train`/condor jobs against the same `data` path (this
+    repo's shared-portal/condor usage makes that a real scenario, not just
+    theoretical) could both `load()` the same base state, merge their own
+    `sections` in independently, and whichever `os.replace()` lands last
+    silently discards the other's freshly-computed section.
     """
     path = sidecar_path(data)
+    lock_path = path.parent / f".{path.name}.lock"
     tmp = path.parent / f".{path.name}.tmp.{os.getpid()}"
     try:
-        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)
+        with open(lock_path, "a") as lock_file:
+            fcntl.flock(lock_file, fcntl.LOCK_EX)
+            try:
+                base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(
+                    files
+                )
+                merged = base.merge(sections)
+                payload = json.dumps(merged.to_json(), separators=(",", ":"))
+                tmp.write_text(payload)
+                os.replace(tmp, path)
+            finally:
+                fcntl.flock(lock_file, fcntl.LOCK_UN)
     except OSError as exc:
         echo(
             f"setup cache: could not write {path} ({exc}) — continuing without caching"

giant/data/transforms.py

@@ -12,7 +12,17 @@ _SIMPLEX_FLOOR = 1e-5
 
 
 def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
-    return np.log(np.asarray(x, dtype=np.float32) + eps)
+    x = np.asarray(x, dtype=np.float32)
+    y = np.log(x + eps)
+    if not np.all(np.isfinite(y)):
+        bad = int(np.sum(~np.isfinite(y)))
+        raise ValueError(
+            f"log_transform: {bad} value(s) produced non-finite output (input "
+            f"< -eps={eps:g}, or already NaN/Inf); every quantity this is "
+            "applied to should be non-negative, so this indicates upstream "
+            "data corruption rather than expected float noise."
+        )
+    return y
 
 
 def inv_log_transform(y: np.ndarray, eps: float = _EPS) -> np.ndarray:
@@ -110,8 +120,22 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
         [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
     )
     axis_norm = np.linalg.norm(axis, axis=1, keepdims=True)  # (N,1)
-    # Replace zero-norm axes with x̂ (the Rodrigues terms that involve the axis
-    # are multiplied by sin_t≈0 and (1-cos_t)≈0, so the choice is irrelevant).
+    # axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
+    # pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
+    # ẑ) vanishes at both. The "choice is irrelevant" claim below only holds
+    # at +ẑ, where sin_t~0 AND (1-cos_t)~0 so every axis-dependent Rodrigues
+    # term vanishes. At -ẑ, sin_t~0 but (1-cos_t)~2 — not negligible — so
+    # snapping to a fixed x̂ there is a genuine (if physically rare)
+    # modeling choice, not a no-op: it picks one representative out of an
+    # inherently ambiguous family of 180°-about-any-transverse-axis
+    # rotations (no single-valued frame convention can be continuous through
+    # this antipode — same obstruction as a sphere's tangent frame having no
+    # continuous choice at a pole). x̂ is still fine to use — it's a fixed,
+    # self-consistent convention that `local_frame_rotation`/
+    # `inv_local_frame_rotation` (same threshold) round-trip correctly
+    # through — but steps whose pre_dir falls in this tiny near-backscatter
+    # cone get a discontinuous "roll" relative to their non-degenerate
+    # neighbors, injecting a small amount of label noise there.
     safe_norm = np.where(axis_norm < 1e-7, 1.0, axis_norm)
     return np.where(axis_norm < 1e-7, np.array([[1.0, 0.0, 0.0]]), axis / safe_norm)
 
@@ -134,8 +158,20 @@ def _validate_unit_pre_dir(pre_dir: np.ndarray) -> np.ndarray:
     drift is corrected silently; a near-zero-norm row has no well-defined
     direction, so it's raised loudly instead of producing a meaningless
     rotation (previously it fell through to an arbitrary axis with no error).
+
+    NaN/Inf rows are also raised on explicitly: `norm < 1e-6` is False for a
+    NaN norm, so without this check a non-finite row would silently pass
+    through and poison everything downstream (e.g. the persisted normalizer
+    stats in `setup_cache`, if the row is swept into a Welford accumulator).
     """
     pre_dir = np.asarray(pre_dir, dtype=np.float32)
+    if not np.all(np.isfinite(pre_dir)):
+        bad = int(np.sum(~np.all(np.isfinite(pre_dir), axis=1)))
+        raise ValueError(
+            f"pre_dir has {bad} row(s) with non-finite (NaN/Inf) components; "
+            "local/inv_local_frame_rotation require a well-defined incoming "
+            "direction for every row."
+        )
     norm = np.linalg.norm(pre_dir, axis=1, keepdims=True)
     if np.any(norm < 1e-6):
         raise ValueError(
@@ -303,12 +339,22 @@ def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
     return sorted_arr[idx] == values
 
 
-def _vectorized_map_lookup(values: np.ndarray, mapping: dict) -> np.ndarray:
+def _vectorized_map_lookup(
+    values: np.ndarray, mapping: dict, strict: bool = True
+) -> np.ndarray:
     """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.
 
     Replaces a per-element Python dict lookup with one `searchsorted` call.
     Raises `KeyError` if any value in `values` isn't a key of `mapping`,
-    matching the dict-comprehension it replaces (never silently misassigns).
+    matching the dict-comprehension it replaces (never silently misassigns)
+    — unless `strict=False`, in which case unmapped values get a dummy index
+    of 0 instead. Only pass `strict=False` where the caller has independently
+    verified the resulting index is never actually read (e.g.
+    `build_cond_features` under `conditioning="physical"`, where
+    `ConditionEncoder` ignores `cond_cat` entirely); it exists so a rollout
+    can be seeded with a species/material outside the training vocab without
+    a spurious `KeyError`, which is the entire point of physical-property
+    conditioning.
     """
     keys = np.asarray(list(mapping.keys()))
     vals = np.asarray(list(mapping.values()), dtype=np.int64)
@@ -319,6 +365,10 @@ def _vectorized_map_lookup(values: np.ndarray, mapping: dict) -> np.ndarray:
     pos = np.clip(pos, 0, len(keys_sorted) - 1)
     found = keys_sorted[pos] == values
     if not found.all():
+        if not strict:
+            out = np.zeros(values.shape, dtype=np.int64)
+            out[found] = vals_sorted[pos[found]]
+            return out
         missing = np.unique(values[~found])
         raise KeyError(f"value(s) not in mapping: {missing[:10].tolist()}")
     return vals_sorted[pos]
@@ -423,11 +473,21 @@ def encode_secondaries(
     else:
         cumsum = np.cumsum(sec_E_list.astype(np.float64), axis=1)
         stick_logits = np.zeros((N, K), dtype=np.float32)
+        # A valid slot whose cumulative secondary energy so far exceeds
+        # e_sec by more than float noise means sec_E_list sums to more than
+        # e_sec — a real upstream data mismatch, not something to paper
+        # over. Flagged once after the loop rather than let `remaining`'s
+        # np.maximum(..., _EPS) floor silently absorb it by saturating that
+        # slot's stick-breaking logit with no signal that anything was off.
+        _SHORTFALL_TOL = 1e-3
+        shortfall_flagged = np.zeros(N, dtype=bool)
         for i in range(K):
             if i == 0:
-                remaining = np.maximum(e_sec, _EPS)
+                remaining_raw = e_sec
             else:
-                remaining = np.maximum(e_sec - cumsum[:, i - 1], _EPS)
+                remaining_raw = e_sec - cumsum[:, i - 1]
+            shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
+            remaining = np.maximum(remaining_raw, _EPS)
             f = np.clip(
                 sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS
             )
@@ -444,6 +504,17 @@ def encode_secondaries(
             )
             stick_logits[:, i] = logit.astype(np.float32)
 
+        if shortfall_flagged.any():
+            n = int(shortfall_flagged.sum())
+            warnings.warn(
+                f"encode_secondaries: {n}/{N} row(s) have sec_E_list summing "
+                "to more than e_sec (beyond float noise) — the overflowing "
+                "slot(s)' stick-breaking logit was saturated instead of "
+                "reflecting a real fraction; check upstream secondary "
+                "energy accounting for these rows.",
+                stacklevel=2,
+            )
+
         # Rotate each slot's direction into the local frame of the primary.
         # pre_dir is broadcast across all K slots.
         dir_local = np.zeros((N, K, 3), dtype=np.float32)
@@ -636,8 +707,15 @@ def build_cond_features(
         [cond_cont, _physical_cond_columns(data, conditioning)]
     ).astype(np.float32)
 
-    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map)
-    mat_idx = _vectorized_map_lookup(data["material"], mat_map)
+    # In "physical" mode cond_cat is only a reporting/router convenience —
+    # ConditionEncoder never reads it (giant/model/network.py) — so a
+    # species/material outside the training vocab (the whole point of
+    # physical-property conditioning) gets a dummy index instead of raising.
+    # In "embedding" mode cond_cat IS the conditioning signal, so an unmapped
+    # value must still raise loudly rather than silently misassign.
+    strict = conditioning == "embedding"
+    pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict)
+    mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict)
     cond_cat = np.column_stack([pdg_idx, mat_idx])
 
     if cond_normalizer is not None:

giant/model/network.py

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

giant/model/wgan.py

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

giant/pipeline.py

@@ -1,3 +1,4 @@
+import os
 from dataclasses import dataclass
 from pathlib import Path
 
@@ -149,7 +150,7 @@ def run_setup_stage(
         cond_norm = entry.cond_norm
         tgt_norm = entry.tgt_norm
         sec_phys_norm = entry.sec_phys_norm
-        energy_sample = entry.energy_reservoir_sample
+        energy_quantiles = entry.energy_quantiles
     else:
         echo("fitting normalizer (streaming) …")
         cond_acc = _WelfordAccumulator(COND_DIM)
@@ -158,12 +159,13 @@ def run_setup_stage(
         # 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.
+        # same pass, not a second scan), then collapse it to a fixed quantile
+        # grid (setup_cache.energy_quantiles_from_sample) so centers can
+        # instead be seeded from actual data quantiles below. Collected
+        # whenever the setup cache is being populated, not only when *this*
+        # run's router is energy-typed, so a later run enabling
+        # --router-type energy against this same (val_fraction, seed,
+        # conditioning) key never needs to rescan just to seed centers.
         collect_energy_sample = energy_router_active or cache is not None
         energy_sampler = (
             _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
@@ -194,24 +196,24 @@ def run_setup_stage(
         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
+        energy_quantiles = (
+            setup_cache.energy_quantiles_from_sample(energy_sampler.sample)
             if energy_sampler is not None
             else np.empty(0, dtype=np.float32)
         )
         if cache is not None:
             cache.normalizers[norm_key] = setup_cache.NormalizerEntry(
-                cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_sample
+                cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles
             )
 
-    if energy_router_active and energy_sample.size > 0:
+    if energy_router_active and energy_quantiles.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()
+        levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
+        raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
+        centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
+            energy_idx
+        ]
+        router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
         echo(
             f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
         )
@@ -255,6 +257,20 @@ def run_train_job(
 
     out_dir = Path(out_dir)
 
+    # Soft warning (never blocks) — CLAUDE.md's Compute environment section
+    # asks that shared portal machines (portal1/deepthought{,2}/bms{1..3})
+    # stay within ~1/4 of CPU/RAM so as not to disturb other users' jobs;
+    # DataLoader's num_workers has no awareness of that on its own.
+    cpu_count = os.cpu_count() or 1
+    quota = max(1, cpu_count // 4)
+    if num_workers > quota:
+        echo(
+            f"warning: --num-workers={num_workers} exceeds ~1/4 of this "
+            f"machine's {cpu_count} CPU(s) ({quota}) — portal machines are "
+            "shared with other users (see CLAUDE.md's Compute environment "
+            "section)"
+        )
+
     router_cfg = m["router"]
     if t["mode"] == "wgan" and router_cfg.get("enabled"):
         raise ValueError(
@@ -415,6 +431,9 @@ def run_train_job(
         lambda_s2=t.get("lambda_s2", 1.0),
         lambda_balance=router_cfg.get("lambda_balance", 0.0),
         lambda_proc=router_cfg.get("lambda_proc", 0.0),
+        lambda_entropy=router_cfg.get("lambda_entropy", 0.0),
+        gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0),
+        gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1),
         normalizer_dict={
             "cond": cond_norm.to_dict(),
             "target": tgt_norm.to_dict(),

giant/rollout.py

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

giant/train.py

@@ -32,6 +32,7 @@ _METRICS_FIELDS = [
     "train_loss_s2",
     "train_loss_balance",
     "train_loss_proc",
+    "train_loss_entropy",
     "train_nsec_acc",
     "d_loss",
     "g_loss",
@@ -43,6 +44,7 @@ _METRICS_FIELDS = [
     "val_loss_s2",
     "val_loss_balance",
     "val_loss_proc",
+    "val_loss_entropy",
     "val_nsec_acc",
     "val_marginal_kl",
     "router_s1_entropy",
@@ -112,6 +114,85 @@ def _update_ema(
         ema_p.mul_(decay).add_(p, alpha=1 - decay)
 
 
+def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -> float:
+    """Linear anneal of the straight-through Gumbel-softmax temperature.
+
+    Deterministic in `step`/`total_steps` alone (no extra state), so it
+    recomputes correctly on `--resume` from a checkpoint's saved `global_step`
+    without needing to persist anything new (see
+    giant.model.network.Router.combine_weights).
+    """
+    progress = min(step / max(total_steps, 1), 1.0)
+    return tau_start + (tau_end - tau_start) * progress
+
+
+def _wandb_run_config(
+    *,
+    mode: str,
+    epochs: int,
+    lr: float,
+    warmup_epochs: int,
+    weight_decay: float,
+    ema_decay: float,
+    lambda_nsec: float,
+    lambda_s2: float,
+    lambda_balance: float,
+    lambda_proc: float,
+    lambda_entropy: float,
+    gumbel_tau_start: float,
+    gumbel_tau_end: float,
+    n_critic: int,
+    gp_weight: float,
+    model_config: dict | None,
+    stage1_params: int,
+    sec_decoder_params: int,
+    critic_params: int,
+    sec_critic_params: int,
+    total_params: int,
+) -> dict:
+    """Build the dict logged as a wandb run's `config`.
+
+    Router-only knobs (`lambda_balance`/`lambda_proc`/`lambda_entropy`/
+    `gumbel_tau_start`/`gumbel_tau_end`) and WGAN-only knobs (`n_critic`/
+    `gp_weight`) are omitted unless actually active, so a run's wandb config
+    doesn't imply hyperparameters from an inactive code path (a disabled
+    router's fine-tuning knobs, or GAN critic settings for a flow/DDPM run).
+    The full `model_config` (including its `router` sub-dict, whatever the
+    router type/state) is always included, so no information is lost — this
+    only trims the flattened top-level convenience duplicates.
+    """
+    router_enabled = bool((model_config or {}).get("router", {}).get("enabled", False))
+    cfg = {
+        "mode": mode,
+        "epochs": epochs,
+        "lr": lr,
+        "warmup_epochs": warmup_epochs,
+        "weight_decay": weight_decay,
+        "ema_decay": ema_decay,
+        "lambda_nsec": lambda_nsec,
+        "lambda_s2": lambda_s2,
+        "model": model_config or {},
+        "stage1_params": stage1_params,
+        "sec_decoder_params": sec_decoder_params,
+        "critic_params": critic_params,
+        "sec_critic_params": sec_critic_params,
+        "total_params": total_params,
+    }
+    if router_enabled:
+        cfg.update(
+            {
+                "lambda_balance": lambda_balance,
+                "lambda_proc": lambda_proc,
+                "lambda_entropy": lambda_entropy,
+                "gumbel_tau_start": gumbel_tau_start,
+                "gumbel_tau_end": gumbel_tau_end,
+            }
+        )
+    if mode == "wgan":
+        cfg.update({"n_critic": n_critic, "gp_weight": gp_weight})
+    return cfg
+
+
 def _compute_losses(
     stage1_model: torch.nn.Module,
     sec_decoder: torch.nn.Module,
@@ -123,6 +204,7 @@ def _compute_losses(
     lambda_s2: float,
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
+    lambda_entropy: float = 0.0,
 ) -> tuple[
     torch.Tensor,
     torch.Tensor,
@@ -131,8 +213,9 @@ def _compute_losses(
     torch.Tensor,
     torch.Tensor,
     torch.Tensor,
+    torch.Tensor,
 ]:
-    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, nsec_acc) for one batch."""
+    """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, L_entropy, nsec_acc) for one batch."""
     cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch
     cond_cont = cond_cont.to(device)
     cond_cat = cond_cat.to(device)
@@ -185,16 +268,25 @@ def _compute_losses(
         l_proc = stage1_model.router.classify_loss(
             cond_cont, cond_cat, proc_idx
         ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx)
+        # Optional entropy-regularization aux loss (see Router.entropy_loss):
+        # penalizes uniform/collapsed gating, a secondary guard against
+        # gate-sharpness collapse that lambda_balance alone can't see.
+        l_entropy = stage1_model.router.entropy_loss(
+            cond_cont, cond_cat
+        ) + sec_decoder.router.entropy_loss(cond_cont, cond_cat)
     else:
         l_balance = torch.zeros((), device=device)
         l_proc = torch.zeros((), device=device)
+        l_entropy = torch.zeros((), device=device)
 
     total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2
     if lambda_balance > 0:
         total = total + lambda_balance * l_balance
     if lambda_proc > 0:
         total = total + lambda_proc * l_proc
-    return total, l_s1, l_nsec, l_s2, l_balance, l_proc, nsec_acc
+    if lambda_entropy > 0:
+        total = total + lambda_entropy * l_entropy
+    return total, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc
 
 
 def _wgan_train_step(
@@ -333,6 +425,9 @@ def train(
     lambda_s2: float = 1.0,
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
+    lambda_entropy: float = 0.0,
+    gumbel_tau_start: float = 1.0,
+    gumbel_tau_end: float = 0.1,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -384,26 +479,29 @@ def train(
             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,
-            },
+            config=_wandb_run_config(
+                mode=mode,
+                epochs=epochs,
+                lr=lr,
+                warmup_epochs=warmup_epochs,
+                weight_decay=weight_decay,
+                ema_decay=ema_decay,
+                lambda_nsec=lambda_nsec,
+                lambda_s2=lambda_s2,
+                lambda_balance=lambda_balance,
+                lambda_proc=lambda_proc,
+                lambda_entropy=lambda_entropy,
+                gumbel_tau_start=gumbel_tau_start,
+                gumbel_tau_end=gumbel_tau_end,
+                n_critic=n_critic,
+                gp_weight=gp_weight,
+                model_config=model_config,
+                stage1_params=stage1_params,
+                sec_decoder_params=sec_decoder_params,
+                critic_params=critic_params,
+                sec_critic_params=sec_critic_params,
+                total_params=total_params,
+            ),
         )
 
     stage1_model = stage1_model.to(device)
@@ -466,7 +564,9 @@ def train(
     # would never finish and cosine decay would barely move.
     steps_per_epoch = max(total_train_batches, 1)
     if mode == "wgan":
-        steps_per_epoch = max(total_train_batches // (n_critic + 1), 1)
+        # Generator steps fire every n_critic-th batch (did_g_step =
+        # step_count % n_critic == 0 in _wgan_train_step), not n_critic + 1.
+        steps_per_epoch = max(total_train_batches // n_critic, 1)
     warmup_steps = warmup_epochs * steps_per_epoch
     total_steps = max(epochs * steps_per_epoch, 1)
 
@@ -481,6 +581,41 @@ def train(
 
     ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None
 
+    def _build_checkpoint(epoch: int, global_step: int, best_val_loss: float) -> dict:
+        ckpt: dict = {
+            "model": stage1_model.state_dict(),
+            "sec_decoder": sec_decoder.state_dict(),
+            "optimizer": optimizer.state_dict(),
+            "lr_sched": lr_sched.state_dict(),
+            "epoch": epoch,
+            "best_val_loss": best_val_loss,
+            "global_step": global_step,
+        }
+        if mode == "wgan":
+            assert (
+                critic is not None
+                and sec_critic is not None
+                and optimizer_d is not None
+            )
+            ckpt["critic"] = critic.state_dict()
+            ckpt["sec_critic"] = sec_critic.state_dict()
+            ckpt["optimizer_d"] = optimizer_d.state_dict()
+        if ema_decay > 0:
+            assert ema_stage1_model is not None and ema_sec_decoder is not None
+            ckpt["model_ema"] = ema_stage1_model.state_dict()
+            ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict()
+        if normalizer_dict is not None:
+            ckpt["normalizer"] = normalizer_dict
+        if pdg_map is not None:
+            ckpt["pdg_map"] = pdg_map
+        if mat_map is not None:
+            ckpt["mat_map"] = mat_map
+        if proc_map is not None:
+            ckpt["proc_map"] = proc_map
+        if model_config is not None:
+            ckpt["model_config"] = model_config
+        return ckpt
+
     start_epoch = 1
     best_val_loss = float("inf")
     resumed_global_step = 0
@@ -503,6 +638,15 @@ def train(
             critic.load_state_dict(ckpt["critic"])
             sec_critic.load_state_dict(ckpt["sec_critic"])
             optimizer_d.load_state_dict(ckpt["optimizer_d"])
+            # Mirrors the `lr` fixup below for the generator optimizer:
+            # optimizer_d.load_state_dict() above restores the checkpoint's
+            # own critic LR, which would otherwise silently override an
+            # explicit `--critic-lr` passed on this resume. optimizer_d has
+            # no LR scheduler (unlike `optimizer`/`lr_sched`), so this is a
+            # flat set rather than a schedule-relative one.
+            resumed_critic_lr = critic_lr if critic_lr is not None else lr
+            for group in optimizer_d.param_groups:
+                group["lr"] = resumed_critic_lr
         optimizer.load_state_dict(ckpt["optimizer"])
         lr_sched.load_state_dict(ckpt["lr_sched"])
         start_epoch = ckpt.get("epoch", 0) + 1
@@ -559,6 +703,7 @@ def train(
             train_s2_sum = 0.0
             train_balance_sum = 0.0
             train_proc_sum = 0.0
+            train_entropy_sum = 0.0
             train_d_sum = 0.0
             train_g_sum = 0.0
             train_wasserstein_sum = 0.0
@@ -580,6 +725,13 @@ def train(
                 dynamic_ncols=True,
             )
             for batch in bar:
+                if has_router:
+                    gumbel_tau = _gumbel_tau(
+                        global_step, total_steps, gumbel_tau_start, gumbel_tau_end
+                    )
+                    stage1_model.router.gumbel_tau = gumbel_tau
+                    sec_decoder.router.gumbel_tau = gumbel_tau
+
                 if mode == "wgan":
                     assert (
                         critic is not None
@@ -626,7 +778,7 @@ def train(
                     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, nsec_acc = (
+                    loss, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc = (
                         _compute_losses(
                             stage1_model,
                             sec_decoder,
@@ -638,6 +790,7 @@ def train(
                             lambda_s2,
                             lambda_balance,
                             lambda_proc,
+                            lambda_entropy,
                         )
                     )
                     optimizer.zero_grad()
@@ -661,6 +814,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_entropy_sum += l_entropy.item() * B
                     train_nsec_acc_sum += nsec_acc.item() * B
 
                 train_n += B
@@ -719,6 +873,7 @@ def train(
                             )
                         log_payload["batch/router_s1_entropy"] = s1_entropy.item()
                         log_payload["batch/router_s2_entropy"] = s2_entropy.item()
+                        log_payload["batch/gumbel_tau"] = gumbel_tau
                     wandb_run.log(log_payload, step=global_step)
 
                 if shutdown.requested:
@@ -726,6 +881,20 @@ def train(
             bar.close()
 
             if shutdown.requested:
+                # Epoch was interrupted mid-loop, so there's no val_loss to
+                # weigh a "best" checkpoint against — save the in-progress
+                # weights as last.pt only, under the last *fully completed*
+                # epoch number so --resume restarts this epoch from scratch
+                # rather than skipping it (weights/optimizer state are still
+                # kept, so those partial-epoch batches aren't wasted work).
+                ckpt = _build_checkpoint(epoch - 1, global_step, best_val_loss)
+                torch.save(ckpt, out_dir / "last.pt")
+                last_completed_epoch = epoch - 1
+                print(
+                    f"saved in-progress weights from partway through epoch "
+                    f"{epoch} to {out_dir / 'last.pt'} "
+                    f"(resume will restart epoch {epoch})"
+                )
                 break
 
             train_loss = train_loss_sum / max(train_n, 1)
@@ -772,7 +941,7 @@ def train(
                 val_loss = val_marginal_kl
                 val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = (
                     val_proc_sum
-                ) = val_nsec_acc_sum = 0.0
+                ) = val_entropy_sum = val_nsec_acc_sum = 0.0
                 val_n = 1
                 val_nsec_acc = 0.0
                 router_s1_entropy = router_s2_entropy = 0.0
@@ -785,6 +954,7 @@ def train(
                 val_s2_sum = 0.0
                 val_balance_sum = 0.0
                 val_proc_sum = 0.0
+                val_entropy_sum = 0.0
                 val_nsec_acc_sum = 0.0
                 val_n = 0
                 if has_router:
@@ -802,19 +972,27 @@ def train(
                     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, nsec_acc = (
-                            _compute_losses(
-                                stage1_model,
-                                sec_decoder,
-                                batch,
-                                mode,
-                                ddpm_schedule,
-                                device,
-                                lambda_nsec,
-                                lambda_s2,
-                                lambda_balance,
-                                lambda_proc,
-                            )
+                        (
+                            loss,
+                            l_s1,
+                            l_nsec,
+                            l_s2,
+                            l_balance,
+                            l_proc,
+                            l_entropy,
+                            nsec_acc,
+                        ) = _compute_losses(
+                            stage1_model,
+                            sec_decoder,
+                            batch,
+                            mode,
+                            ddpm_schedule,
+                            device,
+                            lambda_nsec,
+                            lambda_s2,
+                            lambda_balance,
+                            lambda_proc,
+                            lambda_entropy,
                         )
                         B = batch[0].size(0)
                         val_loss_sum += loss.item() * B
@@ -823,6 +1001,7 @@ def train(
                         val_s2_sum += l_s2.item() * B
                         val_balance_sum += l_balance.item() * B
                         val_proc_sum += l_proc.item() * B
+                        val_entropy_sum += l_entropy.item() * B
                         val_nsec_acc_sum += nsec_acc.item() * B
                         if has_router:
                             cond_cont = batch[0].to(device)
@@ -895,6 +1074,7 @@ def train(
                 f"  s2={train_s2_sum / max(train_n, 1):.3f}"
                 f"  bal={train_balance_sum / max(train_n, 1):.3f}"
                 f"  proc={train_proc_sum / max(train_n, 1):.3f}"
+                f"  entropy={train_entropy_sum / max(train_n, 1):.3f}"
                 f"  d={train_d_sum / max(train_n, 1):.3f}"
                 f"  g={train_g_sum / max(train_n, 1):.3f})"
                 f"  val {val_loss:.4f}"
@@ -909,6 +1089,7 @@ def train(
                 "train_loss_s2": train_s2_sum / max(train_n, 1),
                 "train_loss_balance": train_balance_sum / max(train_n, 1),
                 "train_loss_proc": train_proc_sum / max(train_n, 1),
+                "train_loss_entropy": train_entropy_sum / max(train_n, 1),
                 "train_nsec_acc": train_nsec_acc,
                 "d_loss": train_d_sum / max(train_n, 1),
                 "g_loss": train_g_sum / max(train_n, 1),
@@ -920,6 +1101,7 @@ def train(
                 "val_loss_s2": val_s2_sum / max(val_n, 1),
                 "val_loss_balance": val_balance_sum / max(val_n, 1),
                 "val_loss_proc": val_proc_sum / max(val_n, 1),
+                "val_loss_entropy": val_entropy_sum / max(val_n, 1),
                 "val_nsec_acc": val_nsec_acc,
                 "val_marginal_kl": val_marginal_kl,
                 "router_s1_entropy": router_s1_entropy,
@@ -949,38 +1131,7 @@ def train(
                 # must never decrease.
                 wandb_run.log(metrics_row, step=global_step)
 
-            ckpt: dict = {
-                "model": stage1_model.state_dict(),
-                "sec_decoder": sec_decoder.state_dict(),
-                "optimizer": optimizer.state_dict(),
-                "lr_sched": lr_sched.state_dict(),
-                "epoch": epoch,
-                "best_val_loss": best_val_loss,
-                "global_step": global_step,
-            }
-            if mode == "wgan":
-                assert (
-                    critic is not None
-                    and sec_critic is not None
-                    and optimizer_d is not None
-                )
-                ckpt["critic"] = critic.state_dict()
-                ckpt["sec_critic"] = sec_critic.state_dict()
-                ckpt["optimizer_d"] = optimizer_d.state_dict()
-            if ema_decay > 0:
-                assert ema_stage1_model is not None and ema_sec_decoder is not None
-                ckpt["model_ema"] = ema_stage1_model.state_dict()
-                ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict()
-            if normalizer_dict is not None:
-                ckpt["normalizer"] = normalizer_dict
-            if pdg_map is not None:
-                ckpt["pdg_map"] = pdg_map
-            if mat_map is not None:
-                ckpt["mat_map"] = mat_map
-            if proc_map is not None:
-                ckpt["proc_map"] = proc_map
-            if model_config is not None:
-                ckpt["model_config"] = model_config
+            ckpt = _build_checkpoint(epoch, global_step, best_val_loss)
 
             if val_loss < best_val_loss:
                 best_val_loss = val_loss

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.1.0"
+version = "0.2.0"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"

scripts/bump_dataset_version.py

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

scripts/dwarf.py

@@ -5,6 +5,7 @@ simulation-fanout tools into one Typer app so there's a single command name
 (and `--help`) to remember instead of five differently-hyphenated ones.
 """
 
+import os
 from enum import Enum
 from pathlib import Path
 from typing import Optional
@@ -12,6 +13,7 @@ from typing import Optional
 import typer
 from typing_extensions import Annotated
 
+from giant.config import Conditioning
 from scripts.bump_dataset_version import (
     run_bump_gen,
     run_bump_schema,
@@ -37,6 +39,25 @@ def _main() -> None:
     """dwarf — dataset/tooling CLI (ROOT<->parquet conversion, dataset versioning, sim fanout)."""
 
 
+def _warn_if_exceeds_shared_quota(n: int, flag: str) -> None:
+    """Soft warning (never blocks) when a worker/job count looks likely to
+    grab more than this repo's documented shared-portal-machine etiquette
+    (CLAUDE.md's Compute environment: stay within ~1/4 of CPU/RAM and a
+    single GPU, since portal1/deepthought{,2}/bms{1..3} are shared with
+    other users). Not a hard cap — a legitimate big machine or a
+    deliberately aggressive run is still the caller's call.
+    """
+    cpu_count = os.cpu_count() or 1
+    quota = max(1, cpu_count // 4)
+    if n > quota:
+        typer.echo(
+            f"warning: {flag}={n} exceeds ~1/4 of this machine's "
+            f"{cpu_count} CPU(s) ({quota}) — portal machines are shared "
+            "with other users (see CLAUDE.md's Compute environment section)",
+            err=True,
+        )
+
+
 class Compression(str, Enum):
     snappy = "snappy"
     lz4 = "lz4"
@@ -106,6 +127,7 @@ def convert(
     if jobs < 1:
         typer.echo("error: --jobs must be >= 1", err=True)
         raise typer.Exit(1)
+    _warn_if_exceeds_shared_quota(jobs, "--jobs")
 
     compression_value = (
         "uncompressed" if compression is Compression.none else compression.value
@@ -327,6 +349,10 @@ def create_manifest(
     execute: Annotated[
         bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
     ] = False,
+    force: Annotated[
+        bool,
+        typer.Option("--force", help="Overwrite the manifest if it already exists"),
+    ] = False,
 ) -> None:
     """Create a new manifest from a list of parquet files."""
     run_create_manifest(
@@ -336,6 +362,7 @@ def create_manifest(
         pool=pool,
         type_=type_.value if type_ is not None else None,
         root=str(root),
+        force=force,
     )
 
 
@@ -391,6 +418,7 @@ def make_root(
     ] = False,
 ) -> None:
     """Generate new ROOT shards via a minicalosim executable."""
+    _warn_if_exceeds_shared_quota(jobs, "--jobs")
     run_make_root(
         executable=executable,
         detector=detector,
@@ -472,11 +500,6 @@ def build_geometry_oracle(
     )
 
 
-class Conditioning(str, Enum):
-    physical = "physical"
-    embedding = "embedding"
-
-
 @app.command("warm-cache")
 def warm_cache(
     data: Annotated[

scripts/warm_setup_cache.py

@@ -30,7 +30,7 @@ def run_warm_setup_cache(
     `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
+    energy-router quantile summary is always collected regardless, so a
     later `--router-type energy` run never needs to rescan just to seed
     centers.
     """

tests/test_bump_dataset_version.py

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

tests/test_cli_new_run.py

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

tests/test_condor.py

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

tests/test_config.py

@@ -137,6 +137,16 @@ def test_resolve_expert_dims_missing_keys_also_inherit():
     assert (hidden_dim, n_blocks) == (512, 6)
 
 
+def test_default_config_gumbel_router_defaults_off():
+    # Straight-through Gumbel-softmax combine weights (giant.model.network.
+    # Router.combine_weights) must be opt-in — existing routed configs and
+    # checkpoints should be unaffected unless gumbel is explicitly enabled.
+    router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"]
+    assert router_cfg["gumbel"] is False
+    assert router_cfg["gumbel_tau_start"] == 1.0
+    assert router_cfg["gumbel_tau_end"] == 0.1
+
+
 def test_resolve_expert_dims_explicit_override_wins():
     router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3}
     hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6)
@@ -197,6 +207,64 @@ def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefau
     assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
 
 
+def test_default_out_dir_name_router_gumbel_shown_when_enabled():
+    cfg = _default_cfg(
+        router={"enabled": True, "type": "energy", "n_experts": 8, "gumbel": True}
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_gum"
+
+
+def test_default_out_dir_name_router_gumbel_omitted_when_router_disabled():
+    cfg = _default_cfg(
+        router={"enabled": False, "type": "energy", "n_experts": 8, "gumbel": True}
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430"
+
+
+def test_default_out_dir_name_router_learn_centers_shown_only_when_disabled():
+    cfg_default = _default_cfg(
+        router={"enabled": True, "type": "energy", "n_experts": 8}
+    )
+    assert (
+        gconfig.default_out_dir_name(cfg_default, now=_NOW) == "20260729_1430_r-energy8"
+    )
+
+    cfg_off = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_centers": False,
+        }
+    )
+    assert (
+        gconfig.default_out_dir_name(cfg_off, now=_NOW)
+        == "20260729_1430_r-energy8_nolc"
+    )
+
+
+def test_default_out_dir_name_router_learn_width_and_temperature_shown():
+    cfg = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_width": True,
+        }
+    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_lw"
+
+    cfg2 = _default_cfg(
+        router={
+            "enabled": True,
+            "type": "energy",
+            "n_experts": 8,
+            "learn_temperature": True,
+        }
+    )
+    assert gconfig.default_out_dir_name(cfg2, now=_NOW) == "20260729_1430_r-energy8_lt"
+
+
 def test_default_out_dir_name_mode_shown_bare_no_prefix():
     cfg = _default_cfg(mode="wgan")
     assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan"

tests/test_dataset.py

@@ -30,6 +30,16 @@ def test_make_event_split_no_empty_sets():
     assert len(val_set) > 0
 
 
+def test_make_event_split_val_fraction_zero_holds_out_nothing():
+    """val_fraction=0.0 is an explicit "train on everything" request and
+    must not be silently overridden into holding out 1 event."""
+    rng = np.random.default_rng(3)
+    event_ids = rng.integers(0, 50, size=1000)
+    train_set, val_set = make_event_split(event_ids, val_fraction=0.0)
+    assert val_set == set()
+    assert train_set == set(np.unique(event_ids).tolist())
+
+
 def test_make_event_split_reproducible():
     event_ids = np.arange(100)
     a_tr, a_val = make_event_split(event_ids, val_fraction=0.1, seed=42)

tests/test_dwarf.py

@@ -1,12 +1,35 @@
 from typer.testing import CliRunner
 
+from giant import cli as giant_cli
+from giant.config import Conditioning
 from giant.data import setup_cache
+from scripts import dwarf
 from scripts.dwarf import app
 from test_pipeline import _make_synthetic_steps
 
 runner = CliRunner()
 
 
+def test_conditioning_enum_shared_across_both_clis():
+    """giant.cli and scripts.dwarf must use the one giant.config.Conditioning
+    enum, not independently redefined copies that could silently drift apart
+    on valid --conditioning values."""
+    assert dwarf.Conditioning is Conditioning
+    assert giant_cli.Conditioning is Conditioning
+
+
+def test_warn_if_exceeds_shared_quota_warns_over_quarter_cpu(monkeypatch, capsys):
+    monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8)  # quota = 2
+    dwarf._warn_if_exceeds_shared_quota(3, "--jobs")
+    assert "warning: --jobs=3 exceeds" in capsys.readouterr().err
+
+
+def test_warn_if_exceeds_shared_quota_silent_within_quota(monkeypatch, capsys):
+    monkeypatch.setattr(dwarf.os, "cpu_count", lambda: 8)  # quota = 2
+    dwarf._warn_if_exceeds_shared_quota(2, "--jobs")
+    assert capsys.readouterr().err == ""
+
+
 def test_convert_rejects_jobs_below_one(tmp_path):
     root_file = tmp_path / "shard.root"
     root_file.touch()

tests/test_loader.py

@@ -344,6 +344,15 @@ def test_load_event_ids_applies_offset(tmp_path):
     )
 
 
+def test_load_event_ids_raises_when_event_id_reaches_stride(tmp_path):
+    """A raw event_id >= EVENT_ID_FILE_STRIDE would collide into the next
+    file's offset block if silently allowed through — must raise instead."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"event_id": [0, 1, EVENT_ID_FILE_STRIDE]}).to_parquet(path)
+    with pytest.raises(ValueError, match="EVENT_ID_FILE_STRIDE"):
+        load_event_ids(path)
+
+
 def test_load_steps_applies_offset_to_event_id(tmp_path):
     path = tmp_path / "a.parquet"
     _steps_df([0, 1]).to_parquet(path)

tests/test_pipeline.py

@@ -108,13 +108,13 @@ def _tiny_cfg(**train_overrides):
 
 def _run(data, out_dir, cfg=None, **kwargs):
     echoed: list[str] = []
+    kwargs.setdefault("num_workers", 0)
     run_train_job(
         data=data,
         cfg=cfg or _tiny_cfg(),
         out_dir=out_dir,
         device=torch.device("cpu"),
         shuffle_buffer=64,
-        num_workers=0,
         echo=echoed.append,
         **kwargs,
     )
@@ -143,6 +143,22 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
     assert "normalizer: cache hit" in joined
 
 
+def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
+    tmp_path, data, monkeypatch
+):
+    monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8)  # quota = 2
+    echo = _run(data, tmp_path / "out", num_workers=3)
+    assert any("num-workers=3" in m and "exceeds" in m for m in echo)
+
+
+def test_run_train_job_no_warning_when_num_workers_within_shared_quota(
+    tmp_path, data, monkeypatch
+):
+    monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8)  # quota = 2
+    echo = _run(data, tmp_path / "out", num_workers=2)
+    assert not any("exceeds" in m for m in echo)
+
+
 def test_run_train_job_no_cache_setup_never_writes_sidecar(tmp_path, data):
     _run(data, tmp_path / "out", cache_setup=False)
     assert not setup_cache.sidecar_path(data).exists()

tests/test_rollout.py

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

tests/test_router.py

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

tests/test_setup_cache.py

@@ -98,7 +98,7 @@ def test_save_load_round_trip(tmp_path):
     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])
+    np.testing.assert_allclose(entry.energy_quantiles, [1.0, 2.0, 3.0])
 
 
 def test_load_missing_sidecar_returns_none(tmp_path):
@@ -214,6 +214,74 @@ def test_save_merges_non_colliding_normalizer_keys(tmp_path):
     assert loaded.normalizers["k2"].n_train_steps == 2
 
 
+def test_save_is_serialized_against_concurrent_writers(tmp_path):
+    """Without the flock in setup_cache.save(), two concurrent writers can
+    both load() the same base state and merge their own section in
+    independently, so whichever os.replace() lands last silently drops the
+    other's key — a lost-update race, not a corrupt file. Each of these
+    threads writes a distinct normalizer key many times over; if the
+    load-merge-write critical section isn't actually serialized, at least
+    one thread's key is likely to go missing from the final merged cache."""
+    import threading
+
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    n_writers, n_rounds = 6, 15
+
+    def _writer(idx: int) -> None:
+        for r in range(n_rounds):
+            cache = SetupCache.empty(files)
+            cache.normalizers[f"k{idx}"] = _entry(n_train_steps=r)
+            setup_cache.save(data, files, cache)
+
+    threads = [threading.Thread(target=_writer, args=(i,)) for i in range(n_writers)]
+    for t in threads:
+        t.start()
+    for t in threads:
+        t.join()
+
+    loaded = setup_cache.load(data, files)
+    assert loaded is not None
+    assert set(loaded.normalizers.keys()) == {f"k{i}" for i in range(n_writers)}
+    for i in range(n_writers):
+        assert loaded.normalizers[f"k{i}"].n_train_steps == n_rounds - 1
+
+
+# ── energy_quantiles_from_sample / energy_quantile_at ───────────────────
+
+
+def test_energy_quantiles_from_sample_empty():
+    result = setup_cache.energy_quantiles_from_sample(np.empty(0, dtype=np.float32))
+    assert result.size == 0
+
+
+def test_energy_quantiles_from_sample_has_fixed_grid_size():
+    sample = np.random.default_rng(0).normal(size=5000).astype(np.float32)
+    result = setup_cache.energy_quantiles_from_sample(sample)
+    assert result.shape == (setup_cache.ENERGY_QUANTILE_LEVELS,)
+    assert result[0] == pytest.approx(sample.min(), abs=1e-3)
+    assert result[-1] == pytest.approx(sample.max(), abs=1e-3)
+
+
+def test_energy_quantile_at_matches_direct_quantile_on_stored_grid():
+    sample = np.random.default_rng(1).exponential(size=20_000).astype(np.float32)
+    grid = setup_cache.energy_quantiles_from_sample(sample)
+
+    levels = np.linspace(0.0, 1.0, 5)
+    got = setup_cache.energy_quantile_at(grid, levels)
+    expected = np.quantile(sample, levels)
+
+    np.testing.assert_allclose(got, expected, rtol=0.05)
+
+
+def test_energy_quantile_at_median_of_two_points():
+    grid = np.array([0.0, 10.0], dtype=np.float32)
+    result = setup_cache.energy_quantile_at(grid, np.array([0.0, 0.5, 1.0]))
+    np.testing.assert_allclose(result, [0.0, 5.0, 10.0])
+
+
 # ── n_train_steps_for_split ──────────────────────────────────────────────
 
 

tests/test_train.py

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

tests/test_transforms.py

@@ -1,9 +1,12 @@
+import warnings
+
 import numpy as np
 import pytest
 from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX
 from giant.data.transforms import (
     build_cond_features,
     build_features,
+    encode_secondaries,
     energy_simplex_decode,
     energy_simplex_encode,
     inv_local_frame_rotation,
@@ -24,6 +27,50 @@ def test_log_transform_invertible():
     np.testing.assert_allclose(inv_log_transform(log_transform(x)), x, rtol=1e-5)
 
 
+def test_log_transform_raises_on_input_below_negative_eps():
+    """A meaningfully negative input (upstream data corruption, not float
+    noise near 0) must raise instead of silently returning NaN."""
+    x = np.array([1.0, -5.0], dtype=np.float32)
+    with np.errstate(invalid="ignore"), pytest.raises(ValueError, match="non-finite"):
+        log_transform(x)
+
+
+def test_log_transform_raises_on_nan_input():
+    x = np.array([1.0, np.nan], dtype=np.float32)
+    with pytest.raises(ValueError, match="non-finite"):
+        log_transform(x)
+
+
+def test_encode_secondaries_warns_when_sec_energies_exceed_e_sec():
+    """sec_E_list summing to more than e_sec (before the last slot is even
+    reached) is a real upstream data mismatch — must warn instead of
+    silently saturating the overflowing slot's stick-breaking logit via the
+    _EPS floor. (A single slot alone exceeding what's left of the budget is
+    the normal, expected "last slot takes the remainder" case and must NOT
+    warn — the mismatch here is the *cumulative* sum through an earlier
+    slot already exceeding e_sec.)"""
+    sec_E_list = np.array([[5.0, 4.0, 1.0]], dtype=np.float32)  # sums to 10
+    sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 3, 1)).astype(np.float32)
+    sec_valid = np.array([[True, True, True]])
+    e_sec = np.array([6.0], dtype=np.float32)  # cumsum already 9 by slot 2
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    with pytest.warns(UserWarning, match="sec_E_list summing to more than e_sec"):
+        encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
+
+
+def test_encode_secondaries_no_warning_when_energies_are_consistent():
+    sec_E_list = np.array([[3.0, 2.0]], dtype=np.float32)  # sums to 5
+    sec_dir_list = np.tile([0.0, 0.0, 1.0], (1, 2, 1)).astype(np.float32)
+    sec_valid = np.array([[True, True]])
+    e_sec = np.array([6.0], dtype=np.float32)  # >= 5, no shortfall
+    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir)
+
+
 def test_local_frame_rotation_noop_when_aligned():
     N = 8
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
@@ -68,10 +115,43 @@ def test_local_frame_rotation_rejects_near_zero_pre_dir():
     post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
     with pytest.raises(ValueError, match="near-zero norm"):
         local_frame_rotation(pre_dir, post_dir)
-    with pytest.raises(ValueError, match="near-zero norm"):
+
+
+def test_local_frame_rotation_rejects_nan_pre_dir():
+    """A NaN pre_dir must raise loudly — `norm < 1e-6` is False for NaN, so
+    without an explicit isfinite check this would silently poison the
+    rotation (and any normalizer stats it feeds) instead of erroring."""
+    pre_dir = np.array([[np.nan, 0.0, 1.0], [0.0, 0.0, 1.0]], dtype=np.float32)
+    post_dir = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]], dtype=np.float32)
+    with pytest.raises(ValueError, match="non-finite"):
+        local_frame_rotation(pre_dir, post_dir)
+    with pytest.raises(ValueError, match="non-finite"):
         inv_local_frame_rotation(pre_dir, post_dir)
 
 
+def test_local_frame_rotation_antipodal_pre_dir_uses_x_axis_convention():
+    """pre_dir ~ -ẑ (near-exact backscatter) is a second axis_norm~0
+    degeneracy besides pre_dir ~ +ẑ; unlike the forward case, the Rodrigues
+    axis-dependent terms are NOT negligible there ((1-cos_t)~2), so the x̂
+    fallback is a real (if arbitrary and physically rare) convention choice
+    rather than a no-op. Pin it explicitly — angle-preservation and the
+    round-trip property must still hold even though the "roll" is degenerate.
+    """
+    pre_dir = np.array([[0.0, 0.0, -1.0]], dtype=np.float32)
+    post_dir = np.array([[0.3, 0.4, 0.5]], dtype=np.float32)
+    post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
+
+    rotated = local_frame_rotation(pre_dir, post_dir)
+
+    cos_before = (pre_dir * post_dir).sum(axis=1)
+    cos_after = rotated[:, 2]
+    np.testing.assert_allclose(cos_after, cos_before, atol=1e-5)
+    np.testing.assert_allclose(np.linalg.norm(rotated, axis=1), 1.0, atol=1e-5)
+
+    recovered = inv_local_frame_rotation(pre_dir, rotated)
+    np.testing.assert_allclose(recovered, post_dir, atol=1e-5)
+
+
 def test_local_frame_rotation_normalizes_non_unit_pre_dir():
     """A pre_dir with float32-drift norm (not exactly 1) must still produce the
     same result as its exactly-normalized counterpart, not a skewed frame."""
@@ -494,6 +574,51 @@ def test_vectorized_map_lookup_raises_keyerror_on_missing_value():
         _vectorized_map_lookup(values, mapping)
 
 
+def test_vectorized_map_lookup_strict_false_dummy_indexes_unmapped_values():
+    """strict=False must leave found values untouched and only dummy-index
+    (0) the unmapped ones — never raise, and never disturb a value that IS
+    in the mapping (e.g. one that happens to map to a nonzero index)."""
+    mapping = {1: 5, 2: 7}
+    values = np.array([1, 99, 2, 100])
+    result = _vectorized_map_lookup(values, mapping, strict=False)
+    np.testing.assert_array_equal(result, [5, 0, 7, 0])
+
+
+def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_material():
+    """conditioning="physical" must not KeyError on a pdg/material outside
+    the training-dataset vocab (mat_map/pdg_map) — that's the entire point
+    of the mode (see giant.rollout's known_pdg gate for the paired fix).
+    "embedding" mode must still raise, since cond_cat IS the conditioning
+    signal there. Note this is specifically about the dataset-scoped
+    vocab index, not giant.materials' physical-properties table — a
+    material must still be a real, known Geant4 material (e.g. "G4_Pb",
+    just not one *this* mat_map happened to include) for "physical" mode
+    to derive its Z_eff/A_eff/density/X0/λ_int; a genuinely unknown
+    material name correctly still raises via giant.materials, same as the
+    documented G4_LYSO precedent — that's a separate, intentional guard."""
+    pdg_map = {11: 0, 22: 1}
+    mat_map = {"G4_AIR": 0}
+    data = {
+        "pre_pos": np.zeros((1, 3), dtype=np.float32),
+        "pre_E": np.array([10.0], dtype=np.float32),
+        "pre_dir": np.array([[0.0, 0.0, 1.0]], dtype=np.float32),
+        "layer_id": np.array([0], dtype=np.int32),
+        "pdg": np.array([13], dtype=np.int64),  # not in pdg_map
+        "material": np.array(["G4_Pb"], dtype=object),  # not in mat_map
+        "mass": np.array([105.7], dtype=np.float32),
+        "charge": np.array([-1.0], dtype=np.float32),
+    }
+
+    cond_cont, cond_cat = build_cond_features(
+        data, pdg_map, mat_map, conditioning="physical"
+    )
+    assert cond_cont.shape[-1] == COND_DIM
+    np.testing.assert_array_equal(cond_cat, [[0, 0]])  # dummy indices, no raise
+
+    with pytest.raises(KeyError):
+        build_cond_features(data, pdg_map, mat_map, conditioning="embedding")
+
+
 # ── _WelfordAccumulator ──────────────────────────────────────────────────────
 
 

uv.lock

@@ -534,7 +534,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.1.0"
+version = "0.2.0"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },