Bump version to 0.2.0

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

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

@@ -211,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):

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",
@@ -99,6 +111,19 @@ DEFAULT_CONFIG: dict = {
             # 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
@@ -326,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
@@ -347,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")),

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
@@ -270,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)
@@ -970,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)
@@ -1186,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")
     }
@@ -1204,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]:
@@ -1250,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
 
@@ -256,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(
@@ -417,6 +432,8 @@ def run_train_job(
         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

@@ -114,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,
@@ -347,6 +426,8 @@ def train(
     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,
@@ -398,27 +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,
-                "lambda_entropy": lambda_entropy,
-                "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)
@@ -481,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)
 
@@ -496,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
@@ -518,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
@@ -596,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
@@ -737,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:
@@ -744,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)
@@ -980,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[

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_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
@@ -286,6 +287,110 @@ def test_router_entropy_loss_is_nonnegative_bounded_scalar():
     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 ────────────────────────────────────────────────────────────────
 
 
@@ -380,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

@@ -214,6 +214,41 @@ 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 ───────────────────
 
 

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" },