Merge branch 'master' into feat/router-learnable-width

EnergyRouter's gate sharpness was a single fixed temperature shared by
every expert, with no way for an expert to independently learn how much
of the energy axis it covers. Adds two mutually exclusive, default-off
modes: learn_width (per-expert learnable width) and learn_temperature
(single learnable shared scalar), both bounded via a sigmoid
interpolation warm-started to reproduce today's fixed-temperature gate
exactly at init, to compare against each other without risking the
unbounded-width collapse failure mode. Also promotes gate_stats's
entropy into a generic, optional Router.entropy_loss (lambda_entropy) as
a secondary guard against all experts' widths co-inflating together.

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
 

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")
 
 
@@ -504,17 +528,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 +566,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 +587,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

@@ -77,7 +77,28 @@ 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,
             "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
@@ -372,6 +393,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/setup_cache.py

@@ -31,7 +31,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 +44,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 +107,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 +118,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 +130,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),
         )
 
 

giant/model/network.py

@@ -566,6 +566,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 +639,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 +668,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 +699,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 +738,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")

giant/pipeline.py

@@ -149,7 +149,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 +158,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 +195,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']}"
         )
@@ -415,6 +416,7 @@ 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),
         normalizer_dict={
             "cond": cond_norm.to_dict(),
             "target": tgt_norm.to_dict(),

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",
@@ -123,6 +125,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 +134,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 +189,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 +346,7 @@ def train(
     lambda_s2: float = 1.0,
     lambda_balance: float = 0.0,
     lambda_proc: float = 0.0,
+    lambda_entropy: float = 0.0,
     normalizer_dict: dict | None = None,
     pdg_map: dict | None = None,
     mat_map: dict | None = None,
@@ -395,6 +409,7 @@ def train(
                 "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 {},
@@ -559,6 +574,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
@@ -626,7 +642,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 +654,7 @@ def train(
                             lambda_s2,
                             lambda_balance,
                             lambda_proc,
+                            lambda_entropy,
                         )
                     )
                     optimizer.zero_grad()
@@ -661,6 +678,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
@@ -772,7 +790,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 +803,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 +821,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 +850,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 +923,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 +938,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 +950,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,

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_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_router.py

@@ -141,6 +141,151 @@ 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
+
+
 # ── PdgRouter ────────────────────────────────────────────────────────────────
 
 

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,39 @@ def test_save_merges_non_colliding_normalizer_keys(tmp_path):
     assert loaded.normalizers["k2"].n_train_steps == 2
 
 
+# ── 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 ──────────────────────────────────────────────