Merge pull request 'V0.3.0 stage2 autoregressive' (#27) from v0.3.0-stage2-autoregressive into master

Reviewed-on: #27

.gitea/workflows/ci.yml

@@ -82,7 +82,11 @@ jobs:
           echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
           echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
       - run: uv sync --extra cpu --extra dev
-      - run: uv run pytest
+      - run: uv run pytest --cov --cov-report=term-missing --cov-report=xml
+      - uses: actions/upload-artifact@v3
+        with:
+          name: coverage-report
+          path: coverage.xml
 
   sync-version-on-tag:
     name: Sync project version with tag

.gitignore

@@ -20,3 +20,8 @@ checkpoints/
 
 # giant analyze run directories (shared.json, reduced/, plots/, condor logs)
 /analysis_runs/
+
+# Coverage artifacts
+.coverage
+coverage.xml
+htmlcov/

CLAUDE.md

@@ -22,7 +22,7 @@ giant analyze render <run_dir> --gallery                   # render PDFs + HTML
 dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
                                                   # bump-schema, status, update-manifest, create-manifest,
                                                   # make-root, build-geometry-oracle, warm-cache, hparam-scan
-                                                  # (see scripts/dwarf.py)
+                                                  # (see giant/tools/dwarf.py)
 ```
 
 `cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.
@@ -91,6 +91,6 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
 
 A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
 
-**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time.
+**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`) — v0.2 checkpoint-loading support has **no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
 
 **Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.

README.md

@@ -93,7 +93,7 @@ giant/
 │   │   ├── 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`)
+├── giant/tools/                        # 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,
 │   │                                    #   build-geometry-oracle, warm-cache, hparam-scan

giant/_migration.py

@@ -0,0 +1,70 @@
+"""Shared v0.2 -> v0.3 migration knowledge.
+
+v0.3.0 broke the config format (single `[train]` + `[model]` -> `[conditioning]`/
+`[stage1_model]`/`[stage2_model]`/`[train]`), and that break has to be absorbed by two
+independent migration surfaces: `giant.config.migrate_config` (a v0.2 `config.toml`) and
+`giant.model.network._migrate_legacy_model_config` (a v0.2 checkpoint's flat
+`model_config` dict). Both translate the same v0.2 facts into the same v0.3 shape, so
+the facts live here once rather than as two hand-maintained copies — see issues.md
+Issue 6.
+
+A dependency-free leaf module so neither `config.py` nor `network.py` has to import the
+other to share this.
+"""
+
+# v0.2 model-shaped keys (config.toml's [model] table, or a checkpoint's flat
+# model_config dict — same key names in both) applied identically to both v0.3 stage
+# blocks, because v0.2 had only one trunk shape shared by both stages.
+V02_MODEL_KEY_TO_STAGES: tuple[tuple[str, str], ...] = (
+    ("hidden_dim", "hidden_dim"),
+    ("n_blocks", "n_res_blocks"),
+    ("dropout", "dropout"),
+)
+
+# v0.2 architectural facts that had no corresponding config key at all — always true of
+# a v0.2 model, so both migration surfaces inject them unconditionally. Keyed by dotted
+# path relative to the migrated dict's root. NOTE: conditioning.*.n_layers (2) differs
+# from the v0.3 *default* (1) — not a typo, v0.2's conditioning MLP was always 2 layers
+# deep.
+V02_FIXED_FACTS: dict[str, object] = {
+    "conditioning.out_dim": 128,
+    "conditioning.particle.n_layers": 2,
+    "conditioning.material.n_layers": 2,
+    "stage1_model.active": True,
+    "stage1_model.flow.time_dim": 64,
+    "stage1_model.ddpm.time_dim": 64,
+    "stage2_model.active": True,
+    "stage2_model.flow.time_dim": 64,
+    "stage2_model.ddpm.time_dim": 64,
+    "stage2_model.context_dim": 64,
+    "stage2_model.decoder": "one_shot",
+    "stage2_model.particle_type.target": "physical",
+}
+
+
+def reject_legacy_router_expert_sizing(router_cfg: dict, *, source: str) -> None:
+    """Pop and validate v0.2's per-expert width/depth override, in place.
+
+    v0.3.0 removed per-expert sizing — experts always inherit the stage's
+    hidden_dim/n_res_blocks — so a v0.2 router config/checkpoint that set a non-default
+    `expert_hidden_dim`/`expert_n_blocks` describes experts with a different width/depth
+    than the monolith, and can only be reproduced by v0.2 code. Silently dropping these
+    keys (a router builder's kwarg filtering would do this for free) would resize the
+    experts instead of refusing, so this raises loudly.
+
+    Always pops both keys, whether or not they were non-default, so callers can go on
+    to use the (now-cleaned) `router_cfg` unconditionally. `source` names what's being
+    migrated (e.g. "v0.2 config's model.router" or "this checkpoint's
+    model_config.router") for the error message.
+    """
+    expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0)
+    expert_n_blocks = router_cfg.pop("expert_n_blocks", 0)
+    if not (expert_hidden_dim or expert_n_blocks):
+        return
+    raise ValueError(
+        f"{source} sets expert_hidden_dim/expert_n_blocks to a non-default value "
+        f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed per-expert "
+        "sizing (experts always inherit the stage's hidden_dim/n_res_blocks), so "
+        "this router's experts have a different width/depth than the monolith. "
+        "This checkpoint/config can only be loaded by v0.2 code."
+    )

giant/analysis/catalog.py

@@ -166,9 +166,7 @@ def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]:
     return list(merged.get(str(key), [0] * nbins))
 
 
-def _np_hist_pair(
-    r: np.ndarray, t: np.ndarray, nbins: int
-) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+def _np_hist_pair(r: np.ndarray, t: np.ndarray, nbins: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """Shared-edge histogram of two small per-event arrays (robust range)."""
     both = np.concatenate([r, t]) if (len(r) or len(t)) else np.array([0.0, 1.0])
     lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
@@ -240,9 +238,7 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
 
 def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
     ids, bins = event_energy_bins(lf, edges)
-    return pl.col("event_id").replace_strict(
-        ids, bins, default=-1, return_dtype=pl.Int64
-    )
+    return pl.col("event_id").replace_strict(ids, bins, default=-1, return_dtype=pl.Int64)
 
 
 def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
@@ -265,9 +261,7 @@ def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
     }
 
 
-def _marginal_grouped_finalize(
-    parts: list[dict], ctx: Context, var: str, axis: str
-) -> Reduced:
+def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis: str) -> Reduced:
     label, _ = _var(var)
     edges = _marginal_edges(ctx, var)
     nb = len(edges) - 1
@@ -317,9 +311,7 @@ def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
     return {"r": r.tolist(), "t": t.tolist()}
 
 
-def _event_scalar_finalize(
-    parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str
-) -> Reduced:
+def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str) -> Reduced:
     r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
     t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
     edges, rc, tc = _np_hist_pair(r, t, ctx.n_marginal_bins)
@@ -488,9 +480,7 @@ def _leakage_partial(b: Bundle) -> dict:
 
 def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
     frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts])
-    edges = np.linspace(
-        0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1
-    )
+    edges = np.linspace(0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1)
     counts = np.histogram(frac, edges)[0]
     return Reduced(
         id="leakage_fraction",
@@ -521,18 +511,8 @@ def _sec_frames(b: Bundle):
 
 def _sec_count_per_event_partial(b: Bundle) -> dict:
     r_sec, t_sec = _sec_frames(b)
-    r = (
-        r_sec.group_by("event_id")
-        .agg(pl.len().alias("n"))
-        .collect(engine="streaming")["n"]
-        .to_numpy()
-    )
-    t = (
-        t_sec.group_by("event_id")
-        .agg(pl.len().alias("n"))
-        .collect(engine="streaming")["n"]
-        .to_numpy()
-    )
+    r = r_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
+    t = t_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
     return {"r": r.tolist(), "t": t.tolist()}
 
 
@@ -568,9 +548,7 @@ def _sec_count_per_species_partial(b: Bundle) -> dict:
 def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
     r = sum_merge([p["r"] for p in parts])
     t = sum_merge([p["t"] for p in parts])
-    keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[
-        : len(ctx.top_pdgs)
-    ]
+    keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[: len(ctx.top_pdgs)]
     return Reduced(
         id="sec_count_per_species",
         family="secondaries",
@@ -617,11 +595,9 @@ def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
 
 def _sec_cos_angle_partial(b: Bundle) -> dict:
     edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1)
-    cos = (
-        pl.col("sdx") * pl.col("axis_x")
-        + pl.col("sdy") * pl.col("axis_y")
-        + pl.col("sdz") * pl.col("axis_z")
-    ).clip(-1.0, 1.0)
+    cos = (pl.col("sdx") * pl.col("axis_x") + pl.col("sdy") * pl.col("axis_y") + pl.col("sdz") * pl.col("axis_z")).clip(
+        -1.0, 1.0
+    )
 
     def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]:
         ea = entry_axis(steps_lf)
@@ -659,15 +635,13 @@ _router_gating_partial, _router_gating_finalize = _unchunkable(
     lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys)
 )
 _router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
-    lambda b: compute_router_share_by_pdg(
-        b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs
-    )
+    lambda b: compute_router_share_by_pdg(b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs)
 )
 _router_share_process_partial, _router_share_process_finalize = _unchunkable(
     lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
 )
-_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = (
-    _unchunkable(lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist))
+_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = _unchunkable(
+    lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist)
 )
 
 
@@ -689,9 +663,7 @@ def build_catalog() -> list[PlotSpec]:
                 f"marginal_{var}",
                 "marginals",
                 compute_partial=lambda b, v=var: _marginal_overall_partial(b, v),
-                finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(
-                    parts, ctx, v
-                ),
+                finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(parts, ctx, v),
             )
         )
         for axis in GROUPING_AXES:
@@ -699,12 +671,8 @@ def build_catalog() -> list[PlotSpec]:
                 PlotSpec(
                     f"marginal_{var}_by_{axis}",
                     "marginals",
-                    compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(
-                        b, v, a
-                    ),
-                    finalize=lambda parts, ctx, v=var, a=axis: (
-                        _marginal_grouped_finalize(parts, ctx, v, a)
-                    ),
+                    compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(b, v, a),
+                    finalize=lambda parts, ctx, v=var, a=axis: _marginal_grouped_finalize(parts, ctx, v, a),
                 )
             )
 
@@ -712,9 +680,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "event_total_edep",
             "event",
-            compute_partial=lambda b: _event_scalar_partial(
-                b, "total_edep", use_all=True
-            ),
+            compute_partial=lambda b: _event_scalar_partial(b, "total_edep", use_all=True),
             finalize=lambda parts, ctx: _event_scalar_finalize(
                 parts,
                 ctx,
@@ -732,9 +698,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "event_mean_length",
             "event",
-            compute_partial=lambda b: _event_scalar_partial(
-                b, "mean_length", use_all=False
-            ),
+            compute_partial=lambda b: _event_scalar_partial(b, "mean_length", use_all=False),
             finalize=lambda parts, ctx: _event_scalar_finalize(
                 parts,
                 ctx,
@@ -746,9 +710,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "event_n_steps",
             "event",
-            compute_partial=lambda b: _event_scalar_partial(
-                b, "n_steps", use_all=False
-            ),
+            compute_partial=lambda b: _event_scalar_partial(b, "n_steps", use_all=False),
             finalize=lambda parts, ctx: _event_scalar_finalize(
                 parts,
                 ctx,
@@ -773,9 +735,7 @@ def build_catalog() -> list[PlotSpec]:
         PlotSpec(
             "shower_transverse",
             "shower",
-            compute_partial=lambda b: _profile_partial(
-                b, transverse_expr, "transverse_edges"
-            ),
+            compute_partial=lambda b: _profile_partial(b, transverse_expr, "transverse_edges"),
             finalize=lambda parts, ctx: _profile_finalize(
                 parts,
                 ctx,

giant/analysis/condor.py

@@ -161,9 +161,7 @@ class RunMeta:
         return cls(**json.loads(Path(path).read_text()))
 
 
-def _rows_per_chunk(
-    rollout: str | Path, reference: str | Path, n_chunks: int
-) -> list[int]:
+def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -> list[int]:
     """Rollout+reference row count of each ``event_id % n_chunks`` chunk.
 
     One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
@@ -268,8 +266,7 @@ def compute_reduced(
     effective_n = n_chunks if spec.chunkable else 1
     if not (0 <= chunk_index < effective_n):
         raise ValueError(
-            f"{spec_id}: chunk_index={chunk_index} out of range for "
-            f"n_chunks={effective_n} (chunkable={spec.chunkable})"
+            f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})"
         )
     bundle = Bundle.open(
         rollout,
@@ -327,10 +324,7 @@ def merge_one(spec_id: str, run_dir: str | Path) -> Path:
     effective_n = meta.n_chunks if spec.chunkable else 1
 
     partial_dir = run_path / "reduced_partial"
-    found = {
-        p.chunk: p
-        for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))
-    }
+    found = {p.chunk: p for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))}
     missing = sorted(set(range(effective_n)) - set(found))
     if missing:
         raise FileNotFoundError(
@@ -375,11 +369,7 @@ exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
 
 
 def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
-    reqs_attrs = (
-        "+RemoteJob = True\n"
-        if cfg.remote
-        else "requirements = TARGET.ProvidesETPResources\n"
-    )
+    reqs_attrs = "+RemoteJob = True\n" if cfg.remote else "requirements = TARGET.ProvidesETPResources\n"
     return (
         "universe = docker\n"
         f"docker_image = {cfg.docker_image}\n"
@@ -399,9 +389,7 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> st
     )
 
 
-def _job_walltimes(
-    run_dir: Path, ids: list[str], n_chunks: int
-) -> list[tuple[str, int, int]]:
+def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]:
     """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.
 
     Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
@@ -477,9 +465,7 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
     (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, giant_exe=giant_exe, 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/analysis/context.py

@@ -74,9 +74,7 @@ def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFram
     return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
 
 
-def _combined_quantiles(
-    r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
-) -> tuple[float, float]:
+def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float) -> tuple[float, float]:
     """Robust (lo_q, hi_q) range over the union of two value samples."""
     both = np.concatenate([r_vals, t_vals])
     lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
@@ -104,31 +102,15 @@ def build_context(
 
     # Ranged marginal variables: robust ranges over a shared row subsample.
     exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
-    r_s = (
-        _row_subsample(r_lf, sample_rows, seed)
-        .select(exprs)
-        .collect(engine="streaming")
-    )
-    t_s = (
-        _row_subsample(t_lf, sample_rows, seed)
-        .select(exprs)
-        .collect(engine="streaming")
-    )
+    r_s = _row_subsample(r_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
+    t_s = _row_subsample(t_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
     var_ranges = {
-        name: _combined_quantiles(
-            r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
-        )
-        for name in RANGED_VARS
+        name: _combined_quantiles(r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q) for name in RANGED_VARS
     }
 
     # Energy-bin edges from exact per-event incident energies (cheap group_by).
     def _incident(lf: pl.LazyFrame) -> np.ndarray:
-        return (
-            lf.group_by("event_id")
-            .agg(pl.col("pre_E").max())
-            .collect(engine="streaming")["pre_E"]
-            .to_numpy()
-        )
+        return lf.group_by("event_id").agg(pl.col("pre_E").max()).collect(engine="streaming")["pre_E"].to_numpy()
 
     r_inc, t_inc = _incident(r_lf), _incident(t_lf)
     energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
@@ -145,8 +127,7 @@ def build_context(
     )
     top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
     materials = sorted(
-        set(_counts(r_lf, "material")["material"].to_list())
-        | set(_counts(t_lf, "material")["material"].to_list())
+        set(_counts(r_lf, "material")["material"].to_list()) | set(_counts(t_lf, "material")["material"].to_list())
     )
 
     # Shower depth / transverse ranges from a subsampled proxy.

giant/analysis/grouping.py

@@ -68,9 +68,7 @@ def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
 
 def energy_bin_labels(edges: np.ndarray) -> list[str]:
     """``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
-    return [
-        f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
-    ]
+    return [f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)]
 
 
 def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
@@ -88,9 +86,7 @@ def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
     return idx.clip(0, n_bins - 1)
 
 
-def event_energy_bins(
-    lf: pl.LazyFrame, edges: np.ndarray
-) -> tuple[np.ndarray, np.ndarray]:
+def event_energy_bins(lf: pl.LazyFrame, edges: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
     """Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.
 
     Incident energy is ``max(pre_E)`` per event (the primary). One bounded

giant/analysis/reduce.py

@@ -142,9 +142,7 @@ def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
     """
     ids = entry["event_id"].to_numpy()
     return lf.with_columns(
-        pl.col("event_id")
-        .replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64)
-        .alias(col)
+        pl.col("event_id").replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64).alias(col)
         for col in _ENTRY_AXIS_COLS
     )
 
@@ -254,10 +252,7 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
         lf.group_by("event_id")
         .agg(
             pl.col("edep").sum().alias("deposited"),
-            pl.col("pre_E")
-            .filter(pl.col("termination_reason") == TERM_ESCAPED)
-            .sum()
-            .alias("escaped"),
+            pl.col("pre_E").filter(pl.col("termination_reason") == TERM_ESCAPED).sum().alias("escaped"),
         )
         .collect(engine="streaming")
     )

giant/analysis/render.py

@@ -139,9 +139,7 @@ def _render_overlay(r: Reduced, params: dict):
 def _render_single(r: Reduced, params: dict):
     edges = np.asarray(r.payload["edges"])
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.stairs(
-        _density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
-    )
+    ax.stairs(_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"])
     if r.payload.get("log_y"):
         ax.set_yscale("log")
     if r.payload.get("log_x"):
@@ -187,9 +185,7 @@ def _render_profile(r: Reduced, params: dict):
         mean = np.asarray(r.payload[f"{key}_mean"])
         std = np.asarray(r.payload[f"{key}_std"])
         (line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
-        ax.fill_between(
-            centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
-        )
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=line.get_color())
     ax.set_xlabel(r.xlabel)
     ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
     ps.style_legend(ax, title="source")
@@ -201,9 +197,7 @@ def _render_bar(r: Reduced, params: dict):
     x = np.arange(len(labels))
     width = 0.4
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.bar(
-        x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
-    )
+    ax.bar(x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"])
     ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
     ax.set_xticks(x)
     ax.set_xticklabels(labels, rotation=45, ha="right")
@@ -215,9 +209,7 @@ def _render_bar(r: Reduced, params: dict):
 def _render_router_gating(r: Reduced, params: dict):
     n_experts = r.payload["n_experts"]
     log_x = r.payload.get("log_x", False)
-    fig, axes = ps.new_figure(
-        "slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
-    )
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False)
     flat = axes.ravel()
     for ax, key in zip(flat, ("rollout", "reference")):
         side = r.payload.get(key, {})
@@ -226,9 +218,7 @@ def _render_router_gating(r: Reduced, params: dict):
         if len(centers) and means.size:
             cum = np.zeros(len(centers))
             for i in range(n_experts):
-                ax.fill_between(
-                    centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
-                )
+                ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}")
                 cum = cum + means[:, i]
         if log_x:
             ax.set_xscale("log")
@@ -347,9 +337,7 @@ def render_all(
         families.add(r.family)
         fig = render(r, run_meta)
         ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
-        (family_dir / f"{r.id}.yaml").write_text(
-            yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
-        )
+        (family_dir / f"{r.id}.yaml").write_text(yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False))
         pdfs.append(family_dir / f"{r.id}.pdf")
         import matplotlib.pyplot as plt
 
@@ -370,9 +358,7 @@ def render_all(
     )
     for fam in families:
         (out_dir / fam / "metadata.yaml").write_text(
-            yaml.safe_dump(
-                {"title": fam, "description": f"{fam} plots."}, sort_keys=False
-            )
+            yaml.safe_dump({"title": fam, "description": f"{fam} plots."}, sort_keys=False)
         )
 
     if run_gallery:
@@ -400,6 +386,4 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
         "reference": meta.reference,
         **meta.plot_meta,
     }
-    return render_all(
-        run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
-    )
+    return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)

giant/analysis/router_gating.py

@@ -66,27 +66,11 @@ class _RouterHandle:
     router_type: str
 
 
-def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
-    """(particle_conditioning, material_conditioning) for
-    `giant.data.transforms.build_cond_features` — from either a v0.2
-    checkpoint's flat `model_config["conditioning"]` (one shared string, same
-    for both axes) or a new-format one (independent
-    `model_config["conditioning"]["particle"/"material"]["type"]` — the two
-    axes may differ). Mirrors
-    `giant.cli._conditioning_axes`."""
-    raw = model_cfg.get("conditioning", default)
-    if isinstance(raw, dict):
-        return (
-            raw.get("particle", {}).get("type", default),
-            raw.get("material", {}).get("type", default),
-        )
-    return raw, raw
-
-
 def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     """Load a checkpoint's Stage-1 router, or None if it isn't a MoE checkpoint."""
     import torch
 
+    from giant.checkpoint_io import conditioning_axes
     from giant.data.transforms import Normalizer
     from giant.model.network import build_models
 
@@ -95,9 +79,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     # New nested shape (has a "stage1_model" key) vs. a v0.2 checkpoint's
     # flat model_config.
     router_cfg = (
-        (model_cfg.get("stage1_model") or {}).get("router")
-        if "stage1_model" in model_cfg
-        else model_cfg.get("router")
+        (model_cfg.get("stage1_model") or {}).get("router") if "stage1_model" in model_cfg else model_cfg.get("router")
     )
     if not router_cfg or not router_cfg.get("enabled"):
         return None
@@ -112,7 +94,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     if router is None:
         return None
 
-    particle_conditioning, material_conditioning = _conditioning_axes(model_cfg)
+    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
     return _RouterHandle(
         router=router,
         pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()},
@@ -124,9 +106,7 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None:
     )
 
 
-def _subsample(
-    lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()
-) -> pl.DataFrame:
+def _subsample(lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()) -> pl.DataFrame:
     total = lf.select(pl.len()).collect(engine="streaming").item()
     if total > n:
         threshold = int(n / total * 2**32)
@@ -134,9 +114,7 @@ def _subsample(
     return lf.select(*_COLS, *extra_cols).collect(engine="streaming")
 
 
-def _gate_for_df(
-    handle: _RouterHandle, df: pl.DataFrame
-) -> tuple[pl.DataFrame, np.ndarray]:
+def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame, np.ndarray]:
     """(filtered df, gate_weights) for rows in ``df`` with a known pdg/material.
 
     Rows whose species or material never appeared in the checkpoint's
@@ -160,13 +138,9 @@ def _gate_for_df(
     df = df.filter(pl.Series(known, dtype=pl.Boolean))
 
     data = {
-        "pre_pos": np.column_stack(
-            [df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]
-        ),
+        "pre_pos": np.column_stack([df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]),
         "pre_E": df["pre_E"].to_numpy(),
-        "pre_dir": np.column_stack(
-            [df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]
-        ),
+        "pre_dir": np.column_stack([df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]),
         "layer_id": df["layer_id"].to_numpy(),
         "pdg": df["pdg"].to_numpy(),
         "material": df["material"].to_numpy(),
@@ -180,9 +154,7 @@ def _gate_for_df(
         material_conditioning=handle.material_conditioning,
     )
     with torch.no_grad():
-        gate = handle.router.gate(
-            torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()
-        ).numpy()
+        gate = handle.router.gate(torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()).numpy()
     return df, gate
 
 
@@ -205,9 +177,7 @@ def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict:
     return {"centers": centers[valid].tolist(), "means": means[valid].tolist()}
 
 
-def _top1_shares(
-    categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int
-) -> dict[str, list[float]]:
+def _top1_shares(categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int) -> dict[str, list[float]]:
     """Fraction of each category's rows hard-dispatched to each expert.
 
     Uses `Router.top1` (argmax), not the soft `gate` mean — grouped top-1
@@ -227,10 +197,7 @@ def _top1_shares(
     return shares
 
 
-_NOTE_NOT_MOE = (
-    "checkpoint has no enabled MoE router (model.router.enabled is "
-    "false/absent) — nothing to show"
-)
+_NOTE_NOT_MOE = "checkpoint has no enabled MoE router (model.router.enabled is false/absent) — nothing to show"
 
 _TITLES = {
     "router_gating": "Router gating (mixture-of-experts decision boundaries)",
@@ -266,9 +233,7 @@ def compute_router_gating(
         df = _subsample(lf, _SAMPLE_ROWS, seed)
         df, gate = _gate_for_df(handle, df)
         x = df["pre_E"].to_numpy()
-        sides[name] = (
-            _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
-        )
+        sides[name] = _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
 
     return Reduced(
         id="router_gating",
@@ -304,9 +269,7 @@ def compute_router_share_by_pdg(
         df, gate = _gate_for_df(handle, df)
         if len(df):
             idx = gate.argmax(axis=1)
-            shares = _top1_shares(
-                df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts
-            )
+            shares = _top1_shares(df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts)
         else:
             shares = {str(p): [0.0] * handle.router.n_experts for p in top_pdgs}
         sides[name] = {labels[i]: shares[str(p)] for i, p in enumerate(top_pdgs)}
@@ -351,9 +314,7 @@ def compute_router_share_by_process(
         counts = df["process"].value_counts().sort("count", descending=True)
         order = counts["process"].to_list()[:top_k]
         idx = gate.argmax(axis=1)
-        shares = _top1_shares(
-            df["process"].to_numpy(), idx, order, handle.router.n_experts
-        )
+        shares = _top1_shares(df["process"].to_numpy(), idx, order, handle.router.n_experts)
     else:
         order, shares = [], {}
 

giant/analysis/runtime_estimate.py

@@ -27,7 +27,7 @@ would then wrongly scale up with a bigger dataset. `RUNTIME_SAFETY_MARGIN` is
 deliberately generous (4x total) specifically to absorb that kind of
 contention spike instead. Rerun this calibration (pull fresh
 `condor_history`/`run_meta.json`, refit) if the catalog changes or timings
-drift — a synthetic local rebaseline via `scripts/profile_analysis_costs.py`
+drift — a synthetic local rebaseline via `giant/tools/profile_analysis_costs.py`
 is a reasonable fallback when no real cluster data is available yet, but
 undershoots real wall time badly (it can't see docker pull / `/ceph` I/O
 latency), which is exactly why this file moved off it.
@@ -54,9 +54,7 @@ _FIXED_OVERHEAD_S = 60.0
 # scan. Calibrated from the 3 real router jobs' observed wall times (119, 66,
 # 124s) — max minus _FIXED_OVERHEAD_S, on top of it.
 _ROUTER_FIXED_S = 64.0
-_ROUTER_IDS = frozenset(
-    {"router_gating", "router_share_by_pdg", "router_share_by_process"}
-)
+_ROUTER_IDS = frozenset({"router_gating", "router_share_by_pdg", "router_share_by_process"})
 
 # Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot
 # added after the last calibration run) — the most expensive fitted per-row

giant/analysis/sources.py

@@ -93,10 +93,7 @@ def _check_rollout_metadata(path: Path) -> None:
     metadata = pq.read_schema(path).metadata or {}
     coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
     if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
-        raise ValueError(
-            f"{path} is not a rollout file (coord={coord.decode()!r}); "
-            "expected `giant rollout` output"
-        )
+        raise ValueError(f"{path} is not a rollout file (coord={coord.decode()!r}); expected `giant rollout` output")
 
 
 def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
@@ -121,11 +118,7 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
     else:
         # The reference (a rollout's seed `dataset`) may be a directory of
         # parquet shards rather than a single file — scan them all.
-        lf = (
-            pl.scan_parquet(str(path / "**/*.parquet"))
-            if path.is_dir()
-            else pl.scan_parquet(path)
-        )
+        lf = pl.scan_parquet(str(path / "**/*.parquet")) if path.is_dir() else pl.scan_parquet(path)
     return lf.with_columns(pl.col("pdg").cast(pl.Int64))
 
 
@@ -137,9 +130,7 @@ def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
     """
     if side is Side.reference:
         return lf
-    return lf.filter(
-        ~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS))
-    )
+    return lf.filter(~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS)))
 
 
 def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:

giant/checkpoint_io.py

@@ -0,0 +1,234 @@
+"""Load a trained checkpoint into ready-to-run models (giant.cli's `predict`/`rollout`).
+
+Both commands need the same ~15 steps to go from a checkpoint path to two
+`eval()`-mode models plus their normalizers/vocab maps: load the pickle,
+validate it carries what current code expects, resolve which conditioning
+mode each axis was trained with, restore the top-N vocab maps (if the
+checkpoint used one-hot conditioning), rebuild the normalizers, construct the
+model from `model_config`, and load the requested (raw or EMA) weights. This
+used to be duplicated near-verbatim in both commands (issues.md Issue 5) —
+`load_for_inference` is the single implementation.
+
+This module intentionally has no Typer dependency, so it can be unit-tested
+directly and imported from non-CLI code (`giant.analysis.router_gating`,
+lazily — see that module's docstring for why). Failures raise
+`CheckpointCompatibilityError` with the same wording the CLI has always
+shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import torch
+from torch import nn
+
+from giant import config as gconfig
+from giant.constants import K_MAX
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_from_json
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+
+class CheckpointCompatibilityError(Exception):
+    """Checkpoint is missing something `load_for_inference` needs."""
+
+
+def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
+    """(particle_conditioning, material_conditioning) for
+    `giant.data.transforms.build_cond_features`/`build_features` — from
+    either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
+    shared string, same for both axes) or a new-format one (independent
+    `model_config["conditioning"]["particle"/"material"]["type"]` — the two
+    axes are configured independently and may differ)."""
+    raw = model_cfg.get("conditioning", default)
+    if isinstance(raw, dict):
+        return (
+            raw.get("particle", {}).get("type", default),
+            raw.get("material", {}).get("type", default),
+        )
+    return raw, raw
+
+
+def stage_cfg(model_cfg: dict, stage: str) -> dict:
+    """`model_cfg[f"{stage}_model"]` for a new-format model_config, `{}` for
+    a v0.2 flat one (whose ddpm schedule always used `CosineSchedule`'s own
+    default `T=1000` — never a config key — and which never had
+    `particle_type` at all, so `{}` is the correct fallback for both
+    `ddpm_steps`/`particle_type_other_policy` below)."""
+    val = model_cfg.get(f"{stage}_model")
+    return val if isinstance(val, dict) else {}
+
+
+def ddpm_steps(model_cfg: dict, stage: str) -> int:
+    return stage_cfg(model_cfg, stage).get("ddpm", {}).get("n_steps", 1000)
+
+
+def particle_type_other_policy(model_cfg: dict) -> str:
+    return stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")
+
+
+def load_pdg_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["pdg_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's conditioning/particle_type never needed one (see
+    `giant.pipeline.run_setup_stage`, which only populates it when
+    `conditioning.particle.type` or `stage2_model.particle_type.target` is
+    `"onehot"`)."""
+    raw = ckpt.get("pdg_topn_map")
+    return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+
+
+def load_mat_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
+    this checkpoint's `conditioning.material.type` was never `"onehot"` (see
+    `giant.pipeline.run_setup_stage`)."""
+    raw = ckpt.get("mat_topn_map")
+    return topnmap_from_json(raw, axis="material") if raw is not None else None
+
+
+def load_sec_type_topn_map(ckpt: dict) -> TopNMap | None:
+    """`ckpt["sec_type_topn_map"]` as a `giant.data.loader.TopNMap`, or
+    `None` if this checkpoint's `stage2_model.particle_type.target` was never
+    `"onehot"` (see `giant.pipeline.run_setup_stage`).
+
+    Pre-gitea-#29 checkpoints have no `sec_type_topn_map` key at all — before
+    #29, the secondary-species decode map and the conditioning PDG onehot map
+    were always numerically the same map, saved once under `pdg_topn_map`.
+    For those, fall back to `load_pdg_topn_map` to reproduce that exact
+    behavior; a current checkpoint always has the key (possibly `null`, if
+    `particle_type.target != "onehot"`), so this fallback never fires for one."""
+    if "sec_type_topn_map" in ckpt:
+        raw = ckpt["sec_type_topn_map"]
+        return topnmap_from_json(raw, axis="pdg") if raw is not None else None
+    return load_pdg_topn_map(ckpt)
+
+
+@dataclass(frozen=True)
+class InferenceContext:
+    """Everything needed to run a trained checkpoint forward, resolved once."""
+
+    stage1: nn.Module | None
+    stage2: nn.Module | None
+    cond_norm: Normalizer
+    tgt_norm: Normalizer
+    sec_phys_norm: Normalizer
+    pdg_map: dict[int, int]
+    mat_map: dict[str, int]
+    pdg_topn_map: TopNMap | None
+    mat_topn_map: TopNMap | None
+    sec_type_topn_map: TopNMap | None
+    particle_conditioning: str
+    material_conditioning: str
+    k_max: int
+    stage1_ddpm_steps: int
+    stage2_ddpm_steps: int
+    other_policy: str
+    model_config: dict
+    epoch: int | None
+    best_val_loss: float | None
+
+
+def load_for_inference(
+    checkpoint: Path,
+    device: torch.device,
+    command_name: str,
+    weights: str = "raw",
+    require_stage2: bool = True,
+) -> InferenceContext:
+    """Load *checkpoint* and reconstruct everything `predict`/`rollout` need
+    to run it forward, on *device*, in `eval()` mode.
+
+    *command_name* (e.g. `"predict"`/`"rollout"`) only feeds the "needs both"
+    error message below. *weights* is `"raw"` (the live training weights) or
+    `"ema"` (the EMA shadow copy, see `--ema-decay`). *require_stage2*
+    controls whether a checkpoint with an inactive stage 2
+    (`stage2_model.active = false`) is an error (both current callers need
+    both stages) or an acceptable `stage2 = None` result — kept as a real
+    parameter since `stage{1,2}_model.active` is a real, if currently
+    stage1+stage2-only-in-practice, config option.
+    """
+    ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
+    for key in ("model_config", "sec_decoder"):
+        if key not in ckpt:
+            raise CheckpointCompatibilityError(f"checkpoint has no {key} — retrain with the current code")
+
+    if "sec_phys" not in ckpt.get("normalizer", {}):
+        raise CheckpointCompatibilityError("checkpoint has no normalizer.sec_phys — retrain with the current code")
+
+    gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
+
+    model_cfg = ckpt["model_config"]
+    particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
+    pdg_topn_map = load_pdg_topn_map(ckpt)
+    mat_topn_map = load_mat_topn_map(ckpt)
+    if particle_conditioning == "onehot" and pdg_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's conditioning.particle.type='onehot' but has no pdg_topn_map — retrain with the current code"
+        )
+    if material_conditioning == "onehot" and mat_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's conditioning.material.type='onehot' but has no mat_topn_map — retrain with the current code"
+        )
+    sec_type_topn_map = load_sec_type_topn_map(ckpt)
+    particle_type_target = stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and sec_type_topn_map is None:
+        raise CheckpointCompatibilityError(
+            "checkpoint's stage2_model.particle_type.target='onehot' but has no "
+            "sec_type_topn_map — retrain with the current code"
+        )
+    other_policy = particle_type_other_policy(model_cfg)
+    stage1_ddpm_steps = ddpm_steps(model_cfg, "stage1")
+    stage2_ddpm_steps = ddpm_steps(model_cfg, "stage2")
+    k_max = stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
+    pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
+    mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
+    cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
+    tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
+    sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
+
+    built = build_models(model_cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    if require_stage2 and (stage1 is None or stage2 is None):
+        raise CheckpointCompatibilityError(
+            f"checkpoint has an inactive stage1 or stage2 — {command_name} needs both (see stage{{1,2}}_model.active)"
+        )
+
+    if weights == "raw":
+        model_key, sec_key = "model", "sec_decoder"
+    else:
+        model_key, sec_key = "model_ema", "sec_decoder_ema"
+        if model_key not in ckpt or sec_key not in ckpt:
+            raise CheckpointCompatibilityError(
+                f"{checkpoint} has no EMA weights (trained before --ema-decay, "
+                "or with --ema-decay 0) — use --weights raw"
+            )
+    if stage1 is not None:
+        stage1.load_state_dict(ckpt[model_key])
+        stage1.to(device).eval()
+    if stage2 is not None:
+        stage2.load_state_dict(ckpt[sec_key])
+        stage2.to(device).eval()
+
+    return InferenceContext(
+        stage1=stage1,
+        stage2=stage2,
+        cond_norm=cond_norm,
+        tgt_norm=tgt_norm,
+        sec_phys_norm=sec_phys_norm,
+        pdg_map=pdg_map,
+        mat_map=mat_map,
+        pdg_topn_map=pdg_topn_map,
+        mat_topn_map=mat_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
+        particle_conditioning=particle_conditioning,
+        material_conditioning=material_conditioning,
+        k_max=k_max,
+        stage1_ddpm_steps=stage1_ddpm_steps,
+        stage2_ddpm_steps=stage2_ddpm_steps,
+        other_policy=other_policy,
+        model_config=model_cfg,
+        epoch=ckpt.get("epoch"),
+        best_val_loss=ckpt.get("best_val_loss"),
+    )

giant/data/dataset.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from pathlib import Path
+from typing import NamedTuple
 
 import numpy as np
 import torch
@@ -11,6 +12,37 @@ from giant.data.loader import event_id_offset, iter_file_chunks
 from giant.data.transforms import Normalizer, build_features, sorted_membership
 
 
+class StepBatch(NamedTuple):
+    """One training batch, as yielded by `StreamingStepsDataset`. Field order
+    is load-bearing for existing positional unpacking elsewhere (`trainers.py`,
+    `validate.py`, test fixtures) — append only, never insert or reorder.
+
+    cond_cont:  (B, COND_DIM) float32
+    cond_cat:   (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
+    target_s1:  (B, 9) float32 — normalised Stage-1 primary target
+    n_sec:      (B,) int64 — true secondary count per step
+    sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
+                local_dir, log_mass, charge] per slot (mass/charge
+                normalised iff `sec_phys_normalizer` was given); always
+                computed the same way regardless of
+                stage2_model.particle_type.target, only actually used
+                downstream under target="physical"
+    proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
+                only; zeros when `proc_map` is None)
+    sec_type_idx: (B, k_max) int64 — per-slot class index into
+                `sec_type_class_map`, for particle_type.target in
+                ("onehot", "embedding"); zeros (unused) otherwise
+    """
+
+    cond_cont: torch.Tensor
+    cond_cat: torch.Tensor
+    target_s1: torch.Tensor
+    n_sec: torch.Tensor
+    sec_cont: torch.Tensor
+    proc_idx: torch.Tensor
+    sec_type_idx: torch.Tensor
+
+
 def make_event_split(
     all_event_ids: np.ndarray,
     val_fraction: float = 0.1,
@@ -39,24 +71,7 @@ class StreamingStepsDataset(IterableDataset):
     rather than single rows, so the batch is assembled with vectorized
     numpy slicing instead of a per-row Python loop in the default collate.
 
-    Each batch is a tuple:
-        (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)
-    where:
-        cond_cont:  (B, COND_DIM) float32
-        cond_cat:   (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
-        target_s1:  (B, 9) float32 — normalised Stage-1 primary target
-        n_sec:      (B,) int64 — true secondary count per step
-        sec_cont:   (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
-                    local_dir, log_mass, charge] per slot (mass/charge
-                    normalised iff `sec_phys_normalizer` was given); always
-                    computed the same way regardless of
-                    stage2_model.particle_type.target, only actually used
-                    downstream under target="physical"
-        proc_idx:   (B,) int64 — process-class label (ProcessRouter supervision
-                    only; zeros when `proc_map` is None)
-        sec_type_idx: (B, k_max) int64 — per-slot class index into
-                    `sec_type_class_map`, for particle_type.target in
-                    ("onehot", "embedding"); zeros (unused) otherwise
+    Each batch is a `StepBatch` — see its docstring for field meanings.
 
     `k_max` (constructor arg, default the module constant) should match
     `stage2_model.k_max` — it sets the padded
@@ -123,25 +138,13 @@ class StreamingStepsDataset(IterableDataset):
         buf_n = 0
 
         for path in files:
-            for chunk in iter_file_chunks(
-                path, offset=self._offsets[path], k_max=self.k_max
-            ):
+            for chunk in iter_file_chunks(path, offset=self._offsets[path], k_max=self.k_max):
                 mask = sorted_membership(chunk["event_id"], self._events_arr)
                 if not mask.any():
                     continue
                 chunk = {k: v[mask] for k, v in chunk.items()}
 
-                (
-                    cond_cont,
-                    cond_cat,
-                    target_s1,
-                    n_sec,
-                    sec_cont,
-                    proc_idx,
-                    sec_type_idx,
-                    _,
-                    _,
-                ) = build_features(
+                feats = build_features(
                     chunk,
                     self.pdg_map,
                     self.mat_map,
@@ -157,14 +160,14 @@ class StreamingStepsDataset(IterableDataset):
                     sec_type_class_map=self.sec_type_class_map,
                     k_max=self.k_max,
                 )
-                buf_cont.append(cond_cont)
-                buf_cat.append(cond_cat)
-                buf_tgt.append(target_s1)
-                buf_nsec.append(n_sec)
-                buf_sec.append(sec_cont)
-                buf_proc.append(proc_idx)
-                buf_type.append(sec_type_idx)
-                buf_n += len(cond_cont)
+                buf_cont.append(feats.cond_cont)
+                buf_cat.append(feats.cond_cat)
+                buf_tgt.append(feats.target_s1)
+                buf_nsec.append(feats.n_sec)
+                buf_sec.append(feats.sec_cont)
+                buf_proc.append(feats.proc_idx)
+                buf_type.append(feats.sec_type_idx)
+                buf_n += len(feats.cond_cont)
 
                 if buf_n >= self.shuffle_buffer:
                     (
@@ -228,14 +231,14 @@ class StreamingStepsDataset(IterableDataset):
         n_full = n // bs if not final else (n + bs - 1) // bs
         for start in range(0, n_full * bs, bs):
             end = min(start + bs, n)
-            yield (
-                torch.from_numpy(cont[start:end]).float(),
-                torch.from_numpy(cat[start:end]).long(),
-                torch.from_numpy(tgt[start:end]).float(),
-                torch.from_numpy(nsec[start:end]).long(),
-                torch.from_numpy(sec[start:end]).float(),
-                torch.from_numpy(proc[start:end]).long(),
-                torch.from_numpy(styp[start:end]).long(),
+            yield StepBatch(
+                cond_cont=torch.from_numpy(cont[start:end]).float(),
+                cond_cat=torch.from_numpy(cat[start:end]).long(),
+                target_s1=torch.from_numpy(tgt[start:end]).float(),
+                n_sec=torch.from_numpy(nsec[start:end]).long(),
+                sec_cont=torch.from_numpy(sec[start:end]).float(),
+                proc_idx=torch.from_numpy(proc[start:end]).long(),
+                sec_type_idx=torch.from_numpy(styp[start:end]).long(),
             )
 
         if final:

giant/data/loader.py

@@ -16,7 +16,7 @@ from giant.constants import K_MAX
 MANIFEST_SUFFIX = ".manifest"
 
 # Each input parquet file is a separate Geant4 job converted 1:1 from its own
-# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
+# ROOT file (giant/tools/steps_to_parquet.py), and a job's event_id numbering
 # always restarts from 0 — so when multiple files are loaded together (a
 # directory or .manifest), raw event_id values collide across files even
 # though they refer to unrelated events. Every per-file event_id column gets
@@ -115,9 +115,7 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
     return out
 
 
-def _df_to_dict(
-    df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX
-) -> dict[str, np.ndarray]:
+def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
@@ -136,9 +134,7 @@ def _df_to_dict(
         # / ProcessRouter). Guarded like has_sec_lists: older parquet
         # conversions predating this column still load fine.
         "process": (
-            df["process"].to_numpy(dtype=object)
-            if "process" in df.columns
-            else np.full(len(df), "", dtype=object)
+            df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
         ),
         "step_length": df["step_length"].to_numpy(dtype=np.float32),
         "post_E": df["post_E"].to_numpy(dtype=np.float32),
@@ -151,16 +147,12 @@ def _df_to_dict(
     if has_sec_lists:
         d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
         d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
-        d["sec_dir_list"] = _pad_dir_col(
-            df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max
-        )
+        d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)
 
     return d
 
 
-def load_steps(
-    path: str | Path, offset: int = 0, k_max: int = K_MAX
-) -> dict[str, np.ndarray]:
+def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
     return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
 
 
@@ -170,9 +162,7 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
     return _offset_event_id(ids, offset)
 
 
-def iter_file_chunks(
-    path: str | Path, offset: int = 0, k_max: int = K_MAX
-) -> Iterator[dict[str, np.ndarray]]:
+def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> Iterator[dict[str, np.ndarray]]:
     """Yield one parquet row-group at a time so a large file never fully loads.
 
     `k_max` sets the padded width of the sec_*_list columns (should match
@@ -214,15 +204,11 @@ def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]
     }
 
 
-def iter_cond_chunks(
-    path: str | Path, offset: int = 0
-) -> Iterator[dict[str, np.ndarray]]:
+def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np.ndarray]]:
     """Yield conditioning-only row-groups (no post-step columns read from disk)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _cond_df_to_dict(
-            pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
-        )
+        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)
 
 
 def build_index_maps(
@@ -315,9 +301,7 @@ class TopNMap:
     other_members: dict
 
 
-def build_topn_map_from_files(
-    files: list[Path], column: str, n_classes: int, cast=str
-) -> TopNMap:
+def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
     """Scan `column` and build a frequency-capped value->index map, structurally
     identical to `build_process_map_from_files` (shares its ranking core via
     `_topn_plus_other_map`), generalized over the source column and key type.

giant/data/setup_cache.py

@@ -108,10 +108,7 @@ def normalizer_key(
     # affect which cond_cont columns are computed for real vs. zero-filled
     # (giant.data.transforms._physical_cond_columns), so both must be part of
     # the key or two mixed-axis runs could collide on the same cache entry.
-    return (
-        f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}"
-        f"_mcond={material_conditioning}"
-    )
+    return f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}_mcond={material_conditioning}"
 
 
 # Top-N-map axes: "pdg" keys match pdg_map's int
@@ -126,10 +123,7 @@ def topn_key(axis: str, n_classes: int) -> str:
     sidecar stays reusable across runs with different emb_dim (see the
     dict[int, dict] precedent `proc_maps` sets, keyed by n_experts)."""
     if axis not in _TOPN_AXIS_CASTS:
-        raise ValueError(
-            f"unknown top-N map axis {axis!r}, expected one of "
-            f"{sorted(_TOPN_AXIS_CASTS)}"
-        )
+        raise ValueError(f"unknown top-N map axis {axis!r}, expected one of {sorted(_TOPN_AXIS_CASTS)}")
     return f"{axis}:{n_classes}"
 
 
@@ -165,9 +159,7 @@ 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_quantiles": np.asarray(
-                self.energy_quantiles, dtype=np.float32
-            ).tolist(),
+            "energy_quantiles": np.asarray(self.energy_quantiles, dtype=np.float32).tolist(),
         }
 
     @classmethod
@@ -234,13 +226,8 @@ class SetupCache:
                 np.array(d["event_index"]["counts"], dtype=np.int64),
             )
         proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
-        normalizers = {
-            k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
-        }
-        topn_maps = {
-            k: topnmap_from_json(v, axis=k.split(":", 1)[0])
-            for k, v in d.get("topn_maps", {}).items()
-        }
+        normalizers = {k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()}
+        topn_maps = {k: topnmap_from_json(v, axis=k.split(":", 1)[0]) for k, v in d.get("topn_maps", {}).items()}
         return cls(
             fingerprint=d["fingerprint"],
             git_hash=d.get("git_hash", "unknown"),
@@ -263,18 +250,14 @@ class SetupCache:
             fingerprint=other.fingerprint,
             git_hash=other.git_hash,
             vocab=other.vocab if other.vocab is not None else self.vocab,
-            event_index=(
-                other.event_index if other.event_index is not None else self.event_index
-            ),
+            event_index=(other.event_index if other.event_index is not None else self.event_index),
             proc_maps={**self.proc_maps, **other.proc_maps},
             normalizers={**self.normalizers, **other.normalizers},
             topn_maps={**self.topn_maps, **other.topn_maps},
         )
 
 
-def load(
-    data: str | Path, files: list[Path], echo=lambda *a, **k: None
-) -> SetupCache | None:
+def load(data: str | Path, files: list[Path], echo=lambda *a, **k: None) -> SetupCache | None:
     """Load and validate the sidecar for `data`; `None` on any miss (never raises).
 
     A missing file, corrupt JSON, format-version mismatch, dimension-constant
@@ -298,9 +281,7 @@ def load(
             echo("setup cache: format version changed — ignoring stale cache")
             return None
         if raw.get("dims") != _DIMS:
-            echo(
-                "setup cache: model dimension constants changed — ignoring stale cache"
-            )
+            echo("setup cache: model dimension constants changed — ignoring stale cache")
             return None
         fp = fingerprint_files(files)
         if raw.get("fingerprint") != fp:
@@ -343,9 +324,7 @@ def save(
         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
-                )
+                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)
@@ -353,9 +332,7 @@ def save(
             finally:
                 fcntl.flock(lock_file, fcntl.LOCK_UN)
     except OSError as exc:
-        echo(
-            f"setup cache: could not write {path} ({exc}) — continuing without caching"
-        )
+        echo(f"setup cache: could not write {path} ({exc}) — continuing without caching")
         try:
             tmp.unlink(missing_ok=True)
         except OSError:
@@ -366,16 +343,12 @@ def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.nd
     """Unique event ids + per-event row (step) counts, across all `files`."""
     if not files:
         return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
-    all_ids = np.concatenate(
-        [load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
-    )
+    all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
     unique_ids, counts = np.unique(all_ids, return_counts=True)
     return unique_ids, counts
 
 
-def n_train_steps_for_split(
-    unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
-) -> int:
+def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:
     """Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.
 
     `train_events_arr` must be ascending and duplicate-free (as produced by

giant/data/transforms.py

@@ -1,4 +1,5 @@
 import warnings
+from typing import NamedTuple
 
 import numpy as np
 
@@ -84,9 +85,7 @@ def energy_simplex_encode(
     return z.astype(np.float32)
 
 
-def energy_simplex_decode(
-    z: np.ndarray, pre_E: np.ndarray
-) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+def energy_simplex_decode(z: np.ndarray, pre_E: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
     """Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.
 
     A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
@@ -118,9 +117,7 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
     arbitrary second operand) on every row; profiling on a 114M-row file
     showed `np.cross` as the single hottest call inside this rotation.
     """
-    axis = np.stack(
-        [pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
-    )
+    axis = np.stack([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)
     # axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
     # pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
@@ -199,9 +196,7 @@ def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarra
     kxv = _cross_with_z_axis(axis, post_dir)  # (N,3)
     kdv = (axis * post_dir).sum(axis=1, keepdims=True)  # (N,1)
 
-    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
-        np.float32
-    )
+    return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
 
 
 class Normalizer:
@@ -341,9 +336,7 @@ 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, strict: bool = True, default: int = 0
-) -> np.ndarray:
+def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0) -> 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.
@@ -387,9 +380,7 @@ def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
     disp = post_pos - pre_pos
     norm = np.linalg.norm(disp, axis=1, keepdims=True)
     safe_norm = np.where(norm < 1e-7, 1.0, norm)
-    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
-        np.float32
-    )
+    return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(np.float32)
 
 
 def reconstruct_post_pos(
@@ -408,9 +399,7 @@ def reconstruct_post_pos(
     return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)
 
 
-def inv_local_frame_rotation(
-    pre_dir: np.ndarray, post_dir_local: np.ndarray
-) -> np.ndarray:
+def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> np.ndarray:
     """Inverse of local_frame_rotation: rotate from local frame back to world frame.
 
     Applies R^T (same axis, negative angle) to post_dir_local.
@@ -424,9 +413,7 @@ def inv_local_frame_rotation(
     kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)
 
     # Negative angle: sin_t → -sin_t
-    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
-        np.float32
-    )
+    return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
 
 
 _STICK_LOGIT_CLIP = 10.0  # logit value used for the last valid secondary slot
@@ -491,14 +478,10 @@ def encode_secondaries(
                 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
-            )
+            f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
             logit = np.log(f / (1.0 - f)).astype(np.float32)
             # Last valid slot: give it the full remaining budget
-            is_last = sec_valid[:, i] & ~(
-                sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
-            )
+            is_last = sec_valid[:, i] & ~(sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool))
             logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
             logit = np.where(
                 sec_valid[:, i],
@@ -525,9 +508,7 @@ def encode_secondaries(
             # Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
             valid_mask = sec_valid[:, i]
             if valid_mask.any():
-                dir_local[valid_mask, i] = local_frame_rotation(
-                    pre_dir[valid_mask], sec_dir_list[valid_mask, i]
-                )
+                dir_local[valid_mask, i] = local_frame_rotation(pre_dir[valid_mask], sec_dir_list[valid_mask, i])
 
     if sec_pdg_list is not None:
         from giant.particles import particle_phys_array
@@ -553,9 +534,7 @@ def encode_secondaries(
     return sec_cont.astype(np.float32)
 
 
-def encode_secondary_type_idx(
-    sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict
-) -> np.ndarray:
+def encode_secondary_type_idx(sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict) -> np.ndarray:
     """Per-secondary-slot class index into `class_map` — (N, K_MAX) int64.
 
     `class_map` is either a top-N-plus-other map's `class_map`
@@ -583,9 +562,7 @@ def encode_secondary_type_idx(
     # isn't guaranteed to be a key — an arbitrary present one always is).
     dummy = next(iter(class_map))
     safe_pdg = np.where(sec_valid, sec_pdg_list, dummy)
-    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(
-        N, K
-    )
+    idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(N, K)
     return np.where(sec_valid, idx, 0).astype(np.int64)
 
 
@@ -659,9 +636,7 @@ def decode_secondary_cont(
     for i in range(K):
         valid = sec_valid[:, i]
         if valid.any():
-            sec_dir_world[valid, i] = inv_local_frame_rotation(
-                pre_dir[valid], dir_local[valid, i]
-            )
+            sec_dir_world[valid, i] = inv_local_frame_rotation(pre_dir[valid], dir_local[valid, i])
 
     return sec_E, sec_dir_world, sec_valid
 
@@ -704,9 +679,7 @@ def decode_secondaries(
         sec_cont = sec_cont.copy()
         sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
 
-    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont, n_sec, e_sec, pre_dir)
 
     log_mass = sec_cont[:, :, 4]  # (N, K)
     charge = sec_cont[:, :, 5]  # (N, K)
@@ -760,16 +733,12 @@ def _physical_cond_columns(
     elif particle_conditioning in ("embedding", "onehot"):
         particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
     else:
-        raise ValueError(
-            f"unknown conditioning.particle.type {particle_conditioning!r}"
-        )
+        raise ValueError(f"unknown conditioning.particle.type {particle_conditioning!r}")
 
     if material_conditioning == "physical":
         from giant.materials import material_properties_array
 
-        z_eff, a_eff, density, x0, lambda_int = material_properties_array(
-            data["material"]
-        ).T
+        z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T
         material_cols = np.column_stack(
             [
                 z_eff,
@@ -782,9 +751,7 @@ def _physical_cond_columns(
     elif material_conditioning in ("embedding", "onehot"):
         material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
     else:
-        raise ValueError(
-            f"unknown conditioning.material.type {material_conditioning!r}"
-        )
+        raise ValueError(f"unknown conditioning.material.type {material_conditioning!r}")
 
     return np.column_stack([particle_cols, material_cols]).astype(np.float32)
 
@@ -848,9 +815,7 @@ def build_cond_features(
     cond_cat = np.column_stack(cat_cols)
 
     if cond_normalizer is not None:
-        cond_cont = _cond_normalizer_transform(
-            cond_cont, cond_normalizer, particle_conditioning, material_conditioning
-        )
+        cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, particle_conditioning, material_conditioning)
 
     return cond_cont, cond_cat
 
@@ -896,36 +861,10 @@ def _cond_normalizer_transform(
     return ((cond_cont - mean) / std).astype(np.float32)
 
 
-def build_features(
-    data: dict[str, np.ndarray],
-    pdg_map: dict[int, int],
-    mat_map: dict[str, int],
-    cond_normalizer: Normalizer | None = None,
-    target_normalizer: Normalizer | None = None,
-    sec_phys_normalizer: Normalizer | None = None,
-    fit: bool = False,
-    proc_map: dict[str, int] | None = None,
-    require_secondaries: bool = False,
-    particle_conditioning: str = "embedding",
-    material_conditioning: str = "embedding",
-    sec_phys_only: bool = False,
-    pdg_topn_map: dict[int, int] | None = None,
-    mat_topn_map: dict[str, int] | None = None,
-    sec_type_class_map: dict | None = None,
-    k_max: int = K_MAX,
-) -> tuple[
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    np.ndarray,
-    Normalizer | None,
-    Normalizer | None,
-]:
-    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
-    sec_type_idx) arrays.
+class StepFeatures(NamedTuple):
+    """Output of `build_features`. Field order is load-bearing for existing
+    positional unpacking (tests, `StreamingStepsDataset`) — append only,
+    never insert or reorder.
 
     target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
     n_sec: (N,) integer secondary counts (target for n_sec head)
@@ -944,6 +883,40 @@ def build_features(
               in `("onehot", "embedding")` — see `encode_secondary_type_idx`.
               Zero-filled (and unused) when `sec_type_class_map` is None
               (i.e. `target = "physical"`).
+    """
+
+    cond_cont: np.ndarray
+    cond_cat: np.ndarray
+    target_s1: np.ndarray
+    n_sec: np.ndarray
+    sec_cont: np.ndarray
+    proc_idx: np.ndarray
+    sec_type_idx: np.ndarray
+    cond_normalizer: Normalizer | None
+    target_normalizer: Normalizer | None
+
+
+def build_features(
+    data: dict[str, np.ndarray],
+    pdg_map: dict[int, int],
+    mat_map: dict[str, int],
+    cond_normalizer: Normalizer | None = None,
+    target_normalizer: Normalizer | None = None,
+    sec_phys_normalizer: Normalizer | None = None,
+    fit: bool = False,
+    proc_map: dict[str, int] | None = None,
+    require_secondaries: bool = False,
+    particle_conditioning: str = "embedding",
+    material_conditioning: str = "embedding",
+    sec_phys_only: bool = False,
+    pdg_topn_map: dict[int, int] | None = None,
+    mat_topn_map: dict[str, int] | None = None,
+    sec_type_class_map: dict | None = None,
+    k_max: int = K_MAX,
+) -> StepFeatures:
+    """Assemble a `StepFeatures` of (cond_cont, cond_cat, target_s1, n_sec,
+    sec_cont, proc_idx, sec_type_idx, cond_normalizer, target_normalizer) —
+    see `StepFeatures` for field meanings.
 
     require_secondaries: when True, raise if any step has n_sec > 0 but the
         per-secondary list columns are absent (a mis-converted file that would
@@ -972,13 +945,9 @@ def build_features(
     """
 
     post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
-    travel_dir_local = local_frame_rotation(
-        data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
-    )
+    travel_dir_local = local_frame_rotation(data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"]))
 
-    energy_z = energy_simplex_encode(
-        data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
-    )  # (N, 2)
+    energy_z = energy_simplex_encode(data["edep"], data["e_sec"], data["post_E"], data["pre_E"])  # (N, 2)
 
     target_s1 = np.column_stack(
         [
@@ -1014,9 +983,7 @@ def build_features(
         cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
     cond_cat = np.column_stack(cat_cols)  # (N, 2/3/4)
 
-    n_sec_raw = data["n_sec"].astype(
-        np.int64
-    )  # (N,) unclamped, for the valid-slot mask
+    n_sec_raw = data["n_sec"].astype(np.int64)  # (N,) unclamped, for the valid-slot mask
 
     # Secondary continuous targets
     sec_E_list = data.get("sec_E_list")
@@ -1096,14 +1063,14 @@ def build_features(
     else:
         proc_idx = np.zeros(len(cond_cat), dtype=np.int64)
 
-    return (
-        cond_cont,
-        cond_cat,
-        target_s1,
-        n_sec,
-        sec_cont,
-        proc_idx,
-        sec_type_idx,
-        cond_normalizer,
-        target_normalizer,
+    return StepFeatures(
+        cond_cont=cond_cont,
+        cond_cat=cond_cat,
+        target_s1=target_s1,
+        n_sec=n_sec,
+        sec_cont=sec_cont,
+        proc_idx=proc_idx,
+        sec_type_idx=sec_type_idx,
+        cond_normalizer=cond_normalizer,
+        target_normalizer=target_normalizer,
     )

giant/geometry.py

@@ -31,10 +31,7 @@ import numpy as np
 import pandas as pd
 import pyarrow.parquet as pq
 
-_INSTALL_HINT = (
-    "the geometry oracle needs scikit-learn — install it with "
-    "`uv sync --extra cpu --extra geometry`"
-)
+_INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"
 
 
 def _require_sklearn():
@@ -63,9 +60,7 @@ class _SlabLookup:
     layer_ids: np.ndarray  # (n_segments,) int64, layer_id of each segment
     radius_max: float  # largest transverse radius seen in training data
 
-    def query(
-        self, pos: np.ndarray, margin: float
-    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    def query(self, pos: np.ndarray, margin: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
         other = [i for i in range(3) if i != self.axis]
         z = pos[:, self.axis]
         radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)
@@ -75,11 +70,7 @@ class _SlabLookup:
         material = self.materials[idx]
         layer_id = self.layer_ids[idx]
 
-        escaped = (
-            (z < self.z_edges[0] - margin)
-            | (z > self.z_edges[-1] + margin)
-            | (radius > self.radius_max + margin)
-        )
+        escaped = (z < self.z_edges[0] - margin) | (z > self.z_edges[-1] + margin) | (radius > self.radius_max + margin)
         return material, layer_id, escaped
 
 
@@ -300,16 +291,12 @@ def _fit_slab_lookup(
     """
     other = [i for i in range(3) if i != axis]
     z = pos[:, axis].astype(np.float64)
-    radius = np.sqrt(
-        pos[:, other[0]].astype(np.float64) ** 2
-        + pos[:, other[1]].astype(np.float64) ** 2
-    )
+    radius = np.sqrt(pos[:, other[0]].astype(np.float64) ** 2 + pos[:, other[1]].astype(np.float64) ** 2)
 
     z_min, z_max = float(z.min()), float(z.max())
     if z_min == z_max:
         raise ValueError(
-            "all points share the same depth-axis coordinate — pick a "
-            "different `depth_axis` or use method='knn'/'svm'"
+            "all points share the same depth-axis coordinate — pick a different `depth_axis` or use method='knn'/'svm'"
         )
     edges = np.linspace(z_min, z_max, n_bins + 1)
     bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
@@ -344,12 +331,7 @@ def _fit_slab_lookup(
     bin_layer = bin_layer[fill_from]
 
     # Run-length-encode consecutive bins sharing a label into segments.
-    changed = (
-        np.flatnonzero(
-            (bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
-        )
-        + 1
-    )
+    changed = np.flatnonzero((bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])) + 1
     seg_starts = np.concatenate([[0], changed])
     z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
     materials = bin_material[seg_starts]
@@ -460,11 +442,7 @@ def build_geometry_oracle(
     # Escape threshold from the reference point spacing. Sample a subset for the
     # median 2-NN distance (the 1st neighbour of a training point is itself).
     nn = NearestNeighbors(n_neighbors=2).fit(X)
-    probe = (
-        X
-        if len(X) <= 20_000
-        else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
-    )
+    probe = X if len(X) <= 20_000 else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
     d2, _ = nn.kneighbors(probe, n_neighbors=2)
     median_nn = float(np.median(d2[:, 1]))
     escape_threshold = escape_factor * median_nn

giant/materials.py

@@ -64,18 +64,10 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
     "G4_CESIUM_IODIDE": MaterialProperties(
         z_eff=54.0, a_eff=129.904539, density=4.51, x0=1.860288, lambda_int=39.305990
     ),
-    "G4_Pb": MaterialProperties(
-        z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950
-    ),
-    "G4_W": MaterialProperties(
-        z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580
-    ),
-    "G4_Cu": MaterialProperties(
-        z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940
-    ),
-    "G4_Fe": MaterialProperties(
-        z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300
-    ),
+    "G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950),
+    "G4_W": MaterialProperties(z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580),
+    "G4_Cu": MaterialProperties(z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940),
+    "G4_Fe": MaterialProperties(z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300),
     "G4_BRASS": MaterialProperties(
         z_eff=30.939130,
         a_eff=68.500857,
@@ -83,9 +75,7 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
         x0=1.367465,
         lambda_int=16.947420,
     ),
-    "G4_POLYSTYRENE": MaterialProperties(
-        z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880
-    ),
+    "G4_POLYSTYRENE": MaterialProperties(z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880),
     "G4_PLASTIC_SC_VINYLTOLUENE": MaterialProperties(
         z_eff=3.368421,
         a_eff=6.219791,
@@ -108,20 +98,15 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
         x0=30392.070000,
         lambda_int=71009.500000,
     ),
-    "G4_lAr": MaterialProperties(
-        z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400
-    ),
+    "G4_lAr": MaterialProperties(z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400),
 }
 
 
-def get_material_properties(
-    name: str, table: dict[str, MaterialProperties] | None = None
-) -> MaterialProperties:
+def get_material_properties(name: str, table: dict[str, MaterialProperties] | None = None) -> MaterialProperties:
     t = MATERIAL_PROPERTIES if table is None else table
     if name not in t:
         raise UnknownMaterialError(
-            f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES "
-            f"-- add it (known: {sorted(t)})"
+            f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES -- add it (known: {sorted(t)})"
         )
     props = t[name]
     if any(v is None for v in props):
@@ -134,9 +119,7 @@ def get_material_properties(
     return props
 
 
-def material_properties_array(
-    names: np.ndarray, table: dict[str, MaterialProperties] | None = None
-) -> np.ndarray:
+def material_properties_array(names: np.ndarray, table: dict[str, MaterialProperties] | None = None) -> np.ndarray:
     """(N,) str material names -> (N, 5) float32 [z_eff, a_eff, density, x0, lambda_int]."""
     out = np.array(
         [get_material_properties(str(m), table) for m in np.asarray(names)],

giant/model/_legacy.py

@@ -0,0 +1,114 @@
+"""v0.2 -> v0.3 checkpoint migration: translates a v0.2 checkpoint's flat
+`model_config`/state dicts into the current nested shape (issues.md Issue 8;
+see also `giant._migration` and `giant.config.migrate_config`, the sibling
+config.toml migration surface — issues.md Issue 6)."""
+
+from giant._migration import V02_FIXED_FACTS, reject_legacy_router_expert_sizing
+from giant.constants import EMB_DIM, K_MAX
+
+
+def _migrate_legacy_model_config(model_config: dict) -> dict:
+    """Translate a v0.2 checkpoint's flat `model_config` (giant/pipeline.py's
+    old shape: `hidden_dim`/`n_blocks`/`emb_dim`/`dropout`/`conditioning`/
+    `router`/`mode`/... all at one level) into the nested
+    `{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model",
+    "stage2_model"}` shape `build_models` expects.
+
+    Sets `stage2_model.n_sec.owner = "stage1"` so the n_sec_head weights a v0.2
+    checkpoint carries on its Stage-1 module keep loading there instead of the new
+    default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's
+    own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not
+    just be labeled as such.
+
+    Only the monolithic (non-routed) trunk shape is exercised by the step-2
+    migration test; a routed v0.2 checkpoint still builds correctly here
+    (the router config passes through), but its
+    state dict isn't covered by `migrate_legacy_state_dict` below.
+    """
+    m = model_config
+    conditioning_mode = m.get("conditioning", "embedding")
+    generator = m.get("mode", "flow")
+    hidden_dim = m.get("hidden_dim", 256)
+    n_blocks = m.get("n_blocks", 6)
+    emb_dim = m.get("emb_dim", EMB_DIM)
+    dropout = m.get("dropout", 0.1)
+    k_max = m.get("k_max", K_MAX)
+    noise_dim = m.get("noise_dim", 64)
+    router_cfg = dict(m.get("router") or {})
+    reject_legacy_router_expert_sizing(router_cfg, source="this checkpoint's model_config.router")
+    router_cfg.setdefault("enabled", False)
+
+    F = V02_FIXED_FACTS
+    cond_n_layers = F["conditioning.particle.n_layers"]  # same fact for both axes
+    return {
+        "pdg_vocab": m["pdg_vocab"],
+        "mat_vocab": m["mat_vocab"],
+        "conditioning": {
+            "out_dim": F["conditioning.out_dim"],
+            "share_stages": False,
+            "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
+            "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
+        },
+        "stage1_model": {
+            "active": F["stage1_model.active"],
+            "generator": generator,
+            "hidden_dim": hidden_dim,
+            "n_res_blocks": n_blocks,
+            "dropout": dropout,
+            "flow": {"time_dim": F["stage1_model.flow.time_dim"]},
+            "ddpm": {"time_dim": F["stage1_model.ddpm.time_dim"]},
+            "wgan": {"noise_dim": noise_dim},
+            "router": dict(router_cfg),
+        },
+        "stage2_model": {
+            "active": F["stage2_model.active"],
+            "decoder": F["stage2_model.decoder"],
+            "generator": generator,
+            "hidden_dim": hidden_dim,
+            "n_res_blocks": n_blocks,
+            "dropout": dropout,
+            "k_max": k_max,
+            "context_dim": F["stage2_model.context_dim"],
+            "n_sec": {"mode": "head", "owner": "stage1"},
+            "particle_type": {"target": F["stage2_model.particle_type.target"]},
+            "flow": {"time_dim": F["stage2_model.flow.time_dim"]},
+            "ddpm": {"time_dim": F["stage2_model.ddpm.time_dim"]},
+            "wgan": {"noise_dim": noise_dim},
+            "router": {**router_cfg, "tie_to_stage1": False},
+        },
+    }
+
+
+def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple[dict, dict]:
+    """Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`,
+    `SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new
+    `(Stage1Model, Stage2OneShot)` module structure produced by
+    `build_models(_migrate_legacy_model_config(model_config))`.
+
+    Only the monolithic (non-routed) trunk shape is handled.
+    """
+
+    def _trunk_prefix(k: str) -> str:
+        if k.startswith(("input_proj.", "blocks.", "out_proj.")):
+            return f"trunk.{k}"
+        return k
+
+    new_stage1 = {}
+    for k, v in old_stage1_sd.items():
+        if k.startswith("n_sec_head."):
+            new_stage1[k] = v  # stays top-level (n_sec.owner="stage1")
+        else:
+            new_stage1[_trunk_prefix(k)] = v
+
+    new_stage2 = {}
+    for k, v in old_stage2_sd.items():
+        if k.startswith("cond_enc.base."):
+            new_stage2["cond_enc." + k[len("cond_enc.base.") :]] = v
+        elif k.startswith("cond_enc.stage1_proj."):
+            new_stage2["context_adapter.proj." + k[len("cond_enc.stage1_proj.") :]] = v
+        elif k.startswith("cond_enc.fuse."):
+            new_stage2["fuse." + k[len("cond_enc.fuse.") :]] = v
+        else:
+            new_stage2[_trunk_prefix(k)] = v
+
+    return new_stage1, new_stage2

giant/model/builders.py

@@ -0,0 +1,211 @@
+"""Factories: `build_models`/`build_critics` assemble the top-level stage
+models from a config dict (issues.md Issue 8)."""
+
+import torch.nn as nn
+
+from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig
+from giant.constants import X_DIM
+from giant.model._legacy import _migrate_legacy_model_config
+from giant.model.encoders import ConditionEncoder
+from giant.model.models import (
+    CriticModel,
+    Stage1Model,
+    Stage2Autoregressive,
+    Stage2OneShot,
+    resolve_type_n_classes,
+    stage2_trunk_sec_dim,
+)
+from giant.model.routers import Router, _build_router_from_cfg
+
+# ---------------------------------------------------------------------------
+# Factories
+# ---------------------------------------------------------------------------
+
+
+def build_models(model_config: dict) -> dict[str, nn.Module | None]:
+    """Construct `{"stage1": ..., "stage2": ...}` from a config dict — either
+    the new nested shape (has a `"stage1_model"` key, plus `"pdg_vocab"`/
+    `"mat_vocab"`/`"conditioning"` at the top level) or a v0.2 checkpoint's
+    flat `model_config`, auto-migrated via `_migrate_legacy_model_config`.
+
+    A stage is `None` in the result when that stage's `active = False`.
+    `stage2_model.router.tie_to_stage1` shares stage 1's literal `Router`
+    instance rather than building a second, independently-parameterized one
+    (v0.2's actual — probably accidental — behaviour: two routers built from
+    one config with no semantic relationship between them).
+
+    `conditioning.share_stages = true` builds one `ConditionEncoder`
+    instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/
+    `Stage2Autoregressive`'s `cond_enc` param), instead of each stage
+    building its own — halving the conditioning parameter count and forcing a
+    common representation. `false` (default) keeps v0.2 behaviour:
+    independent instances with identical config but independent weights.
+    """
+    cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
+    pdg_vocab = cfg["pdg_vocab"]
+    mat_vocab = cfg["mat_vocab"]
+    conditioning = cfg["conditioning"]
+    particle_cfg = conditioning["particle"]
+    material_cfg = conditioning["material"]
+    particle_conditioning = particle_cfg["type"]
+    conditioning_cfg = ConditioningConfig.from_dict(conditioning)
+    s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
+    s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
+    cond_out_dim = conditioning_cfg.out_dim
+    shared_cond_enc: ConditionEncoder | None = None
+    if conditioning_cfg.share_stages:
+        shared_cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+
+    result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
+
+    stage1_router: Router | None = None
+    if s1_spec.active:
+        router_cfg = cfg["stage1_model"].get("router") or {}
+        if s1_spec.router.enabled:
+            stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
+        generator = s1_spec.generator
+        # wgan has no time_dim concept (no diffusion/flow time variable) —
+        # matches the pre-dataclass .get("time_dim", 64) fallback, which
+        # always hit its default for a wgan sub-block too.
+        time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64
+        n_sec_owner = s2_spec.n_sec.owner
+        n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None
+        result["stage1"] = Stage1Model(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            hidden_dim=s1_spec.hidden_dim,
+            n_res_blocks=s1_spec.n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            dropout=s1_spec.dropout,
+            generator=generator,
+            time_dim=time_dim,
+            noise_dim=s1_spec.wgan.noise_dim,
+            router=stage1_router,
+            n_sec_head_k_max=n_sec_head_k_max,
+            cond_enc=shared_cond_enc,
+        )
+
+    if s2_spec.active:
+        decoder = s2_spec.decoder
+        router_cfg = cfg["stage2_model"].get("router") or {}
+        stage2_router: Router | None = None
+        if s2_spec.router.enabled:
+            if s2_spec.router.tie_to_stage1 and stage1_router is not None:
+                stage2_router = stage1_router
+            else:
+                stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
+        generator = s2_spec.generator
+        # wgan has no time_dim concept — see the matching comment in stage 1
+        # above.
+        time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64
+        n_sec_owner = s2_spec.n_sec.owner
+        k_max = s2_spec.k_max
+        particle_type_cfg = s2_spec.particle_type.to_dict()
+
+        if decoder == "autoregressive":
+            ar_cfg = s2_spec.autoregressive
+            result["stage2"] = Stage2Autoregressive(
+                pdg_vocab=pdg_vocab,
+                mat_vocab=mat_vocab,
+                particle_cfg=particle_cfg,
+                material_cfg=material_cfg,
+                hidden_dim=s2_spec.hidden_dim,
+                n_res_blocks=s2_spec.n_res_blocks,
+                cond_out_dim=cond_out_dim,
+                context_dim=s2_spec.context_dim,
+                dropout=s2_spec.dropout,
+                generator=generator,
+                time_dim=time_dim,
+                noise_dim=s2_spec.wgan.noise_dim,
+                k_max=k_max,
+                router=stage2_router,
+                build_n_sec_head=n_sec_owner != "stage1",
+                particle_type_cfg=particle_type_cfg,
+                history=ar_cfg.history,
+                attn_n_heads=ar_cfg.attn_n_heads,
+                attn_n_layers=ar_cfg.attn_n_layers,
+                cond_enc=shared_cond_enc,
+            )
+        else:
+            sec_dim = stage2_trunk_sec_dim(
+                particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+            )
+            result["stage2"] = Stage2OneShot(
+                pdg_vocab=pdg_vocab,
+                mat_vocab=mat_vocab,
+                particle_cfg=particle_cfg,
+                material_cfg=material_cfg,
+                hidden_dim=s2_spec.hidden_dim,
+                n_res_blocks=s2_spec.n_res_blocks,
+                cond_out_dim=cond_out_dim,
+                context_dim=s2_spec.context_dim,
+                sec_dim=sec_dim,
+                dropout=s2_spec.dropout,
+                generator=generator,
+                time_dim=time_dim,
+                noise_dim=s2_spec.wgan.noise_dim,
+                k_max=k_max,
+                router=stage2_router,
+                build_n_sec_head=n_sec_owner != "stage1",
+                particle_type_cfg=particle_type_cfg,
+                cond_enc=shared_cond_enc,
+            )
+
+    return result
+
+
+def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
+    """Construct `{"stage1": ..., "stage2": ...}` critics for `generator =
+    "wgan"` training. Training-only — never persisted for inference the way
+    `build_models`'s pair is. `None` for a stage that's inactive or not
+    WGAN."""
+    cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
+    pdg_vocab = cfg["pdg_vocab"]
+    mat_vocab = cfg["mat_vocab"]
+    conditioning = cfg["conditioning"]
+    particle_cfg = conditioning["particle"]
+    material_cfg = conditioning["material"]
+    conditioning_cfg = ConditioningConfig.from_dict(conditioning)
+    cond_out_dim = conditioning_cfg.out_dim
+    s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
+    s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
+
+    result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
+
+    if s1_spec.active and s1_spec.generator == "wgan":
+        result["stage1"] = CriticModel(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            in_dim=X_DIM,
+            hidden_dim=s1_spec.wgan.critic_hidden_dim or s1_spec.hidden_dim,
+            n_res_blocks=s1_spec.wgan.critic_n_res_blocks or s1_spec.n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            dropout=s1_spec.dropout,
+            stage="stage1",
+        )
+
+    if s2_spec.active and s2_spec.generator == "wgan":
+        k_max = s2_spec.k_max
+        particle_type_cfg = s2_spec.particle_type.to_dict()
+        in_dim = stage2_trunk_sec_dim(
+            particle_type_cfg, "wgan", k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+        )
+        result["stage2"] = CriticModel(
+            pdg_vocab=pdg_vocab,
+            mat_vocab=mat_vocab,
+            particle_cfg=particle_cfg,
+            material_cfg=material_cfg,
+            in_dim=in_dim,
+            hidden_dim=s2_spec.wgan.critic_hidden_dim or s2_spec.hidden_dim,
+            n_res_blocks=s2_spec.wgan.critic_n_res_blocks or s2_spec.n_res_blocks,
+            cond_out_dim=cond_out_dim,
+            dropout=s2_spec.dropout,
+            stage="stage2",
+            context_dim=s2_spec.context_dim,
+        )
+
+    return result

giant/model/encoders.py

@@ -0,0 +1,122 @@
+"""Conditioning encoder — fuses continuous conditioning with particle/material
+identity (issues.md Issue 8)."""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
+from giant.model.layers import _make_axis_mlp
+
+
+def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]:
+    """`cond_cat` column indices for each axis's top-N-onehot index, or
+    `None` if that axis isn't `"onehot"`.
+
+    Columns 0/1 are always the dense pdg/material vocab index. The particle
+    top-N column (if any) comes next, then the material top-N column (if
+    any) — `giant.data.transforms.build_cond_features`/`build_features`
+    append columns in this same order, so the two sides must never drift
+    apart.
+    """
+    col = 2
+    particle_col = None
+    if particle_type == "onehot":
+        particle_col = col
+        col += 1
+    material_col = None
+    if material_type == "onehot":
+        material_col = col
+        col += 1
+    return particle_col, material_col
+
+
+class ConditionEncoder(nn.Module):
+    """Fuses continuous conditioning with particle/material identity.
+
+    The particle and material axes are configured independently
+    (`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`)
+    and may mix freely, e.g. material "physical" with particle "embedding".
+    Three modes per axis:
+      - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s
+        dense training-vocab index. Memorizes the training menu.
+      - "physical": an `n_layers`-deep MLP over the axis's raw physical
+        properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
+        giant.data.transforms.build_features), computable for any PDG code /
+        material name rather than only ones seen in training.
+      - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other
+        class map (`giant.data.loader.build_topn_map_from_files`/
+        `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra
+        top-N-index column(s) — see `_cat_col_layout`.
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        cont_dim: int = COND_DIM,
+        out_dim: int = 128,
+    ) -> None:
+        super().__init__()
+        self.particle_cfg = dict(particle_cfg)
+        self.material_cfg = dict(material_cfg)
+        self._particle_topn_col, self._material_topn_col = cat_col_layout(particle_cfg["type"], material_cfg["type"])
+
+        p_type = particle_cfg["type"]
+        p_emb_dim = particle_cfg["emb_dim"]
+        if p_type == "embedding":
+            self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim)
+        elif p_type == "physical":
+            self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1))
+        elif p_type != "onehot":
+            raise ValueError(f"unknown conditioning.particle.type {p_type!r}")
+
+        m_type = material_cfg["type"]
+        m_emb_dim = material_cfg["emb_dim"]
+        if m_type == "embedding":
+            self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim)
+        elif m_type == "physical":
+            self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1))
+        elif m_type != "onehot":
+            raise ValueError(f"unknown conditioning.material.type {m_type!r}")
+
+        in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim
+        self.mlp = nn.Sequential(
+            nn.Linear(in_dim, out_dim),
+            nn.SiLU(),
+            nn.Linear(out_dim, out_dim),
+        )
+
+    def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
+        p_type = self.particle_cfg["type"]
+        if p_type == "embedding":
+            return self.pdg_emb(cond_cat[:, 0])
+        if p_type == "physical":
+            particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
+            return self.particle_mlp(particle_phys)
+        assert self._particle_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self._particle_topn_col],
+            num_classes=self.particle_cfg["emb_dim"],
+        ).float()
+
+    def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
+        m_type = self.material_cfg["type"]
+        if m_type == "embedding":
+            return self.mat_emb(cond_cat[:, 1])
+        if m_type == "physical":
+            material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
+            return self.material_mlp(material_phys)
+        assert self._material_topn_col is not None
+        return F.one_hot(
+            cond_cat[:, self._material_topn_col],
+            num_classes=self.material_cfg["emb_dim"],
+        ).float()
+
+    def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        pdg_e = self._particle_embed(cond_cont, cond_cat)
+        mat_e = self._material_embed(cond_cont, cond_cat)
+        x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
+        return self.mlp(x)

giant/model/history.py

@@ -0,0 +1,154 @@
+"""History encoders — stage-2 autoregressive only. Self-contained, no
+dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+
+import torch
+import torch.nn as nn
+
+
+class HistoryEncoder(nn.Module):
+    """Interface for stage-2 autoregressive per-token history summaries:
+    `forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over
+    a full (teacher-forced) token sequence — used by training. `MarkovHistory`
+    and `AttentionHistory` are the two implementations. Inference
+    (`giant/sample.py`) generates one token at a
+    time and cannot afford `forward`'s per-step cost to be O(K) (attention
+    would then be O(K^2) over a rollout's k_max loop); encoders that need
+    incremental state for that path additionally implement `init_cache`/
+    `step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since
+    its per-step cost is already O(1) (it only ever looks at the previous
+    token, not the full prefix)."""
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class MarkovHistory(HistoryEncoder):
+    """Summarizes the previous secondary's own `(energy_fraction, direction,
+    type_representation)` through one small MLP — the "markov" history:
+    token i+1 only ever sees token i plus the running scalars
+    (`remaining_frac`/`slot_idx`, fused in separately by
+    `Stage2Autoregressive._token_cond`), not the full prefix.
+
+    At slot 0 (`has_prev` False) substitutes a learned start vector rather
+    than zeros — a reasonable default.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.start = nn.Parameter(torch.zeros(in_dim))
+        self.mlp = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU())
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        start = self.start.view(1, 1, -1).expand_as(feat)
+        x = torch.where(has_prev.unsqueeze(-1), feat, start)
+        return self.mlp(x)
+
+
+class _CausalAttnBlock(nn.Module):
+    """One pre-norm causal self-attention block for `AttentionHistory`.
+
+    Exposes two forward paths that must agree (see
+    `test_attention_history_step_matches_forward` in `tests/test_network.py`):
+    `forward` — the full-sequence, causally-masked pass used for training;
+    `step` — an incremental pass for inference, given the *pre-attention*
+    normalized hidden states of every earlier position (`kv_cache`, i.e.
+    `norm1(x)` for positions `< t`, not `x` itself). Caching `norm1(x)` rather
+    than raw `x` is what makes `step` correct: this block's attention needs
+    exactly that quantity as keys/values, and `LayerNorm` has no cross-position
+    interaction, so recomputing it per position instead of caching it would
+    still be correct but pointlessly repeat work. The *next* block's cache is
+    built from a different sequence (this block's output), so each block owns
+    an independent cache entry.
+    """
+
+    def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))
+
+    def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
+        h = self.norm1(x)
+        attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
+        x = x + attn_out
+        x = x + self.mlp(self.norm2(x))
+        return x
+
+    def step(self, x_new: torch.Tensor, kv_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
+        """`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None`
+        (first position) or `(B, T, dim)` — `norm1(x)` of every earlier
+        position at this same block. Returns `(out, new_kv_cache)`, `out`
+        being this position's block output (`(B, 1, dim)`, to feed the next
+        block's `step`), `new_kv_cache` the same cache extended by this
+        position (to reuse at this block's *next* `step` call)."""
+        h_new = self.norm1(x_new)
+        kv = h_new if kv_cache is None else torch.cat([kv_cache, h_new], dim=1)
+        attn_out, _ = self.attn(h_new, kv, kv, need_weights=False)
+        x = x_new + attn_out
+        x = x + self.mlp(self.norm2(x))
+        return x, kv
+
+
+class AttentionHistory(HistoryEncoder):
+    """Causal self-attention over the emitted-token prefix — the more
+    expressive alternative to `MarkovHistory`'s fixed previous-token-only
+    summary. `feat`/`has_prev`
+    follow the same shifted-by-one convention `MarkovHistory` and
+    `Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s
+    own `(energy_fraction, direction, type_representation)`, with a learned
+    start vector substituted at `has_prev == False` positions (only slot 0 in
+    practice — see `giant.training.stage2_inputs._ar_has_prev`). Causal masking then makes
+    position `i`'s output a function of `feat[:, 1:i+1]` — i.e. tokens
+    `0..i-1` — exactly the prefix available when predicting token `i`.
+
+    `forward` is the parallel training path (one pass over the whole
+    teacher-forced sequence); `init_cache`/`step` are the incremental
+    inference path `giant/sample.py` uses, one new token per call, to avoid
+    re-encoding the whole prefix from scratch every slot — `step` must be
+    called exactly once per slot (its cache-extension is not idempotent),
+    so a slot's output must be reused for
+    every model call within that slot (`forward`'s ODE substeps, or a separate
+    `predict_type` call) rather than re-derived — see
+    `Stage2Autoregressive.history_step`.
+    """
+
+    def __init__(self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2) -> None:
+        super().__init__()
+        self.start = nn.Parameter(torch.zeros(in_dim))
+        self.in_proj = nn.Linear(in_dim, out_dim)
+        self.blocks = nn.ModuleList([_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)])
+
+    def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        start = self.start.view(1, 1, -1).expand_as(feat)
+        x = torch.where(has_prev.unsqueeze(-1), feat, start)
+        return self.in_proj(x)
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        B, K, _ = feat.shape
+        x = self._embed(feat, has_prev)
+        mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
+        for block in self.blocks:
+            x = block(x, mask)
+        return x
+
+    def init_cache(self) -> list[torch.Tensor | None]:
+        return [None for _ in self.blocks]
+
+    def step(
+        self,
+        token_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        cache: list[torch.Tensor | None],
+    ) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
+        """`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest
+        token's own features (what would be `feat[:, k]` in `forward`).
+        Advances every block's cache by this position and returns this
+        position's output (`(B, 1, out_dim)`, the correct history summary for
+        the NEXT slot) plus the updated cache."""
+        x = self._embed(token_feat, has_prev)
+        new_cache: list[torch.Tensor | None] = []
+        for block, kv in zip(self.blocks, cache):
+            x, kv_new = block.step(x, kv)
+            new_cache.append(kv_new)
+        return x, new_cache

giant/model/layers.py

@@ -0,0 +1,76 @@
+"""Small stateless-ish building blocks shared across encoders/trunks/models —
+no dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+
+import math
+
+import torch
+import torch.nn as nn
+
+
+class SinusoidalEmbedding(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        assert dim % 2 == 0, "dim must be even"
+        half = dim // 2
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
+        self.register_buffer("freqs", freqs)
+
+    def forward(self, t: torch.Tensor) -> torch.Tensor:
+        t = t.reshape(-1, 1).float()
+        args = t * self.freqs.unsqueeze(0)  # (B, half)
+        return torch.cat([args.sin(), args.cos()], dim=-1)  # (B, dim)
+
+
+def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential:
+    """`n_layers`-deep MLP producing an `emb_dim`-wide vector from `in_dim`
+    physical properties (`conditioning.{particle,material}.n_layers`).
+
+    `n_layers=1` (the v0.3.0 default): a single `Linear`, no hidden
+    activation. `n_layers=2` reproduces v0.2's hardcoded depth exactly —
+    `Linear -> SiLU -> Linear` — which is why `migrate_config` back-fills
+    `n_layers=2` for migrated configs rather than the v0.3 default of 1 (see
+    its docstring).
+    """
+    if n_layers < 1:
+        raise ValueError(f"n_layers must be >= 1, got {n_layers}")
+    if n_layers == 1:
+        return nn.Sequential(nn.Linear(in_dim, emb_dim))
+    layers: list[nn.Module] = [nn.Linear(in_dim, emb_dim), nn.SiLU()]
+    for _ in range(n_layers - 2):
+        layers += [nn.Linear(emb_dim, emb_dim), nn.SiLU()]
+    layers.append(nn.Linear(emb_dim, emb_dim))
+    return nn.Sequential(*layers)
+
+
+class ContextAdapter(nn.Module):
+    """Projects a stage's outcome (e.g. Stage 1's 9D target) down to a
+    fixed-width context vector for a downstream stage's conditioning —
+    `stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj`
+    (+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since
+    `SecondaryConditionEncoder` as a wrapper class disappears."""
+
+    def __init__(self, in_dim: int, context_dim: int) -> None:
+        super().__init__()
+        self.proj = nn.Linear(in_dim, context_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.tanh(self.proj(x))
+
+
+class ResBlock(nn.Module):
+    def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
+        super().__init__()
+        self.norm = nn.LayerNorm(dim)
+        self.linear1 = nn.Linear(dim, dim)
+        self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
+        self.act = nn.SiLU()
+        self.dropout = nn.Dropout(dropout)
+        self.linear2 = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        h = self.linear1(h) + self.cond_proj(cond)
+        h = self.act(h)
+        h = self.dropout(h)
+        h = self.linear2(h)
+        return x + h

giant/model/models.py

@@ -0,0 +1,593 @@
+"""Top-level stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`,
+`CriticModel` — composed from encoders/trunks/history (issues.md Issue 8)."""
+
+import torch
+import torch.nn as nn
+
+from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model.encoders import ConditionEncoder
+from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory
+from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding
+from giant.model.routers import Router
+from giant.model.trunks import build_trunk
+
+# ---------------------------------------------------------------------------
+# Stage models
+# ---------------------------------------------------------------------------
+
+
+def resolve_type_n_classes(particle_type_cfg: dict, particle_emb_dim: int) -> int:
+    """Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in
+    place of a bare `conditioning.particle.emb_dim` read. Under
+    `target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 =
+    inherit `conditioning.particle.emb_dim`) — see gitea #29, which decoupled
+    the secondary-species vocabulary size from the unrelated
+    physical-conditioning MLP's output width. Under `target = "embedding"`
+    (or `"physical"`, which ignores this value entirely) `n_classes` doesn't
+    apply — the width stays `conditioning.particle.emb_dim`, the embedding
+    table's own dimensionality (`validate_config` requires
+    `conditioning.particle.type = "embedding"` here)."""
+    if particle_type_cfg.get("target", "physical") == "onehot":
+        return particle_type_cfg.get("n_classes", 0) or particle_emb_dim
+    return particle_emb_dim
+
+
+def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
+    """Width of a single secondary slot's type slice —
+    `PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
+    `emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
+    this many classes/dims wide — callers resolve `emb_dim` via
+    `resolve_type_n_classes` first)."""
+    target = particle_type_cfg.get("target", "physical")
+    return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
+
+
+def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int:
+    """`Stage2OneShot`'s trunk output width.
+
+    `target = "physical"` is untouched from v0.2/today:
+    `k_max * SEC_SLOT_DIM`, the type slice folded into the same
+    flow-matched/WGAN vector as the continuous stick/dir slots.
+
+    `target` in `("onehot", "embedding")`: under `generator == "wgan"` the
+    type slice is still folded in (adversarial for onehot via ST-Gumbel,
+    already-continuous for embedding), just `emb_dim` wide instead of
+    `PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under
+    `generator in ("flow", "ddpm")` the type slice isn't part of this vector
+    at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk
+    only covers `k_max * CONT_SLOT_DIM`.
+    """
+    target = particle_type_cfg.get("target", "physical")
+    if target == "physical":
+        return k_max * SEC_SLOT_DIM
+    if generator == "wgan":
+        return k_max * (CONT_SLOT_DIM + emb_dim)
+    return k_max * CONT_SLOT_DIM
+
+
+class Stage1Model(nn.Module):
+    """Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
+    move it to stage 2, except for a migrated v0.2 checkpoint
+    (`n_sec_head_k_max` given), where it stays attached here
+    since that's where its weights live and what conditioning it was trained
+    against (see `_migrate_legacy_model_config`).
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — `conditioning.share_stages = true`: `build_models`
+    constructs one shared instance and passes it to both stages, halving the
+    conditioning parameter count and forcing a common representation."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        x_dim: int = X_DIM,
+        dropout: float = 0.0,
+        generator: str = "flow",
+        time_dim: int = 64,
+        noise_dim: int = 64,
+        router: Router | None = None,
+        n_sec_head_k_max: int | None = None,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+        )
+        has_time = generator in ("flow", "ddpm")
+        self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
+        merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
+        in_dim = noise_dim if generator == "wgan" else x_dim
+        self.trunk = build_trunk(router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
+        self.n_sec_head = None
+        if n_sec_head_k_max is not None:
+            self.n_sec_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, n_sec_head_k_max + 1),
+            )
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        t: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
+        return self.trunk(x_t, cond, cond_cont, cond_cat)
+
+    def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """Return n_sec logits (B, K_MAX+1) from conditioning alone. Only
+        valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
+        configs predict n_sec from Stage2OneShot instead."""
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                "this Stage1Model has no n_sec_head — n_sec now lives on "
+                "stage 2 by default; this method only exists "
+                "for a migrated v0.2 checkpoint (n_sec.owner='stage1')"
+            )
+        c_emb = self.cond_enc(cond_cont, cond_cat)
+        return self.n_sec_head(c_emb)
+
+
+class Stage2OneShot(nn.Module):
+    """Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
+    reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`,
+    step 4/5, not implemented yet).
+
+    Owns `n_sec_head` by default unless `build_n_sec_head=False`
+    (a migrated v0.2 checkpoint, whose n_sec_head instead attaches to
+    Stage1Model — see `_migrate_legacy_model_config`).
+
+    `particle_type_cfg["target"]` (default `"physical"`) selects the
+    secondary-type mechanism: `"physical"` keeps the type slice folded into
+    the trunk's own
+    flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by
+    the caller via `stage2_trunk_sec_dim` — already reflects this). Under
+    `"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type
+    slice is predicted by a separate `type_head` instead (same shape pattern
+    as `n_sec_head`) — `sec_dim` then covers only the continuous
+    stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
+    Under `generator == "wgan"` the type slice stays folded into `sec_dim`
+    (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
+    unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        context_dim: int = 64,
+        sec_dim: int = SEC_DIM,
+        x_dim: int = X_DIM,
+        dropout: float = 0.0,
+        generator: str = "wgan",
+        time_dim: int = 64,
+        noise_dim: int = 64,
+        k_max: int = K_MAX,
+        router: Router | None = None,
+        build_n_sec_head: bool = True,
+        particle_type_cfg: dict | None = None,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.k_max = k_max
+        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
+        self.type_dim = stage2_type_dim(
+            self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
+        )
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+        )
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+        has_time = generator in ("flow", "ddpm")
+        self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
+        merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
+        in_dim = noise_dim if generator == "wgan" else sec_dim
+        self.trunk = build_trunk(router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
+        self.n_sec_head = None
+        if build_n_sec_head:
+            self.n_sec_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max + 1),
+            )
+        self.type_head = None
+        target = self.particle_type_cfg.get("target", "physical")
+        if target != "physical" and generator in ("flow", "ddpm"):
+            emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
+            self.type_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max * emb_dim),
+            )
+            self._type_k_max = k_max
+            self._type_emb_dim = emb_dim
+
+    def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        base = self.cond_enc(cond_cont, cond_cat)
+        ctx = self.context_adapter(stage1_out)
+        return self.fuse(torch.cat([base, ctx], dim=-1))
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        t: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
+        return self.trunk(x_t, cond, cond_cont, cond_cat)
+
+    def predict_n_sec(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                "this Stage2OneShot has no n_sec_head — it belongs to a "
+                "migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
+                "stage1.predict_n_sec(cond_cont, cond_cat) instead"
+            )
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        return self.n_sec_head(c_emb)
+
+    def predict_type(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+    ) -> torch.Tensor:
+        """`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or
+        vectors (`target="embedding"`) — only under `generator in ("flow",
+        "ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s
+        own output instead (see class docstring)."""
+        if self.type_head is None:
+            raise RuntimeError(
+                "this Stage2OneShot has no type_head — either "
+                "particle_type.target='physical' (the type slice is part of "
+                "forward()'s own output) or generator='wgan' (the WGAN "
+                "trainer reads the type slice out of forward()'s output "
+                "directly instead)"
+            )
+        c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
+        return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim)
+
+
+class Stage2Autoregressive(nn.Module):
+    """Emits secondaries one at a time in descending-energy order, instead
+    of `Stage2OneShot`'s simultaneous
+    k_max-slot prediction. `history` selects `MarkovHistory` or
+    `AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only).
+    `teacher_forcing` handling lives entirely in the trainer
+    (`giant/train.py`), since it only affects how training inputs are
+    assembled, not this module's architecture.
+
+    Under teacher forcing every token's conditioning is built from ground
+    truth, so a whole K-token sequence trains in one parallel batched pass:
+    `forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded
+    to `k_max`) — this also means a future one-token-at-a-time inference loop
+    (`K=1` per call, step 6) needs no interface change here.
+
+    Two independent conditioning paths, mirroring `Stage2OneShot`'s
+    `_cond_embed` but split in two: `_base_cond` (`cond_enc` +
+    `context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend
+    on token position; `_token_cond` additionally fuses in the history
+    encoding and two running scalars (remaining energy-budget fraction,
+    normalized slot index), and feeds `forward`/`predict_type`/the trunk.
+
+    `cond_enc`, if given, is used in place of building a fresh
+    `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
+    """
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        context_dim: int = 64,
+        x_dim: int = X_DIM,
+        dropout: float = 0.0,
+        generator: str = "wgan",
+        time_dim: int = 64,
+        noise_dim: int = 64,
+        k_max: int = K_MAX,
+        router: Router | None = None,
+        build_n_sec_head: bool = True,
+        particle_type_cfg: dict | None = None,
+        history: str = "markov",
+        attn_n_heads: int = 4,
+        attn_n_layers: int = 2,
+        cond_enc: ConditionEncoder | None = None,
+    ) -> None:
+        super().__init__()
+        if history not in ("markov", "attention"):
+            raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'")
+        self.history_kind = history
+        self.generator_kind = generator
+        self.noise_dim = noise_dim
+        self.k_max = k_max
+        self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
+        emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
+        self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim)
+
+        self.cond_enc = (
+            cond_enc
+            if cond_enc is not None
+            else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+        )
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.base_fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+
+        # Reuses conditioning.out_dim for the history encoder's own output
+        # width — there's no dedicated stage2_model.autoregressive key for
+        # this, a reasonable default rather than a design-doc-specified value.
+        history_dim = cond_out_dim
+        hist_in_dim = CONT_SLOT_DIM + self.type_dim
+        self.history_encoder: HistoryEncoder = (
+            AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers)
+            if history == "attention"
+            else MarkovHistory(hist_in_dim, history_dim)
+        )
+        token_fuse_in = cond_out_dim + context_dim + history_dim + 2  # +2: remaining_frac, slot_idx
+        self.token_fuse = nn.Sequential(
+            nn.Linear(token_fuse_in, cond_out_dim),
+            nn.SiLU(),
+        )
+
+        has_time = generator in ("flow", "ddpm")
+        self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
+        merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
+        token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim)
+        in_dim = noise_dim if generator == "wgan" else token_dim
+        self.trunk = build_trunk(
+            router,
+            in_dim,
+            token_dim,
+            hidden_dim,
+            n_res_blocks,
+            merged_cond_dim,
+            dropout,
+        )
+
+        self.n_sec_head = None
+        if build_n_sec_head:
+            self.n_sec_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, k_max + 1),
+            )
+        self.type_head = None
+        target = self.particle_type_cfg.get("target", "physical")
+        if target != "physical" and generator in ("flow", "ddpm"):
+            self.type_head = nn.Sequential(
+                nn.Linear(cond_out_dim, hidden_dim // 2),
+                nn.SiLU(),
+                nn.Linear(hidden_dim // 2, self.type_dim),
+            )
+
+    def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        base = self.cond_enc(cond_cont, cond_cat)
+        ctx = self.context_adapter(stage1_out)
+        return self.base_fuse(torch.cat([base, ctx], dim=-1))
+
+    def _token_cond(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """`hist`, if given, overrides recomputing `self.history_encoder`
+        from `history_feat`/`has_prev` — the inference-time KV-cache path
+        (`Stage2Autoregressive.history_step`) precomputes it once per slot and
+        passes it in here so a slot's (possibly several) model calls — an ODE
+        loop's substeps, or a separate `predict_type` call — read the same
+        cached history instead of each re-deriving (and, under attention,
+        re-appending to the cache — see `AttentionHistory.step`'s docstring)."""
+        K = history_feat.size(1)
+        base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1)
+        ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1)
+        if hist is None:
+            hist = self.history_encoder(history_feat, has_prev)
+        scalars = torch.stack([remaining_frac, slot_idx], dim=-1)
+        return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1))
+
+    def init_history_cache(self):
+        """Inference-only incremental-decoding state for `self.history_encoder`
+        (`giant/sample.py`'s AR loop): `None` under `history="markov"` (its
+        per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or
+        `AttentionHistory.init_cache()` under `history="attention"`."""
+        if isinstance(self.history_encoder, AttentionHistory):
+            return self.history_encoder.init_cache()
+        return None
+
+    def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
+        """One inference slot's worth of history encoding: advances `cache`
+        (from `init_history_cache`, or a previous `history_step` call) by
+        `token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
+        token, same convention `giant.sample.sample_secondaries_ar` already
+        threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is
+        this slot's history summary (pass it as `_token_cond`'s `hist=` to
+        every model call made for this slot), `new_cache` is what to pass into
+        the *next* slot's `history_step`. Must be called exactly once per
+        slot — see `AttentionHistory.step`'s docstring."""
+        if isinstance(self.history_encoder, AttentionHistory):
+            return self.history_encoder.step(token_feat, has_prev, cache)
+        return self.history_encoder(token_feat, has_prev), cache
+
+    def forward(
+        self,
+        x_t: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        t: torch.Tensor | None = None,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        B, K = x_t.shape[0], x_t.shape[1]
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+            hist=hist,
+        )
+        if self.time_emb is not None:
+            assert t is not None
+            t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1)
+            cond = torch.cat([t_emb, c_emb], dim=-1)
+        else:
+            cond = c_emb
+        x_flat = x_t.reshape(B * K, -1)
+        cond_flat = cond.reshape(B * K, -1)
+        cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
+        cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
+        out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
+        return out.view(B, K, -1)
+
+    def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        if self.n_sec_head is None:
+            raise RuntimeError(
+                "this Stage2Autoregressive has no n_sec_head — it belongs to "
+                "a migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
+                "stage1.predict_n_sec(cond_cont, cond_cat) instead"
+            )
+        return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out))
+
+    def predict_type(
+        self,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor,
+        history_feat: torch.Tensor,
+        has_prev: torch.Tensor,
+        remaining_frac: torch.Tensor,
+        slot_idx: torch.Tensor,
+        hist: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        if self.type_head is None:
+            raise RuntimeError(
+                "this Stage2Autoregressive has no type_head — either "
+                "particle_type.target='physical' (the type slice is part of "
+                "forward()'s own output) or generator='wgan' (the WGAN "
+                "trainer reads the type slice out of forward()'s output "
+                "directly instead)"
+            )
+        c_emb = self._token_cond(
+            cond_cont,
+            cond_cat,
+            stage1_out,
+            history_feat,
+            has_prev,
+            remaining_frac,
+            slot_idx,
+            hist=hist,
+        )
+        B, K, _ = c_emb.shape
+        return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
+
+
+class CriticModel(nn.Module):
+    """Generator-agnostic WGAN-GP critic body: a scalar realism score, for
+    either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"`
+    mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as
+    `Stage2OneShot`). Used only when that stage's `generator == "wgan"`."""
+
+    def __init__(
+        self,
+        pdg_vocab: int,
+        mat_vocab: int,
+        particle_cfg: dict,
+        material_cfg: dict,
+        in_dim: int,
+        hidden_dim: int = 256,
+        n_res_blocks: int = 6,
+        cond_out_dim: int = 128,
+        dropout: float = 0.0,
+        stage: str = "stage1",
+        context_dim: int = 64,
+        context_in_dim: int = X_DIM,
+    ) -> None:
+        super().__init__()
+        if stage not in ("stage1", "stage2"):
+            raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
+        self.stage = stage
+        self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
+        if stage == "stage2":
+            self.context_adapter = ContextAdapter(context_in_dim, context_dim)
+            self.fuse = nn.Sequential(
+                nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+                nn.SiLU(),
+            )
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)])
+        self.out_norm = nn.LayerNorm(hidden_dim)
+        self.out_proj = nn.Linear(hidden_dim, 1)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        stage1_out: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        base = self.cond_enc(cond_cont, cond_cat)
+        if self.stage == "stage2":
+            ctx = self.context_adapter(stage1_out)
+            cond = self.fuse(torch.cat([base, ctx], dim=-1))
+        else:
+            cond = base
+        h = self.input_proj(x)
+        for block in self.blocks:
+            h = block(h, cond)
+        return self.out_proj(self.out_norm(h)).squeeze(-1)

giant/model/routers.py

@@ -0,0 +1,375 @@
+"""Mixture-of-experts routing: `Router` base + registry, the four concrete
+router types, and composed/config-driven construction — self-contained, no
+dependency on any other `giant.model` submodule (issues.md Issue 8)."""
+
+import inspect
+import math
+import re
+from collections.abc import Sequence
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from giant.constants import COND_DIM
+
+# ---------------------------------------------------------------------------
+# Routers — carried over unchanged from v0.2
+# ---------------------------------------------------------------------------
+
+
+class Router(nn.Module):
+    """Contract for a pluggable mixture-of-experts routing axis.
+
+    Subclasses implement `gate` (soft partition-of-unity weights over
+    experts, used in train mode for a fully differentiable mixture);
+    `top1` and `balance_loss` have working defaults so a new routing axis
+    is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
+    """
+
+    def __init__(self, n_experts: int) -> None:
+        super().__init__()
+        self.n_experts = n_experts
+        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()`. Opt-in
+        straight-through Gumbel-softmax (`gumbel=True`, train mode only):
+        hardens the forward pass to a one-hot sample (matching eval-time
+        top-1 dispatch) while keeping the soft sample's gradient on backward.
+        """
+        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)
+
+    def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
+        importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
+        return (importance.std() / (importance.mean() + 1e-8)) ** 2
+
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        """Optional supervised auxiliary loss shaping the router's own belief.
+
+        Default: none (a scalar 0). Routers gating on an unobservable
+        pre-step quantity (e.g. ProcessRouter) override this.
+        """
+        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."""
+        norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
+        return norm_entropy
+
+    def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for
+        the full explanation, unchanged in v0.3.0."""
+        gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
+        row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
+        norm_entropy = row_entropy.mean() / math.log(self.n_experts)
+        importance = gate.sum(dim=0)  # (n_experts,)
+        return norm_entropy, importance
+
+
+ROUTER_REGISTRY: dict[str, type[Router]] = {}
+
+
+def register_router(name: str):
+    def decorator(cls: type[Router]) -> type[Router]:
+        ROUTER_REGISTRY[name] = cls
+        return cls
+
+    return decorator
+
+
+def build_router(name: str, n_experts: int, **kwargs) -> Router:
+    """Factory: look up a `Router` subclass by name from the registry.
+
+    Every registered router type is fed the same `router` config dict;
+    kwargs not declared by that type's constructor are silently dropped, so
+    per-type hyperparameters (e.g. EnergyRouter's `temperature`) can coexist
+    in one config without special-casing.
+    """
+    if name not in ROUTER_REGISTRY:
+        raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
+    cls = ROUTER_REGISTRY[name]
+    accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
+    filtered = {k: v for k, v in kwargs.items() if k in accepted}
+    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 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 the initial effective width/temperature exactly equals `value`."""
+    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.
+
+    Reads `cond_cont[:, energy_idx]` (ignores cond_cat). `gate(e) =
+    softmax_i(-(e - c_i)^2 / tau)`; as tau -> 0 this hardens to
+    nearest-center (Voronoi) selection, exactly what `top1` uses at eval.
+    """
+
+    def __init__(
+        self,
+        n_experts: int = 4,
+        temperature: float = 0.5,
+        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 ({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:
+            if len(centers_init) != n_experts:
+                raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
+            centers = torch.tensor(list(centers_init), dtype=torch.float32)
+        if learn_centers:
+            self.centers = nn.Parameter(centers)
+        else:
+            self.register_buffer("centers", centers)
+
+    def effective_width(self) -> torch.Tensor | float:
+        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.effective_width(), dim=-1)
+
+
+@register_router("pdg")
+class PdgRouter(Router):
+    """Soft turn-on gate over a learned PDG embedding (own table, separate
+    from the trunk's `ConditionEncoder`). No supervision needed — PDG code
+    is already known at pre-step time."""
+
+    def __init__(
+        self,
+        n_experts: int,
+        pdg_vocab: int,
+        emb_dim: int = 8,
+        temperature: float = 0.5,
+        learn_centers: bool = True,
+    ) -> None:
+        super().__init__(n_experts)
+        self.temperature = temperature
+        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
+        centers = torch.randn(n_experts, emb_dim) * 0.1
+        if learn_centers:
+            self.centers = nn.Parameter(centers)
+        else:
+            self.register_buffer("centers", centers)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
+        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1)  # (B, n_experts)
+        return torch.softmax(-d2 / self.temperature, dim=-1)
+
+
+@register_router("process")
+class ProcessRouter(Router):
+    """Routes on the physics process expected to end the step — a post-step
+    outcome, so a small classifier over pre-step conditioning predicts it
+    (own pdg/material embeddings, separate from the trunk's ConditionEncoder).
+    `n_experts` doubles as the number of process classes. Supervised via
+    `classify_loss` against the true `process` label at train time only;
+    `gate`/`top1` never see it."""
+
+    def __init__(
+        self,
+        n_experts: int,
+        pdg_vocab: int,
+        mat_vocab: int,
+        emb_dim: int = 8,
+        hidden_dim: int = 64,
+    ) -> None:
+        super().__init__(n_experts)
+        self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
+        self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
+        self.classifier = nn.Sequential(
+            nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim),
+            nn.SiLU(),
+            nn.Linear(hidden_dim, n_experts),
+        )
+
+    def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        pdg_e = self.pdg_emb(cond_cat[:, 0])
+        mat_e = self.mat_emb(cond_cat[:, 1])
+        h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
+        return self.classifier(h)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
+
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
+
+
+class ComposedRouter(Router):
+    """Joint router over independent axes (e.g. energy x pdg), outer-product
+    gated. Not registered in `ROUTER_REGISTRY`; use `build_composed_router`."""
+
+    def __init__(self, routers: list[Router]) -> None:
+        if not routers:
+            raise ValueError("ComposedRouter needs at least one sub-router")
+        n_experts = 1
+        for r in routers:
+            n_experts *= r.n_experts
+        super().__init__(n_experts)
+        self.routers = nn.ModuleList(routers)
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
+        for router in self.routers[1:]:
+            g = router.gate(cond_cont, cond_cat)  # (B, n_i)
+            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1)  # (B, prod so far)
+        return joint
+
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+        total = torch.zeros((), device=cond_cont.device)
+        for router in self.routers:
+            total = total + router.classify_loss(cond_cont, cond_cat, labels)
+        return total
+
+
+def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
+    """Build a `ComposedRouter` from a list of per-axis router specs — see
+    `_parse_composed_axes`."""
+    routers = [
+        build_router(
+            spec["type"],
+            spec["n_experts"],
+            **{
+                **shared_kwargs,
+                **{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
+            },
+        )
+        for spec in specs
+    ]
+    return ComposedRouter(routers)
+
+
+_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")
+
+
+def _parse_composed_axes(router_cfg: dict) -> list[dict]:
+    """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.
+
+    e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`,
+    `axis1_n_experts = 3`, `axis1_emb_dim = 8`. Axis indices must be
+    contiguous from 0.
+    """
+    axes: dict[int, dict] = {}
+    for key, value in router_cfg.items():
+        m = _AXIS_KEY_RE.match(key)
+        if m is None:
+            continue
+        idx, field = int(m.group(1)), m.group(2)
+        axes.setdefault(idx, {})[field] = value
+    missing = set(range(len(axes))) - axes.keys()
+    if missing:
+        raise ValueError(f"composed router config has gaps at axis indices {missing}")
+    return [axes[i] for i in range(len(axes))]
+
+
+# Router types that read cond_cat's pdg index through their own
+# nn.Embedding(pdg_vocab, ...), regardless of the trunk's particle
+# conditioning mode — see _check_router_conditioning_compat.
+_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
+
+
+def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None:
+    """Reject a router axis that reintroduces a training-vocab PDG lookup
+    under `conditioning.particle.type = "physical"`.
+
+    `PdgRouter`/`ProcessRouter` always build their own dataset-scoped
+    `nn.Embedding(pdg_vocab, ...)`, independent of `ConditionEncoder`'s
+    particle mode. Pairing either with `"physical"` would silently
+    reintroduce a training-menu-scoped lookup at the routing layer,
+    defeating the point of physical-property conditioning. Raised loudly at
+    model-build time.
+    """
+    bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
+    if bad and particle_conditioning == "physical":
+        raise ValueError(
+            f"router type(s) {bad} always use a training-vocab PDG embedding, "
+            "which is incompatible with conditioning.particle.type='physical' "
+            "(whose whole point is generalizing beyond that vocab) — pick a "
+            "different router type (e.g. 'energy') or use "
+            "conditioning.particle.type='embedding'."
+        )
+
+
+def _build_router_from_cfg(
+    router_cfg: dict,
+    pdg_vocab: int,
+    mat_vocab: int,
+    particle_conditioning: str = "embedding",
+) -> Router:
+    """Resolve one stage's `router` config into a `Router`, single-axis or
+    composed. `gumbel` is set as a post-construction attribute (shared by
+    every router type, not a per-type constructor kwarg)."""
+    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    if router_cfg["type"] == "composed":
+        axes = _parse_composed_axes(router_cfg)
+        _check_router_conditioning_compat([a["type"] for a in axes], particle_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"]], particle_conditioning)
+    router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
+    router_kwargs.setdefault("pdg_vocab", pdg_vocab)
+    router_kwargs.setdefault("mat_vocab", mat_vocab)
+    router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
+    router.gumbel = bool(router_cfg.get("gumbel", False))
+    return router

giant/model/schedule.py

@@ -11,9 +11,7 @@ class CosineSchedule:
         steps = np.arange(T + 1, dtype=np.float64)
         f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
         alpha_bars = (f / f[0]).astype(np.float32)
-        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(
-            np.float32
-        )
+        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)
 
         self.betas = torch.from_numpy(betas)
         self.alphas = torch.from_numpy((1.0 - betas))

giant/model/trunks.py

@@ -0,0 +1,155 @@
+"""Trunks: everything downstream of the fused conditioning vector — monolithic
+or expert-routed (issues.md Issue 8)."""
+
+import torch
+import torch.nn as nn
+
+from giant.model.layers import ResBlock
+from giant.model.routers import Router
+
+
+class ExpertTrunk(nn.Module):
+    """One small expert: `input_proj -> ResBlock stack -> out_proj`.
+
+    Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR
+    tokens later (`noise_dim` in, `4 + type_dim` out), even though every
+    step-2/3 caller still has `in_dim == out_dim`.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)])
+        self.out_proj = nn.Linear(hidden_dim, out_dim)
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
+        x = self.input_proj(x)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+
+def _route_forward(
+    experts: nn.ModuleList,
+    router: Router,
+    x: torch.Tensor,
+    cond: torch.Tensor,
+    cond_cont: torch.Tensor,
+    cond_cat: torch.Tensor,
+    training: bool,
+) -> torch.Tensor:
+    """Shared dispatch for `RoutedTrunk`.
+
+    Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
+    N-expert dense compute, fully differentiable (`weight` is
+    `router.combine_weights`). Eval mode: grouped top-1 dispatch — each row
+    runs exactly one expert, the actual source of the per-call speedup.
+    """
+    if training:
+        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts)
+        out = torch.zeros(x.shape[0], experts[0].out_proj.out_features, device=x.device)
+        for i, expert in enumerate(experts):
+            out = out + weights[:, i : i + 1] * expert(x, cond)
+        return out
+
+    idx = router.top1(cond_cont, cond_cat)  # (B,)
+    out_dim = experts[0].out_proj.out_features
+    out = torch.zeros(x.shape[0], out_dim, device=x.device)
+    for i, expert in enumerate(experts):
+        mask = idx == i
+        if mask.any():
+            out[mask] = expert(x[mask], cond[mask])
+    return out
+
+
+class Trunk(nn.Module):
+    """Interface implemented by `MonolithicTrunk`/`RoutedTrunk`: everything
+    downstream of the fused conditioning vector, i.e. the actual generative
+    trunk of a stage (`input_proj -> blocks -> out_proj`, monolithic or
+    expert-routed)."""
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        raise NotImplementedError
+
+
+class MonolithicTrunk(Trunk):
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_res_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.input_proj = nn.Linear(in_dim, hidden_dim)
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_res_blocks)])
+        self.out_proj = nn.Linear(hidden_dim, out_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        x = self.input_proj(x)
+        for block in self.blocks:
+            x = block(x, cond)
+        return self.out_proj(x)
+
+
+class RoutedTrunk(Trunk):
+    def __init__(
+        self,
+        router: Router,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_res_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.router = router
+        self.experts = nn.ModuleList(
+            [ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) for _ in range(router.n_experts)]
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+    ) -> torch.Tensor:
+        return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training)
+
+
+def build_trunk(
+    router: Router | None,
+    in_dim: int,
+    out_dim: int,
+    hidden_dim: int,
+    n_res_blocks: int,
+    cond_dim: int,
+    dropout: float = 0.0,
+) -> Trunk:
+    if router is not None:
+        return RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
+    return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)

giant/particles.py

@@ -58,9 +58,7 @@ def particle_mass_charge(pdg: int) -> tuple[float, float]:
         if _pdgid.is_nucleus(pdg):
             z, a = _pdgid.Z(pdg), _pdgid.A(pdg)
             if z is None or a is None:
-                raise ValueError(
-                    f"PDG {pdg}: is_nucleus but Z/A decode failed"
-                ) from None
+                raise ValueError(f"PDG {pdg}: is_nucleus but Z/A decode failed") from None
             return float(a) * _AMU_MEV, float(z)
         raise ValueError(
             f"PDG code {pdg} could not be resolved via the `particle` package "
@@ -109,9 +107,7 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd
     if len(resolved) == 0:
         raise ValueError("nearest_known_pdg: no resolvable candidates")
     codes = np.array([r[0] for r in resolved], dtype=np.int64)
-    table_log_mass = np.log(
-        np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS
-    )
+    table_log_mass = np.log(np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS)
     table_charge = np.array([r[2] for r in resolved], dtype=np.float64)
 
     mass = np.asarray(mass, dtype=np.float64)
@@ -144,9 +140,8 @@ def decode_topn_class(
     other_policy: str = "sample",
     rng: np.random.Generator | None = None,
 ) -> np.ndarray:
-    """`conditioning.particle.type` / `stage2_model.particle_type.target =
-    "onehot"` inference decode: per-row top-N class index -> concrete PDG
-    code.
+    """`stage2_model.particle_type.target = "onehot"` inference decode:
+    per-row top-N class index -> concrete secondary-species PDG code.
 
     class_idx: int array, any shape, values in `[0, n_classes)`.
     topn_map: the `TopNMap` (`giant.data.loader.build_pdg_topn_map_from_files`)
@@ -154,8 +149,11 @@ def decode_topn_class(
         except at the shared "other" index) plus `other_members` (the
         empirical within-"other" distribution, needed for `other_policy =
         "sample"`/`"modal"`).
-    n_classes: `conditioning.particle.emb_dim` — the class count; the "other"
-        bucket is index `n_classes - 1` by construction
+    n_classes: the resolved secondary-species class count
+        (`giant.model.models.resolve_type_n_classes` —
+        `stage2_model.particle_type.n_classes`, 0 = inherit
+        `conditioning.particle.emb_dim`; see gitea #29); the "other" bucket
+        is index `n_classes - 1` by construction
         (`giant.data.loader._topn_plus_other_map`).
     other_policy: `"sample"` draws from `other_members`' empirical frequency;
         `"modal"` always the single most common "other" member; `"drop"`
@@ -180,10 +178,7 @@ def decode_topn_class(
     n_other = int(other_mask.sum())
     if n_other:
         if not topn_map.other_members:
-            raise ValueError(
-                "decode_topn_class: 'other' class predicted but "
-                "topn_map.other_members is empty"
-            )
+            raise ValueError("decode_topn_class: 'other' class predicted but topn_map.other_members is empty")
         members = np.array(list(topn_map.other_members.keys()), dtype=np.int64)
         counts = np.array(list(topn_map.other_members.values()), dtype=np.float64)
         if other_policy == "drop":

giant/pipeline.py

@@ -31,7 +31,7 @@ from giant.data.transforms import (
     sorted_membership,
 )
 from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import build_models, build_critics
+from giant.model.network import build_models, build_critics, resolve_type_n_classes
 from giant.training import train as run_training
 
 
@@ -49,6 +49,7 @@ class SetupStageResult:
     mat_map: dict[str, int]
     proc_map: dict[str, int] | None
     pdg_topn_map: TopNMap | None
+    sec_type_topn_map: TopNMap | None
     mat_topn_map: TopNMap | None
     cond_norm: Normalizer
     tgt_norm: Normalizer
@@ -73,21 +74,14 @@ def _seed_energy_router(
     if not active:
         return
     if energy_quantiles.size == 0:
-        echo(
-            "  warning: no energy samples collected — EnergyRouter falls back to "
-            "default centers"
-        )
+        echo("  warning: no energy samples collected — EnergyRouter falls back to default centers")
         return
     assert cond_norm.mean is not None and cond_norm.std is not None
     levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
     raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
-    centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
-        energy_idx
-    ]
+    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']}"
-    )
+    echo(f"  seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}")
 
 
 def run_setup_stage(
@@ -142,23 +136,14 @@ def run_setup_stage(
         if cache is not None:
             cache.event_index = (unique_ids, counts)
 
-    train_events, val_events = make_event_split(
-        unique_ids, val_fraction=val_fraction, seed=seed
-    )
+    train_events, val_events = make_event_split(unique_ids, val_fraction=val_fraction, seed=seed)
     events_arr = np.array(sorted(train_events))
     n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
-    echo(
-        f"  {int(counts.sum()):,} steps  |  "
-        f"{len(train_events)} train events  |  "
-        f"{len(val_events)} val events"
-    )
+    echo(f"  {int(counts.sum()):,} steps  |  {len(train_events)} train events  |  {len(val_events)} val events")
 
     if cache is not None and cache.vocab is not None:
         pdg_map, mat_map = cache.vocab
-        echo(
-            f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
-            f"{len(mat_map)} materials)"
-        )
+        echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
     else:
         echo("building vocabulary maps …")
         pdg_map, mat_map = build_index_maps_from_files(files)
@@ -174,11 +159,7 @@ def run_setup_stage(
     # need that generality.
     proc_map: dict[str, int] | None = None
     process_router_cfg = next(
-        (
-            r
-            for r in (stage1_router, stage2_router)
-            if r.get("enabled") and r.get("type") == "process"
-        ),
+        (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
         None,
     )
     if process_router_cfg is not None:
@@ -186,10 +167,7 @@ def run_setup_stage(
         cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
         if cached_proc_map is not None:
             proc_map = cached_proc_map
-            echo(
-                f"process vocabulary: cache hit ({len(proc_map)} labels, "
-                f"{n_experts} experts)"
-            )
+            echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
         else:
             echo("building process vocabulary …")
             proc_map = build_process_map_from_files(files, n_experts=n_experts)
@@ -198,34 +176,42 @@ def run_setup_stage(
                 cache.proc_maps[n_experts] = proc_map
 
     # Top-N-plus-other maps for onehot conditioning/type axes.
-    # The PDG axis is shared by
-    # conditioning.particle.type="onehot" and
-    # stage2_model.particle_type.target="onehot" (both key off
-    # conditioning.particle.emb_dim), so at most one PDG scan is needed even
-    # if both consumers are active. The material axis is independent.
+    # The PDG axis is used independently by conditioning.particle.type="onehot"
+    # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
+    # (secondary-species decode) — their class counts can now differ (gitea
+    # #29: stage2_model.particle_type.n_classes, 0 = inherit
+    # conditioning.particle.emb_dim), so each is resolved and built
+    # independently via _pdg_topn below. cache.topn_maps is keyed by
+    # (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
+    # the same N the second call is a cache hit against the first — no extra
+    # scan in the common case where they still match. The material axis is
+    # independent of both.
     particle_cfg = cfg["conditioning"]["particle"]
     material_cfg = cfg["conditioning"]["material"]
-    particle_type_target = cfg["stage2_model"].get("particle_type", {}).get("target")
+    particle_type_cfg_dict = cfg["stage2_model"].get("particle_type") or {}
+    particle_type_target = config.ParticleTypeConfig.from_dict(particle_type_cfg_dict).target
 
-    pdg_topn_map: TopNMap | None = None
-    if particle_cfg["type"] == "onehot" or particle_type_target == "onehot":
-        n_classes = particle_cfg["emb_dim"]
+    def _pdg_topn(n_classes: int) -> TopNMap:
         cache_key = setup_cache.topn_key("pdg", n_classes)
         cached = cache.topn_maps.get(cache_key) if cache is not None else None
         if cached is not None:
-            pdg_topn_map = cached
-            echo(
-                f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, "
-                f"{n_classes} classes)"
-            )
-        else:
-            echo("building pdg top-N map …")
-            pdg_topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
-            echo(
-                f"  {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes"
-            )
-            if cache is not None:
-                cache.topn_maps[cache_key] = pdg_topn_map
+            echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
+            return cached
+        echo("building pdg top-N map …")
+        topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
+        echo(f"  {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
+        if cache is not None:
+            cache.topn_maps[cache_key] = topn_map
+        return topn_map
+
+    pdg_topn_map: TopNMap | None = None
+    if particle_cfg["type"] == "onehot":
+        pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
+
+    sec_type_topn_map: TopNMap | None = None
+    if particle_type_target == "onehot":
+        sec_type_n_classes = resolve_type_n_classes(particle_type_cfg_dict, particle_cfg["emb_dim"])
+        sec_type_topn_map = _pdg_topn(sec_type_n_classes)
 
     mat_topn_map: TopNMap | None = None
     if material_cfg["type"] == "onehot":
@@ -234,29 +220,17 @@ def run_setup_stage(
         cached = cache.topn_maps.get(cache_key) if cache is not None else None
         if cached is not None:
             mat_topn_map = cached
-            echo(
-                f"material top-N map: cache hit ({len(mat_topn_map.class_map)} "
-                f"materials, {n_classes} classes)"
-            )
+            echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
         else:
             echo("building material top-N map …")
-            mat_topn_map = build_topn_map_from_files(
-                files, "material", n_classes=n_classes, cast=str
-            )
-            echo(
-                f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes"
-            )
+            mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
+            echo(f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
             if cache is not None:
                 cache.topn_maps[cache_key] = mat_topn_map
 
-    energy_router_active = any(
-        r.get("enabled") and r.get("type") == "energy"
-        for r in (stage1_router, stage2_router)
-    )
+    energy_router_active = any(r.get("enabled") and r.get("type") == "energy" for r in (stage1_router, stage2_router))
     energy_idx = 3
-    norm_key = setup_cache.normalizer_key(
-        val_fraction, seed, particle_conditioning, material_conditioning
-    )
+    norm_key = setup_cache.normalizer_key(val_fraction, seed, particle_conditioning, material_conditioning)
     entry = cache.normalizers.get(norm_key) if cache is not None else None
 
     if entry is not None:
@@ -281,28 +255,28 @@ def run_setup_stage(
         # router 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
-        )
+        energy_sampler = _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
         for i, path in enumerate(files):
             for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
                 mask = sorted_membership(chunk["event_id"], events_arr)
                 if not mask.any():
                     continue
                 chunk_tr = {k: v[mask] for k, v in chunk.items()}
-                cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _, _ = (
-                    build_features(
-                        chunk_tr,
-                        pdg_map,
-                        mat_map,
-                        proc_map=proc_map,
-                        require_secondaries=True,
-                        particle_conditioning=particle_conditioning,
-                        material_conditioning=material_conditioning,
-                        sec_phys_only=True,
-                        k_max=k_max,
-                    )
+                feats = build_features(
+                    chunk_tr,
+                    pdg_map,
+                    mat_map,
+                    proc_map=proc_map,
+                    require_secondaries=True,
+                    particle_conditioning=particle_conditioning,
+                    material_conditioning=material_conditioning,
+                    sec_phys_only=True,
+                    k_max=k_max,
                 )
+                cond_cont = feats.cond_cont
+                target_s1 = feats.target_s1
+                n_sec = feats.n_sec
+                sec_cont = feats.sec_cont
                 cond_acc.update(cond_cont)
                 tgt_acc.update(target_s1)
                 if energy_sampler is not None:
@@ -336,6 +310,7 @@ def run_setup_stage(
         mat_map=mat_map,
         proc_map=proc_map,
         pdg_topn_map=pdg_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         mat_topn_map=mat_topn_map,
         cond_norm=cond_norm,
         tgt_norm=tgt_norm,
@@ -423,12 +398,10 @@ def run_train_job(
     # The secondary type-index map depends on stage2_model.particle_type.target,
     # independently of conditioning's own onehot/embedding choice above
     # (physical stays untouched/None).
-    particle_type_target = (
-        cfg["stage2_model"].get("particle_type", {}).get("target", "physical")
-    )
+    particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target
     if particle_type_target == "onehot":
-        assert setup.pdg_topn_map is not None
-        sec_type_class_map = setup.pdg_topn_map.class_map
+        assert setup.sec_type_topn_map is not None
+        sec_type_class_map = setup.sec_type_topn_map.class_map
     elif particle_type_target == "embedding":
         sec_type_class_map = pdg_map
     else:
@@ -532,6 +505,7 @@ def run_train_job(
         mat_map={str(k): v for k, v in mat_map.items()},
         proc_map=proc_map,
         pdg_topn_map=setup.pdg_topn_map,
+        sec_type_topn_map=setup.sec_type_topn_map,
         mat_topn_map=setup.mat_topn_map,
         model_config=model_config,
         resume_path=resume,

giant/rollout.py

@@ -117,12 +117,10 @@ def decode_secondary_identity(
     pre_dir: np.ndarray,
     sec_phys_norm: Normalizer,
     pdg_map: dict[int, int],
-    pdg_topn_map: "TopNMap | None",
+    sec_type_topn_map: "TopNMap | None",
     other_policy: str,
     rng: np.random.Generator | None,
-) -> tuple[
-    np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None
-]:
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
     """Decode Stage 2's raw (sec_cont, sec_type) output into physical
     secondary attributes, branching on `sec_decoder.particle_type_cfg`:
 
@@ -150,27 +148,25 @@ def decode_secondary_identity(
         sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
             sec_full, n_sec_np, e_sec, pre_dir, sec_phys_normalizer=sec_phys_norm
         )
-        sec_pdg = nearest_known_pdg(
-            sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
-        ).reshape(sec_mass.shape)
+        sec_pdg = nearest_known_pdg(sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()).reshape(
+            sec_mass.shape
+        )
         return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, None
 
-    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(
-        sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir
-    )
+    sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir)
     sec_type_np = sec_type.cpu().numpy()
     l1_dist = None
 
     if target == "onehot":
-        if pdg_topn_map is None:
+        if sec_type_topn_map is None:
             raise RuntimeError(
-                "particle_type.target='onehot' rollout needs pdg_topn_map "
-                "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
+                "particle_type.target='onehot' rollout needs sec_type_topn_map "
+                "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
             )
         class_idx = sec_type_np.argmax(axis=-1)
         sec_pdg = decode_topn_class(
             class_idx,
-            pdg_topn_map,
+            sec_type_topn_map,
             n_classes=sec_decoder.type_dim,
             other_policy=other_policy,
             rng=rng,
@@ -182,12 +178,8 @@ def decode_secondary_identity(
         l1_dist = np.where(sec_valid, l1_dist, 0.0).astype(np.float32)
 
     sec_mass, sec_charge = particle_phys_array(sec_pdg.reshape(-1)).T
-    sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(
-        np.float32
-    )
-    sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(
-        np.float32
-    )
+    sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(np.float32)
+    sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(np.float32)
     sec_pdg = np.where(sec_valid, sec_pdg, 0).astype(np.int64)
     return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, l1_dist
 
@@ -299,13 +291,9 @@ class _Recorder:
     RAM until the very end.
     """
 
-    def __init__(
-        self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
-    ) -> None:
+    def __init__(self, sink: Callable[[dict[str, np.ndarray]], None] | None = None) -> None:
         self._sink = sink
-        self._cols: dict[str, list] | None = (
-            None if sink is not None else {k: [] for k in _RECORD_KEYS}
-        )
+        self._cols: dict[str, list] | None = None if sink is not None else {k: [] for k in _RECORD_KEYS}
         self.n_rows = 0
         self.termination_reason_counts: Counter[str] = Counter()
 
@@ -313,10 +301,7 @@ class _Recorder:
         n = len(cols["event_id"])
         if n == 0:
             return
-        row = {
-            k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
-            for k in _RECORD_KEYS
-        }
+        row = {k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) for k in _RECORD_KEYS}
         self.n_rows += n
         reasons = row["termination_reason"]
         nonempty = reasons[reasons != ""]
@@ -333,8 +318,7 @@ class _Recorder:
 
     def to_dict(self) -> dict[str, np.ndarray]:
         assert self._cols is not None, (
-            "to_dict() is unavailable when streaming to a sink — use "
-            "n_rows/termination_reason_counts instead"
+            "to_dict() is unavailable when streaming to a sink — use n_rows/termination_reason_counts instead"
         )
         out = {}
         for k, chunks in self._cols.items():
@@ -460,6 +444,7 @@ def rollout(
     material_conditioning: str = "embedding",
     pdg_topn_map: "TopNMap | None" = None,
     mat_topn_map: "TopNMap | None" = None,
+    sec_type_topn_map: "TopNMap | None" = None,
     other_policy: str = "sample",
     seed: int | None = None,
     stage1_ddpm_steps: int = 1000,
@@ -485,14 +470,17 @@ def rollout(
     autoregressive) is inferred from `sec_decoder`'s own class — see
     `sample_stage1`/`sample_stage2` (giant.sample).
 
-    `pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen
-    to share `pdg_topn_map` (one PDG map, not two): they're required
-    whenever `particle_conditioning`/`material_conditioning` is `"onehot"`
-    (feeds `build_cond_features`'s extra `cond_cat` top-N columns), and
-    `pdg_topn_map`/`other_policy` are additionally read under
+    `pdg_topn_map`/`mat_topn_map`/`sec_type_topn_map` serve three independent
+    purposes, no longer required to share one map (see gitea #29):
+    `pdg_topn_map`/`mat_topn_map` are required whenever
+    `particle_conditioning`/`material_conditioning` is `"onehot"` (feeds
+    `build_cond_features`'s extra `cond_cat` top-N columns); `sec_type_topn_map`/
+    `other_policy` are required instead under
     `stage2_model.particle_type.target = "onehot"` (secondary-species
-    decode). `seed` seeds the `other_policy = "sample"` draw only
-    (torch/numpy sampling itself is seeded by the caller, same as today).
+    decode) — its class count (`stage2_model.particle_type.n_classes`) may
+    differ from `pdg_topn_map`'s. `seed` seeds the `other_policy = "sample"`
+    draw only (torch/numpy sampling itself is seeded by the caller, same as
+    today).
 
     `l1_dist_collector`, if given, accumulates the embedding-distance
     diagnostic across the whole run — see `L1DistCollector`. Only populated
@@ -503,6 +491,11 @@ def rollout(
             "conditioning.particle.type='onehot' rollout needs pdg_topn_map "
             "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
         )
+    if sec_decoder.particle_type_cfg.get("target") == "onehot" and sec_type_topn_map is None:
+        raise RuntimeError(
+            "stage2_model.particle_type.target='onehot' rollout needs sec_type_topn_map "
+            "(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
+        )
     if material_conditioning == "onehot" and mat_topn_map is None:
         raise RuntimeError(
             "conditioning.material.type='onehot' rollout needs mat_topn_map "
@@ -552,6 +545,7 @@ def rollout(
                     material_conditioning,
                     pdg_topn_map,
                     mat_topn_map,
+                    sec_type_topn_map,
                     other_policy,
                     rng,
                     stage1_ddpm_steps,
@@ -590,6 +584,7 @@ def _step_chunk(
     material_conditioning,
     pdg_topn_map,
     mat_topn_map,
+    sec_type_topn_map,
     other_policy,
     rng,
     stage1_ddpm_steps,
@@ -619,33 +614,19 @@ def _step_chunk(
     # --- Pre-step termination gates (in priority order; each track picks one) ---
     stop = np.zeros(n, dtype=bool)
     escaped_sel = escaped & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
-        )
-    )
+    rec.add(**_terminal_rows(tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))))
     stop |= escaped_sel
 
     unknown_sel = ~known_pdg & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
-        )
-    )
+    rec.add(**_terminal_rows(tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]))
     stop |= unknown_sel
 
     cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
-    rec.add(
-        **_terminal_rows(
-            tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
-        )
-    )
+    rec.add(**_terminal_rows(tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]))
     stop |= cutoff_sel
 
     maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
-    rec.add(
-        **_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
-    )
+    rec.add(**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel]))
     stop |= maxstep_sel
 
     active = ~stop
@@ -679,41 +660,27 @@ def _step_chunk(
         cond_norm,
         particle_conditioning=particle_conditioning,
         material_conditioning=material_conditioning,
-        pdg_topn_map=pdg_topn_map.class_map
-        if particle_conditioning == "onehot"
-        else None,
-        mat_topn_map=mat_topn_map.class_map
-        if material_conditioning == "onehot"
-        else None,
+        pdg_topn_map=pdg_topn_map.class_map if particle_conditioning == "onehot" else None,
+        mat_topn_map=mat_topn_map.class_map if material_conditioning == "onehot" else None,
     )
     cc = torch.from_numpy(cond_cont).float().to(device)
     ck = torch.from_numpy(cond_cat).long().to(device)
 
-    stage1_norm, n_sec_pred_stage1 = sample_stage1(
-        stage1_model, cc, ck, steps, stage1_ddpm_steps
-    )
+    stage1_norm, n_sec_pred_stage1 = sample_stage1(stage1_model, cc, ck, steps, stage1_ddpm_steps)
     raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
 
     step_length = inv_log_transform(raw[:, 0])
     edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])
 
     post_dir_local = raw[:, 3:6].copy()
-    post_dir_local /= np.clip(
-        np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
-    )
+    post_dir_local /= np.clip(np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None)
     post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)
 
     travel_dir_local = raw[:, 6:9].copy()
-    travel_dir_local /= np.clip(
-        np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
-    )
-    post_pos = reconstruct_post_pos(
-        tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
-    )
+    travel_dir_local /= np.clip(np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None)
+    post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local)
 
-    n_sec_pred = resolve_n_sec(
-        stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1
-    )
+    n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)
     n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
 
     # --- Secondaries ---
@@ -724,23 +691,19 @@ def _step_chunk(
     # decode_secondary_identity's docstring for how each
     # particle_type.target differs on whether PDG resolution is a real
     # identity decision or just a reporting label.
-    sec_cont, sec_type, _valid = sample_stage2(
-        sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps
-    )
-    sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = (
-        decode_secondary_identity(
-            sec_decoder,
-            sec_cont,
-            sec_type,
-            n_sec_np,
-            e_sec,
-            tr["pre_dir"],
-            sec_phys_norm,
-            pdg_map,
-            pdg_topn_map,
-            other_policy,
-            rng,
-        )
+    sec_cont, sec_type, _valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
+    sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
+        sec_decoder,
+        sec_cont,
+        sec_type,
+        n_sec_np,
+        e_sec,
+        tr["pre_dir"],
+        sec_phys_norm,
+        pdg_map,
+        sec_type_topn_map,
+        other_policy,
+        rng,
     )
     sec_valid = np.arange(sec_E.shape[1])[None, :] < n_sec_np[:, None]
     if l1_dist_collector is not None and sec_type_l1_dist is not None:

giant/sample.py

@@ -67,9 +67,7 @@ def sample_ddpm(
         alpha = schedule.alphas[i]
         alpha_bar = schedule.alpha_bars[i]
         z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
-        x = (1.0 / alpha.sqrt()) * (
-            x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
-        ) + beta.sqrt() * z
+        x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
     return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
 
 
@@ -173,9 +171,7 @@ def _decode_stage2_flat(
     else:
         sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
         sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
     return sec_cont, sec_type, sec_valid
 
 
@@ -206,9 +202,7 @@ def sample_secondaries(
         v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
         x = x + v * dt
 
-    return _decode_stage2_flat(
-        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
 
 
 @torch.no_grad()
@@ -226,9 +220,7 @@ def sample_secondaries_wgan(
     B = cond_cont.size(0)
     z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
     x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
-    return _decode_stage2_flat(
-        sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
 
 
 @torch.no_grad()
@@ -306,12 +298,8 @@ def sample_secondaries_ar(
         has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
         history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
         remaining_frac = remaining.unsqueeze(1)  # (B, 1)
-        slot_idx = torch.full(
-            (B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32
-        )
-        hist, history_cache = sec_decoder.history_step(
-            history_feat, has_prev, history_cache
-        )
+        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
         if generator == "wgan":
             z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
@@ -366,21 +354,15 @@ def sample_secondaries_ar(
         sec_type[:, k] = type_k
 
         if target == "onehot":
-            type_for_history = F.one_hot(
-                type_k.argmax(dim=-1), num_classes=type_dim
-            ).float()
+            type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
         else:
             type_for_history = type_k
 
         stick_fraction = torch.sigmoid(cont_k[:, 0])
-        prev_repr = torch.cat(
-            [stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1
-        )
+        prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
         remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
 
-    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
-        1
-    )
+    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
     return sec_cont, sec_type, sec_valid
 
 
@@ -426,16 +408,10 @@ def sample_stage2(
     case, so there's nothing to dispatch to here.
     """
     if isinstance(sec_decoder, Stage2Autoregressive):
-        return sample_secondaries_ar(
-            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
-        )
+        return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
     if sec_decoder.generator_kind == "wgan":
-        return sample_secondaries_wgan(
-            sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
-        )
-    return sample_secondaries(
-        sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps
-    )
+        return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
+    return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
 
 
 def resolve_n_sec(

giant/tools/bump_dataset_version.py

@@ -1,5 +1,5 @@
 """Cut a new raw generation or processed schema version for the geant_steps
-dataset tree (see scripts/migrate_geant_steps.py for the layout):
+dataset tree (see giant/tools/migrate_geant_steps.py for the layout):
 
     raw/<kind>/<gen>/<detector>/shard-NNN.root
     processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
@@ -82,9 +82,7 @@ def _max_index(parent: Path, pattern: re.Pattern) -> int:
 
 def _git_user_name() -> str | None:
     try:
-        out = subprocess.run(
-            ["git", "config", "user.name"], capture_output=True, text=True, timeout=2
-        )
+        out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2)
     except (OSError, subprocess.SubprocessError):
         # OSError (e.g. git not on PATH) and subprocess.SubprocessError
         # (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
@@ -142,9 +140,7 @@ def plan_bump_schema(
     raw_gen_dir = root / "raw" / kind / gen_tag
     processed_gen_dir = root / "processed" / kind / gen_tag
     if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
-        raise SystemExit(
-            f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
-        )
+        raise SystemExit(f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first")
     if target is not None:
         if not SCHEMA_RE.match(target):
             raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
@@ -154,9 +150,7 @@ def plan_bump_schema(
         schema_tag = f"schema{next_schema}"
     new_dirs = [processed_gen_dir / schema_tag]
     by_suffix = f" ({by})" if by else ""
-    log_line = (
-        f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
-    )
+    log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
     return new_dirs, log_line
 
 
@@ -212,9 +206,7 @@ def _manifest_referenced_files(pools_root: Path) -> set[Path]:
     return referenced
 
 
-def _referenced_root_count(
-    raw_gen_dir: Path, processed_gen_dir: Path
-) -> tuple[int, int]:
+def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
     """(total .root files, count with a same-named .parquet under any schema) for one gen."""
     if not raw_gen_dir.is_dir():
         return 0, 0
@@ -243,9 +235,7 @@ def _referenced_root_count(
     return total, referenced
 
 
-def _referenced_parquet_count(
-    schema_dir: Path, manifest_referenced: set[Path]
-) -> tuple[int, int]:
+def _referenced_parquet_count(schema_dir: Path, manifest_referenced: set[Path]) -> tuple[int, int]:
     """(total .parquet files, count listed in at least one manifest) for one schema dir."""
     if not schema_dir.is_dir():
         return 0, 0
@@ -342,11 +332,7 @@ def print_status(root: Path) -> None:
     grand_files = 0
     for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
         kind = kind_dir.name
-        gens = sorted(
-            int(m.group(1))
-            for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
-            if m
-        )
+        gens = sorted(int(m.group(1)) for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir()) if m)
         print(_colorize(f"{kind}/", "kind"))
         kind_total = 0
         kind_files = 0
@@ -359,25 +345,18 @@ def print_status(root: Path) -> None:
             schemas = sorted(
                 int(m.group(1))
                 for m in (
-                    SCHEMA_RE.match(p.name)
-                    for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
-                    if p.is_dir()
+                    SCHEMA_RE.match(p.name) for p in (schema_dir.iterdir() if schema_dir.is_dir() else []) if p.is_dir()
                 )
                 if m
             )
             schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
             schema_counts = {
-                s: _referenced_parquet_count(
-                    schema_dir / f"schema{s}", manifest_referenced
-                )
-                for s in schemas
+                s: _referenced_parquet_count(schema_dir / f"schema{s}", manifest_referenced) for s in schemas
             }
             processed_size = sum(schema_sizes.values())
             processed_files = sum(c[0] for c in schema_counts.values())
             processed_referenced = sum(c[1] for c in schema_counts.values())
-            raw_files, raw_referenced = _referenced_root_count(
-                raw_gen_dir, processed_gen_dir
-            )
+            raw_files, raw_referenced = _referenced_root_count(raw_gen_dir, processed_gen_dir)
             gen_total = raw_size + processed_size
             gen_files = raw_files + processed_files
             kind_total += gen_total
@@ -425,9 +404,7 @@ def print_status(root: Path) -> None:
                         print(_reason_line(schema_reason, indent=4))
             else:
                 print(_colorize("      (none)", "schema"))
-        print(
-            _row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files)
-        )
+        print(_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files))
         print()
         grand_total += kind_total
         grand_files += kind_files
@@ -526,13 +503,8 @@ def plan_update_manifest(
     return result, missing
 
 
-def apply_update_manifest(
-    manifest_path: Path, lines: list[tuple[str, str | None]]
-) -> None:
-    out = [
-        replacement if replacement is not None else original
-        for original, replacement in lines
-    ]
+def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
+    out = [replacement if replacement is not None else original for original, replacement in lines]
     manifest_path.write_text("\n".join(out) + "\n")
 
 
@@ -552,9 +524,7 @@ def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
     return files
 
 
-def plan_create_manifest(
-    output_path: Path, parquet_files: list[Path]
-) -> tuple[list[str], list[Path], list[Path]]:
+def plan_create_manifest(output_path: Path, parquet_files: list[Path]) -> tuple[list[str], list[Path], list[Path]]:
     """Return (relative_lines, missing_files, resolved_abs_paths)."""
     manifest_dir = output_path.resolve().parent
     lines: list[str] = []
@@ -569,9 +539,7 @@ def plan_create_manifest(
     return lines, missing, resolved
 
 
-def check_holdout_overlap(
-    output_path: Path, resolved_new_files: list[Path]
-) -> list[tuple[str, Path]]:
+def check_holdout_overlap(output_path: Path, resolved_new_files: list[Path]) -> list[tuple[str, Path]]:
     """Return (other_manifest_name, file) pairs where new files clash with existing manifests.
 
     The check is triggered when output_path is (or will be) holdout.manifest, or when a
@@ -611,7 +579,7 @@ def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
 
 
 # ---------------------------------------------------------------------------
-# CLI entry points (called from scripts/dwarf.py)
+# CLI entry points (called from giant/tools/dwarf.py)
 # ---------------------------------------------------------------------------
 
 
@@ -641,9 +609,7 @@ def _run_bump(
     if gen is None:
         new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
     else:
-        new_dirs, log_line = plan_bump_schema(
-            root_path, kind, gen, reason, by, date, to
-        )
+        new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date, to)
 
     print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
     print("new directories:")

giant/tools/check_migration_v02_v03.py

@@ -10,9 +10,9 @@ machine against an actual trained checkpoint before merging
 
 Usage (from the repo root, on a portal machine):
 
-    uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt
-    uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --ema
-    uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1
+    uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt
+    uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --ema
+    uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1
 
 Run it once against a flow (or ddpm) checkpoint and once against a wgan
 checkpoint ("one flow checkpoint and one WGAN checkpoint").
@@ -84,19 +84,12 @@ def main() -> int:
     mode = model_config.get("mode", "flow")
     routed = bool((model_config.get("router") or {}).get("enabled"))
     print(f"checkpoint: {args.checkpoint}")
-    print(
-        f"  mode={mode!r}  conditioning={model_config.get('conditioning')!r}  "
-        f"routed={routed}  ema={args.ema}"
-    )
+    print(f"  mode={mode!r}  conditioning={model_config.get('conditioning')!r}  routed={routed}  ema={args.ema}")
 
     stage1_key = "model_ema" if args.ema and "model_ema" in ckpt else "model"
-    stage2_key = (
-        "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
-    )
+    stage2_key = "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
     if args.ema and stage1_key == "model":
-        print(
-            "  warning: --ema requested but no model_ema in checkpoint, using raw weights"
-        )
+        print("  warning: --ema requested but no model_ema in checkpoint, using raw weights")
 
     # --- old side: the frozen v0.2 snapshot, loaded with the checkpoint's own weights ---
     old_stage1, old_stage2 = legacy.build_models(model_config)
@@ -119,9 +112,7 @@ def main() -> int:
         print("PASS (construction only, routed checkpoint)")
         return 0
 
-    remapped1, remapped2 = net.migrate_legacy_state_dict(
-        ckpt[stage1_key], ckpt[stage2_key]
-    )
+    remapped1, remapped2 = net.migrate_legacy_state_dict(ckpt[stage1_key], ckpt[stage2_key])
     missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
     missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
     if missing1 or unexpected1 or missing2 or unexpected2:
@@ -132,9 +123,7 @@ def main() -> int:
     new_stage1.eval()
     new_stage2.eval()
 
-    cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(
-        model_config, args.batch, args.seed
-    )
+    cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(model_config, args.batch, args.seed)
 
     ok = True
     with torch.no_grad():

giant/tools/create_root_files.py

@@ -32,7 +32,7 @@ from dataclasses import dataclass
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
-# Must match scripts/bump_dataset_version.py's GEN_RE.
+# Must match giant/tools/bump_dataset_version.py's GEN_RE.
 GEN_RE = re.compile(r"^gen\d+$")
 SHARD_RE = re.compile(r"^shard-(\d+)\.root$")
 
@@ -64,9 +64,7 @@ def parse_detector_spec(spec: str) -> tuple[str, str | None]:
     if ":" in spec:
         label, config = spec.split(":", 1)
         if not label or not config:
-            raise PlanError(
-                f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG"
-            )
+            raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
         return label, config
     return spec, None
 
@@ -94,9 +92,7 @@ def plan_jobs(
         raise PlanError(f"--gen must look like 'genN', got {gen!r}")
     gen_dir = dataset_root / "raw" / kind / gen
     if not gen_dir.is_dir():
-        raise PlanError(
-            f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first"
-        )
+        raise PlanError(f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first")
 
     jobs = []
     for spec in detector_specs:
@@ -124,9 +120,7 @@ def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
     return zlib.crc32(key.encode()) & 0x7FFFFFFF
 
 
-def build_cmd(
-    executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None
-) -> list[str]:
+def build_cmd(executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None) -> list[str]:
     """minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
     cmd = [str(executable)]
     if job.config:
@@ -147,10 +141,7 @@ def run_job(
     gen: str,
     tmp_root: Path,
 ) -> JobResult:
-    workdir = (
-        tmp_root
-        / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
-    )
+    workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
     workdir.mkdir(parents=True)
 
     cmd = build_cmd(executable, job, events_per_file, energy_gev)
@@ -174,20 +165,12 @@ def run_job(
             job,
             False,
             None,
-            f"expected exactly one .root output in {workdir}, found {len(produced)}: "
-            f"{[p.name for p in produced]}",
+            f"expected exactly one .root output in {workdir}, found {len(produced)}: {[p.name for p in produced]}",
             result.stdout,
             result.stderr,
         )
 
-    dest = (
-        dataset_root
-        / "raw"
-        / kind
-        / gen
-        / job.detector
-        / f"shard-{job.shard_index:03d}.root"
-    )
+    dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
     if dest.exists():
         return JobResult(
             job,
@@ -279,14 +262,7 @@ def run_make_root(
     print(f"executable: {executable}")
     for job in planned_jobs:
         cmd = build_cmd(executable, job, events_per_file, energy_gev)
-        dest = (
-            dataset_root_path
-            / "raw"
-            / kind
-            / gen
-            / job.detector
-            / f"shard-{job.shard_index:03d}.root"
-        )
+        dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
         seed = job_seed(kind, gen, job, energy_gev)
         print(f"  MINICALOSIM_SEED={seed} {' '.join(cmd)}  ->  {dest}")
 

giant/tools/dwarf.py

@@ -1,6 +1,6 @@
 """dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.
 
-Unifies the standalone scripts/*.py conversion, migration, versioning, and
+Unifies the standalone giant/tools/*.py conversion, migration, versioning, and
 simulation-fanout tools into one Typer app so there's a single command name
 (and `--help`) to remember instead of five differently-hyphenated ones.
 """
@@ -14,20 +14,20 @@ import typer
 from typing_extensions import Annotated
 
 from giant.config import Conditioning
-from scripts.bump_dataset_version import (
+from giant.tools.bump_dataset_version import (
     run_bump_gen,
     run_bump_schema,
     run_create_manifest,
     run_status,
     run_update_manifest,
 )
-from scripts.create_root_files import run_make_root
-from scripts.geometry_oracle import run_build_geometry_oracle
-from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
-from scripts.migrate_geant_steps import run_migration
-from scripts.steps_to_parquet import convert_steps_to_parquet
-from scripts.steps_to_parquet_parallel import run_parallel_job
-from scripts.warm_setup_cache import run_warm_setup_cache
+from giant.tools.create_root_files import run_make_root
+from giant.tools.geometry_oracle import run_build_geometry_oracle
+from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
+from giant.tools.migrate_geant_steps import run_migration
+from giant.tools.steps_to_parquet import convert_steps_to_parquet
+from giant.tools.steps_to_parquet_parallel import run_parallel_job
+from giant.tools.warm_setup_cache import run_warm_setup_cache
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -80,8 +80,7 @@ def convert(
         typer.Option(
             "--output",
             "-o",
-            help="Output Parquet file (default: <input>.parquet). Only valid "
-            "with a single input file and --jobs 1.",
+            help="Output Parquet file (default: <input>.parquet). Only valid with a single input file and --jobs 1.",
         ),
     ] = None,
     batch_size: Annotated[
@@ -91,9 +90,7 @@ def convert(
             help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
         ),
     ] = "100 MB",
-    tree: Annotated[
-        str, typer.Option("--tree", help="Tree name inside the ROOT file")
-    ] = "Steps",
+    tree: Annotated[str, typer.Option("--tree", help="Tree name inside the ROOT file")] = "Steps",
     compression: Annotated[
         Compression, typer.Option("--compression", help="Parquet compression codec")
     ] = Compression.snappy,
@@ -129,15 +126,11 @@ def convert(
         raise typer.Exit(1)
     _warn_if_exceeds_shared_quota(jobs, "--jobs")
 
-    compression_value = (
-        "uncompressed" if compression is Compression.none else compression.value
-    )
+    compression_value = "uncompressed" if compression is Compression.none else compression.value
 
     if jobs == 1:
         if output is not None and len(root_files) > 1:
-            typer.echo(
-                "error: --output can only be used with a single input file", err=True
-            )
+            typer.echo("error: --output can only be used with a single input file", err=True)
             raise typer.Exit(1)
         total_orphaned = 0
         for root_file in root_files:
@@ -150,10 +143,7 @@ def convert(
             )
             total_orphaned += n_orphaned
         if total_orphaned:
-            typer.echo(
-                f"\n{total_orphaned} orphaned child track(s) dropped across "
-                f"{len(root_files)} file(s)."
-            )
+            typer.echo(f"\n{total_orphaned} orphaned child track(s) dropped across {len(root_files)} file(s).")
         return
 
     if output is not None:
@@ -176,9 +166,7 @@ def convert(
 
 @app.command()
 def migrate(
-    root: Annotated[
-        Path, typer.Argument(help="Dataset root to migrate in place")
-    ] = _DATASET_ROOT_DEFAULT,
+    root: Annotated[Path, typer.Argument(help="Dataset root to migrate in place")] = _DATASET_ROOT_DEFAULT,
     execute: Annotated[
         bool,
         typer.Option(
@@ -190,8 +178,7 @@ def migrate(
         bool,
         typer.Option(
             "--copy",
-            help="Copy instead of move, leaving the originals in place "
-            "(e.g. if another process is still reading them)",
+            help="Copy instead of move, leaving the originals in place (e.g. if another process is still reading them)",
         ),
     ] = False,
 ) -> None:
@@ -202,12 +189,8 @@ def migrate(
 @app.command("bump-gen")
 def bump_gen(
     reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    by: Annotated[
-        Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
-    ] = None,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
     date: Annotated[
         Optional[str],
         typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -220,12 +203,8 @@ def bump_gen(
             help="Target gen tag (default: one past the current highest)",
         ),
     ] = None,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Apply (default: dry run)")
-    ] = False,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new raw generation."""
     run_bump_gen(
@@ -243,12 +222,8 @@ def bump_gen(
 def bump_schema(
     gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
     reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    by: Annotated[
-        Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
-    ] = None,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
     date: Annotated[
         Optional[str],
         typer.Option("--date", help="Override date (default: today, ISO)"),
@@ -261,12 +236,8 @@ def bump_schema(
             help="Target schema tag (default: one past the current highest)",
         ),
     ] = None,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Apply (default: dry run)")
-    ] = False,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new schema within a gen."""
     run_bump_schema(
@@ -283,9 +254,7 @@ def bump_schema(
 
 @app.command()
 def status(
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """List existing gens/schemas per kind."""
     run_status(str(root))
@@ -293,9 +262,7 @@ def status(
 
 @app.command("update-manifest")
 def update_manifest(
-    manifests: Annotated[
-        list[Path], typer.Argument(help="One or more .manifest files to update")
-    ],
+    manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
     schema: Annotated[
         Optional[str],
         typer.Option(
@@ -306,9 +273,7 @@ def update_manifest(
     ] = None,
     gen: Annotated[
         Optional[str],
-        typer.Option(
-            "--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"
-        ),
+        typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
     ] = None,
     execute: Annotated[
         bool,
@@ -316,9 +281,7 @@ def update_manifest(
     ] = False,
 ) -> None:
     """Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
-    run_update_manifest(
-        [str(m) for m in manifests], schema=schema, execute=execute, gen=gen
-    )
+    run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)
 
 
 @app.command("create-manifest")
@@ -333,22 +296,15 @@ def create_manifest(
         typer.Option(
             "--pool",
             metavar="DETECTOR",
-            help="Detector name; combined with --type and --root to form "
-            "<root>/pools/<detector>/<type>.manifest",
+            help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
         ),
     ] = None,
     type_: Annotated[
         Optional[PoolType],
-        typer.Option(
-            "--type", help="Pool type — full, holdout, or dev (required with --pool)"
-        ),
+        typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
     ] = None,
-    root: Annotated[
-        Path, typer.Option("--root", help="Dataset root (used with --pool)")
-    ] = _DATASET_ROOT_DEFAULT,
-    execute: Annotated[
-        bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
-    ] = False,
+    root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
+    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"),
@@ -368,9 +324,7 @@ def create_manifest(
 
 @app.command("make-root")
 def make_root(
-    executable: Annotated[
-        Path, typer.Option("--executable", help="Built minicalosim run_* executable")
-    ],
+    executable: Annotated[Path, typer.Option("--executable", help="Built minicalosim run_* executable")],
     detector: Annotated[
         list[str],
         typer.Option(
@@ -382,15 +336,9 @@ def make_root(
             "Repeatable.",
         ),
     ],
-    num_files: Annotated[
-        int, typer.Option("--num-files", help="New shards to create per detector")
-    ],
-    events_per_file: Annotated[
-        int, typer.Option("--events-per-file", help="nEvents passed to the executable")
-    ],
-    gen: Annotated[
-        str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
-    ],
+    num_files: Annotated[int, typer.Option("--num-files", help="New shards to create per detector")],
+    events_per_file: Annotated[int, typer.Option("--events-per-file", help="nEvents passed to the executable")],
+    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")],
     energy_gev: Annotated[
         float | None,
         typer.Option(
@@ -401,20 +349,12 @@ def make_root(
             "to name the dataset accordingly.",
         ),
     ] = None,
-    kind: Annotated[
-        str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
-    ] = "steps",
-    dataset_root: Annotated[
-        Path, typer.Option("--dataset-root", help="Dataset root")
-    ] = _DATASET_ROOT_DEFAULT,
-    jobs: Annotated[
-        int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
-    ] = 4,
+    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
+    dataset_root: Annotated[Path, typer.Option("--dataset-root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
+    jobs: Annotated[int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")] = 4,
     execute: Annotated[
         bool,
-        typer.Option(
-            "--execute", help="Actually run jobs (default: dry run / print plan)"
-        ),
+        typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)"),
     ] = False,
 ) -> None:
     """Generate new ROOT shards via a minicalosim executable."""
@@ -441,9 +381,7 @@ class OracleMethod(str, Enum):
 
 @app.command("build-geometry-oracle")
 def build_geometry_oracle(
-    data: Annotated[
-        Path, typer.Argument(help="Steps parquet file or directory of steps files")
-    ],
+    data: Annotated[Path, typer.Argument(help="Steps parquet file or directory of steps files")],
     out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
     method: Annotated[
         OracleMethod,
@@ -456,9 +394,7 @@ def build_geometry_oracle(
             ),
         ),
     ] = OracleMethod.slab,
-    k: Annotated[
-        int, typer.Option("--k", help="Neighbours for the knn classifier")
-    ] = 1,
+    k: Annotated[int, typer.Option("--k", help="Neighbours for the knn classifier")] = 1,
     subsample: Annotated[
         int,
         typer.Option("--subsample", help="Max reference points sampled from the data"),
@@ -504,9 +440,7 @@ def build_geometry_oracle(
 def warm_cache(
     data: Annotated[
         Path,
-        typer.Argument(
-            help="Parquet file, directory, or .manifest — same as `giant train`'s"
-        ),
+        typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
     ],
     val_fraction: Annotated[
         float,
@@ -518,16 +452,13 @@ def warm_cache(
     ] = 0.1,
     seed: Annotated[
         int,
-        typer.Option(
-            "--seed", "-s", help="Must match the `giant train` run(s) to warm for"
-        ),
+        typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
     ] = 0,
     particle_conditioning: Annotated[
         Conditioning,
         typer.Option(
             "--particle-conditioning",
-            help="Must match the `giant train` run(s)' conditioning.particle.type "
-            "to warm for",
+            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for",
         ),
     ] = Conditioning.physical,
     material_conditioning: Annotated[
@@ -543,21 +474,14 @@ def warm_cache(
         bool,
         typer.Option(
             "--router/--no-router",
-            help="Warm the process vocabulary too (only takes effect with "
-            "--router-type process)",
+            help="Warm the process vocabulary too (only takes effect with --router-type process)",
         ),
     ] = False,
-    router_type: Annotated[
-        str, typer.Option("--router-type", help="Router implementation name")
-    ] = "energy",
-    n_experts: Annotated[
-        int, typer.Option("--n-experts", help="Number of routed experts")
-    ] = 4,
+    router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
+    n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
     rebuild: Annotated[
         bool,
-        typer.Option(
-            "--rebuild", help="Ignore any existing sidecar and recompute every section"
-        ),
+        typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
     ] = False,
 ) -> None:
     """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.

giant/tools/geometry_oracle.py

@@ -38,9 +38,7 @@ def run_build_geometry_oracle(
         n_bins=n_bins,
     )
 
-    print(
-        f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}"
-    )
+    print(f"method: {method}  reference points: {oracle.metadata['n_reference_points']:,}")
     print("classes (material, layer_id):")
     for material, layer_id in oracle.classes:
         print(f"  {material:<12} layer_id={layer_id}")

giant/tools/hparam_scan.py

@@ -154,9 +154,7 @@ def run_hparam_scan(
         wall_time_s = time.monotonic() - start
 
         if metrics_path.exists():
-            epochs_completed, final_val_loss, best_val_loss = final_metrics(
-                metrics_path
-            )
+            epochs_completed, final_val_loss, best_val_loss = final_metrics(metrics_path)
             append_summary(
                 summary_path,
                 {
@@ -171,11 +169,6 @@ def run_hparam_scan(
                     "wall_time_s": round(wall_time_s, 1),
                 },
             )
-            print(
-                f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} "
-                f"({wall_time_s:.1f}s)"
-            )
+            print(f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} ({wall_time_s:.1f}s)")
         else:
-            print(
-                f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
-            )
+            print(f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log")

giant/tools/migrate_geant_steps.py

@@ -50,12 +50,8 @@ PREDICTED_RE = re.compile(
     r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)"
     r"_predicted(?P<local>_local)?\.parquet$"
 )
-SHARD_RE = re.compile(
-    r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$"
-)
-LEGACY_PREDICTED_RE = re.compile(
-    r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$"
-)
+SHARD_RE = re.compile(r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$")
+LEGACY_PREDICTED_RE = re.compile(r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$")
 LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?P<ext>root|parquet)$")
 
 
@@ -110,24 +106,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
         if m:
             detector, shard, ext = m["detector"], int(m["shard"]), m["ext"]
             if ext == "root":
-                dst = (
-                    src_root
-                    / "raw"
-                    / "steps"
-                    / GEN
-                    / detector
-                    / f"shard-{shard:03d}.root"
-                )
+                dst = src_root / "raw" / "steps" / GEN / detector / f"shard-{shard:03d}.root"
             else:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "steps"
-                    / GEN
-                    / SCHEMA
-                    / detector
-                    / f"shard-{shard:03d}.parquet"
-                )
+                dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
             moves.append((path, dst))
             continue
 
@@ -135,19 +116,9 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
         if m:
             ext = m["ext"]
             if ext == "root":
-                dst = (
-                    src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
-                )
+                dst = src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
             else:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "hits"
-                    / LEGACY_GEN
-                    / LEGACY_SCHEMA
-                    / "pbwo4"
-                    / "shard-000.parquet"
-                )
+                dst = src_root / "processed" / "hits" / LEGACY_GEN / LEGACY_SCHEMA / "pbwo4" / "shard-000.parquet"
             moves.append((path, dst))
             continue
 
@@ -165,20 +136,9 @@ def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
         for pool, shards in rules.items():
             manifest_path = manifest_dir / f"{pool}{MANIFEST_SUFFIX}"
             for shard in shards:
-                dst = (
-                    src_root
-                    / "processed"
-                    / "steps"
-                    / GEN
-                    / SCHEMA
-                    / detector
-                    / f"shard-{shard:03d}.parquet"
-                )
+                dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
                 manifests[manifest_path].append((shard, dst))
-    return {
-        k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)]
-        for k, v in manifests.items()
-    }
+    return {k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)] for k, v in manifests.items()}
 
 
 def run_migration(root: str, execute: bool, copy: bool) -> None:

giant/tools/profile_analysis_costs.py

@@ -13,7 +13,7 @@ real checkpoint) they short-circuit almost instantly and are excluded here —
 see `runtime_estimate.py`'s `_ROUTER_FIXED_S` for how those are handled
 instead.
 
-Usage: ``uv run python scripts/profile_analysis_costs.py``
+Usage: ``uv run python giant/tools/profile_analysis_costs.py``
 """
 
 from __future__ import annotations
@@ -94,9 +94,7 @@ def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
             "post_dx": post_dir[:, 0],
             "post_dy": post_dir[:, 1],
             "post_dz": post_dir[:, 2],
-            "edep": np.where(
-                is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep
-            ),
+            "edep": np.where(is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep),
             "step_length": np.where(is_synthetic, 0.0, step_length),
             "material": rng.choice(_MATERIALS, size=n),
             "layer_id": rng.integers(0, 30, size=n),
@@ -163,9 +161,7 @@ def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
     )
 
 
-def _time(
-    spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path
-) -> float:
+def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path) -> float:
     t0 = time.perf_counter()
     compute_reduced(
         spec_id,

giant/tools/steps_to_parquet_parallel.py

@@ -2,11 +2,11 @@
 
 A single `dwarf convert` call converts a list of files one at a time; this
 module runs up to --jobs conversions concurrently, each as its own `dwarf
-convert` subprocess (invoked via `python -m scripts.dwarf`, so it picks up
+convert` subprocess (invoked via `python -m giant.tools.dwarf`, so it picks up
 the active venv/uv environment automatically).
 
 Inputs must live under <dataset-root>/raw/<kind>/<gen>/<detector>/<file>.root
-(see scripts/migrate_geant_steps.py) — each is written to the matching
+(see giant/tools/migrate_geant_steps.py) — each is written to the matching
 processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet, where <schema>
 defaults to the highest schemaN already under processed/<kind>/<gen>/ (pass
 --schema to pick a specific one, e.g. one just created by `dwarf bump-schema`).
@@ -22,7 +22,7 @@ import sys
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
-# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
+# Must match giant/tools/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
 GEN_RE = re.compile(r"^gen\d+$")
 SCHEMA_RE = re.compile(r"^schema(\d+)$")
 
@@ -49,9 +49,7 @@ def latest_schema_tag(processed_gen_dir: Path) -> str | None:
     return best_tag
 
 
-def resolve_destination(
-    root_file: Path, dataset_root: Path, schema_override: str | None
-) -> Path:
+def resolve_destination(root_file: Path, dataset_root: Path, schema_override: str | None) -> Path:
     """Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
     to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.
 
@@ -66,12 +64,7 @@ def resolve_destination(
         raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")
 
     parts = rel.parts
-    if (
-        len(parts) != 5
-        or parts[0] != "raw"
-        or not GEN_RE.match(parts[2])
-        or not parts[4].endswith(".root")
-    ):
+    if len(parts) != 5 or parts[0] != "raw" or not GEN_RE.match(parts[2]) or not parts[4].endswith(".root"):
         raise DestinationError(
             f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
             f"under {dataset_root} (got relative path: {rel})"
@@ -89,7 +82,7 @@ def resolve_destination(
     return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"
 
 
-_DWARF_CONVERT_CMD = [sys.executable, "-m", "scripts.dwarf", "convert"]
+_DWARF_CONVERT_CMD = [sys.executable, "-m", "giant.tools.dwarf", "convert"]
 
 
 def _convert_one(
@@ -134,7 +127,7 @@ def run_parallel(
     written next to the input .root).
 
     *cmd_prefix* overrides the subprocess command run per file (defaults to
-    `python -m scripts.dwarf convert`) — used by tests to substitute a fake
+    `python -m giant.tools.dwarf convert`) — used by tests to substitute a fake
     conversion script.
 
     Returns one (root_file, returncode, stdout, stderr) tuple per file, in
@@ -216,13 +209,7 @@ def run_parallel_job(
             print(f"  {root_file}", file=sys.stderr)
         raise SystemExit(1)
 
-    total_orphaned = sum(
-        int(m.group(1))
-        for _, _, stdout, _ in results
-        for m in _ORPHAN_RE.finditer(stdout)
-    )
+    total_orphaned = sum(int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout))
     if total_orphaned:
-        print(
-            f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)."
-        )
+        print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
     print(f"\nAll {len(results)} conversion(s) completed.")

giant/tools/warm_setup_cache.py

@@ -9,6 +9,7 @@ for the sidecar itself.
 
 from pathlib import Path
 
+from giant import config as gconfig
 from giant.constants import K_MAX
 from giant.pipeline import run_setup_stage
 
@@ -43,19 +44,25 @@ def run_warm_setup_cache(
         "type": router_type,
         "n_experts": n_experts,
     }
-    # A minimal v0.3 cfg — only the keys run_setup_stage actually reads
-    # (conditioning.{particle,material}.type, stage{1,2}_model.router). This
-    # CLI only ever configures one router (matching today's single
+    # Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so
+    # run_setup_stage always sees every key it might read (e.g.
+    # conditioning.particle.emb_dim, stage2_model.particle_type.target) at
+    # its real default, not silently missing/None — see issues.md Issue 1.
+    # This CLI only ever configures one router (matching today's single
     # --router-type flag), so it's placed on stage1_model; stage2_model's
     # stays disabled.
-    cfg = {
-        "conditioning": {
-            "particle": {"type": particle_conditioning},
-            "material": {"type": material_conditioning},
+    cfg = gconfig.merge_cli_overrides(
+        gconfig.DEFAULT_CONFIG,
+        None,
+        {
+            "conditioning": {
+                "particle": {"type": particle_conditioning},
+                "material": {"type": material_conditioning},
+            },
+            "stage1_model": {"router": router_cfg},
+            "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
         },
-        "stage1_model": {"router": router_cfg},
-        "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
-    }
+    )
     run_setup_stage(
         Path(data),
         val_fraction=val_fraction,

giant/training/loop.py

@@ -78,9 +78,7 @@ def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs)
     return validate_marginals(model, val_loader, device=device, **kwargs)
 
 
-def _marginal_kl(
-    trainers: dict[str, StageTrainer], val_loader, device, **kwargs
-) -> float:
+def _marginal_kl(trainers: dict[str, StageTrainer], val_loader, device, **kwargs) -> float:
     """Mean marginal KL over the stage-1 sampling chain, or NaN when stage 1
     is inactive or `validate_marginals` declined to produce a result."""
     stage1 = trainers.get("stage1")
@@ -90,9 +88,7 @@ def _marginal_kl(
         stage1,
         val_loader,
         device,
-        sec_decoder=trainers["stage2"].sampling_model()
-        if "stage2" in trainers
-        else None,
+        sec_decoder=trainers["stage2"].sampling_model() if "stage2" in trainers else None,
         **kwargs,
     )
     if result is None:
@@ -113,6 +109,7 @@ def train(
     mat_map: dict | None = None,
     proc_map: dict | None = None,
     pdg_topn_map: TopNMap | None = None,
+    sec_type_topn_map: TopNMap | None = None,
     mat_topn_map: TopNMap | None = None,
     model_config: dict | None = None,
     resume_path: str | Path | None = None,
@@ -137,9 +134,7 @@ def train(
 
     trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches)
     if not trainers:
-        raise ValueError(
-            "no active stage — stage1_model.active and stage2_model.active are both false"
-        )
+        raise ValueError("no active stage — stage1_model.active and stage2_model.active are both false")
     has_adversarial = any(not tr.supports_val_loss for tr in trainers.values())
 
     checkpoint_extras = {
@@ -147,12 +142,9 @@ def train(
         "pdg_map": pdg_map,
         "mat_map": mat_map,
         "proc_map": proc_map,
-        "pdg_topn_map": topnmap_to_json(pdg_topn_map)
-        if pdg_topn_map is not None
-        else None,
-        "mat_topn_map": topnmap_to_json(mat_topn_map)
-        if mat_topn_map is not None
-        else None,
+        "pdg_topn_map": topnmap_to_json(pdg_topn_map) if pdg_topn_map is not None else None,
+        "sec_type_topn_map": topnmap_to_json(sec_type_topn_map) if sec_type_topn_map is not None else None,
+        "mat_topn_map": topnmap_to_json(mat_topn_map) if mat_topn_map is not None else None,
         "model_config": model_config,
     }
 
@@ -166,10 +158,7 @@ def train(
         best_val_loss = ckpt.get("best_val_loss", float("inf"))
         global_step = ckpt.get("global_step", 0)
         if start_epoch > epochs:
-            print(
-                f"checkpoint already completed epoch {start_epoch - 1} "
-                f"(>= --epochs {epochs}) — nothing to train"
-            )
+            print(f"checkpoint already completed epoch {start_epoch - 1} (>= --epochs {epochs}) — nothing to train")
             return
 
     collector = MetricsCollector.create(
@@ -206,10 +195,7 @@ def train(
             for batch in bar:
                 B = batch[0].size(0)
                 collector.add_train_batch(
-                    {
-                        name: trainer.step(batch, device, global_step)
-                        for name, trainer in trainers.items()
-                    },
+                    {name: trainer.step(batch, device, global_step) for name, trainer in trainers.items()},
                     B,
                 )
                 bar.set_postfix_str(collector.postfix(), refresh=False)
@@ -220,9 +206,7 @@ def train(
             bar.close()
 
             if shutdown.requested:
-                ckpt = build_checkpoint(
-                    trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras
-                )
+                ckpt = build_checkpoint(trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras)
                 torch.save(ckpt, out_dir / "last.pt")
                 last_completed_epoch = epoch - 1
                 print(
@@ -244,15 +228,10 @@ def train(
                             break
                         B = batch[0].size(0)
                         collector.add_val_batch(
-                            {
-                                name: tr.val_loss(batch, device)
-                                for name, tr in scored.items()
-                            },
+                            {name: tr.val_loss(batch, device) for name, tr in scored.items()},
                             B,
                         )
-                        collector.observe_routers(
-                            batch[0].to(device), batch[1].to(device), B
-                        )
+                        collector.observe_routers(batch[0].to(device), batch[1].to(device), B)
 
             # An adversarial stage has no averageable validation loss, so it
             # needs the marginal-KL signal every epoch to pick a best
@@ -264,10 +243,7 @@ def train(
             elif validate_every > 0 and epoch % validate_every == 0:
                 stage1 = trainers.get("stage1")
                 ddpm_steps = 1000
-                if (
-                    isinstance(stage1, FlowDDPMStageTrainer)
-                    and stage1.ddpm_schedule is not None
-                ):
+                if isinstance(stage1, FlowDDPMStageTrainer) and stage1.ddpm_schedule is not None:
                     ddpm_steps = stage1.ddpm_schedule.T
                 marginal_kl = _marginal_kl(
                     trainers,
@@ -292,22 +268,16 @@ def train(
             collector.set("val/marginal_kl", marginal_kl)
             collector.set(
                 "gpu_mem_mb",
-                torch.cuda.max_memory_allocated(device) / (1024 * 1024)
-                if device.type == "cuda"
-                else 0.0,
-            )
-            collector.set(
-                "samples_per_sec", collector.train_samples / max(epoch_time, 1e-8)
+                torch.cuda.max_memory_allocated(device) / (1024 * 1024) if device.type == "cuda" else 0.0,
             )
+            collector.set("samples_per_sec", collector.train_samples / max(epoch_time, 1e-8))
             collector.set("is_best", int(is_best))
             collector.set("epoch_time_s", epoch_time)
 
             print(collector.summary_line(val_loss, epoch_time, is_best))
             collector.write_epoch(global_step)
 
-            ckpt = build_checkpoint(
-                trainers, epoch, global_step, best_val_loss, checkpoint_extras
-            )
+            ckpt = build_checkpoint(trainers, epoch, global_step, best_val_loss, checkpoint_extras)
             if is_best:
                 best_val_loss = val_loss
                 ckpt["best_val_loss"] = best_val_loss

giant/training/metrics.py

@@ -116,11 +116,7 @@ class _RouterAccumulator:
 
     def add(self, entropy: torch.Tensor, importance: torch.Tensor, n: int) -> None:
         self.entropy += entropy.item() * n
-        self.importance = (
-            importance.clone()
-            if self.importance is None
-            else self.importance + importance
-        )
+        self.importance = importance.clone() if self.importance is None else self.importance + importance
         self.n += n
 
     def stats(self) -> dict[str, float]:
@@ -168,22 +164,18 @@ class MetricsCollector:
         self._train = {name: _Accumulator() for name in trainers}
         self._val = {name: _Accumulator() for name in trainers}
         self._routers = {
-            name: _RouterAccumulator(tr.router.n_experts)
-            for name, tr in trainers.items()
-            if tr.router is not None
+            name: _RouterAccumulator(tr.router.n_experts) for name, tr in trainers.items() if tr.router is not None
         }
         # Only "mean" specs need summing; "last" specs are read straight off
         # the accumulator's most recent stats dict. "grad_norm" is always
         # summed — it feeds the run-level `grad_norm` column whether or not
         # a trainer reports it per stage.
         self._train_keys = {
-            name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"}
-            | {"grad_norm"}
+            name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"} | {"grad_norm"}
             for name, tr in trainers.items()
         }
         self._val_keys = {
-            name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"}
-            for name, tr in trainers.items()
+            name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"} for name, tr in trainers.items()
         }
         self._run_values: dict[str, float] = {}
         self._epoch = 0
@@ -224,13 +216,9 @@ class MetricsCollector:
                 import wandb
             except ImportError as exc:
                 raise RuntimeError(
-                    "train.wandb = true (--wandb) requires the 'wandb' package — "
-                    "install it via `uv sync --extra wandb`"
+                    "train.wandb = true (--wandb) requires the 'wandb' package — install it via `uv sync --extra wandb`"
                 ) from exc
-            param_counts = {
-                name: sum(p.numel() for p in tr.model.parameters())
-                for name, tr in trainers.items()
-            }
+            param_counts = {name: sum(p.numel() for p in tr.model.parameters()) for name, tr in trainers.items()}
             param_counts["total"] = sum(param_counts.values())
             wandb_run = wandb.init(
                 project=wandb_project,
@@ -290,9 +278,7 @@ class MetricsCollector:
             self._batch_grad_norm += stage_stats.get("grad_norm", 0.0)
         if self._ema_seeded:
             self._ema_loss += _EMA_ALPHA * (self._batch_loss - self._ema_loss)
-            self._ema_grad_norm += _EMA_ALPHA * (
-                self._batch_grad_norm - self._ema_grad_norm
-            )
+            self._ema_grad_norm += _EMA_ALPHA * (self._batch_grad_norm - self._ema_grad_norm)
         else:
             self._ema_loss = self._batch_loss
             self._ema_grad_norm = self._batch_grad_norm
@@ -303,9 +289,7 @@ class MetricsCollector:
             self._val[name].add(stage_stats, self._val_keys[name], batch_size)
 
     @torch.no_grad()
-    def observe_routers(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int
-    ) -> None:
+    def observe_routers(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int) -> None:
         """Record gate diagnostics for every routed stage on this batch.
 
         Called from the validation pass only (as in v0.2/v0.3.0), so a stage
@@ -339,9 +323,7 @@ class MetricsCollector:
             payload[f"batch/{name}/lr"] = trainer.optimizer.param_groups[0]["lr"]
             if trainer.router is not None:
                 with torch.no_grad():
-                    entropy, _ = trainer.router.gate_stats(
-                        batch[0].to(device), batch[1].to(device)
-                    )
+                    entropy, _ = trainer.router.gate_stats(batch[0].to(device), batch[1].to(device))
                 payload[f"batch/{name}/router/entropy"] = entropy.item()
         self.wandb_run.log(payload, step=global_step)
 
@@ -367,10 +349,7 @@ class MetricsCollector:
         self._run_values[column] = value
 
     def summary_line(self, val_loss: float, epoch_time: float, is_best: bool) -> str:
-        bits = [
-            trainer.summary(self.train_means(name))
-            for name, trainer in self.trainers.items()
-        ]
+        bits = [trainer.summary(self.train_means(name)) for name, trainer in self.trainers.items()]
         marker = "  [best]" if is_best else ""
         return (
             f"epoch {self._epoch:{self.epoch_width}d}/{self.epochs}  "

giant/training/stage2_inputs.py

@@ -100,9 +100,9 @@ def _assemble_stage2_real(
     `Stage2Autoregressive`'s per-token target also uses; the two must stay in
     lockstep. See `_assemble_stage2_ar_target`'s docstring for the
     (target, generator) width rules."""
-    return _assemble_stage2_ar_target(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    ).flatten(1)
+    return _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim).flatten(
+        1
+    )
 
 
 def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor:
@@ -137,18 +137,14 @@ def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
     return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
 
 
-def _ar_meta(
-    k_max: int, batch: int, device: torch.device, fraction: torch.Tensor
-) -> dict[str, torch.Tensor]:
+def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]:
     """`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR
     conditioning tensors that don't depend on *which* history representation
     (ground truth vs. the scheduled-sampling mix) produced `fraction`.
     Shared by `_assemble_stage2_ar_inputs` and
     `_assemble_stage2_ar_inputs_scheduled`, which differ only in
     `history_feat`."""
-    slot_idx = (
-        torch.arange(k_max, device=device).float() / max(k_max - 1, 1)
-    ).unsqueeze(0)
+    slot_idx = (torch.arange(k_max, device=device).float() / max(k_max - 1, 1)).unsqueeze(0)
     return {
         "has_prev": _ar_has_prev(k_max, device).expand(batch, -1),
         "remaining_frac": _remaining_energy_fraction(fraction),
@@ -182,9 +178,7 @@ def _assemble_stage2_ar_inputs(
     return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)}
 
 
-def _stage2_tf_prob(
-    mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int
-) -> float:
+def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
     """P(condition slot k+1 on the TRUE token k rather than the model's own
     prediction), for the current epoch
     (`stage2_model.autoregressive.teacher_forcing`).
@@ -260,9 +254,7 @@ def _assemble_stage2_ar_inputs_scheduled(
     device = sec_cont.device
     B, K = sec_cont.shape[0], sec_cont.shape[1]
     if p_tf >= 1.0:
-        return _assemble_stage2_ar_inputs(
-            sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim
-        )
+        return _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim)
 
     was_training = model.training
     sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
@@ -273,9 +265,7 @@ def _assemble_stage2_ar_inputs_scheduled(
 
     fraction_gt = _stick_fraction(sec_cont)
     dir_gt = sec_cont[..., 1:CONT_SLOT_DIM]
-    type_repr_gt = _type_repr(
-        sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim
-    )
+    type_repr_gt = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
     fraction_pred, dir_pred, type_repr_pred = _history_repr_from_ar_sample(
         sec_cont_pred, sec_type_pred, particle_type_cfg
     )
@@ -323,8 +313,6 @@ def _relax_onehot_type_slice(
     cont, type_logits = x[..., :cont_dim], x[..., cont_dim:]
     if grad_probe is not None:
         cont.register_hook(lambda g: grad_probe.__setitem__("cont", g.norm().item()))
-        type_logits.register_hook(
-            lambda g: grad_probe.__setitem__("type", g.norm().item())
-        )
+        type_logits.register_hook(lambda g: grad_probe.__setitem__("type", g.norm().item()))
     type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1)
     return torch.cat([cont, type_soft], dim=-1).reshape(B, -1)

giant/training/trainers.py

@@ -16,13 +16,16 @@ adversarial and non-adversarial stages identically.
 import copy
 import math
 from dataclasses import dataclass, field
+from typing import NamedTuple
 
 import torch
 import torch.nn.functional as F
 import torch.optim as optim
 
+from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
 from giant.constants import CONT_SLOT_DIM
-from giant.model.network import Router, stage2_type_dim
+from giant.data.dataset import StepBatch
+from giant.model.network import Router, resolve_type_n_classes, stage2_type_dim
 from giant.model.schedule import (
     CosineSchedule,
     flow_matching_loss,
@@ -41,9 +44,7 @@ from giant.training.stage2_inputs import (
 
 
 @torch.no_grad()
-def _update_ema(
-    ema_model: torch.nn.Module, model: torch.nn.Module, decay: float
-) -> None:
+def _update_ema(ema_model: torch.nn.Module, model: torch.nn.Module, decay: float) -> None:
     for ema_p, p in zip(ema_model.parameters(), model.parameters()):
         ema_p.mul_(decay).add_(p, alpha=1 - decay)
 
@@ -73,8 +74,8 @@ def _cosine_warmup_lambda(warmup_steps: int, total_steps: int):
     return _lr_lambda
 
 
-def _batch_to_device(batch: tuple, device: torch.device) -> tuple:
-    return tuple(t.to(device) for t in batch)
+def _batch_to_device(batch: StepBatch, device: torch.device) -> StepBatch:
+    return type(batch)(*(t.to(device) for t in batch))
 
 
 @dataclass(frozen=True)
@@ -96,8 +97,8 @@ class StageSpec:
     n_sec_lambda: float = 0.1
 
     # particle-type target (stage 2 only)
-    particle_type: dict = field(default_factory=lambda: {"target": "physical"})
-    particle_type_emb_dim: int = 16
+    particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
+    particle_type_n_classes: int = 16
 
     # optimization
     lr: float = 3e-4
@@ -129,63 +130,69 @@ class StageSpec:
     type_gumbel_tau_end: float = 0.1
 
     @classmethod
-    def from_config(
-        cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int
-    ) -> "StageSpec":
-        t = cfg["train"]
-        stage_cfg = cfg[f"{name}_model"]
-        router_cfg = stage_cfg.get("router") or {}
-        wgan_cfg = stage_cfg.get("wgan") or {}
-        ar_cfg = (stage_cfg.get("autoregressive") or {}) if is_stage2 else {}
+    def from_config(cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int) -> "StageSpec":
+        t = TrainConfig.from_dict(cfg["train"])
+        # n_sec/particle_type/decoder/autoregressive/wgan's gumbel_tau_* are
+        # stage-2-only concepts, always read off s2_spec (guarded by
+        # is_stage2 where the stage-1 StageSpec needs a different value) —
+        # historically n_sec/particle_type were read from stage2_model
+        # unconditionally even for the stage-1 StageSpec, preserved here for
+        # behavioral parity. stage_spec covers the fields both stage configs
+        # share structurally (generator, lambda, router, ddpm, and wgan's
+        # base fields — Stage2ModelConfig's sub-configs all subclass
+        # stage 1's).
+        s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
+        stage_spec = s2_spec if is_stage2 else Stage1ModelConfig.from_dict(cfg["stage1_model"])
         return cls(
             name=name,
             is_stage2=is_stage2,
-            generator=stage_cfg["generator"],
-            decoder=stage_cfg.get("decoder", "one_shot") if is_stage2 else "one_shot",
-            lambda_weight=stage_cfg.get("lambda", 1.0),
-            n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1),
-            particle_type=cfg["stage2_model"].get("particle_type")
-            or {"target": "physical"},
-            particle_type_emb_dim=cfg["conditioning"]["particle"]["emb_dim"],
+            generator=stage_spec.generator,
+            decoder=s2_spec.decoder if is_stage2 else "one_shot",
+            lambda_weight=stage_spec.lambda_weight,
+            n_sec_lambda=s2_spec.n_sec.lambda_weight,
+            particle_type=s2_spec.particle_type,
+            particle_type_n_classes=resolve_type_n_classes(
+                s2_spec.particle_type.to_dict(), cfg["conditioning"]["particle"]["emb_dim"]
+            ),
             # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
-            # (giant/config.py), so they read directly; the field defaults
-            # below exist only for tests that construct StageSpec by hand.
-            lr=t["lr"],
-            weight_decay=t["weight_decay"],
-            ema_decay=t["ema_decay"],
-            warmup_epochs=t["warmup_epochs"],
-            epochs=t["epochs"],
+            # (giant/config.py), so TrainConfig.from_dict never has to fall
+            # back to a literal here; the field defaults below exist only
+            # for tests that construct StageSpec by hand.
+            lr=t.lr,
+            weight_decay=t.weight_decay,
+            ema_decay=t.ema_decay,
+            warmup_epochs=t.warmup_epochs,
+            epochs=t.epochs,
             steps_per_epoch=max(steps_per_epoch, 1),
-            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),
-            teacher_forcing=ar_cfg.get("teacher_forcing", "always"),
-            tf_p_start=ar_cfg.get("tf_p_start", 1.0),
-            tf_p_end=ar_cfg.get("tf_p_end", 1.0),
+            lambda_balance=stage_spec.router.lambda_balance,
+            lambda_proc=stage_spec.router.lambda_proc,
+            lambda_entropy=stage_spec.router.lambda_entropy,
+            gumbel_tau_start=stage_spec.router.gumbel_tau_start,
+            gumbel_tau_end=stage_spec.router.gumbel_tau_end,
+            teacher_forcing=s2_spec.autoregressive.teacher_forcing if is_stage2 else cls.teacher_forcing,
+            tf_p_start=s2_spec.autoregressive.tf_p_start if is_stage2 else cls.tf_p_start,
+            tf_p_end=s2_spec.autoregressive.tf_p_end if is_stage2 else cls.tf_p_end,
             # AR self-sampling under scheduled/never teacher forcing reuses
             # train.validate_steps as its flow-matching ODE step count — no
             # dedicated config key for this (the autoregressive config lists
             # tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only).
-            ar_sample_steps=t["validate_steps"],
-            ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000),
-            n_critic=wgan_cfg.get("n_critic", 5),
-            gp_weight=wgan_cfg.get("gp_weight", 10.0),
-            critic_lr=wgan_cfg.get("critic_lr", 0.0),
-            type_gumbel_tau_start=wgan_cfg.get("gumbel_tau_start", 1.0),
-            type_gumbel_tau_end=wgan_cfg.get("gumbel_tau_end", 0.1),
+            ar_sample_steps=t.validate_steps,
+            ddpm_n_steps=stage_spec.ddpm.n_steps,
+            n_critic=stage_spec.wgan.n_critic,
+            gp_weight=stage_spec.wgan.gp_weight,
+            critic_lr=stage_spec.wgan.critic_lr,
+            type_gumbel_tau_start=s2_spec.wgan.gumbel_tau_start if is_stage2 else cls.type_gumbel_tau_start,
+            type_gumbel_tau_end=s2_spec.wgan.gumbel_tau_end if is_stage2 else cls.type_gumbel_tau_end,
         )
 
 
 class StageTrainer:
     """One active stage's optimizer(s), EMA, and per-batch step.
 
-    Reads only the shared batch tuple `(cond_cont, cond_cat, x1_s1, n_sec,
-    sec_cont, proc_idx, sec_type_idx)` — stage 2 always conditions on the
-    ground-truth `x1_s1` (`stage2_model.stage1_context = "truth"`,
-    stage-level teacher forcing; `"sampled"` is not implemented), so stage
-    trainers never need each other's output at train time. This means
+    Reads only the shared `StepBatch` (`giant.data.dataset`) — stage 2 always
+    conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context =
+    "truth"`, stage-level teacher forcing; `"sampled"` is not implemented),
+    so stage trainers never need each other's output at train time. This means
     "stage-2-only training is a cheap ablation, not new plumbing" falls out
     for free: a trainer only exists for active stages, and inactive stages
     are simply never constructed.
@@ -232,8 +239,8 @@ class StageTrainer:
         self.router = _stage_router(self.model)
         self._modules = (self.model, *extra_modules)
 
-        self.particle_type_cfg = dict(spec.particle_type or {"target": "physical"})
-        self.particle_type_emb_dim = spec.particle_type_emb_dim
+        self.particle_type_cfg = spec.particle_type.to_dict()
+        self.particle_type_n_classes = spec.particle_type_n_classes
         self.ema_decay = spec.ema_decay
 
         self.ema_model: torch.nn.Module | None = None
@@ -244,19 +251,17 @@ class StageTrainer:
 
     # --- schedule -------------------------------------------------------
 
-    def _init_lr_schedule(
-        self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int
-    ) -> None:
+    def _init_lr_schedule(self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int) -> None:
         self._lr_lambda = _cosine_warmup_lambda(warmup_steps, total_steps)
         self.total_steps = total_steps
         self.lr_sched = optim.lr_scheduler.LambdaLR(optimizer, self._lr_lambda)
 
     # --- per-batch (subclass responsibility) ----------------------------
 
-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         raise NotImplementedError
 
-    def val_loss(self, batch: tuple, device: torch.device) -> dict:
+    def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
         raise NotImplementedError
 
     # --- reporting hooks ------------------------------------------------
@@ -270,9 +275,7 @@ class StageTrainer:
         """This stage's fragment of the end-of-epoch console line."""
         raise NotImplementedError
 
-    def val_objective(
-        self, train_means: dict, val_means: dict, marginal_kl: float
-    ) -> float:
+    def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
         """This stage's contribution to the best-checkpoint selection score."""
         raise NotImplementedError
 
@@ -328,7 +331,7 @@ class StageTrainer:
             n_sec,
             self.particle_type_cfg,
             self.model.cond_enc,
-            self.particle_type_emb_dim,
+            self.particle_type_n_classes,
             p_tf,
             self.spec.ar_sample_steps,
         )
@@ -354,14 +357,12 @@ class StageTrainer:
             self.particle_type_cfg,
             generator,
             self.model.cond_enc,
-            self.particle_type_emb_dim,
+            self.particle_type_n_classes,
         )
         return target.flatten(1) if flatten else target
 
     @staticmethod
-    def _sec_mask(
-        n_sec: torch.Tensor, k_max: int, device: torch.device
-    ) -> torch.Tensor:
+    def _sec_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> torch.Tensor:
         """`(B, K_MAX)` bool prefix mask: slot k is valid iff `k < n_sec`."""
         return torch.arange(k_max, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
 
@@ -397,9 +398,7 @@ class StageTrainer:
         return l_nsec, nsec_acc
 
     @staticmethod
-    def _step_optimizer(
-        optimizer: optim.Optimizer, loss: torch.Tensor, params: list
-    ) -> float:
+    def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
         """`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
         the pre-clip grad norm. The one place the grad-clip constant lives."""
         optimizer.zero_grad()
@@ -450,9 +449,7 @@ class StageTrainer:
 class FlowDDPMStageTrainer(StageTrainer):
     """flow or ddpm generator for a single stage."""
 
-    def __init__(
-        self, spec: StageSpec, model: torch.nn.Module, device: torch.device
-    ) -> None:
+    def __init__(self, spec: StageSpec, model: torch.nn.Module, device: torch.device) -> None:
         if spec.is_stage2 and spec.generator not in ("flow",):
             raise NotImplementedError(
                 f"stage2_model.generator={spec.generator!r} is accepted by the "
@@ -466,26 +463,16 @@ class FlowDDPMStageTrainer(StageTrainer):
         # NotImplementedError above): "physical" keeps it folded in
         # (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding"
         # pull it out into model.type_head instead (0 here).
-        self._flow_type_dim = (
-            None
-            if self.particle_type_cfg.get("target", "physical") == "physical"
-            else 0
-        )
+        self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0
 
         self.params = list(self.model.parameters())
-        self.optimizer = optim.AdamW(
-            self.params, lr=spec.lr, weight_decay=spec.weight_decay
-        )
+        self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
         self._init_lr_schedule(
             self.optimizer,
             warmup_steps=spec.warmup_epochs * spec.steps_per_epoch,
             total_steps=max(spec.epochs * spec.steps_per_epoch, 1),
         )
-        self.ddpm_schedule = (
-            CosineSchedule(T=spec.ddpm_n_steps).to(device)
-            if spec.generator == "ddpm"
-            else None
-        )
+        self.ddpm_schedule = CosineSchedule(T=spec.ddpm_n_steps).to(device) if spec.generator == "ddpm" else None
 
         self.train_metrics = [
             train_metric(key)
@@ -515,9 +502,7 @@ class FlowDDPMStageTrainer(StageTrainer):
         ]
         self.stage_metrics = [stage_metric("lr")]
 
-    def _generator_loss(
-        self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None
-    ):
+    def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None):
         if not self.is_stage2:
             if self.generator == "flow":
                 return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
@@ -584,22 +569,16 @@ class FlowDDPMStageTrainer(StageTrainer):
         mask = sec_mask.float()
         denom = mask.sum().clamp(min=1)
         if self.particle_type_cfg.get("target") == "onehot":
-            ce = F.cross_entropy(
-                type_out.transpose(1, 2), sec_type_idx, reduction="none"
-            )
+            ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none")
             l_type = (ce * mask).sum() / denom
-            type_acc = (
-                (type_out.argmax(-1) == sec_type_idx).float() * mask
-            ).sum() / denom
+            type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom
         else:  # "embedding"
             target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
             se = ((type_out - target_vec) ** 2).mean(-1)
             l_type = (se * mask).sum() / denom
         return l_type, type_acc
 
-    def _compute(
-        self, batch: tuple, device: torch.device, epoch: int | None = None
-    ) -> dict:
+    def _compute(self, batch: StepBatch, device: torch.device, epoch: int | None = None) -> dict:
         """`epoch=None` (the `val_loss` path) always uses full teacher
         forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation
         should stay a stable, non-stochastic ground-truth comparison; only
@@ -619,23 +598,13 @@ class FlowDDPMStageTrainer(StageTrainer):
         x1_s2 = None
         ar_inputs = None
         if self.is_stage2 and self.decoder == "autoregressive":
-            ar_inputs = self._ar_inputs(
-                cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
-            )
-            x1_s2 = self._sec_target(
-                sec_cont, sec_type_idx, self.generator, flatten=False
-            )
+            ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
+            x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False)
         elif self.is_stage2:
-            x1_s2 = self._sec_target(
-                sec_cont, sec_type_idx, self.generator, flatten=True
-            )
+            x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True)
 
-        l_gen = self._generator_loss(
-            cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs
-        )
-        l_nsec, nsec_acc = self._n_sec_loss(
-            cond_cont, cond_cat, stage1_ctx, n_sec, device
-        )
+        l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs)
+        l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
 
         l_type, type_acc = self._type_loss(
             cond_cont,
@@ -649,15 +618,14 @@ class FlowDDPMStageTrainer(StageTrainer):
 
         l_balance = l_proc = l_entropy = torch.zeros((), device=device)
         if self.router is not None:
-            l_balance = self.router.balance_loss(cond_cont, cond_cat)
-            l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
-            l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
+            if self.spec.lambda_balance > 0:
+                l_balance = self.router.balance_loss(cond_cont, cond_cat)
+            if self.spec.lambda_proc > 0:
+                l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
+            if self.spec.lambda_entropy > 0:
+                l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
 
-        total = (
-            self.spec.lambda_weight * l_gen
-            + self.spec.n_sec_lambda * l_nsec
-            + self.particle_type_lambda * l_type
-        )
+        total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type
         if self.spec.lambda_balance > 0:
             total = total + self.spec.lambda_balance * l_balance
         if self.spec.lambda_proc > 0:
@@ -677,7 +645,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             "nsec_acc": nsec_acc,
         }
 
-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         if self.router is not None:
             self.router.gumbel_tau = _gumbel_tau(
                 global_step,
@@ -697,10 +665,8 @@ class FlowDDPMStageTrainer(StageTrainer):
         return stats
 
     @torch.no_grad()
-    def val_loss(self, batch: tuple, device: torch.device) -> dict:
-        return {
-            key: value.item() for key, value in self._compute(batch, device).items()
-        }
+    def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
+        return {key: value.item() for key, value in self._compute(batch, device).items()}
 
     # --- reporting ------------------------------------------------------
 
@@ -710,12 +676,20 @@ class FlowDDPMStageTrainer(StageTrainer):
     def summary(self, means: dict) -> str:
         return f"{self.name}[loss={means.get('loss', 0.0):.3f}]"
 
-    def val_objective(
-        self, train_means: dict, val_means: dict, marginal_kl: float
-    ) -> float:
+    def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
         return val_means.get("loss", 0.0)
 
 
+class _Stage2RealFakeBatch(NamedTuple):
+    """Subset of `StepBatch` that `_stage2_real_and_fake` needs."""
+
+    cond_cont: torch.Tensor
+    cond_cat: torch.Tensor
+    n_sec: torch.Tensor
+    sec_cont: torch.Tensor
+    sec_type_idx: torch.Tensor
+
+
 class WGANStageTrainer(StageTrainer):
     """WGAN-GP generator+critic for a single stage (see giant/model/wgan.py).
 
@@ -779,7 +753,7 @@ class WGANStageTrainer(StageTrainer):
         self.val_metrics = []
         self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")]
 
-    def _stage2_real_and_fake(self, batch_tensors, stage1_ctx, global_step, device):
+    def _stage2_real_and_fake(self, batch_tensors: _Stage2RealFakeBatch, stage1_ctx, global_step, device):
         """Build `(real, fake_raw, mask, critic_fn)` for stage 2, covering
         both decoders and all three particle-type targets. `fake_raw` still
         needs the caller's straight-through relaxation under
@@ -787,7 +761,7 @@ class WGANStageTrainer(StageTrainer):
         multiplied on the fake side yet."""
         cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
         B = cond_cont.size(0)
-        type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim)
+        type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes)
         slot_width = CONT_SLOT_DIM + type_dim
         k_max = sec_cont.size(1)
 
@@ -799,15 +773,8 @@ class WGANStageTrainer(StageTrainer):
 
         if self.decoder == "autoregressive":
             epoch = global_step // self.spec.steps_per_epoch
-            ar = self._ar_inputs(
-                cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch
-            )
-            real = (
-                self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(
-                    B, -1
-                )
-                * mask
-            )
+            ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
+            real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(B, -1) * mask
             z = torch.randn(B, k_max, self.model.noise_dim, device=device)
             fake_raw = self.model(
                 z,
@@ -826,7 +793,7 @@ class WGANStageTrainer(StageTrainer):
 
         return real, fake_raw, mask, critic_fn
 
-    def step(self, batch: tuple, device: torch.device, global_step: int) -> dict:
+    def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
         (
             cond_cont,
             cond_cat,
@@ -851,7 +818,7 @@ class WGANStageTrainer(StageTrainer):
             mask = None
         else:
             real, fake_raw, mask, critic_fn = self._stage2_real_and_fake(
-                (cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
+                _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
                 stage1_ctx,
                 global_step,
                 device,
@@ -873,7 +840,7 @@ class WGANStageTrainer(StageTrainer):
                     fake_raw,
                     sec_cont.size(1),
                     CONT_SLOT_DIM,
-                    stage2_type_dim(self.particle_type_cfg, self.particle_type_emb_dim),
+                    stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes),
                     tau,
                     grad_probe=grad_probe,
                 )
@@ -891,9 +858,7 @@ class WGANStageTrainer(StageTrainer):
 
         # --- generator (+ n_sec) step ---
         did_g_step = global_step % self.n_critic == 0
-        l_nsec, nsec_acc = self._n_sec_loss(
-            cond_cont, cond_cat, stage1_ctx, n_sec, device
-        )
+        l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
 
         # On a non-generator-step batch with no n_sec_head on this stage
         # (n_sec now defaults to stage 2), there's nothing for
@@ -902,9 +867,7 @@ class WGANStageTrainer(StageTrainer):
         skip_g_step = not did_g_step and self.model.n_sec_head is None
         if did_g_step:
             g_loss_adv = generator_loss(critic_fn, fake)
-            g_loss = (
-                self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
-            )
+            g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
         else:
             g_loss_adv = torch.zeros((), device=device)
             g_loss = self.spec.n_sec_lambda * l_nsec
@@ -941,14 +904,9 @@ class WGANStageTrainer(StageTrainer):
         return stats["d_loss"] + stats["g_loss"]
 
     def summary(self, means: dict) -> str:
-        return (
-            f"{self.name}[d={means.get('d_loss', 0.0):.3f} "
-            f"g={means.get('g_loss', 0.0):.3f}]"
-        )
+        return f"{self.name}[d={means.get('d_loss', 0.0):.3f} g={means.get('g_loss', 0.0):.3f}]"
 
-    def val_objective(
-        self, train_means: dict, val_means: dict, marginal_kl: float
-    ) -> float:
+    def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
         """No monotone per-batch WGAN loss fit for averaging, so
         best-checkpoint selection uses the real marginal-KL signal when
         `validate_marginals` produced one, and falls back to this epoch's own
@@ -1002,8 +960,7 @@ def build_stage_trainers(
         if spec.generator == "wgan":
             critic = critics.get(name)
             assert critic is not None, (
-                f"{name}_model.generator='wgan' requires a critic (see "
-                "giant.model.network.build_critics)"
+                f"{name}_model.generator='wgan' requires a critic (see giant.model.network.build_critics)"
             )
             trainers[name] = WGANStageTrainer(spec, model, critic, device)
         else:

giant/validate.py

@@ -8,9 +8,7 @@ from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 _SEC_PHYS_NAMES = ["log_mass", "charge"]
 
 
-def _histogram_kl(
-    p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
-) -> float:
+def _histogram_kl(p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8) -> float:
     """KL(P || Q) between two 1D samples, estimated via a shared histogram."""
     lo = min(p_samples.min(), q_samples.min())
     hi = max(p_samples.max(), q_samples.max())
@@ -33,9 +31,7 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
     return counts / total if total > 0 else counts
 
 
-def _categorical_kl(
-    real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8
-) -> float:
+def _categorical_kl(real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8) -> float:
     """KL(P_real || Q_gen) between two class-index samples over `n_classes`
     categories, estimated from bincount fractions. NaN if either side has no
     valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
@@ -48,9 +44,7 @@ def _categorical_kl(
     return float(np.sum(p * np.log(p / q)))
 
 
-def _embedding_nearest_class(
-    vectors: torch.Tensor, emb_weight: torch.Tensor
-) -> np.ndarray:
+def _embedding_nearest_class(vectors: torch.Tensor, emb_weight: torch.Tensor) -> np.ndarray:
     """Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
     (vocab, emb_dim) — same computation as
     `giant.particles.decode_embedding_nearest`, but returning the raw class
@@ -109,11 +103,7 @@ def validate_marginals(
         sec_decoder.eval()
 
     k_max = sec_decoder.k_max if sec_decoder is not None else 0
-    target = (
-        sec_decoder.particle_type_cfg.get("target", "physical")
-        if sec_decoder is not None
-        else "physical"
-    )
+    target = sec_decoder.particle_type_cfg.get("target", "physical") if sec_decoder is not None else "physical"
 
     all_real, all_gen = [], []
     all_n_sec_real, all_n_sec_pred = [], []
@@ -125,15 +115,12 @@ def validate_marginals(
     for i, batch in enumerate(val_loader):
         if n_batches is not None and i >= n_batches:
             break
-        # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
-        # sec_type_idx).
-        cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch
-        cond_cont = cond_cont.to(device)
-        cond_cat = cond_cat.to(device)
+        # batch is a StepBatch (giant.data.dataset).
+        x1, n_sec, sec_cont, sec_type_idx = batch.target_s1, batch.n_sec, batch.sec_cont, batch.sec_type_idx
+        cond_cont = batch.cond_cont.to(device)
+        cond_cat = batch.cond_cat.to(device)
 
-        gen, n_sec_pred = sample_stage1(
-            stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps
-        )
+        gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)
 
         all_real.append(x1.numpy())
         all_gen.append(gen.cpu().numpy())
@@ -141,9 +128,7 @@ def validate_marginals(
         if sec_decoder is None:
             continue
 
-        n_sec_pred = resolve_n_sec(
-            stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
-        )
+        n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
         n_sec_pred_np = n_sec_pred.cpu().numpy()
         n_sec_np = n_sec.numpy()
         all_n_sec_real.append(n_sec_np)
@@ -155,9 +140,7 @@ def validate_marginals(
         sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
             sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
         )
-        gen_frac = 1.0 / (
-            1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
-        )
+        gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
         gen_valid = sec_valid_pred.cpu().numpy()
 
         if target == "physical":
@@ -182,17 +165,9 @@ def validate_marginals(
     real = np.concatenate(all_real, axis=0)
     generated = np.concatenate(all_gen, axis=0)
 
-    kl_divergence = np.array(
-        [
-            _histogram_kl(real[:, j], generated[:, j], bins=kl_bins)
-            for j in range(real.shape[1])
-        ]
-    )
+    kl_divergence = np.array([_histogram_kl(real[:, j], generated[:, j], bins=kl_bins) for j in range(real.shape[1])])
 
-    header = (
-        f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} "
-        f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
-    )
+    header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
     print(f"\n{header}")
     print("-" * len(header))
     for j, name in enumerate(LOCAL_TARGET_NAMES):
@@ -215,10 +190,7 @@ def validate_marginals(
     n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
 
     energy_fraction_kl = np.full(k_max, np.nan)
-    print(
-        f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  "
-        f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
-    )
+    print(f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f}  mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}")
     n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
     print(n_sec_dist_header)
     print("-" * len(n_sec_dist_header))
@@ -240,10 +212,7 @@ def validate_marginals(
             continue
         kl = _histogram_kl(r, g, bins=kl_bins)
         energy_fraction_kl[j] = kl
-        print(
-            f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} "
-            f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}"
-        )
+        print(f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}")
 
     result.update(
         {
@@ -259,12 +228,7 @@ def validate_marginals(
         phys_gen = np.concatenate(all_phys_gen, axis=0)  # (M, 2)
 
         if len(phys_real) > 0 and len(phys_gen) > 0:
-            phys_kl = np.array(
-                [
-                    _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
-                    for j in range(2)
-                ]
-            )
+            phys_kl = np.array([_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)])
         else:
             phys_kl = np.full(2, np.nan)
 
@@ -278,21 +242,14 @@ def validate_marginals(
             if len(r) == 0 or len(g) == 0:
                 continue
             print(
-                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
-                f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
+                f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
             )
 
-        result.update(
-            {"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl}
-        )
+        result.update({"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl})
     else:
         type_class_real = np.concatenate(all_type_class_real, axis=0)
         type_class_gen = np.concatenate(all_type_class_gen, axis=0)
-        n_classes = (
-            sec_decoder.type_dim
-            if target == "onehot"
-            else sec_decoder.cond_enc.pdg_emb.weight.size(0)
-        )
+        n_classes = sec_decoder.type_dim if target == "onehot" else sec_decoder.cond_enc.pdg_emb.weight.size(0)
         type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)
 
         print(

pyproject.toml

@@ -23,6 +23,7 @@ cuda = [
 ]
 dev = [
     "pytest>=8,<10",
+    "pytest-cov>=5,<8",
     "ruff>=0.15,<1",
     "ty>=0.0.50,<0.1",
     "giant[convert,analysis,geometry,wandb]",
@@ -49,14 +50,27 @@ analysis = [
 
 [project.scripts]
 giant = "giant.cli:app"
-dwarf = "scripts.dwarf:app"
+dwarf = "giant.tools.dwarf:app"
+
+[tool.ruff]
+line-length = 120
+
+[tool.coverage.run]
+source = ["giant"]
+omit = ["*/legacy/*"]
+
+[tool.coverage.report]
+exclude_also = [
+    "if TYPE_CHECKING:",
+    "raise NotImplementedError",
+]
 
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"
 
 [tool.hatch.build.targets.wheel]
-packages = ["giant", "scripts"]
+packages = ["giant"]
 
 [tool.uv]
 conflicts = [

tests/legacy/network_v02_snapshot.py

@@ -37,11 +37,7 @@ class SinusoidalEmbedding(nn.Module):
         super().__init__()
         assert dim % 2 == 0, "dim must be even"
         half = dim // 2
-        freqs = torch.exp(
-            -math.log(10000)
-            * torch.arange(half, dtype=torch.float32)
-            / max(half - 1, 1)
-        )
+        freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
         self.register_buffer("freqs", freqs)
 
     def forward(self, t: torch.Tensor) -> torch.Tensor:
@@ -107,9 +103,7 @@ class ConditionEncoder(nn.Module):
             pdg_e = self.pdg_emb(cond_cat[:, 0])
             mat_e = self.mat_emb(cond_cat[:, 1])
         else:
-            particle_phys = cond_cont[
-                :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM
-            ]
+            particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
             material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
             pdg_e = self.particle_mlp(particle_phys)
             mat_e = self.material_mlp(material_phys)
@@ -168,12 +162,7 @@ class DenoisingMLP(nn.Module):
         )
         merged_cond_dim = time_dim + cond_out_dim
         self.input_proj = nn.Linear(x_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, x_dim)
         # Predicts n_sec as classification over {0, 1, ..., k_max}.
         # Applied to the condition encoding (not the diffused latent).
@@ -286,12 +275,7 @@ class SecondaryDecoder(nn.Module):
         )
         merged_cond_dim = time_dim + cond_out_dim
         self.input_proj = nn.Linear(sec_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, sec_dim)
 
     def forward(
@@ -345,12 +329,7 @@ class WGANGenerator(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(noise_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, x_dim)
         self.n_sec_head = nn.Sequential(
             nn.Linear(cond_out_dim, hidden_dim // 2),
@@ -410,12 +389,7 @@ class Critic(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(x_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_norm = nn.LayerNorm(hidden_dim)
         self.out_proj = nn.Linear(hidden_dim, 1)
 
@@ -466,12 +440,7 @@ class WGANSecondaryGenerator(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(noise_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, sec_dim)
 
     def forward(
@@ -515,12 +484,7 @@ class SecondaryCritic(nn.Module):
             conditioning=conditioning,
         )
         self.input_proj = nn.Linear(sec_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, cond_out_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_norm = nn.LayerNorm(hidden_dim)
         self.out_proj = nn.Linear(hidden_dim, 1)
 
@@ -550,9 +514,7 @@ class Router(nn.Module):
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         raise NotImplementedError
 
-    def combine_weights(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         probs = self.gate(cond_cont, cond_cat)
         if not (self.gumbel and self.training):
             return probs
@@ -562,26 +524,18 @@ class Router(nn.Module):
     def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         return self.gate(cond_cont, cond_cat).argmax(dim=-1)
 
-    def balance_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         importance = self.gate(cond_cont, cond_cat).sum(dim=0)  # (n_experts,)
         return (importance.std() / (importance.mean() + 1e-8)) ** 2
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         return torch.zeros((), device=cond_cont.device)
 
-    def entropy_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor
-    ) -> torch.Tensor:
+    def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         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]:
+    def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         gate = self.gate(cond_cont, cond_cat)  # (B, n_experts)
         row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1)  # (B,)
         norm_entropy = row_entropy.mean() / math.log(self.n_experts)
@@ -602,9 +556,7 @@ def register_router(name: str):
 
 def build_router(name: str, n_experts: int, **kwargs) -> Router:
     if name not in ROUTER_REGISTRY:
-        raise ValueError(
-            f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}"
-        )
+        raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
     cls = ROUTER_REGISTRY[name]
     accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
     filtered = {k: v for k, v in kwargs.items() if k in accepted}
@@ -644,8 +596,7 @@ class EnergyRouter(Router):
         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"
+                    f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
                 )
             self._width_lo = width_min_ratio * temperature
             self._width_hi = width_max_ratio * temperature
@@ -658,10 +609,7 @@ class EnergyRouter(Router):
             centers = torch.linspace(-2.0, 2.0, n_experts)
         else:
             if len(centers_init) != n_experts:
-                raise ValueError(
-                    f"centers_init has {len(centers_init)} values, "
-                    f"expected n_experts={n_experts}"
-                )
+                raise ValueError(f"centers_init has {len(centers_init)} values, expected n_experts={n_experts}")
             centers = torch.tensor(list(centers_init), dtype=torch.float32)
         if learn_centers:
             self.centers = nn.Parameter(centers)
@@ -702,9 +650,7 @@ class PdgRouter(Router):
 
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         e = self.pdg_emb(cond_cat[:, 0])  # (B, emb_dim)
-        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(
-            -1
-        )  # (B, n_experts)
+        d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1)  # (B, n_experts)
         return torch.softmax(-d2 / self.temperature, dim=-1)
 
 
@@ -736,9 +682,7 @@ class ProcessRouter(Router):
     def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
         return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
 
 
@@ -756,14 +700,10 @@ class ComposedRouter(Router):
         joint = self.routers[0].gate(cond_cont, cond_cat)  # (B, n_0)
         for router in self.routers[1:]:
             g = router.gate(cond_cont, cond_cat)  # (B, n_i)
-            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(
-                1
-            )  # (B, prod so far)
+            joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1)  # (B, prod so far)
         return joint
 
-    def classify_loss(
-        self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor
-    ) -> torch.Tensor:
+    def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
         total = torch.zeros((), device=cond_cont.device)
         for router in self.routers:
             total = total + router.classify_loss(cond_cont, cond_cat, labels)
@@ -796,12 +736,7 @@ class ExpertTrunk(nn.Module):
     ) -> None:
         super().__init__()
         self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList(
-            [
-                ResBlock(hidden_dim, merged_cond_dim, dropout=dropout)
-                for _ in range(n_blocks)
-            ]
-        )
+        self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)])
         self.out_proj = nn.Linear(hidden_dim, in_dim)
 
     def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
@@ -891,9 +826,7 @@ class RoutedDenoisingMLP(nn.Module):
         t_emb = self.time_emb(t)
         c_emb = self.cond_enc(cond_cont, cond_cat)
         cond = torch.cat([t_emb, c_emb], dim=-1)
-        return _route_forward(
-            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
-        )
+        return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)
 
     def predict_n_sec(
         self,
@@ -957,9 +890,7 @@ class RoutedSecondaryDecoder(nn.Module):
         t_emb = self.time_emb(t)
         c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out)
         cond = torch.cat([t_emb, c_emb], dim=-1)
-        return _route_forward(
-            self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training
-        )
+        return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training)
 
 
 _STAGE1_MODEL_KEYS = {
@@ -1006,9 +937,7 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]:
 _VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
 
 
-def _check_router_conditioning_compat(
-    router_types: list[str], conditioning: str
-) -> None:
+def _check_router_conditioning_compat(router_types: list[str], conditioning: str) -> None:
     bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
     if bad and conditioning == "physical":
         raise ValueError(
@@ -1019,9 +948,7 @@ def _check_router_conditioning_compat(
         )
 
 
-def _build_router_from_cfg(
-    router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding"
-) -> Router:
+def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding") -> Router:
     shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
     if router_cfg["type"] == "composed":
         axes = _parse_composed_axes(router_cfg)
@@ -1030,9 +957,7 @@ def _build_router_from_cfg(
         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")
-    }
+    router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
     router_kwargs.setdefault("pdg_vocab", pdg_vocab)
     router_kwargs.setdefault("mat_vocab", mat_vocab)
     router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
@@ -1042,15 +967,9 @@ def _build_router_from_cfg(
 
 def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
     if model_config.get("mode") == "wgan":
-        stage1 = WGANGenerator(
-            **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}
-        )
+        stage1 = WGANGenerator(**{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS})
         sec_decoder = WGANSecondaryGenerator(
-            **{
-                k: v
-                for k, v in model_config.items()
-                if k in _WGAN_SEC_GENERATOR_MODEL_KEYS
-            }
+            **{k: v for k, v in model_config.items() if k in _WGAN_SEC_GENERATOR_MODEL_KEYS}
         )
         return stage1, sec_decoder
 
@@ -1061,44 +980,30 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
         shared = dict(
             pdg_vocab=pdg_vocab,
             mat_vocab=mat_vocab,
-            expert_hidden_dim=model_config.get("expert_hidden_dim")
-            or model_config.get("hidden_dim", 128),
-            expert_n_blocks=model_config.get("expert_n_blocks")
-            or model_config.get("n_blocks", 3),
+            expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
+            expert_n_blocks=model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
             emb_dim=model_config.get("emb_dim", EMB_DIM),
             dropout=model_config.get("dropout", 0.1),
             conditioning=model_config.get("conditioning", "embedding"),
         )
         conditioning = shared["conditioning"]
         stage1 = RoutedDenoisingMLP(
-            router=_build_router_from_cfg(
-                router_cfg, pdg_vocab, mat_vocab, conditioning
-            ),
+            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, conditioning
-            ),
+            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
             **shared,
         )
         return stage1, sec_decoder
 
-    stage1 = DenoisingMLP(
-        **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}
-    )
-    sec_decoder = SecondaryDecoder(
-        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
+    stage1 = DenoisingMLP(**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS})
+    sec_decoder = SecondaryDecoder(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
     return stage1, sec_decoder
 
 
 def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]:
-    critic = Critic(
-        **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}
-    )
-    sec_critic = SecondaryCritic(
-        **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}
-    )
+    critic = Critic(**{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS})
+    sec_critic = SecondaryCritic(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS})
     return critic, sec_critic

tests/test_bump_dataset_version.py

@@ -1,7 +1,7 @@
 import os
 import subprocess
 
-from scripts import bump_dataset_version
+from giant.tools import bump_dataset_version
 
 plan_bump_gen = bump_dataset_version.plan_bump_gen
 plan_bump_schema = bump_dataset_version.plan_bump_schema
@@ -22,9 +22,7 @@ def test_git_user_name_returns_none_on_timeout(monkeypatch):
 
 
 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"
-    )
+    dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
     assert dirs == [
         tmp_path / "raw" / "steps" / "gen1",
         tmp_path / "processed" / "steps" / "gen1" / "schema1",
@@ -56,9 +54,7 @@ def test_bump_gen_kinds_are_independent(tmp_path):
 
 def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
-    dirs, log_line = plan_bump_schema(
-        tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01"
-    )
+    dirs, log_line = plan_bump_schema(tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01")
     assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"]
     assert "`gen1`/`schema1`" in log_line
 
@@ -66,18 +62,14 @@ def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
 def test_bump_schema_increments_within_its_gen(tmp_path):
     (tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
     (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(
-        tmp_path, "steps", "gen1", "next schema", None, "2026-01-01"
-    )
+    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
     assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]
 
 
 def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
     (tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
     (tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
-    dirs, _ = plan_bump_schema(
-        tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01"
-    )
+    dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
     assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
 
 
@@ -91,9 +83,7 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
 
 def test_bump_gen_to_specific_tag(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
-    dirs, log_line = plan_bump_gen(
-        tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5"
-    )
+    dirs, log_line = plan_bump_gen(tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5")
     assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
     assert "`gen5`" in log_line
 
@@ -125,9 +115,7 @@ def test_bump_schema_to_specific_tag(tmp_path):
 def test_bump_schema_rejects_invalid_to_tag(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
     try:
-        plan_bump_schema(
-            tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
-        )
+        plan_bump_schema(tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3")
         assert False, "expected SystemExit"
     except SystemExit:
         pass
@@ -164,15 +152,7 @@ def _make_parquet(path):
 
 
 def test_update_manifest_bumps_to_specified_schema(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -194,15 +174,7 @@ def test_update_manifest_auto_detects_highest_schema(tmp_path):
     for schema in ("schema1", "schema2", "schema3"):
         d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
         d.mkdir(parents=True)
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
     parquet.touch()
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -231,15 +203,7 @@ def test_update_manifest_reports_missing_targets(tmp_path):
 
 
 def test_update_manifest_skips_already_at_target(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -254,15 +218,7 @@ def test_update_manifest_skips_already_at_target(tmp_path):
 
 
 def test_update_manifest_preserves_comments_and_blanks(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -278,15 +234,7 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
 
 
 def test_update_manifest_bumps_gen(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen2"
-        / "schema1"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -303,15 +251,7 @@ def test_update_manifest_bumps_gen(tmp_path):
 
 
 def test_update_manifest_bumps_gen_and_schema(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen2"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -328,15 +268,7 @@ def test_update_manifest_bumps_gen_and_schema(tmp_path):
 
 
 def test_apply_update_manifest_writes_file(tmp_path):
-    parquet = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
     _make_parquet(parquet)
 
     manifest_dir = tmp_path / "pools" / "pbwo4"
@@ -358,24 +290,8 @@ def test_apply_update_manifest_writes_file(tmp_path):
 
 
 def test_create_manifest_writes_relative_paths(tmp_path):
-    pq1 = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
-    pq2 = (
-        tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema2"
-        / "pbwo4"
-        / "shard-001.parquet"
-    )
+    pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
+    pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
     _make_parquet(pq1)
     _make_parquet(pq2)
 
@@ -419,9 +335,7 @@ def test_run_create_manifest_refuses_to_overwrite_existing_output(tmp_path):
     output.write_text("original contents\n")
 
     try:
-        bump_dataset_version.run_create_manifest(
-            [str(pq)], execute=True, output=str(output)
-        )
+        bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output))
         assert False, "expected SystemExit"
     except SystemExit:
         pass
@@ -435,9 +349,7 @@ def test_run_create_manifest_force_overwrites_existing_output(tmp_path):
     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
-    )
+    bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output), force=True)
     assert output.read_text() != "original contents\n"
 
 

tests/test_catalog.py

@@ -13,9 +13,7 @@ from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 def _build_ctx() -> Context:
     r, t = _rollout_frame(), _reference_frame()
-    return build_context(
-        r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
-    )
+    return build_context(r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
 
 
 @pytest.fixture(scope="module")
@@ -142,10 +140,7 @@ def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
 
     # 4 chunks over only 2 distinct event_ids also exercises empty chunks.
     n_chunks = 4 if spec.chunkable else 1
-    parts = [
-        spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks)))
-        for k in range(n_chunks)
-    ]
+    parts = [spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
     chunked = spec.finalize(parts, ctx)
 
     assert chunked.id == unchunked.id

tests/test_checkpoint_io.py

@@ -0,0 +1,280 @@
+"""Tests for giant.checkpoint_io.load_for_inference (issues.md Issue 5) —
+the shared bootstrap `giant predict`/`giant rollout` use to go from a
+checkpoint path to ready-to-run models."""
+
+from __future__ import annotations
+
+import copy
+
+import numpy as np
+import pytest
+import torch
+
+from giant import config as gconfig
+from giant.checkpoint_io import (
+    CheckpointCompatibilityError,
+    InferenceContext,
+    conditioning_axes,
+    load_for_inference,
+    stage_cfg,
+)
+from giant.data.loader import TopNMap
+from giant.data.setup_cache import topnmap_to_json
+from giant.data.transforms import Normalizer
+from giant.model.network import build_models
+
+PDG_MAP = {11: 0, 22: 1, -11: 2}
+MAT_MAP = {"G4_PbWO4": 0, "G4_AIR": 1}
+
+
+def _model_cfg(stage2_active: bool = True) -> dict:
+    """DEFAULT_CONFIG-derived, shrunk for speed — same pattern as
+    tests/test_network.py::_minimal_model_config. Default `conditioning`
+    (both axes "physical") needs no top-N vocab map, so this is a cheap,
+    fully self-contained happy-path config."""
+    cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
+    cfg["conditioning"]["particle"]["emb_dim"] = 4
+    cfg["conditioning"]["material"]["emb_dim"] = 4
+    cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
+    cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
+    cfg["stage2_model"]["active"] = stage2_active
+    return {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def _norms() -> tuple[Normalizer, Normalizer, Normalizer]:
+    rng = np.random.default_rng(0)
+    cond = Normalizer().fit(rng.standard_normal((100, 15)).astype(np.float32))
+    tgt = Normalizer().fit(rng.standard_normal((100, 9)).astype(np.float32))
+    sec_phys = Normalizer().fit(rng.standard_normal((100, 2)).astype(np.float32))
+    return cond, tgt, sec_phys
+
+
+def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overrides):
+    cfg = model_cfg if model_cfg is not None else _model_cfg()
+    built = build_models(cfg)
+    stage1, stage2 = built["stage1"], built["stage2"]
+    cond, tgt, sec_phys = _norms()
+    ckpt: dict = {
+        "model_config": cfg,
+        "model": stage1.state_dict() if stage1 is not None else {},
+        "sec_decoder": stage2.state_dict() if stage2 is not None else {},
+        "pdg_map": PDG_MAP,
+        "mat_map": MAT_MAP,
+        "normalizer": {"cond": cond.to_dict(), "target": tgt.to_dict(), "sec_phys": sec_phys.to_dict()},
+        "epoch": 3,
+        "best_val_loss": 0.5,
+    }
+    if ema:
+        ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
+        ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
+    # DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
+    # "onehot", and giant train's pipeline (gitea #29) now always writes a
+    # sec_type_topn_map in that case — default one in here too, unless a
+    # test explicitly overrides it, so fixtures represent a real, loadable
+    # checkpoint by default rather than exercising the "missing" guard by
+    # accident.
+    particle_type_target = cfg.get("stage2_model", {}).get("particle_type", {}).get("target", "onehot")
+    if particle_type_target == "onehot" and "sec_type_topn_map" not in ckpt_overrides:
+        default_sec_type_topn = TopNMap(class_map=dict(zip(PDG_MAP, range(len(PDG_MAP)))), other_members={})
+        ckpt["sec_type_topn_map"] = topnmap_to_json(default_sec_type_topn)
+    ckpt.update(ckpt_overrides)
+    path = tmp_path / "ckpt.pt"
+    torch.save(ckpt, path)
+    return path
+
+
+def _onehot_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"
+    return cfg
+
+
+# ---------------------------------------------------------------------------
+# Happy path
+# ---------------------------------------------------------------------------
+
+
+def test_happy_path_returns_populated_context(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert isinstance(ctx, InferenceContext)
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+    assert not ctx.stage1.training
+    assert not ctx.stage2.training
+    assert next(ctx.stage1.parameters()).device == torch.device("cpu")
+    assert ctx.pdg_map == PDG_MAP
+    assert ctx.mat_map == MAT_MAP
+    assert all(isinstance(k, int) for k in ctx.pdg_map)
+    assert all(isinstance(k, str) for k in ctx.mat_map)
+    assert ctx.particle_conditioning == "physical"
+    assert ctx.material_conditioning == "physical"
+    assert ctx.k_max == 3
+    assert ctx.epoch == 3
+    assert ctx.best_val_loss == 0.5
+    assert ctx.model_config["stage1_model"]["hidden_dim"] == 8
+
+
+def test_happy_path_normalizer_values_round_trip(tmp_path):
+    cond, tgt, sec_phys = _norms()
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+    assert ctx.cond_norm.mean is not None and cond.mean is not None
+    assert ctx.tgt_norm.mean is not None and tgt.mean is not None
+    assert ctx.sec_phys_norm.mean is not None and sec_phys.mean is not None
+    np.testing.assert_allclose(ctx.cond_norm.mean, cond.mean)
+    np.testing.assert_allclose(ctx.tgt_norm.mean, tgt.mean)
+    np.testing.assert_allclose(ctx.sec_phys_norm.mean, sec_phys.mean)
+
+
+# ---------------------------------------------------------------------------
+# Guards
+# ---------------------------------------------------------------------------
+
+
+def test_missing_model_config_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["model_config"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no model_config"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_decoder_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_decoder"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no sec_decoder"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_missing_sec_phys_normalizer_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["normalizer"]["sec_phys"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no normalizer.sec_phys"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_without_topn_map_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_onehot_model_cfg())
+
+    with pytest.raises(CheckpointCompatibilityError, match="pdg_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
+    topn = TopNMap(class_map={11: 0, 22: 1}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+    )
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.particle_conditioning == "onehot"
+    assert ctx.pdg_topn_map is not None
+    assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}
+
+
+def test_onehot_particle_type_target_without_sec_type_topn_map_raises(tmp_path):
+    """DEFAULT_CONFIG's stage2_model.particle_type.target="onehot" needs a
+    sec_type_topn_map (gitea #29) — a checkpoint with neither key at all
+    (not even the pre-#29 pdg_topn_map to fall back to) must fail loudly."""
+    checkpoint = _write_checkpoint(tmp_path, sec_type_topn_map=None)
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    with pytest.raises(CheckpointCompatibilityError, match="sec_type_topn_map"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict")
+
+
+def test_pre_gitea_29_checkpoint_falls_back_to_pdg_topn_map_for_sec_type(tmp_path):
+    """A checkpoint written before gitea #29 has no sec_type_topn_map key at
+    all — conditioning and secondary-type onehot maps were always the same
+    map, saved once under pdg_topn_map. load_for_inference must reproduce
+    that exact pre-#29 behavior for such a checkpoint."""
+    topn = TopNMap(class_map={11: 0, 22: 1, -11: 2}, other_members={})
+    checkpoint = _write_checkpoint(
+        tmp_path,
+        model_cfg=_onehot_model_cfg(),
+        pdg_topn_map=topnmap_to_json(topn),
+        sec_type_topn_map=None,
+    )
+    ckpt = torch.load(checkpoint, weights_only=False)
+    del ckpt["sec_type_topn_map"]
+    torch.save(ckpt, checkpoint)
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.sec_type_topn_map is not None
+    assert ctx.sec_type_topn_map.class_map == {11: 0, 22: 1, -11: 2}
+
+
+def test_ema_weights_requested_but_missing_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=False)
+
+    with pytest.raises(CheckpointCompatibilityError, match="no EMA weights"):
+        load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+
+
+def test_ema_weights_requested_and_present_succeeds(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, ema=True)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+
+
+@pytest.mark.parametrize("command_name", ["predict", "rollout"])
+def test_inactive_stage_with_require_stage2_raises_with_command_name(tmp_path, command_name):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    with pytest.raises(CheckpointCompatibilityError, match=f"{command_name} needs both"):
+        load_for_inference(checkpoint, torch.device("cpu"), command_name)
+
+
+def test_inactive_stage_with_require_stage2_false_succeeds_with_stage2_none(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
+
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", require_stage2=False)
+    assert ctx.stage1 is not None
+    assert ctx.stage2 is None
+
+
+# ---------------------------------------------------------------------------
+# conditioning_axes / stage_cfg
+# ---------------------------------------------------------------------------
+
+
+def test_conditioning_axes_v02_flat_string_applies_to_both_axes():
+    assert conditioning_axes({"conditioning": "embedding"}) == ("embedding", "embedding")
+
+
+def test_conditioning_axes_v03_nested_dict_independent_per_axis():
+    model_cfg = {"conditioning": {"particle": {"type": "onehot"}, "material": {"type": "physical"}}}
+    assert conditioning_axes(model_cfg) == ("onehot", "physical")
+
+
+def test_conditioning_axes_missing_key_uses_default():
+    assert conditioning_axes({}, default="embedding") == ("embedding", "embedding")
+
+
+def test_stage_cfg_new_shape_returns_subdict():
+    model_cfg = {"stage2_model": {"k_max": 7}}
+    assert stage_cfg(model_cfg, "stage2") == {"k_max": 7}
+
+
+def test_stage_cfg_v02_flat_shape_returns_empty_dict():
+    model_cfg = {"hidden_dim": 32, "n_blocks": 4}
+    assert stage_cfg(model_cfg, "stage2") == {}

tests/test_cli_new_run.py

@@ -101,9 +101,7 @@ def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
     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"]
-    )
+    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()
 

tests/test_cli_predict.py

@@ -1,13 +1,18 @@
 import uuid
 
+import torch
 import yaml
+from typer.testing import CliRunner
 
 from giant.cli import (
     _CEPH_PREDICTIONS,
     _resolve_prediction_output,
     _write_prediction_ref,
+    app,
 )
 
+runner = CliRunner()
+
 
 # ---------------------------------------------------------------------------
 # _resolve_prediction_output
@@ -116,9 +121,7 @@ def test_ref_yaml_includes_comment_when_provided(tmp_path):
     dataset = tmp_path / "full.manifest"
     pred_uuid = str(uuid.uuid4())
 
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3")
     data = yaml.safe_load(ref_path.read_text())
 
     assert data["comment"] == "baseline sweep run 3"
@@ -133,9 +136,7 @@ def test_ref_timestamp_is_iso_format(tmp_path):
     checkpoint.touch()
 
     pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
     data = yaml.safe_load(ref_path.read_text())
 
     # Must parse without error and be timezone-aware (UTC).
@@ -150,9 +151,24 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
     checkpoint.touch()
 
     pred_uuid = str(uuid.uuid4())
-    ref_path = _write_prediction_ref(
-        checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
-    )
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
     data = yaml.safe_load(ref_path.read_text())
 
     assert data["checkpoint"].startswith("/")
+
+
+# ---------------------------------------------------------------------------
+# Bootstrap failure surfaces via the CLI (issues.md Issue 5 — confirms
+# CheckpointCompatibilityError -> typer.Exit(1) actually wires up end-to-end,
+# not just at the giant.checkpoint_io unit level).
+# ---------------------------------------------------------------------------
+
+
+def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
+    checkpoint = tmp_path / "bad.pt"
+    torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
+
+    result = runner.invoke(app, ["predict", "dummy.parquet", "--checkpoint", str(checkpoint)])
+
+    assert result.exit_code == 1
+    assert "checkpoint has no model_config" in result.output

tests/test_cli_rollout.py

@@ -0,0 +1,33 @@
+"""Thin CLI smoke coverage for `giant rollout` (issues.md Issue 5) — confirms
+the CheckpointCompatibilityError raised by giant.checkpoint_io.load_for_inference
+surfaces as a clean typer.Exit(1) with the expected message, end-to-end
+through the CLI, not just at the giant.checkpoint_io unit level."""
+
+from __future__ import annotations
+
+import torch
+from typer.testing import CliRunner
+
+from giant.cli import app
+
+runner = CliRunner()
+
+
+def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
+    checkpoint = tmp_path / "bad.pt"
+    torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "checkpoint has no model_config" in result.output

tests/test_cli_train_overrides.py

@@ -41,6 +41,10 @@ def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):
 
 
 def test_stage2_only_knobs(monkeypatch, tmp_path):
+    # --stage2-stage1-context is exercised separately at the overrides-dict
+    # level (test_overrides_from_flags_stage2_only_knobs in test_config.py):
+    # its only non-default value, "sampled", is rejected by validate_config
+    # (issues.md Issue 1), so it can't appear in a full CLI invocation here.
     cfg = _invoke_and_capture_cfg(
         monkeypatch,
         tmp_path,
@@ -53,15 +57,12 @@ def test_stage2_only_knobs(monkeypatch, tmp_path):
             "32",
             "--stage2-context-dim",
             "16",
-            "--stage2-stage1-context",
-            "sampled",
         ],
     )
     assert cfg["stage2_model"]["decoder"] == "one_shot"
     assert cfg["stage2_model"]["k_max"] == 8
     assert cfg["stage2_model"]["hidden_dim"] == 32
     assert cfg["stage2_model"]["context_dim"] == 16
-    assert cfg["stage2_model"]["stage1_context"] == "sampled"
     # untouched stage1 defaults
     assert cfg["stage1_model"]["hidden_dim"] == 256
 
@@ -94,3 +95,83 @@ def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
     assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
     assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
     assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
+
+
+def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
+    monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
+    )
+    assert result.exit_code == 1
+    assert "--batch-size must be an integer or 'auto'" in result.output
+
+
+def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    resume_dir = tmp_path / "resumed_run"
+    resume_dir.mkdir()
+    (resume_dir / "last.pt").touch()
+    explicit_out = tmp_path / "explicit_run"
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(explicit_out), "--resume", str(resume_dir / "last.pt")],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == explicit_out
+
+
+def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+
+    resume_dir = tmp_path / "resumed_run"
+    resume_dir.mkdir()
+    (resume_dir / "last.pt").touch()
+
+    result = runner.invoke(cli.app, ["train", "dummy.parquet", "--resume", str(resume_dir / "last.pt")])
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == resume_dir
+
+
+def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
+        captured["out_dir"] = out_dir
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.chdir(tmp_path)
+
+    result = runner.invoke(cli.app, ["train", "dummy.parquet"])
+    assert result.exit_code == 0, result.output
+    assert captured["out_dir"] == Path("checkpoints") / cli.gconfig.default_out_dir_name(cli.gconfig.DEFAULT_CONFIG)
+
+
+def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
+    captured: dict = {}
+
+    def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
+        captured["batch_size"] = cfg["train"]["batch_size"]
+
+    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
+
+    result = runner.invoke(
+        cli.app,
+        ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "auto"],
+    )
+    assert result.exit_code == 0, result.output
+    assert captured["batch_size"] == 123
+    assert "batch_size: 123 (auto-estimated from free GPU memory)" in result.output

tests/test_condor.py

@@ -62,9 +62,7 @@ def _fake_venv(repo_dir: Path) -> None:
     giant.chmod(0o755)
 
 
-def _prep(
-    rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1
-) -> Path:
+def _prep(rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1) -> Path:
     """``prep`` with small test-sized context bins/sampling."""
     return prep(
         rollout_yaml,
@@ -224,16 +222,12 @@ 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, monkeypatch: pytest.MonkeyPatch
-):
+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")
-    )
+    monkeypatch.setattr(sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python"))
     with pytest.raises(FileNotFoundError, match="uv sync"):
         write_submit(cfg)
 
@@ -241,9 +235,7 @@ def test_write_submit_requires_synced_venv(
 def test_write_submit_remote_flag(tmp_path: Path):
     run_dir = _prep(_write_inputs(tmp_path))
     _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True
-    )
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True)
     txt = write_submit(cfg).read_text()
     assert "+RemoteJob = True" in txt
     assert "ProvidesETPResources" not in txt
@@ -253,9 +245,7 @@ def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
     assert get_spec("router_gating").chunkable is False
     run_dir = _prep(_write_inputs(tmp_path), chunks=4)
     _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
-    )
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
     write_submit(cfg)
     jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
     counts: dict[str, int] = {}
@@ -272,9 +262,7 @@ def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
     _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
-    )
+    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)
 
@@ -297,16 +285,9 @@ def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
     run_dir = _prep(_write_inputs(tmp_path), chunks=2)
     meta = RunMeta.load(run_dir / "run_meta.json")
     _fake_venv(tmp_path)
-    cfg = SubmitConfig(
-        run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2
-    )
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2)
     write_submit(cfg)
-    jobs = {
-        (i, int(k)): int(w)
-        for i, k, w in (
-            line.split(",") for line in (run_dir / "jobs.txt").read_text().split()
-        )
-    }
+    jobs = {(i, int(k)): int(w) for i, k, w in (line.split(",") for line in (run_dir / "jobs.txt").read_text().split())}
     for chunk in range(2):
         expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
         assert jobs[("marginal_edep", chunk)] == expected

tests/test_config.py

@@ -27,6 +27,95 @@ def test_conditioning_enum_has_onehot():
     }
 
 
+# ---------------------------------------------------------------------------
+# Config dataclasses (issues.md Issue 1)
+# ---------------------------------------------------------------------------
+
+
+def test_giant_config_to_dict_matches_default_config():
+    """DEFAULT_CONFIG is generated from GiantConfig().to_dict() (not
+    hand-maintained), so the two cannot structurally drift apart — but this
+    pins the *equality* too, catching e.g. a stray in-place mutation of
+    DEFAULT_CONFIG added elsewhere after import."""
+    assert gconfig.GiantConfig().to_dict() == gconfig.DEFAULT_CONFIG
+
+
+@pytest.mark.parametrize(
+    "cls",
+    [
+        gconfig.ConditioningAxisConfig,
+        gconfig.ConditioningConfig,
+        gconfig.FlowConfig,
+        gconfig.DdpmConfig,
+        gconfig.Stage1WganConfig,
+        gconfig.Stage2WganConfig,
+        gconfig.RouterConfig,
+        gconfig.Stage2RouterConfig,
+        gconfig.NSecConfig,
+        gconfig.ParticleTypeConfig,
+        gconfig.AutoregressiveConfig,
+        gconfig.Stage1ModelConfig,
+        gconfig.Stage2ModelConfig,
+        gconfig.TrainConfig,
+        gconfig.GiantConfig,
+    ],
+)
+def test_config_dataclass_from_dict_round_trips_through_to_dict(cls):
+    assert cls.from_dict(cls().to_dict()) == cls()
+    assert cls.from_dict(None) == cls()
+
+
+def test_stage2_model_config_defaults_match_documented_v030_intent():
+    """The two keys issues.md Issue 1 found drifted between DEFAULT_CONFIG
+    and build_models/StageSpec.from_config's own .get(key, default)
+    fallbacks — pinned directly against the dataclass that is now their
+    shared single source of truth."""
+    spec = gconfig.Stage2ModelConfig()
+    assert spec.decoder == "autoregressive"
+    assert spec.particle_type.target == "onehot"
+
+
+def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
+    """gitea #29: n_classes=0 means "inherit conditioning.particle.emb_dim"
+    — the default must stay 0 so an existing config.toml with no
+    stage2_model.particle_type.n_classes key reproduces pre-#29 behavior."""
+    assert gconfig.ParticleTypeConfig().n_classes == 0
+    spec = gconfig.ParticleTypeConfig.from_dict({"n_classes": 32})
+    assert spec.n_classes == 32
+    assert spec.to_dict()["n_classes"] == 32
+
+
+def test_router_config_extra_round_trips_composed_axis_keys():
+    d = {"enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4}
+    router = gconfig.RouterConfig.from_dict(d)
+    assert router.enabled is True
+    assert router.extra == {"axis0_type": "energy", "axis0_n_experts": 4}
+    assert router.to_dict()["axis0_type"] == "energy"
+
+
+def test_stage2_router_config_tie_to_stage1_not_leaked_into_extra():
+    router = gconfig.Stage2RouterConfig.from_dict({"tie_to_stage1": True})
+    assert router.tie_to_stage1 is True
+    assert "tie_to_stage1" not in router.extra
+
+
+def test_stage1_router_config_has_no_tie_to_stage1_key():
+    """Stage 1's router schema must not gain stage 2's tie_to_stage1 key —
+    that would change every future run's saved config.toml shape."""
+    assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict()
+
+
+def test_n_sec_config_owner_defaults_to_stage2():
+    n_sec = gconfig.NSecConfig()
+    assert n_sec.owner == "stage2"
+
+
+def test_n_sec_config_owner_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
+    assert n_sec.owner == "stage1"
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1"}
+
+
 # ---------------------------------------------------------------------------
 # _deep_merge
 # ---------------------------------------------------------------------------
@@ -112,9 +201,7 @@ def test_migrate_config_lambda_nsec_and_lambda_s2():
 
 
 def test_migrate_config_wgan_knobs_map_to_both_stages():
-    new = gconfig.migrate_config(
-        {"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}}
-    )
+    new = gconfig.migrate_config({"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}})
     for stage in ("stage1_model", "stage2_model"):
         assert new[stage]["wgan"]["n_critic"] == 3
         assert new[stage]["wgan"]["gp_weight"] == 5.0
@@ -122,9 +209,7 @@ def test_migrate_config_wgan_knobs_map_to_both_stages():
 
 
 def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
-    new = gconfig.migrate_config(
-        {"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}}
-    )
+    new = gconfig.migrate_config({"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}})
     for stage in ("stage1_model", "stage2_model"):
         assert new[stage]["hidden_dim"] == 128
         assert new[stage]["n_res_blocks"] == 4
@@ -132,9 +217,7 @@ def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages():
 
 
 def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes():
-    new = gconfig.migrate_config(
-        {"model": {"emb_dim": 32, "conditioning": "embedding"}}
-    )
+    new = gconfig.migrate_config({"model": {"emb_dim": 32, "conditioning": "embedding"}})
     for axis in ("particle", "material"):
         assert new["conditioning"][axis]["emb_dim"] == 32
         assert new["conditioning"][axis]["type"] == "embedding"
@@ -181,11 +264,7 @@ def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false():
 
 
 def test_migrate_config_router_nonzero_expert_dims_raises():
-    cfg = {
-        "model": {
-            "router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}
-        }
-    }
+    cfg = {"model": {"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}}}
     try:
         gconfig.migrate_config(cfg)
         assert False, "expected ValueError"
@@ -267,9 +346,7 @@ def test_merge_cli_overrides_nested_override_keeps_siblings():
     assert cfg["stage1_model"]["hidden_dim"] == 256  # untouched sibling section
 
 
-def test_merge_cli_overrides_file_then_explicit_override_precedence(
-    tmp_path, monkeypatch
-):
+def test_merge_cli_overrides_file_then_explicit_override_precedence(tmp_path, monkeypatch):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123")
     path = tmp_path / "config.toml"
     _write_toml(
@@ -317,9 +394,7 @@ def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, c
     assert "current999" in captured.err
 
 
-def test_merge_cli_overrides_no_warning_on_matching_git_hash(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_on_matching_git_hash(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
     path = tmp_path / "config.toml"
     _write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n")
@@ -328,9 +403,7 @@ def test_merge_cli_overrides_no_warning_on_matching_git_hash(
     assert capsys.readouterr().err == ""
 
 
-def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_when_git_hash_unknown(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown")
     path = tmp_path / "config.toml"
     _write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n")
@@ -339,9 +412,7 @@ def test_merge_cli_overrides_no_warning_when_git_hash_unknown(
     assert capsys.readouterr().err == ""
 
 
-def test_merge_cli_overrides_no_warning_when_meta_section_absent(
-    tmp_path, monkeypatch, capsys
-):
+def test_merge_cli_overrides_no_warning_when_meta_section_absent(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
     path = tmp_path / "config.toml"
     path.write_text("[train]\nepochs = 5\n")
@@ -351,12 +422,8 @@ def test_merge_cli_overrides_no_warning_when_meta_section_absent(
 
 
 def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch):
-    monkeypatch.setattr(
-        gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243"
-    )
-    cfg = gconfig.merge_cli_overrides(
-        gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {}
-    )
+    monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {})
     assert cfg["stage1_model"]["generator"] == "flow"
     assert cfg["stage1_model"]["hidden_dim"] == 256
     assert cfg["stage2_model"]["hidden_dim"] == 256
@@ -431,10 +498,7 @@ def _cfg_with(**dotted_overrides):
 
 
 def test_default_out_dir_name_all_defaults_is_just_the_timestamp():
-    assert (
-        gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW)
-        == "20260729_1430"
-    )
+    assert gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) == "20260729_1430"
 
 
 def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix():
@@ -499,9 +563,7 @@ def test_default_out_dir_name_router_gumbel_shown_when_enabled():
             "stage1_model.router.gumbel": True,
         }
     )
-    assert (
-        gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
-    )
+    assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum"
 
 
 def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
@@ -524,9 +586,7 @@ def test_default_out_dir_name_overflow_caps_and_hashes_remainder():
     # First 6 by priority: stage1_generator, stage2_generator, stage2_decoder,
     # stage2_history, particle_type_target, stage1_router — stage2_router
     # overflows into the hash suffix.
-    assert name.startswith(
-        "20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+"
-    )
+    assert name.startswith("20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+")
 
 
 def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive():
@@ -637,6 +697,15 @@ def test_validate_config_stop_token_not_implemented():
         assert "stop_token" in str(e)
 
 
+def test_validate_config_stage1_context_sampled_not_implemented():
+    cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e)
+
+
 def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
     """'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint' —
     both stages active means giant rollout
@@ -688,6 +757,31 @@ def test_validate_config_ar_default_markov_always_passes():
     gconfig.validate_config(cfg)  # must not raise
 
 
+def test_validate_config_ar_order_energy_desc_passes():
+    """'energy_desc' is the only implemented order — must not raise."""
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.order": "energy_desc",
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_ar_order_invalid_value_rejected():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.decoder": "autoregressive",
+            "stage2_model.autoregressive.order": "energy_asc",
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "order" in str(e)
+
+
 def test_validate_config_ar_history_attention_passes():
     """v0.3.0 step 7 implements history='attention' — must not raise."""
     cfg = _cfg_with(
@@ -741,11 +835,12 @@ def test_validate_config_ar_teacher_forcing_invalid_value_rejected():
 
 
 def test_validate_config_ar_checks_skipped_under_one_shot():
-    """history/teacher_forcing values that would fail under AR are irrelevant
-    (and unchecked) when decoder='one_shot'."""
+    """order/history/teacher_forcing values that would fail under AR are
+    irrelevant (and unchecked) when decoder='one_shot'."""
     cfg = _cfg_with(
         **{
             "stage2_model.decoder": "one_shot",
+            "stage2_model.autoregressive.order": "bogus",
             "stage2_model.autoregressive.history": "attention",
             "stage2_model.autoregressive.teacher_forcing": "scheduled",
         }
@@ -753,20 +848,243 @@ def test_validate_config_ar_checks_skipped_under_one_shot():
     gconfig.validate_config(cfg)  # must not raise
 
 
+# ---------------------------------------------------------------------------
+# validate_config_keys / merge_cli_overrides unknown-key rejection
+# ---------------------------------------------------------------------------
+
+
+def test_validate_config_keys_default_config_passes():
+    gconfig.validate_config_keys(gconfig.DEFAULT_CONFIG)  # must not raise
+
+
+def test_validate_config_keys_rejects_unknown_top_level_key():
+    cfg = _cfg_with(**{"bogus_section.foo": 1})
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "bogus_section" in str(e)
+
+
+def test_validate_config_keys_rejects_unknown_nested_key_with_close_match_hint():
+    cfg = _cfg_with(**{"stage1_model.n_res_block": 12})
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.n_res_block" in str(e)
+        assert "n_res_blocks" in str(e)
+
+
+def test_validate_config_keys_allows_composed_router_axis_keys():
+    cfg = _cfg_with(
+        **{
+            "stage1_model.router.enabled": True,
+            "stage1_model.router.type": "composed",
+            "stage1_model.router.axis0_type": "energy",
+            "stage1_model.router.axis0_n_experts": 4,
+            "stage1_model.router.axis1_type": "pdg",
+            "stage1_model.router.axis1_emb_dim": 8,
+        }
+    )
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_allows_centers_init():
+    cfg = _cfg_with(**{"stage1_model.router.centers_init": [-1.0, 0.0, 1.0]})
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_validate_config_keys_rejects_unrelated_unknown_router_key():
+    cfg = _cfg_with(**{"stage1_model.router.n_expert": 4})  # typo for n_experts
+    try:
+        gconfig.validate_config_keys(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_model.router.n_expert" in str(e)
+        assert "n_experts" in str(e)
+
+
+def test_validate_config_keys_skips_meta_section():
+    cfg = _cfg_with()
+    cfg["meta"] = {"config_version": 3, "git_hash": "abc123"}
+    gconfig.validate_config_keys(cfg)  # must not raise
+
+
+def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path):
+    path = tmp_path / "config.toml"
+    path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n")
+    try:
+        gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {})
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "n_res_block" in str(e)
+
+
+def test_merge_cli_overrides_rejects_typo_in_cli_overrides():
+    try:
+        gconfig.merge_cli_overrides(
+            gconfig.DEFAULT_CONFIG,
+            None,
+            {"stage1_model": {"n_res_block": 12}},
+        )
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "n_res_block" in str(e)
+
+
+@pytest.mark.parametrize("fixture_name", ["default.toml", "wgan_h128_b4_physical.toml"])
+def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_name, monkeypatch):
+    monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243")
+    gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / fixture_name, {})  # must not raise
+
+
+# ---------------------------------------------------------------------------
+# overrides_from_flags (issues.md Issue 3): the flag -> config-path table
+# shared by `giant train`/`giant new-run`. Each test below pins one
+# precedence rule directly, without CliRunner — see also
+# tests/test_cli_train_overrides.py for the thin end-to-end smoke coverage.
+# ---------------------------------------------------------------------------
+
+
+def test_overrides_from_flags_empty_values_yield_empty_overrides():
+    assert gconfig.overrides_from_flags({}) == {}
+    assert gconfig.overrides_from_flags({"epochs": None, "hidden_dim": None}) == {}
+
+
+def test_overrides_from_flags_train_block_passthrough():
+    overrides = gconfig.overrides_from_flags({"epochs": 5, "lr": 1e-3, "hidden_dim": None})
+    assert overrides == {"train": {"epochs": 5, "lr": 1e-3}}
+
+
+@pytest.mark.parametrize(
+    ("shorthand", "explicit", "path_key"),
+    [
+        ("hidden_dim", "stage1_hidden_dim", "hidden_dim"),
+        ("n_blocks", "stage1_n_res_blocks", "n_res_blocks"),
+        ("dropout", "stage1_dropout", "dropout"),
+    ],
+)
+def test_overrides_from_flags_stage1_explicit_overrides_shorthand(shorthand, explicit, path_key):
+    overrides = gconfig.overrides_from_flags({shorthand: 1, explicit: 2})
+    assert overrides["stage1_model"][path_key] == 2
+
+
+@pytest.mark.parametrize(
+    ("shorthand", "path_key"),
+    [("hidden_dim", "hidden_dim"), ("n_blocks", "n_res_blocks"), ("dropout", "dropout")],
+)
+def test_overrides_from_flags_stage1_shorthand_alone(shorthand, path_key):
+    overrides = gconfig.overrides_from_flags({shorthand: 7})
+    assert overrides["stage1_model"][path_key] == 7
+
+
+def test_overrides_from_flags_stage2_only_knobs():
+    overrides = gconfig.overrides_from_flags(
+        {
+            "stage2_hidden_dim": 32,
+            "stage2_n_res_blocks": 4,
+            "stage2_dropout": 0.1,
+            "stage2_decoder": "one_shot",
+            "stage2_k_max": 8,
+            "stage2_context_dim": 16,
+            "stage2_stage1_context": "sampled",
+        }
+    )
+    assert overrides["stage2_model"] == {
+        "hidden_dim": 32,
+        "n_res_blocks": 4,
+        "dropout": 0.1,
+        "decoder": "one_shot",
+        "k_max": 8,
+        "context_dim": 16,
+        "stage1_context": "sampled",
+    }
+    assert "stage1_model" not in overrides
+
+
+def test_overrides_from_flags_mode_fans_to_both_stages():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan"})
+    assert overrides["stage1_model"]["generator"] == "wgan"
+    assert overrides["stage2_model"]["generator"] == "wgan"
+
+
+def test_overrides_from_flags_stage1_generator_overrides_mode_for_stage1_only():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage1_generator": "flow"})
+    assert overrides["stage1_model"]["generator"] == "flow"
+    assert overrides["stage2_model"]["generator"] == "wgan"
+
+
+def test_overrides_from_flags_stage2_generator_overrides_mode_for_stage2_only():
+    overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage2_generator": "flow"})
+    assert overrides["stage1_model"]["generator"] == "wgan"
+    assert overrides["stage2_model"]["generator"] == "flow"
+
+
+def test_overrides_from_flags_emb_dim_sets_both_conditioning_axes():
+    overrides = gconfig.overrides_from_flags({"emb_dim": 24})
+    assert overrides["conditioning"]["particle"]["emb_dim"] == 24
+    assert overrides["conditioning"]["material"]["emb_dim"] == 24
+
+
+def test_overrides_from_flags_conditioning_sets_both_axes_type():
+    overrides = gconfig.overrides_from_flags({"conditioning": "onehot"})
+    assert overrides["conditioning"]["particle"]["type"] == "onehot"
+    assert overrides["conditioning"]["material"]["type"] == "onehot"
+
+
+def test_overrides_from_flags_router_config_only_touches_stage1():
+    overrides = gconfig.overrides_from_flags({"router_config": {"enabled": True, "type": "energy"}})
+    assert overrides["stage1_model"]["router"] == {"enabled": True, "type": "energy"}
+    assert "stage2_model" not in overrides
+
+
+@pytest.mark.parametrize(
+    ("shared", "stage1_specific", "stage2_specific", "path_key"),
+    [
+        ("n_critic", "stage1_n_critic", "stage2_n_critic", "n_critic"),
+        ("gp_weight", "stage1_gp_weight", "stage2_gp_weight", "gp_weight"),
+        ("noise_dim", "stage1_noise_dim", "stage2_noise_dim", "noise_dim"),
+        ("critic_lr", "stage1_critic_lr", "stage2_critic_lr", "critic_lr"),
+    ],
+)
+def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific, stage2_specific, path_key):
+    overrides = gconfig.overrides_from_flags({shared: 5.0, stage1_specific: 3.0})
+    assert overrides["stage1_model"]["wgan"][path_key] == 3.0
+    assert overrides["stage2_model"]["wgan"][path_key] == 5.0
+
+    overrides = gconfig.overrides_from_flags({shared: 5.0, stage2_specific: 2.5})
+    assert overrides["stage1_model"]["wgan"][path_key] == 5.0
+    assert overrides["stage2_model"]["wgan"][path_key] == 2.5
+
+
+@pytest.mark.parametrize(
+    ("stage_flag", "stage_model", "path_key"),
+    [
+        ("stage1_critic_hidden_dim", "stage1_model", "critic_hidden_dim"),
+        ("stage1_critic_n_res_blocks", "stage1_model", "critic_n_res_blocks"),
+        ("stage2_critic_hidden_dim", "stage2_model", "critic_hidden_dim"),
+        ("stage2_critic_n_res_blocks", "stage2_model", "critic_n_res_blocks"),
+    ],
+)
+def test_overrides_from_flags_critic_sizing_is_stage_scoped_only(stage_flag, stage_model, path_key):
+    """critic_hidden_dim/critic_n_res_blocks are architectural per-stage
+    knobs (gitea #28) — unlike n_critic/gp_weight/noise_dim/critic_lr above,
+    there is deliberately no shared alias that fans out to both stages."""
+    overrides = gconfig.overrides_from_flags({stage_flag: 32})
+    assert overrides == {stage_model: {"wgan": {path_key: 32}}}
+
+
 # ---------------------------------------------------------------------------
 # checkpoint config-mismatch warnings (unchanged surface, still exercised)
 # ---------------------------------------------------------------------------
 
 
-def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
     ckpt_path = tmp_path / "best.pt"
     ckpt_path.write_bytes(b"")
-    _write_toml(
-        tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n"
-    )
+    _write_toml(tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n")
 
     gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
 
@@ -776,9 +1094,7 @@ def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(
     assert "current999" in captured.err
 
 
-def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "current999")
     ckpt_path = tmp_path / "best.pt"
     ckpt_path.write_bytes(b"")
@@ -787,15 +1103,11 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(
     assert capsys.readouterr().err == ""
 
 
-def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(
-    tmp_path, monkeypatch, capsys
-):
+def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(tmp_path, monkeypatch, capsys):
     monkeypatch.setattr(gconfig, "git_hash", lambda: "same123")
     ckpt_path = tmp_path / "best.pt"
     ckpt_path.write_bytes(b"")
-    _write_toml(
-        tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n"
-    )
+    _write_toml(tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n")
 
     gconfig.warn_if_checkpoint_config_mismatch(ckpt_path)
     assert capsys.readouterr().err == ""

tests/test_config_consumed_keys.py

@@ -0,0 +1,156 @@
+"""Consumed-keys audit (issues.md Issue 5).
+
+`validate_config_keys` (`giant/config.py`) only checks that a config key is
+*declared* — present somewhere in `DEFAULT_CONFIG`, which is generated from
+the frozen dataclasses. It says nothing about whether anything actually
+*reads* the value once parsed. Issues 1, 2 and 4 are three keys that slipped
+through exactly that gap: declared, round-tripped, silently ignored. This
+module walks every leaf path in `DEFAULT_CONFIG` and asserts each is either
+genuinely consumed by the model-building/training/rollout code, or explicitly
+recorded in `_KNOWN_UNUSED` with a reason.
+
+"Consumed" is approximated by static analysis rather than true call-graph
+reachability: for each leaf path's field name, does it appear anywhere in a
+fixed whitelist of source files as a real attribute access, a dict-key-shaped
+string constant, or a function/constructor parameter name (the last of these
+because `Router` subclasses receive their config via `**kwargs` filtered by
+signature — see `giant.model.routers.build_router`)? Docstrings are excluded
+from the string-constant scan so prose mentioning a dotted config path in
+passing can't masquerade as a read of it. This whitelist-based approach is
+deliberately narrower than "anywhere in `giant/`": scanning the whole package
+produces false negatives from unrelated identifier collisions (e.g.
+`giant/analysis/router_gating.py`'s `_top1_shares(..., order: list, ...)`
+parameter would otherwise make `stage2_model.autoregressive.order` read as
+"consumed").
+"""
+
+import ast
+from pathlib import Path
+
+from giant.config import DEFAULT_CONFIG
+
+_REPO_ROOT = Path(__file__).resolve().parents[1]
+
+# Files that legitimately consume model_config / training config at
+# build/train/rollout time. Not `giant/cli.py` (a CLI flag existing is not
+# consumption — that's precisely how Issue 1 slipped through), not
+# `giant/config.py` itself (declaring/parsing a field is not reading it), and
+# not `giant/model/_legacy.py` (the protected v0.2 migration surface, which
+# intentionally re-derives old flat keys under old names).
+_CONSUMER_ROOTS = ("giant/model", "giant/training")
+_CONSUMER_FILES = (
+    "giant/sample.py",
+    "giant/pipeline.py",
+    "giant/rollout.py",
+    "giant/checkpoint_io.py",
+    "giant/particles.py",
+    "giant/materials.py",
+)
+_EXCLUDED_FILES = ("giant/model/_legacy.py",)
+
+# Leaf DEFAULT_CONFIG paths that are declared but not (yet) read anywhere in
+# the consumer whitelist above. Each entry must name the issue that tracks
+# it. If a key here starts showing up as consumed, the fix landed and this
+# entry is stale — see test_known_unused_allow_list_has_no_stale_entries.
+_KNOWN_UNUSED = {
+    "stage2_model.stage1_context": (
+        "issues.md Issue 1 — trainers.py hardcodes stage1_ctx to the "
+        "ground-truth stage-1 output; 'sampled' is now rejected loudly by "
+        "validate_config (not silently accepted), but the key still isn't "
+        "read by any build/train consumer file since only 'truth' can pass "
+        "validation — see Issue 16 for the real implementation"
+    ),
+    "stage2_model.autoregressive.order": (
+        "gitea #30 — validate_config now checks order is 'energy_desc', but "
+        "nothing in the build/train/rollout consumer whitelist reads the "
+        "value itself since it's still single-valued"
+    ),
+}
+
+# "lambda" is a Python keyword, so the dataclasses expose the dict key
+# "lambda" as the field `lambda_weight` (giant/config.py:49-50).
+_FIELD_NAME_OVERRIDES = {"lambda": "lambda_weight"}
+
+
+def _leaf_paths(node: dict, prefix: str = "") -> list[str]:
+    paths = []
+    for key, value in node.items():
+        if prefix == "" and key == "meta":
+            continue
+        path = f"{prefix}.{key}" if prefix else key
+        if isinstance(value, dict):
+            paths.extend(_leaf_paths(value, path))
+        else:
+            paths.append(path)
+    return paths
+
+
+def _field_name(leaf_path: str) -> str:
+    name = leaf_path.rsplit(".", 1)[-1]
+    return _FIELD_NAME_OVERRIDES.get(name, name)
+
+
+def _is_docstring_expr(expr: ast.Expr) -> bool:
+    return isinstance(expr.value, ast.Constant) and isinstance(expr.value.value, str)
+
+
+def _collect_names(source: str, filename: str) -> set[str]:
+    tree = ast.parse(source, filename=filename)
+    docstring_ids = set()
+    for node in ast.walk(tree):
+        if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
+            body = getattr(node, "body", [])
+            if body and isinstance(body[0], ast.Expr) and _is_docstring_expr(body[0]):
+                docstring_ids.add(id(body[0].value))
+
+    names: set[str] = set()
+    for node in ast.walk(tree):
+        if isinstance(node, ast.Attribute):
+            names.add(node.attr)
+        elif isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstring_ids:
+            names.add(node.value)
+        elif isinstance(node, ast.arg):
+            names.add(node.arg)
+        elif isinstance(node, ast.keyword) and node.arg is not None:
+            names.add(node.arg)
+    return names
+
+
+def _consumer_files() -> list[Path]:
+    files: set[Path] = {_REPO_ROOT / f for f in _CONSUMER_FILES}
+    for root in _CONSUMER_ROOTS:
+        files |= set((_REPO_ROOT / root).rglob("*.py"))
+    files -= {_REPO_ROOT / f for f in _EXCLUDED_FILES}
+    return sorted(files)
+
+
+def _consumed_names() -> set[str]:
+    names: set[str] = set()
+    for path in _consumer_files():
+        names |= _collect_names(path.read_text(), str(path))
+    return names
+
+
+def test_every_config_key_is_consumed_or_allow_listed():
+    consumed = _consumed_names()
+    unconsumed = {p for p in _leaf_paths(DEFAULT_CONFIG) if _field_name(p) not in consumed}
+    unexplained = unconsumed - _KNOWN_UNUSED.keys()
+    assert not unexplained, (
+        f"config key(s) {sorted(unexplained)} are declared in DEFAULT_CONFIG "
+        "but not read anywhere in the build/train/rollout consumer files "
+        f"({[str(f.relative_to(_REPO_ROOT)) for f in _consumer_files()]}) — "
+        "either wire the key up, or add it to _KNOWN_UNUSED with a reason "
+        "(see issues.md Issue 5)"
+    )
+
+
+def test_known_unused_allow_list_has_no_stale_entries():
+    consumed = _consumed_names()
+    all_paths = set(_leaf_paths(DEFAULT_CONFIG))
+    stale = {p for p in _KNOWN_UNUSED if p not in all_paths or _field_name(p) in consumed}
+    assert not stale, (
+        f"_KNOWN_UNUSED entry/entries {sorted(stale)} no longer belong on the "
+        "allow-list — either the key was removed from DEFAULT_CONFIG, or it "
+        "is now consumed (the underlying issue was fixed). Remove the stale "
+        "entry/entries."
+    )

tests/test_create_root_files.py

@@ -4,7 +4,7 @@ from pathlib import Path
 
 import pytest
 
-from scripts import create_root_files
+from giant.tools import create_root_files
 
 parse_detector_spec = create_root_files.parse_detector_spec
 next_shard_index = create_root_files.next_shard_index
@@ -16,9 +16,7 @@ SimJob = create_root_files.SimJob
 PlanError = create_root_files.PlanError
 
 
-def _write_fake_executable(
-    path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0
-) -> Path:
+def _write_fake_executable(path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0) -> Path:
     """Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files
     into its own cwd (so callers can verify each job gets an isolated workdir
     and that the workdir ends up holding *only* the .root output, matching
@@ -89,17 +87,13 @@ def test_next_shard_index_continues_past_existing(tmp_path):
 
 def test_plan_jobs_rejects_missing_gen(tmp_path):
     with pytest.raises(PlanError):
-        plan_jobs(
-            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
-        )
+        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")
 
 
 def test_plan_jobs_rejects_malformed_gen(tmp_path):
     (tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
     with pytest.raises(PlanError):
-        plan_jobs(
-            ["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
-        )
+        plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")
 
 
 def test_plan_jobs_continues_from_existing_shards(tmp_path):
@@ -108,9 +102,7 @@ def test_plan_jobs_continues_from_existing_shards(tmp_path):
     (gen_dir / "pbwo4" / "shard-000.root").touch()
     (gen_dir / "pbwo4" / "shard-001.root").touch()
 
-    jobs = plan_jobs(
-        ["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
-    )
+    jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")
 
     assert [j.shard_index for j in jobs] == [2, 3, 4]
     assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
@@ -320,9 +312,7 @@ def test_run_all_caps_concurrency(tmp_path):
     assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}
 
     intervals = [json.loads(d.read_text()) for d in dests]
-    events = sorted(
-        [(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
-    )
+    events = sorted([(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals])
     concurrent = 0
     peak = 0
     for _, delta in events:

tests/test_dataset.py

@@ -95,7 +95,7 @@ def _dummy_normalizer(width):
 
 def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
     """Two files that each restart event_id from 0 (one Geant4 job per file,
-    see scripts/steps_to_parquet.py) must not have their same-numbered events
+    see giant/tools/steps_to_parquet.py) must not have their same-numbered events
     collapsed together: every row from every file must show up in exactly one
     of train/val, and the number of distinct events must be the sum across
     files, not the union of raw ids."""

tests/test_dwarf.py

@@ -3,15 +3,15 @@ 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 giant.tools import dwarf
+from giant.tools.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
+    """giant.cli and giant.tools.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
@@ -51,9 +51,7 @@ def test_convert_rejects_output_with_multiple_files(tmp_path):
 def test_convert_rejects_output_with_parallel_jobs(tmp_path):
     root_file = tmp_path / "shard.root"
     root_file.touch()
-    result = runner.invoke(
-        app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
-    )
+    result = runner.invoke(app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"])
     assert result.exit_code != 0
     assert "--output cannot be combined with --jobs > 1" in result.output
 

tests/test_geometry.py

@@ -4,6 +4,7 @@ from pathlib import Path
 from unittest.mock import patch
 
 import numpy as np
+import pandas as pd
 import pytest
 
 from giant import geometry as g
@@ -12,6 +13,70 @@ from giant import geometry as g
 pytest.importorskip("sklearn")
 
 
+def _steps_frame(n=5, with_post=True):
+    rng = np.random.default_rng(0)
+    data = {
+        "pre_x": rng.uniform(-10, 10, n),
+        "pre_y": rng.uniform(-10, 10, n),
+        "pre_z": rng.uniform(-10, 10, n),
+        "material": ["G4_AIR"] * n,
+        "layer_id": np.arange(n, dtype=np.int64),
+    }
+    if with_post:
+        data["post_x"] = rng.uniform(-10, 10, n)
+        data["post_y"] = rng.uniform(-10, 10, n)
+        data["post_z"] = rng.uniform(-10, 10, n)
+    return pd.DataFrame(data)
+
+
+def test_iter_point_batches_missing_columns_raises(tmp_path):
+    path = tmp_path / "steps.parquet"
+    pd.DataFrame({"pre_x": [0.0]}).to_parquet(path)
+    with pytest.raises(ValueError, match="missing columns"):
+        next(g._iter_point_batches(path))
+
+
+def test_iter_point_batches_without_post_columns_yields_pre_only(tmp_path):
+    path = tmp_path / "steps.parquet"
+    df = _steps_frame(n=5, with_post=False)
+    df.to_parquet(path)
+
+    (pos, mat, lay) = next(g._iter_point_batches(path))
+
+    assert pos.shape == (5, 3)
+    np.testing.assert_allclose(pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
+    assert list(mat) == ["G4_AIR"] * 5
+    np.testing.assert_array_equal(lay, np.arange(5))
+
+
+def test_iter_point_batches_with_post_columns_doubles_and_concatenates_points(
+    tmp_path,
+):
+    path = tmp_path / "steps.parquet"
+    df = _steps_frame(n=5, with_post=True)
+    df.to_parquet(path)
+
+    (pos, mat, lay) = next(g._iter_point_batches(path))
+
+    # Every step contributes both its pre_pos and post_pos, sharing the
+    # step's material/layer_id label — so batches double in length.
+    assert pos.shape == (10, 3)
+    np.testing.assert_allclose(pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
+    np.testing.assert_allclose(pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32))
+    assert list(mat) == ["G4_AIR"] * 10
+    np.testing.assert_array_equal(lay, np.concatenate([np.arange(5), np.arange(5)]))
+
+
+def test_iter_point_batches_respects_batch_size(tmp_path):
+    path = tmp_path / "steps.parquet"
+    df = _steps_frame(n=10, with_post=False)
+    df.to_parquet(path, row_group_size=10)
+
+    batches = list(g._iter_point_batches(path, batch_size=4))
+
+    assert [len(pos) for pos, _, _ in batches] == [4, 4, 2]
+
+
 def _box_batch(n, rng):
     """A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1."""
     pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32)
@@ -137,9 +202,7 @@ def test_slab_classes_discovered():
 
 def test_slab_query_labels_by_depth():
     orc = _build_slab()
-    pos = np.array(
-        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
-    )  # layer 0, gap, layer 1
+    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]])  # layer 0, gap, layer 1
     material, layer_id, escaped = orc.query(pos)
     assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
     assert list(layer_id) == [0, -1, 1]
@@ -168,9 +231,7 @@ def test_slab_save_load_roundtrip(tmp_path):
     orc.save(p)
     loaded = g.GeometryOracle.load(p)
 
-    pos = np.array(
-        [[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
-    )
+    pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]])
     m0, l0, e0 = orc.query(pos)
     m1, l1, e1 = loaded.query(pos)
     assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()

tests/test_loader.py

@@ -315,9 +315,7 @@ def test_build_index_maps_from_files_ordering_independent_of_file_order(tmp_path
 
 def test_build_index_maps_from_files_numeric_sort_for_nuclear_codes(tmp_path):
     path = tmp_path / "a.parquet"
-    pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(
-        path
-    )
+    pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(path)
 
     pdg_map, _ = build_index_maps_from_files([path])
     assert list(pdg_map.keys()) == [11, 22, 1000060120]
@@ -398,9 +396,7 @@ def test_load_event_ids_applies_offset(tmp_path):
     path = tmp_path / "a.parquet"
     pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path)
     offset = event_id_offset(1)
-    np.testing.assert_array_equal(
-        load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2]
-    )
+    np.testing.assert_array_equal(load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2])
 
 
 def test_load_event_ids_raises_when_event_id_reaches_stride(tmp_path):

tests/test_materials.py

@@ -23,11 +23,7 @@ def test_get_material_properties_unfilled_entry_raises():
 
 
 def test_get_material_properties_returns_filled_entry_from_injected_table():
-    table = {
-        "G4_Pb": MaterialProperties(
-            z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59
-        )
-    }
+    table = {"G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59)}
     props = get_material_properties("G4_Pb", table)
     assert props.z_eff == 82.0
     assert props.a_eff == 207.2

tests/test_migration_v02_v03.py

@@ -56,9 +56,7 @@ def _random_batch(seed: int):
 
 def _assert_bit_identical(a: torch.Tensor, b: torch.Tensor, label: str) -> None:
     assert a.shape == b.shape, f"{label}: shape mismatch {a.shape} vs {b.shape}"
-    assert torch.equal(a, b), (
-        f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
-    )
+    assert torch.equal(a, b), f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
 
 
 def _run_migration_check(mode: str, conditioning: str) -> None:
@@ -132,14 +130,12 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
     new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"]
     assert isinstance(new_stage1, net.Stage1Model)
     assert isinstance(new_stage2, net.Stage2OneShot)
-    # legacy_owner="stage1": n_sec lives on stage1, not stage2, for a
+    # n_sec.owner="stage1": n_sec lives on stage1, not stage2, for a
     # migrated v0.2 checkpoint.
     assert new_stage1.n_sec_head is not None
     assert new_stage2.n_sec_head is None
 
-    remapped1, remapped2 = net.migrate_legacy_state_dict(
-        old_stage1.state_dict(), old_stage2.state_dict()
-    )
+    remapped1, remapped2 = net.migrate_legacy_state_dict(old_stage1.state_dict(), old_stage2.state_dict())
     missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
     missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
     assert not missing1 and not unexpected1
@@ -159,9 +155,7 @@ def _run_migration_check(mode: str, conditioning: str) -> None:
             new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t)
 
     _assert_bit_identical(old_out1, new_out1, f"stage1 output ({mode}, {conditioning})")
-    _assert_bit_identical(
-        old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})"
-    )
+    _assert_bit_identical(old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})")
     _assert_bit_identical(old_out2, new_out2, f"stage2 output ({mode}, {conditioning})")
 
 
@@ -183,7 +177,7 @@ def test_migration_wgan_physical():
 
 def test_migrate_legacy_model_config_shape():
     """_migrate_legacy_model_config produces the nested shape build_models
-    expects, with the legacy_owner marker set so build_models routes the
+    expects, with the n_sec.owner marker set so build_models routes the
     n_sec head back onto stage 1."""
     legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
     migrated = net._migrate_legacy_model_config(legacy_cfg)
@@ -193,7 +187,7 @@ def test_migrate_legacy_model_config_shape():
     assert migrated["conditioning"]["particle"]["n_layers"] == 2
     assert migrated["conditioning"]["material"]["n_layers"] == 2
     assert migrated["stage1_model"]["hidden_dim"] == HIDDEN_DIM
-    assert migrated["stage2_model"]["n_sec"]["legacy_owner"] == "stage1"
+    assert migrated["stage2_model"]["n_sec"]["owner"] == "stage1"
     assert migrated["stage2_model"]["decoder"] == "one_shot"
 
 
@@ -284,6 +278,6 @@ def test_build_models_accepts_new_nested_shape_unchanged():
     models = net.build_models(cfg)
     assert isinstance(models["stage1"], net.Stage1Model)
     assert isinstance(models["stage2"], net.Stage2OneShot)
-    # Fresh v0.3.0 config, no legacy_owner: n_sec lives on stage 2.
+    # Fresh v0.3.0 config, n_sec.owner defaults to "stage2": n_sec lives on stage 2.
     assert models["stage1"].n_sec_head is None
     assert models["stage2"].n_sec_head is not None

tests/test_network.py

@@ -12,6 +12,7 @@ from giant.model.network import (
     Stage1Model,
     Stage2Autoregressive,
     Stage2OneShot,
+    build_critics,
     build_models,
     cat_col_layout,
     stage2_trunk_sec_dim,
@@ -85,9 +86,7 @@ def test_stage1_model_gradients_flow():
 def test_stage1_model_no_n_sec_head_by_default():
     """Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
     it moves to stage 2."""
-    model = Stage1Model(
-        pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG
-    )
+    model = Stage1Model(pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG)
     assert model.n_sec_head is None
 
 
@@ -121,29 +120,18 @@ def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
 
 def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
     k_max = 15
-    assert (
-        stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16)
-        == k_max * SEC_SLOT_DIM
-    )
-    assert (
-        stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16)
-        == k_max * SEC_SLOT_DIM
-    )
+    assert stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
+    assert stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
 
 
 def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
     k_max = 15
-    assert stage2_trunk_sec_dim(
-        {"target": "onehot"}, "wgan", k_max, emb_dim=16
-    ) == k_max * (CONT_SLOT_DIM + 16)
+    assert stage2_trunk_sec_dim({"target": "onehot"}, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)
 
 
 def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
     k_max = 15
-    assert (
-        stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16)
-        == k_max * CONT_SLOT_DIM
-    )
+    assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
 
 
 # --- ConditionEncoder onehot mode -------------------------------------------
@@ -300,6 +288,33 @@ def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
     assert out.shape == (B, k_max * CONT_SLOT_DIM)
 
 
+def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29: stage2_model.particle_type.n_classes, not
+    conditioning.particle.emb_dim, sizes the onehot type_head/type_dim when
+    explicitly set — the two used to be silently the same number."""
+    k_max = 5
+    particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20)
+    model = Stage2OneShot(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        sec_dim=sec_dim,
+        generator="flow",
+        k_max=k_max,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == k_max * 20
+
+
 # --- MarkovHistory -----------------------------------------------------------
 
 
@@ -435,21 +450,40 @@ def test_stage2_autoregressive_history_invalid_raises():
         _build_stage2_ar("onehot", "wgan", history="bogus")
 
 
+def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29, Stage2Autoregressive side — see the Stage2OneShot version
+    of this test for the full rationale."""
+    particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
+    model = Stage2Autoregressive(
+        pdg_vocab=5,
+        mat_vocab=3,
+        particle_cfg=particle_cfg,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=1,
+        cond_out_dim=16,
+        context_dim=8,
+        generator="flow",
+        k_max=5,
+        particle_type_cfg=particle_type_cfg,
+    )
+    assert model.type_dim == 20  # not particle_cfg["emb_dim"] == 6
+    assert model.type_head is not None
+    assert model.type_head[-1].out_features == 20
+
+
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
 def test_stage2_autoregressive_forward_shape(target, generator, history):
     B, K, emb_dim = 4, 5, 6
-    model = _build_stage2_ar(
-        target, generator, emb_dim=emb_dim, k_max=K, history=history
-    )
+    model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K, history=history)
     cond_cont = torch.randn(B, COND_DIM)
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": target}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
     if generator == "wgan":
         x_t = torch.randn(B, K, model.noise_dim)
@@ -488,9 +522,7 @@ def test_stage2_autoregressive_predict_type_shape():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     out = model.predict_type(
         cond_cont,
         cond_cat,
@@ -511,9 +543,7 @@ def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, gen
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": target}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     with pytest.raises(RuntimeError):
         model.predict_type(
             cond_cont,
@@ -533,9 +563,7 @@ def test_stage2_autoregressive_gradients_flow_wgan_onehot():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     z = torch.randn(B, K, model.noise_dim)
     gen_out = model(
         z,
@@ -560,9 +588,7 @@ def test_stage2_autoregressive_gradients_flow_onehot():
     cond_cat = torch.zeros(B, 2, dtype=torch.long)
     stage1_out = torch.randn(B, 9)
     type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(
-        B, K, CONT_SLOT_DIM + type_dim
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
     token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
     x_t = torch.randn(B, K, token_dim)
     t = torch.rand(B, K)
@@ -601,17 +627,13 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
     itself rather than `AttentionHistory` in isolation
     (`test_attention_history_step_matches_forward` covers that lower layer)."""
     B, K, emb_dim = 3, 6, 6
-    model = _build_stage2_ar(
-        "physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention"
-    )
+    model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention")
     model.eval()
     type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
     hist_in_dim = CONT_SLOT_DIM + type_dim
     own_feat = torch.randn(B, K, hist_in_dim)  # token i's own raw feature
     has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
-    history_feat = torch.cat(
-        [torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1
-    )
+    history_feat = torch.cat([torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1)
 
     with torch.no_grad():
         expected = model.history_encoder(history_feat, has_prev_full)
@@ -681,3 +703,126 @@ def test_build_models_share_stages_true_shared_params_are_in_both_stage_paramete
     assert shared_ids
     assert shared_ids <= {id(p) for p in stage1.parameters()}
     assert shared_ids <= {id(p) for p in stage2.parameters()}
+
+
+def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim():
+    """gitea #29 end-to-end through build_models: setting
+    stage2_model.particle_type.n_classes independently of
+    conditioning.particle.emb_dim actually resizes the built stage2 model,
+    not just the two lower-level unit tests above."""
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
+    built = build_models(cfg)
+    assert built["stage2"] is not None
+    assert built["stage2"].type_dim == 11
+
+
+def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
+    cfg = _minimal_model_config(share_stages=False)  # conditioning.particle.emb_dim = 4
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 4}
+    default_n_classes_critic = build_critics(cfg)["stage2"]
+    assert default_n_classes_critic is not None
+
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 11
+    wider_critic = build_critics(cfg)["stage2"]
+    assert wider_critic is not None
+    # k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
+    # widening n_classes alone (emb_dim stays 4) must widen the critic input.
+    assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features
+
+
+# ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
+
+
+def _partial_model_config() -> dict:
+    """A hand-built model_config that omits stage2_model.decoder and
+    stage2_model.particle_type — deliberately not derived from
+    DEFAULT_CONFIG, unlike _minimal_model_config above. Regression fixture
+    for issues.md Issue 1: build_models/build_critics/StageSpec.from_config's
+    own fallback defaults for these two keys must equal DEFAULT_CONFIG's
+    ("autoregressive" / "onehot"), not the old, now-wrong v0.2-shaped
+    ("one_shot" / "physical") literals that used to live in three separate
+    .get(key, default) call sites."""
+    return {
+        "pdg_vocab": 3,
+        "mat_vocab": 2,
+        "conditioning": {
+            "particle": {"type": "physical", "emb_dim": 4, "n_layers": 1},
+            "material": {"type": "physical", "emb_dim": 4, "n_layers": 1},
+        },
+        "stage1_model": {"active": False},
+        "stage2_model": {
+            "generator": "wgan",
+            "hidden_dim": 8,
+            "n_res_blocks": 1,
+            "k_max": 3,
+            # decoder and particle_type deliberately omitted
+        },
+    }
+
+
+def test_build_models_omitted_decoder_and_particle_type_match_default_config():
+    built = build_models(_partial_model_config())
+    assert isinstance(built["stage2"], Stage2Autoregressive)
+    assert built["stage2"].particle_type_cfg["target"] == "onehot"
+
+
+def test_build_critics_omitted_particle_type_matches_default_config():
+    cfg = _partial_model_config()
+    cfg["stage2_model"]["generator"] = "wgan"
+    onehot_critic = build_critics(cfg)["stage2"]
+    assert onehot_critic is not None
+    onehot_in_dim = onehot_critic.input_proj.in_features
+
+    cfg["stage2_model"]["particle_type"] = {"target": "physical"}
+    physical_critic = build_critics(cfg)["stage2"]
+    assert physical_critic is not None
+    physical_in_dim = physical_critic.input_proj.in_features
+
+    # onehot's per-slot type width is emb_dim classes vs. physical's fixed
+    # (log-mass, charge) pair — different unless emb_dim happens to be 2, so
+    # this also confirms the critic was actually built in onehot mode by
+    # default, not silently falling back to physical.
+    assert onehot_in_dim != physical_in_dim
+
+
+# ── build_critics: critic_hidden_dim/critic_n_res_blocks honoured (gitea #28) ─
+
+
+def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_generator_size():
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage1_model"]["hidden_dim"] = 8
+    cfg["stage1_model"]["n_res_blocks"] = 1
+
+    inherited = build_critics(cfg)["stage1"]
+    assert inherited is not None
+    assert inherited.input_proj.out_features == 8
+    assert len(inherited.blocks) == 1
+
+    cfg["stage1_model"]["wgan"]["critic_hidden_dim"] = 16
+    cfg["stage1_model"]["wgan"]["critic_n_res_blocks"] = 3
+    overridden = build_critics(cfg)["stage1"]
+    assert overridden is not None
+    assert overridden.input_proj.out_features == 16
+    assert len(overridden.blocks) == 3
+
+
+def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["hidden_dim"] = 8
+    cfg["stage2_model"]["n_res_blocks"] = 1
+
+    inherited = build_critics(cfg)["stage2"]
+    assert inherited is not None
+    assert inherited.input_proj.out_features == 8
+    assert len(inherited.blocks) == 1
+
+    cfg["stage2_model"]["wgan"]["critic_hidden_dim"] = 16
+    cfg["stage2_model"]["wgan"]["critic_n_res_blocks"] = 3
+    overridden = build_critics(cfg)["stage2"]
+    assert overridden is not None
+    assert overridden.input_proj.out_features == 16
+    assert len(overridden.blocks) == 3

tests/test_particles.py

@@ -49,9 +49,7 @@ def test_ground_state_nucleus_resolved_via_particle_package():
     """He-4 (Z=2, A=4) is a common nuclide in `particle`'s ground-state table."""
     mass, charge = particle_mass_charge(1000020040)
     assert charge == pytest.approx(2.0)
-    assert mass == pytest.approx(
-        4 * 931.494, rel=0.05
-    )  # near A*amu, binding-energy-corrected
+    assert mass == pytest.approx(4 * 931.494, rel=0.05)  # near A*amu, binding-energy-corrected
 
 
 def test_nuclear_isomer_falls_back_to_z_a_decode():
@@ -174,9 +172,7 @@ def test_decode_topn_class_other_drop_returns_zero_sentinel():
 def test_decode_topn_class_other_sample_stays_within_members():
     topn_map, n_classes = _topn_fixture()
     rng = np.random.default_rng(0)
-    out = decode_topn_class(
-        np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng
-    )
+    out = decode_topn_class(np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng)
     assert set(out.tolist()) <= {2212, 2112}
 
 

tests/test_phase2.py

@@ -57,9 +57,7 @@ def _sec_decoder(pdg=3, mat=2, conditioning="embedding"):
 
 def _cond(B=8, pdg=3, mat=2):
     cond_cont = torch.randn(B, COND_DIM)
-    cond_cat = torch.stack(
-        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
-    )
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
     return cond_cont, cond_cat
 
 
@@ -154,9 +152,7 @@ def test_flow_matching_loss_secondary_scalar():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
-    loss = flow_matching_loss_secondary(
-        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
-    )
+    loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
     assert loss.shape == ()
     assert loss.item() >= 0.0
 
@@ -169,9 +165,7 @@ def test_flow_matching_loss_secondary_mask_zeros_padding():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
-    loss = flow_matching_loss_secondary(
-        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
-    )
+    loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
     assert loss.item() == pytest.approx(0.0, abs=1e-6)
 
 
@@ -182,9 +176,7 @@ def test_flow_matching_loss_secondary_has_grad():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
-    flow_matching_loss_secondary(
-        decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
-    ).backward()
+    flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask).backward()
     assert any(p.grad is not None for p in decoder.parameters())
 
 
@@ -220,9 +212,7 @@ def test_flow_matching_loss_secondary_ar_scalar():
     x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
-        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     sec_mask = torch.ones(B, K, dtype=torch.bool)
     loss = flow_matching_loss_secondary_ar(
         decoder,
@@ -246,9 +236,7 @@ def test_flow_matching_loss_secondary_ar_mask_zeros_padding():
     x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
-        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     sec_mask = torch.zeros(B, K, dtype=torch.bool)
     loss = flow_matching_loss_secondary_ar(
         decoder,
@@ -271,9 +259,7 @@ def test_flow_matching_loss_secondary_ar_has_grad():
     x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
-    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(
-        B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM
-    )
+    history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     sec_mask = torch.ones(B, K, dtype=torch.bool)
     flow_matching_loss_secondary_ar(
         decoder,
@@ -299,9 +285,7 @@ def test_sample_secondaries_shapes():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
-    sec_cont, sec_phys, sec_valid = sample_secondaries(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3
-    )
+    sec_cont, sec_phys, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3)
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM)
     assert sec_valid.shape == (B, K_MAX)
@@ -314,9 +298,7 @@ def test_sample_secondaries_valid_mask_matches_n_sec():
     cond_cont, cond_cat = _cond(B, pdg, mat)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
-    _, _, sec_valid = sample_secondaries(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    _, _, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     for i, n in enumerate(n_sec_pred.tolist()):
         assert sec_valid[i, :n].all()
         assert not sec_valid[i, n:].any()
@@ -350,9 +332,7 @@ def test_encode_secondaries_energy_conservation():
     pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
     pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
 
-    sec_cont = encode_secondaries(
-        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
-    )
+    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
     assert sec_cont.shape == (N, K_MAX, 6)
     assert np.isfinite(sec_cont).all()
 
@@ -399,9 +379,7 @@ def test_encode_secondaries_stick_logits_match_naive_reference():
             logit = np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP)
             expected[row, i] = logit
 
-    np.testing.assert_allclose(
-        stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4
-    )
+    np.testing.assert_allclose(stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4)
 
 
 def test_encode_secondaries_direction_encoding():
@@ -513,9 +491,7 @@ def test_encode_secondaries_physical_columns_match_ground_truth_pdg():
     sec_valid[0, 0] = True
     pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
 
-    sec_cont = encode_secondaries(
-        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
-    )
+    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
     mass, charge = particle_mass_charge(11)
     assert sec_cont[0, 0, 4] == pytest.approx(log_transform(np.array([mass]))[0])
     assert sec_cont[0, 0, 5] == pytest.approx(charge)
@@ -549,9 +525,7 @@ def test_decode_secondaries_valid_slots_sum_to_e_sec():
     e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
 
-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
 
     valid_sum = (sec_E * sec_valid).sum(axis=1)
     has_secondaries = n_sec > 0
@@ -574,9 +548,7 @@ def test_decode_secondaries_zero_n_sec_has_zero_energy():
     e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
 
-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
 
     assert not sec_valid.any()
     np.testing.assert_allclose(sec_E, 0.0)
@@ -596,9 +568,7 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
     e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
     pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
 
-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
-        sec_cont, n_sec, e_sec, pre_dir
-    )
+    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
 
     for i, k in enumerate(n_sec):
         if k == 0:
@@ -622,12 +592,8 @@ def test_decode_secondaries_rescale_preserves_relative_shares():
     n_sec = np.array([4])
     pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
 
-    sec_E_small, _, _, _, sec_valid = decode_secondaries(
-        sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir
-    )
-    sec_E_large, _, _, _, _ = decode_secondaries(
-        sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir
-    )
+    sec_E_small, _, _, _, sec_valid = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
+    sec_E_large, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir)
 
     ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
     ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
@@ -649,20 +615,14 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
     sec_valid[0, 0] = True
     pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
 
-    sec_cont = encode_secondaries(
-        sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
-    )
+    sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
     norm = Normalizer()
     norm.mean = np.array([-2.0, 0.5], dtype=np.float32)
     norm.std = np.array([3.0, 1.5], dtype=np.float32)
     sec_cont_normed = sec_cont.copy()
-    sec_cont_normed[:, :, 4:6] = norm.transform(
-        sec_cont[:, :, 4:6].reshape(-1, 2)
-    ).reshape(N, K_MAX, 2)
+    sec_cont_normed[:, :, 4:6] = norm.transform(sec_cont[:, :, 4:6].reshape(-1, 2)).reshape(N, K_MAX, 2)
 
     n_sec = np.array([1])
-    _, _, sec_mass, sec_charge, _ = decode_secondaries(
-        sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm
-    )
+    _, _, sec_mass, sec_charge, _ = decode_secondaries(sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm)
     assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
     assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)

tests/test_pipeline.py

@@ -6,9 +6,10 @@ import pytest
 import torch
 
 from giant import config as gconfig
+from giant.constants import COND_DIM
 from giant.data import setup_cache
 from giant.data.transforms import Normalizer
-from giant.pipeline import run_train_job
+from giant.pipeline import _seed_energy_router, run_train_job
 
 
 def _unit(v):
@@ -47,9 +48,7 @@ def _make_synthetic_steps(path, n_events=20, seed=0):
             pre_dir = np.array([0.0, 0.0, 1.0])
             post_dir = _unit(rng.normal(size=3))
             post_pos = pre_pos + step_length * pre_dir
-            sec_energies = (
-                list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
-            )
+            sec_energies = list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
             sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)]
             sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)]
             rows.append(
@@ -153,28 +152,83 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
     assert "normalizer: cache hit" in joined
 
 
-def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data):
+def test_run_train_job_builds_caches_and_persists_sec_type_topn_map(tmp_path, data):
     """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
-    "onehot" — a plain _tiny_cfg() run must
-    build the shared pdg top-N map, cache it in the setup-cache sidecar, and
-    persist it into the checkpoint, with no extra config needed."""
+    "onehot" while conditioning.particle.type stays "physical" — a plain
+    _tiny_cfg() run must build the secondary-type-only pdg top-N map (gitea
+    #29: no longer shared with any conditioning-side onehot map), cache it in
+    the setup-cache sidecar, and persist it into the checkpoint's
+    sec_type_topn_map key, with no extra config needed. pdg_topn_map
+    (conditioning-only) stays unbuilt since conditioning.particle.type is
+    "physical" here."""
     echo1 = _run(data, tmp_path / "out1")
     assert any("building pdg top-N map" in m for m in echo1)
 
     loaded = setup_cache.load(data, [data])
     assert loaded is not None
-    key = setup_cache.topn_key("pdg", 4)  # conditioning.particle.emb_dim = 4
+    # stage2_model.particle_type.n_classes = 0 -> conditioning.particle.emb_dim = 4
+    key = setup_cache.topn_key("pdg", 4)
     assert key in loaded.topn_maps
     assert set(loaded.topn_maps[key].class_map.keys()) >= {11, 22}
 
     ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
-    assert "pdg_topn_map" in ckpt
-    assert set(ckpt["pdg_topn_map"]["class_map"].keys()) >= {"11", "22"}
+    assert ckpt.get("pdg_topn_map") is None
+    assert "sec_type_topn_map" in ckpt
+    assert set(ckpt["sec_type_topn_map"]["class_map"].keys()) >= {"11", "22"}
 
     echo2 = _run(data, tmp_path / "out2")
     assert any("pdg top-N map: cache hit" in m for m in echo2)
 
 
+def test_run_train_job_independent_cond_and_sec_type_topn_maps(tmp_path, data):
+    """conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" with different class counts
+    (gitea #29's fix: stage2_model.particle_type.n_classes decouples the two)
+    build two distinct top-N maps, cached under their own (axis, n_classes)
+    key and persisted under two distinct checkpoint keys — no longer forced
+    to share conditioning.particle.emb_dim."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["particle"]["type"] = "onehot"  # emb_dim = 4, from _tiny_cfg
+    cfg["stage2_model"]["particle_type"]["n_classes"] = 3
+    echo = _run(data, tmp_path / "out", cfg=cfg)
+    assert any("mapped to 4 classes" in m for m in echo)
+    assert any("mapped to 3 classes" in m for m in echo)
+
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    cond_key = setup_cache.topn_key("pdg", 4)
+    type_key = setup_cache.topn_key("pdg", 3)
+    assert cond_key in loaded.topn_maps
+    assert type_key in loaded.topn_maps
+
+    ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
+    assert ckpt.get("pdg_topn_map") is not None
+    assert ckpt.get("sec_type_topn_map") is not None
+
+
+def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data):
+    """conditioning.material.type="onehot" is an independent axis from the
+    pdg one above, with its own build/cache-hit branch in run_setup_stage —
+    exercise both here the same way the pdg test above does."""
+    cfg = _tiny_cfg()
+    cfg["conditioning"]["material"]["type"] = "onehot"
+    echo1 = _run(data, tmp_path / "out1", cfg=cfg)
+    assert any("building material top-N map" in m for m in echo1)
+
+    loaded = setup_cache.load(data, [data])
+    assert loaded is not None
+    key = setup_cache.topn_key("material", 4)  # conditioning.material.emb_dim = 4
+    assert key in loaded.topn_maps
+    assert set(loaded.topn_maps[key].class_map.keys()) >= {"G4_AIR", "G4_Fe"}
+
+    ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
+    assert "mat_topn_map" in ckpt
+    assert set(ckpt["mat_topn_map"]["class_map"].keys()) >= {"G4_AIR", "G4_Fe"}
+
+    echo2 = _run(data, tmp_path / "out2", cfg=cfg)
+    assert any("material top-N map: cache hit" in m for m in echo2)
+
+
 def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
     cfg = _tiny_cfg()
     cfg["stage2_model"]["particle_type"] = {"target": "physical", "lambda": 1.0}
@@ -187,9 +241,7 @@ def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
 
 
 @pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
-def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
-    tmp_path, data, monkeypatch
-):
+def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(tmp_path, data, monkeypatch):
     # num_workers>0 makes DataLoader actually fork worker subprocesses
     # (unlike every other test here, which runs with num_workers=0) — pytest
     # itself is multi-threaded, hence Python's fork-safety warning below.
@@ -199,9 +251,7 @@ def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
 
 
 @pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
-def test_run_train_job_no_warning_when_num_workers_within_shared_quota(
-    tmp_path, data, monkeypatch
-):
+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)
@@ -275,12 +325,8 @@ def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path,
     from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
 
     assert cond_norm.mean is not None and cond_norm.std is not None
-    np.testing.assert_allclose(
-        cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0
-    )
-    np.testing.assert_allclose(
-        cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0
-    )
+    np.testing.assert_allclose(cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0)
+    np.testing.assert_allclose(cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0)
     material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
     assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
 
@@ -318,11 +364,63 @@ def test_run_train_job_matches_uncached_output(tmp_path, data):
     cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False)
 
     for key in ("cond", "target", "sec_phys"):
-        np.testing.assert_allclose(
-            uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]
-        )
-        np.testing.assert_allclose(
-            uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]
-        )
+        np.testing.assert_allclose(uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"])
+        np.testing.assert_allclose(uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"])
     assert uncached["pdg_map"] == cached["pdg_map"]
     assert uncached["mat_map"] == cached["mat_map"]
+
+
+def _fitted_cond_norm(seed=0):
+    rng = np.random.default_rng(seed)
+    return Normalizer().fit(rng.normal(size=(64, COND_DIM)).astype(np.float32))
+
+
+@pytest.mark.parametrize(
+    "router_cfg",
+    [
+        {"enabled": False, "type": "energy", "n_experts": 4},
+        {"enabled": True, "type": "pdg", "n_experts": 4},
+    ],
+)
+def test_seed_energy_router_noop_when_not_an_enabled_energy_router(router_cfg):
+    cond_norm = _fitted_cond_norm()
+    echoed = []
+    _seed_energy_router(router_cfg, cond_norm, np.array([1.0, 2.0]), 3, echoed.append)
+    assert "centers_init" not in router_cfg
+    assert echoed == []
+
+
+def test_seed_energy_router_falls_back_to_default_and_warns_when_no_samples():
+    router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
+    cond_norm = _fitted_cond_norm()
+    echoed = []
+    _seed_energy_router(router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append)
+    assert "centers_init" not in router_cfg
+    assert len(echoed) == 1
+    assert "falls back to default centers" in echoed[0]
+
+
+def test_seed_energy_router_seeds_centers_from_data_quantiles():
+    router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
+    cond_norm = _fitted_cond_norm()
+    energy_idx = 3
+    # A grid of "raw" quantile values as setup_cache.energy_quantiles_from_sample
+    # would produce them: monotonically increasing, in the same (log-energy)
+    # units as the conditioning column being normalized against.
+    energy_quantiles = np.linspace(1.0, 10.0, 33).astype(np.float32)
+    echoed = []
+    _seed_energy_router(router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append)
+
+    assert "centers_init" in router_cfg
+    centers = np.asarray(router_cfg["centers_init"], dtype=np.float32)
+    assert centers.shape == (router_cfg["n_experts"],)
+
+    levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
+    raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
+    expected = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
+    np.testing.assert_allclose(centers, expected, rtol=1e-5)
+    # Quantile levels are increasing, and the normalizer's std is positive, so
+    # the seeded centers must preserve that order rather than e.g. reversing it.
+    assert np.all(np.diff(centers) > 0)
+    assert len(echoed) == 1
+    assert "seeded EnergyRouter centers" in echoed[0]

tests/test_render.py

@@ -4,10 +4,12 @@ from __future__ import annotations
 
 from pathlib import Path
 
+import numpy as np
 import pytest
 
 pytest.importorskip("plotstyle")
 
+from giant.analysis import render as render_mod  # noqa: E402
 from giant.analysis.reduced import Reduced  # noqa: E402
 
 
@@ -19,6 +21,152 @@ def _try_render(reduced: list[Reduced], out: Path):
     return render_all(out / "reduced", out / "plots")
 
 
+def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
+    reduced = [
+        Reduced(
+            "rg",
+            "router",
+            "router_gating",
+            "Router gating",
+            "pre-step energy [MeV]",
+            {
+                "n_experts": 2,
+                "log_x": True,
+                "router_type": "energy",
+                "rollout": {
+                    "centers": [1.0, 10.0, 100.0],
+                    "means": [[0.6, 0.4], [0.5, 0.5], [0.4, 0.6]],
+                },
+                "reference": {
+                    "centers": [1.0, 10.0, 100.0],
+                    "means": [[0.55, 0.45], [0.5, 0.5], [0.45, 0.55]],
+                },
+            },
+        ),
+        Reduced(
+            "rs",
+            "router",
+            "router_share",
+            "Router share",
+            "species",
+            {
+                "categories": ["e-", "gamma"],
+                "n_experts": 2,
+                "router_type": "energy",
+                "rollout": {"e-": [0.7, 0.3], "gamma": [0.2, 0.8]},
+                "reference": {"e-": [0.6, 0.4], "gamma": [0.3, 0.7]},
+            },
+        ),
+        Reduced(
+            "ru",
+            "router",
+            "unavailable",
+            "Router unavailable",
+            "x",
+            {"note": "router diagnostics unavailable: no router in this run"},
+        ),
+        Reduced(
+            "g4",
+            "marginals",
+            "grouped_hist",
+            "Grouped (4)",
+            "x",
+            {
+                "edges": [0, 1, 2],
+                "groups": {lbl: {"rollout": [1, 2], "reference": [2, 1]} for lbl in ("a", "b", "c", "d")},
+                "log_y": True,
+            },
+        ),
+        Reduced(
+            "sl",
+            "species",
+            "single_hist",
+            "Single (log-x)",
+            "x",
+            {"edges": [1, 10, 100], "rollout": [5, 1], "log_x": True, "log_y": True},
+        ),
+    ]
+    try:
+        pdfs = _try_render(reduced, tmp_path)
+    except RuntimeError as e:  # LaTeX missing at render time
+        pytest.skip(f"LaTeX rendering unavailable: {e}")
+    assert len(pdfs) == len(reduced)
+    assert all(p.exists() for p in pdfs)
+
+
+def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
+    # render_mod.subprocess *is* the stdlib subprocess module, so a blanket
+    # patch of .run would also swallow the real subprocess.run calls
+    # matplotlib's texmanager makes to compile LaTeX during savefig — only
+    # intercept the "gallery generate" call itself and pass everything else
+    # (LaTeX included) through to the real subprocess.run.
+    calls = []
+    real_run = render_mod.subprocess.run
+
+    def fake_run(*a, **k):
+        if a and a[0] and a[0][0] == "gallery":
+            calls.append((a, k))
+            return None
+        return real_run(*a, **k)
+
+    monkeypatch.setattr(render_mod.subprocess, "run", fake_run)
+    reduced = [
+        Reduced(
+            "s",
+            "species",
+            "single_hist",
+            "Single",
+            "x",
+            {"edges": [0, 1, 2], "rollout": [5, 1]},
+        )
+    ]
+    for r in reduced:
+        r.save(tmp_path / "reduced" / f"{r.id}.json")
+    try:
+        render_mod.render_all(tmp_path / "reduced", tmp_path / "plots", run_gallery=True)
+    except RuntimeError as e:
+        pytest.skip(f"LaTeX rendering unavailable: {e}")
+
+    assert len(calls) == 1
+    args, kwargs = calls[0]
+    assert args[0] == ["gallery", "generate", "--source", str(tmp_path / "plots")]
+    assert kwargs == {"check": True}
+
+
+def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkeypatch):
+    from giant.analysis import condor as condor_mod
+
+    run_dir = tmp_path / "run"
+    (run_dir / "reduced").mkdir(parents=True)
+
+    merge_calls = []
+    monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
+    meta = condor_mod.RunMeta(
+        rollout="rollout.parquet",
+        reference="reference.parquet",
+        run_dir=str(run_dir),
+        title="my-run",
+        plot_meta={"checkpoint": "ckpt/best.pt"},
+    )
+    monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))
+
+    Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]}).save(
+        run_dir / "reduced" / "s.json"
+    )
+
+    try:
+        pdfs = render_mod.render_run(run_dir)
+    except RuntimeError as e:
+        pytest.skip(f"LaTeX rendering unavailable: {e}")
+
+    assert merge_calls == [run_dir]
+    assert len(pdfs) == 1
+    plot_meta = (run_dir / "plots" / "species" / "s.yaml").read_text()
+    assert "checkpoint" in plot_meta
+    root_meta = (run_dir / "plots" / "metadata.yaml").read_text()
+    assert "my-run" in root_meta
+
+
 def test_render_one_of_each_kind(tmp_path: Path):
     reduced = [
         Reduced(
@@ -90,3 +238,133 @@ def test_render_one_of_each_kind(tmp_path: Path):
     assert len(pdfs) == len(reduced)
     assert all(p.exists() for p in pdfs)
     assert (tmp_path / "plots" / "metadata.yaml").exists()
+
+
+# ── pure-function helpers: no matplotlib figure needed ──────────────────
+
+
+def test_density_zero_total_returns_counts_unchanged():
+    counts = np.array([0.0, 0.0, 0.0])
+    out = render_mod._density(counts, np.array([0.0, 1.0, 2.0, 3.0]))
+    np.testing.assert_array_equal(out, counts)
+
+
+def test_density_normalizes_by_total_and_bin_width():
+    counts = [1, 3]
+    edges = np.array([0.0, 2.0, 4.0])  # bin width 2
+    out = render_mod._density(counts, edges)
+    np.testing.assert_allclose(out, np.array([1, 3]) / (4 * 2))
+
+
+def test_router_summary_disabled_is_off():
+    assert render_mod._router_summary({"enabled": False, "type": "energy"}) == "off"
+    assert render_mod._router_summary({}) == "off"
+
+
+def test_router_summary_enabled_formats_type_and_n_experts():
+    cfg = {"enabled": True, "type": "energy", "n_experts": 8}
+    assert render_mod._router_summary(cfg) == "energy×8"
+
+
+def test_figure_params_v2_basics_and_router_and_epoch():
+    mc = {
+        "stage1_model": {
+            "hidden_dim": 256,
+            "n_res_blocks": 4,
+            "generator": "flow",
+            "router": {"enabled": True, "type": "energy", "n_experts": 4},
+        },
+        "conditioning": {"particle": {"type": "physical"}},
+    }
+    run_meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10}
+    params = render_mod._figure_params(run_meta | {"model_config": mc})
+    assert params == {
+        "hidden_dim": 256,
+        "n_res_blocks": 4,
+        "mode": "flow",
+        "conditioning": "physical",
+        "router": "energy×4",
+        "epoch": 12,
+        "best_val_loss": 0.1235,
+        "steps": 10,
+    }
+
+
+def test_figure_params_v2_wgan_reports_noise_dim_not_steps():
+    mc = {
+        "stage1_model": {
+            "generator": "wgan",
+            "wgan": {"noise_dim": 32},
+        },
+    }
+    run_meta = {"model_config": mc, "steps": 10}
+    params = render_mod._figure_params(run_meta)
+    assert params["mode"] == "wgan"
+    assert params["noise_dim"] == 32
+    assert "steps" not in params
+
+
+def test_figure_params_v2_reports_mode_s2_only_when_it_differs():
+    same = {
+        "stage1_model": {"generator": "flow"},
+        "stage2_model": {"generator": "flow"},
+    }
+    assert "mode_s2" not in render_mod._figure_params({"model_config": same})
+
+    mixed = {
+        "stage1_model": {"generator": "flow"},
+        "stage2_model": {"generator": "wgan"},
+    }
+    params = render_mod._figure_params({"model_config": mixed})
+    assert params["mode_s2"] == "wgan"
+
+
+def test_figure_params_old_shape_basics():
+    run_meta = {
+        "model_config": {
+            "hidden_dim": 128,
+            "n_blocks": 3,
+            "mode": "ddpm",
+            "conditioning": "embedding",
+            "router": {"enabled": False},
+        },
+        "training_epoch": 5,
+        "best_val_loss": 0.5,
+        "steps": 20,
+    }
+    params = render_mod._figure_params(run_meta)
+    assert params == {
+        "hidden_dim": 128,
+        "n_blocks": 3,
+        "mode": "ddpm",
+        "conditioning": "embedding",
+        "router": "off",
+        "epoch": 5,
+        "best_val_loss": 0.5,
+        "steps": 20,
+    }
+
+
+def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps():
+    run_meta = {
+        "model_config": {"mode": "wgan", "noise_dim": 16},
+        "steps": 20,
+    }
+    params = render_mod._figure_params(run_meta)
+    assert params["noise_dim"] == 16
+    assert "steps" not in params
+
+
+def test_plot_metadata_includes_note_and_run_meta_parameters():
+    r = Reduced("u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"})
+    meta = render_mod._plot_metadata(r, {"title": "run-1", "checkpoint": "ckpt.pt"})
+    assert meta["note"] == "no router data"
+    assert meta["parameters"] == {"checkpoint": "ckpt.pt"}
+    assert "title" not in meta["parameters"]
+
+
+def test_plot_metadata_omits_parameters_when_run_meta_empty():
+    r = Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1]})
+    meta = render_mod._plot_metadata(r, {})
+    assert "parameters" not in meta
+    assert "note" not in meta

tests/test_rollout.py

@@ -121,12 +121,8 @@ def fake_material_props(monkeypatch):
     import giant.materials as gm
 
     fake = {
-        "G4_AIR": gm.MaterialProperties(
-            z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5
-        ),
-        "G4_PbWO4": gm.MaterialProperties(
-            z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
-        ),
+        "G4_AIR": gm.MaterialProperties(z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5),
+        "G4_PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7),
     }
     monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
     return fake
@@ -434,7 +430,7 @@ def _run_v3(
     max_tracks_per_event=100,
     seeds=None,
     conditioning="physical",
-    pdg_topn_map=None,
+    sec_type_topn_map=None,
     other_policy="sample",
     seed=0,
     stage1_ddpm_steps=1000,
@@ -461,7 +457,7 @@ def _run_v3(
         escape_threshold=escape_threshold,
         particle_conditioning=conditioning,
         material_conditioning=conditioning,
-        pdg_topn_map=pdg_topn_map,
+        sec_type_topn_map=sec_type_topn_map,
         other_policy=other_policy,
         seed=seed,
         stage1_ddpm_steps=stage1_ddpm_steps,
@@ -496,9 +492,7 @@ def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
 
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 @pytest.mark.parametrize("generator2", ["flow", "wgan"])
-def test_rollout_physical_target_decoder_generator_matrix(
-    fake_material_props, decoder, generator2
-):
+def test_rollout_physical_target_decoder_generator_matrix(fake_material_props, decoder, generator2):
     """Every (decoder, stage2 generator) combination under
     particle_type.target="physical" must run to completion and conserve
     energy."""
@@ -538,21 +532,19 @@ def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
     (giant.particles.particle_phys_array) become the secondary's identity —
     unlike "physical", not just a reporting label."""
     s1, s2 = _models_v3(decoder=decoder, target="onehot", emb_dim=3)
-    rec = _run_v3(s1, s2, pdg_topn_map=PDG_TOPN_MAP, other_policy="modal")
+    rec = _run_v3(s1, s2, sec_type_topn_map=PDG_TOPN_MAP, other_policy="modal")
     assert len(rec["event_id"]) > 0
     # Every spawned secondary's nominal pdg must be one decode_topn_class can
     # actually produce (the topn map's known classes + its "other" members).
-    possible = set(PDG_TOPN_MAP.class_map.keys()) | set(
-        PDG_TOPN_MAP.other_members.keys()
-    )
+    possible = set(PDG_TOPN_MAP.class_map.keys()) | set(PDG_TOPN_MAP.other_members.keys())
     secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
     assert secondary_pdgs <= possible
 
 
 def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
     s1, s2 = _models_v3(target="onehot", emb_dim=3)
-    with pytest.raises(RuntimeError, match="pdg_topn_map"):
-        _run_v3(s1, s2, pdg_topn_map=None)
+    with pytest.raises(RuntimeError, match="sec_type_topn_map"):
+        _run_v3(s1, s2, sec_type_topn_map=None)
 
 
 # --- conditioning.{particle,material}.type = "onehot" — a separate axis from
@@ -589,9 +581,7 @@ def _onehot_conditioning_models():
     return s1.eval(), s2.eval()
 
 
-def _run_onehot_conditioning(
-    pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP
-):
+def _run_onehot_conditioning(pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP):
     s1, s2 = _onehot_conditioning_models()
     cond, tgt, sec_phys = _norms()
     return rollout(
@@ -637,15 +627,93 @@ def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
         _run_onehot_conditioning(mat_topn_map=None)
 
 
+SEC_TYPE_TOPN_MAP_DIFFERENT_N = TopNMap(class_map={22: 0, 11: 1, -11: 2, 13: 3}, other_members={2112: 3, 2212: 1})
+
+
+def _run_conditioning_and_type_onehot_different_n_classes():
+    """Both conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" active at once, with
+    stage2_model.particle_type.n_classes deliberately different from
+    conditioning.particle.emb_dim (gitea #29)."""
+    cond_emb_dim = len(PDG_MAP)  # 3
+    type_n_classes = 5  # deliberately different from cond_emb_dim
+    particle_cfg = {"type": "onehot", "emb_dim": cond_emb_dim, "n_layers": 1}
+    material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
+    particle_type_cfg = {"target": "onehot", "n_classes": type_n_classes}
+
+    s1 = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+    ).eval()
+    sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", K_MAX, type_n_classes)
+    s2 = Stage2OneShot(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=particle_cfg,
+        material_cfg=material_cfg,
+        hidden_dim=32,
+        n_res_blocks=2,
+        sec_dim=sec_dim,
+        generator="flow",
+        time_dim=16,
+        k_max=K_MAX,
+        particle_type_cfg=particle_type_cfg,
+    ).eval()
+    # Sanity: the model's own type_dim followed n_classes, not cond_emb_dim.
+    assert s2.type_dim == type_n_classes
+
+    cond, tgt, sec_phys = _norms()
+    return rollout(
+        s1,
+        s2,
+        _oracle(),
+        _seeds(),
+        cond,
+        tgt,
+        sec_phys,
+        PDG_MAP,
+        MAT_MAP,
+        energy_cutoff=1.0,
+        max_steps=15,
+        steps=3,
+        batch_size=128,
+        max_tracks_per_event=100,
+        escape_threshold=1e9,
+        particle_conditioning="onehot",
+        material_conditioning="onehot",
+        pdg_topn_map=COND_PDG_TOPN_MAP,
+        mat_topn_map=COND_MAT_TOPN_MAP,
+        sec_type_topn_map=SEC_TYPE_TOPN_MAP_DIFFERENT_N,
+        other_policy="modal",
+    )
+
+
+def test_rollout_conditioning_and_type_onehot_with_different_n_classes(fake_material_props):
+    """gitea #29 end-to-end: conditioning.particle.type="onehot" and
+    stage2_model.particle_type.target="onehot" now use independently sized
+    top-N maps (stage2_model.particle_type.n_classes != conditioning.particle
+    .emb_dim), and rollout must decode secondaries using the type-side map,
+    not silently reuse the conditioning-side one (the pre-#29 bug)."""
+    rec = _run_conditioning_and_type_onehot_different_n_classes()
+    assert len(rec["event_id"]) > 0
+    possible = set(SEC_TYPE_TOPN_MAP_DIFFERENT_N.class_map.keys()) | set(
+        SEC_TYPE_TOPN_MAP_DIFFERENT_N.other_members.keys()
+    )
+    secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
+    assert secondary_pdgs <= possible
+
+
 @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
 def test_rollout_embedding_target_end_to_end(decoder):
     """particle_type.target="embedding" L1-snaps to the nearest row of the
     conditioning's own particle embedding table, so every resolved PDG must
     be a real member of the dense training vocab (pdg_map) — unlike
     "onehot", there is no "other" bucket to fall outside of."""
-    s1, s2 = _models_v3(
-        conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4
-    )
+    s1, s2 = _models_v3(conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4)
     rec = _run_v3(s1, s2, conditioning="embedding")
     assert len(rec["event_id"]) > 0
     secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())

tests/test_router.py

@@ -25,9 +25,7 @@ MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
 
 def _cond(B=8, pdg=3, mat=2):
     cond_cont = torch.randn(B, COND_DIM)
-    cond_cat = torch.stack(
-        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
-    )
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
     return cond_cont, cond_cat
 
 
@@ -78,9 +76,7 @@ def test_energy_router_gate_partition_of_unity():
 def test_energy_router_top1_matches_gate_argmax():
     router = EnergyRouter(n_experts=4)
     cond_cont, cond_cat = _cond(16)
-    assert torch.equal(
-        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
-    )
+    assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
 
 
 def test_energy_router_hardens_as_temperature_shrinks():
@@ -128,12 +124,8 @@ def test_energy_router_centers_init_wrong_length_raises():
 
 
 def test_energy_router_centers_init_respects_learn_centers_flag():
-    learned = EnergyRouter(
-        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True
-    )
-    fixed = EnergyRouter(
-        n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False
-    )
+    learned = EnergyRouter(n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True)
+    fixed = EnergyRouter(n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False)
     assert isinstance(learned.centers, torch.nn.Parameter)
     assert not isinstance(fixed.centers, torch.nn.Parameter)
 
@@ -160,9 +152,7 @@ def test_energy_router_learn_width_matches_fixed_temperature_at_init():
     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
-    )
+    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),
@@ -195,21 +185,15 @@ def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
         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"
-    )
+    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
-        )
+        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"
-    )
+    raise AssertionError("expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0")
 
 
 def test_energy_router_effective_width_stays_within_bounds():
@@ -246,9 +230,7 @@ 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
-    )
+    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)
@@ -264,9 +246,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
     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]
-    )
+    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
 
@@ -280,9 +260,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
 
 
 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
-    )
+    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)
@@ -419,9 +397,7 @@ def test_pdg_router_gate_partition_of_unity():
 def test_pdg_router_top1_matches_gate_argmax():
     router = PdgRouter(n_experts=4, pdg_vocab=3)
     cond_cont, cond_cat = _cond(16)
-    assert torch.equal(
-        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
-    )
+    assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
 
 
 def test_pdg_router_hardens_as_temperature_shrinks():
@@ -586,9 +562,7 @@ def test_process_router_gate_partition_of_unity():
 def test_process_router_top1_matches_gate_argmax():
     router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
     cond_cont, cond_cat = _cond(16)
-    assert torch.equal(
-        router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
-    )
+    assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
 
 
 def test_process_router_balance_loss_is_nonnegative_scalar():
@@ -663,16 +637,12 @@ def test_build_models_routed_with_process_router():
 
 
 def test_composed_router_n_experts_is_product():
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
     assert router.n_experts == 12
 
 
 def test_composed_router_gate_partition_of_unity():
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
     cond_cont, cond_cat = _cond(16, pdg=5)
     g = router.gate(cond_cont, cond_cat)
     assert g.shape == (16, 12)
@@ -709,9 +679,7 @@ def test_composed_router_top1_factors_into_per_axis_argmax():
 
 
 def test_composed_router_supports_different_expert_counts_per_axis():
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)])
     assert router.n_experts == 10
     cond_cont, cond_cat = _cond(8, pdg=5)
     assert router.gate(cond_cont, cond_cat).shape == (8, 10)
@@ -719,9 +687,7 @@ def test_composed_router_supports_different_expert_counts_per_axis():
 
 def test_composed_router_classify_loss_sums_sub_router_losses():
     """energy/pdg both default to zero, so the composed loss should too."""
-    router = ComposedRouter(
-        [EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
-    )
+    router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
     cond_cont, cond_cat = _cond(16, pdg=5)
     labels = torch.randint(0, 4, (16,))
     loss = router.classify_loss(cond_cont, cond_cat, labels)
@@ -964,9 +930,7 @@ def test_routed_stage1_eval_dispatch_matches_manual_grouping():
         idx = model.trunk.router.top1(cond_cont, cond_cat)
         manual = torch.zeros_like(x_t)
         for i in range(B):
-            manual[i] = model.trunk.experts[int(idx[i])](
-                x_t[i : i + 1], cond[i : i + 1]
-            )[0]
+            manual[i] = model.trunk.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
 
     torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
 

tests/test_sample.py

@@ -30,9 +30,7 @@ def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]
 
 def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
     cond_cont = torch.randn(B, COND_DIM)
-    cond_cat = torch.stack(
-        [torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
-    )
+    cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
     return cond_cont, cond_cat
 
 
@@ -43,12 +41,8 @@ def _conditioning_for(target: str) -> str:
     return "embedding" if target == "embedding" else "physical"
 
 
-def _stage2_oneshot(
-    target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2
-) -> Stage2OneShot:
-    particle_cfg, material_cfg = _particle_material_cfg(
-        _conditioning_for(target), emb_dim
-    )
+def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot:
+    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     particle_type_cfg = {"target": target}
     # build_models (giant/model/network.py) computes sec_dim this same way
     # before constructing Stage2OneShot — its own default (SEC_DIM, the
@@ -78,9 +72,7 @@ def _stage2_ar(
     k_max: int = 5,
     history: str = "markov",
 ) -> Stage2Autoregressive:
-    particle_cfg, material_cfg = _particle_material_cfg(
-        _conditioning_for(target), emb_dim
-    )
+    particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     return Stage2Autoregressive(
         pdg_vocab=pdg,
         mat_vocab=mat,
@@ -163,9 +155,7 @@ def test_sample_secondaries_flow_shapes_by_target(target):
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
-    sec_cont, sec_type, sec_valid = sample_secondaries(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
     assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
     assert sec_valid.shape == (B, K_MAX)
@@ -180,9 +170,7 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
-    sec_cont, sec_type, sec_valid = sample_secondaries_wgan(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries_wgan(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
     assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
     assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
     assert sec_valid.shape == (B, K_MAX)
@@ -196,15 +184,11 @@ def test_sample_secondaries_wgan_shapes_by_target(target):
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 def test_sample_secondaries_ar_shapes(target, generator, history):
     B, k_max, emb_dim = 4, 5, 6
-    decoder = _stage2_ar(
-        target, generator, emb_dim=emb_dim, k_max=k_max, history=history
-    )
+    decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max, history=history)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, k_max + 1, (B,))
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, k_max, CONT_SLOT_DIM)
     assert sec_type.shape == (B, k_max, _expected_type_dim(target, emb_dim))
     assert sec_valid.shape == (B, k_max)
@@ -220,9 +204,7 @@ def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.tensor([0, 2, k_max])
-    _, _, sec_valid = sample_secondaries_ar(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     for i, n in enumerate(n_sec_pred.tolist()):
         assert sec_valid[i, :n].all()
         assert not sec_valid[i, n:].any()
@@ -237,8 +219,6 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.tensor([0, 1, 1])
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
-        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
-    )
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
     assert sec_valid.tolist() == [[False], [True], [True]]

tests/test_setup_cache.py

@@ -14,9 +14,7 @@ from giant.data.transforms import Normalizer
 
 def _touch_parquet(path, n=1):
     path.parent.mkdir(parents=True, exist_ok=True)
-    pd.DataFrame(
-        {"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}
-    ).to_parquet(path)
+    pd.DataFrame({"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}).to_parquet(path)
     return path
 
 
@@ -29,9 +27,7 @@ def _normalizer(width=3):
 
 def _entry(n_train_steps=100, sample=None):
     sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) if sample is None else sample
-    return NormalizerEntry(
-        _normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample
-    )
+    return NormalizerEntry(_normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample)
 
 
 # ── sidecar_path ─────────────────────────────────────────────────────────
@@ -39,9 +35,7 @@ def _entry(n_train_steps=100, sample=None):
 
 def test_sidecar_path_single_file(tmp_path):
     f = tmp_path / "shard.parquet"
-    assert (
-        setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
-    )
+    assert setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
 
 
 def test_sidecar_path_directory(tmp_path):
@@ -51,10 +45,7 @@ def test_sidecar_path_directory(tmp_path):
 
 def test_sidecar_path_manifest(tmp_path):
     m = tmp_path / "pools" / "full.manifest"
-    assert (
-        setup_cache.sidecar_path(m)
-        == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
-    )
+    assert setup_cache.sidecar_path(m) == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
 
 
 # ── fingerprint_files ────────────────────────────────────────────────────
@@ -128,6 +119,11 @@ def test_save_load_round_trip_topn_maps(tmp_path):
     assert mat_m.class_map == {"G4_AIR": 0, "PbWO4": 1}
 
 
+def test_topn_key_unknown_axis_raises():
+    with pytest.raises(ValueError, match="unknown top-N map axis"):
+        setup_cache.topn_key("process", 4)
+
+
 def test_load_missing_sidecar_returns_none(tmp_path):
     data = _touch_parquet(tmp_path / "shard.parquet")
     assert setup_cache.load(data, [data]) is None
@@ -176,6 +172,25 @@ def test_load_invalidates_on_file_content_change(tmp_path):
     assert setup_cache.load(data, files) is None
 
 
+def test_load_returns_none_on_malformed_cache_body(tmp_path):
+    """format_version/dims/fingerprint all check out, but the cache body
+    itself doesn't match SetupCache.from_json's expected shape (e.g. hand-
+    edited or written by a version that changed a nested key) — a clean
+    miss, not a crash."""
+    data = _touch_parquet(tmp_path / "shard.parquet")
+    files = [data]
+    setup_cache.save(data, files, SetupCache.empty(files))
+
+    path = setup_cache.sidecar_path(data)
+    raw = json.loads(path.read_text())
+    raw["vocab"] = {"pdg_map": {"11": 0}}  # missing required "mat_map" key
+    path.write_text(json.dumps(raw))
+
+    echoed = []
+    assert setup_cache.load(data, files, echo=echoed.append) is None
+    assert any("malformed" in m for m in echoed)
+
+
 def test_load_soft_warns_on_git_hash_mismatch_but_still_hits(tmp_path, capsys):
     data = _touch_parquet(tmp_path / "shard.parquet")
     files = [data]
@@ -352,3 +367,11 @@ def test_compute_event_index_from_files_single_file_unaffected(tmp_path):
 
     np.testing.assert_array_equal(unique_ids, [5, 7])
     np.testing.assert_array_equal(counts, [2, 1])
+
+
+def test_compute_event_index_from_files_empty_file_list():
+    unique_ids, counts = setup_cache.compute_event_index_from_files([])
+    assert unique_ids.size == 0
+    assert counts.size == 0
+    assert unique_ids.dtype == np.int64
+    assert counts.dtype == np.int64

tests/test_steps_to_parquet.py

@@ -1,6 +1,6 @@
 import polars as pl
 
-from scripts import steps_to_parquet
+from giant.tools import steps_to_parquet
 
 
 def _frame() -> pl.DataFrame:
@@ -24,18 +24,14 @@ def _frame() -> pl.DataFrame:
 def test_e_sec_sums_child_first_step_energy():
     out, n_orphaned = steps_to_parquet._add_secondary_attributes(_frame())
     assert n_orphaned == 0
-    e_sec = dict(
-        zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
-    )
+    e_sec = dict(zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"]))
     assert e_sec[(1, 0, 0)] == 15.0  # one child, first-step pre_E 15
     assert e_sec[(1, 0, 1)] == 50.0  # two children, 20 + 30
 
 
 def test_e_sec_zero_when_no_children():
     out, _ = steps_to_parquet._add_secondary_attributes(_frame())
-    childless = out.filter(
-        (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
-    )
+    childless = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1))
     assert childless["e_sec"].item() == 0.0
 
 
@@ -69,9 +65,7 @@ def test_orphaned_child_track_is_dropped_not_nulled():
         }
     )
     out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
-    row = out.filter(
-        (pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)
-    )
+    row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
 
     assert n_orphaned == 1
     assert row["child_track_ids"].to_list() == [[2]]

tests/test_steps_to_parquet_parallel.py

@@ -2,7 +2,7 @@ import json
 import sys
 from pathlib import Path
 
-from scripts import steps_to_parquet_parallel
+from giant.tools import steps_to_parquet_parallel
 
 run_parallel = steps_to_parquet_parallel.run_parallel
 resolve_destination = steps_to_parquet_parallel.resolve_destination
@@ -124,31 +124,13 @@ def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
 def test_resolve_destination_uses_latest_schema(tmp_path):
     root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
     dest = resolve_destination(root_file, tmp_path, schema_override=None)
-    assert (
-        dest
-        == tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema3"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
 
 
 def test_resolve_destination_schema_override_wins(tmp_path):
     root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
     dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
-    assert (
-        dest
-        == tmp_path
-        / "processed"
-        / "steps"
-        / "gen1"
-        / "schema9"
-        / "pbwo4"
-        / "shard-000.parquet"
-    )
+    assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"
 
 
 def test_resolve_destination_errors_without_any_schema(tmp_path):

tests/test_train.py

@@ -5,6 +5,7 @@ import csv
 import math
 import tempfile
 from pathlib import Path
+from unittest.mock import MagicMock
 
 import pytest
 import torch
@@ -17,6 +18,7 @@ from giant.constants import (
     SEC_SLOT_DIM,
     X_DIM,
 )
+from giant.data.dataset import StepBatch
 from giant.model.network import build_critics, build_models
 from giant.training import (
     FlowDDPMStageTrainer,
@@ -74,9 +76,7 @@ def test_wandb_run_config_includes_full_cfg_and_param_counts():
         "stage1_model": {"generator": "flow"},
         "stage2_model": {"generator": "wgan"},
     }
-    wcfg = _wandb_run_config(
-        cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100}
-    )
+    wcfg = _wandb_run_config(cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100})
     assert wcfg["train"] == {"lr": 3e-4}
     assert wcfg["stage1_model"] == {"generator": "flow"}
     assert wcfg["stage2_model"] == {"generator": "wgan"}
@@ -162,9 +162,7 @@ def test_type_repr_shapes_and_values(target):
     expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim
     assert repr_.shape == (B, K, expected_width)
     if target == "physical":
-        assert torch.equal(
-            repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
-        )
+        assert torch.equal(repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM])
     if target == "onehot":
         assert torch.all(repr_.sum(-1) == 1.0)
 
@@ -180,9 +178,7 @@ def test_type_repr_shapes_and_values(target):
         ("embedding", "wgan"),
     ],
 )
-def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
-    target, generator
-):
+def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target, generator):
     """Regression test tying the refactor together: _assemble_stage2_real is
     now defined as _assemble_stage2_ar_target(...).flatten(1)."""
     B, emb_dim = 4, 6
@@ -192,12 +188,8 @@ def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(
     if target == "embedding":
         cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
     particle_type_cfg = {"target": target}
-    flat = _assemble_stage2_real(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    )
-    unflat = _assemble_stage2_ar_target(
-        sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim
-    )
+    flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
+    unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
     assert torch.equal(unflat.flatten(1), flat)
 
 
@@ -206,9 +198,7 @@ def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width():
     sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
     sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
     cond_enc = torch.nn.Module()
-    out = _assemble_stage2_ar_inputs(
-        sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim
-    )
+    out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim)
     assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
     assert out["has_prev"].shape == (B, K_MAX)
     assert out["remaining_frac"].shape == (B, K_MAX)
@@ -221,9 +211,7 @@ def test_relax_onehot_type_slice_grad_probe_populates_both_norms():
     B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
     x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
     grad_probe: dict[str, float] = {}
-    out = _relax_onehot_type_slice(
-        x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe
-    )
+    out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe)
     out.sum().backward()
     assert grad_probe["cont"] >= 0.0
     assert grad_probe["type"] >= 0.0
@@ -279,6 +267,12 @@ def _base_cfg():
             "k_max": K_MAX,
             "context_dim": 16,
             "n_sec": {"mode": "head", "lambda": 0.1},
+            # Explicit, not relying on the fallback default (which is
+            # "onehot", matching DEFAULT_CONFIG — see issues.md Issue 1):
+            # the "physical"-labelled cases below (and this fixture's own
+            # comment history) intend this as the base "physical" case,
+            # with "*_onehot"/"*_embedding" cases opting in explicitly.
+            "particle_type": {"target": "physical", "lambda": 1.0},
             "flow": {"time_dim": 16},
             "ddpm": {"time_dim": 16, "n_steps": 50},
             "wgan": {
@@ -324,9 +318,7 @@ def _fake_batches(n_batches, batch_size, seed=0):
         sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
         proc_idx = torch.zeros(batch_size, dtype=torch.long)
         sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
-        batches.append(
-            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
-        )
+        batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches
 
 
@@ -409,17 +401,13 @@ def _run_train(cfg, out_dir, resume_path=None):
         ),
         (
             "stage2_onehot_target_wgan",
-            lambda cfg: cfg["stage2_model"].__setitem__(
-                "particle_type", {"target": "onehot", "lambda": 1.0}
-            ),
+            lambda cfg: cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
         ),
         (
             "stage2_onehot_target_flow",
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -427,9 +415,7 @@ def _run_train(cfg, out_dir, resume_path=None):
             lambda cfg: (
                 cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
@@ -438,18 +424,14 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
             "ar_wgan_onehot",
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -461,9 +443,7 @@ def _run_train(cfg, out_dir, resume_path=None):
             lambda cfg: (
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
         (
@@ -473,9 +453,7 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["conditioning"]["material"].__setitem__("type", "embedding"),
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "embedding", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
             ),
         ),
         (
@@ -491,9 +469,7 @@ def _run_train(cfg, out_dir, resume_path=None):
                 cfg["stage1_model"].__setitem__("generator", "wgan"),
                 cfg["stage2_model"].__setitem__("generator", "flow"),
                 cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
-                cfg["stage2_model"].__setitem__(
-                    "particle_type", {"target": "onehot", "lambda": 1.0}
-                ),
+                cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
             ),
         ),
     ],
@@ -585,9 +561,7 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
         ema_decay=0.0,
         steps_per_epoch=4,
     )
-    trainer = WGANStageTrainer(
-        spec, models["stage1"], critics["stage1"], torch.device("cpu")
-    )
+    trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu"))
     assert trainer.model.n_sec_head is None
     batch = _fake_batches(1, 8)[0]
     stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -595,22 +569,86 @@ def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
 
 
 def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
-    spec = StageSpec(
-        name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0
-    )
+    spec = StageSpec(name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0)
     with pytest.raises(NotImplementedError):
         FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
 
 
+def test_stage_spec_from_config_omitted_decoder_and_particle_type_match_default_config():
+    """Regression for issues.md Issue 1: StageSpec.from_config's own fallback
+    defaults for stage2_model.decoder/particle_type must equal
+    DEFAULT_CONFIG's ("autoregressive" / "onehot"), not the old, now-wrong
+    ("one_shot" / "physical") literals a .get(key, default) call used to
+    supply when a hand-built cfg omitted these keys."""
+    cfg = _base_cfg()
+    del cfg["stage2_model"]["decoder"]
+    del cfg["stage2_model"]["particle_type"]
+    spec = StageSpec.from_config(cfg, "stage2", is_stage2=True, steps_per_epoch=1)
+    assert spec.decoder == "autoregressive"
+    assert spec.particle_type.target == "onehot"
+
+
+def _routed_stage1_trainer(lambda_balance, lambda_proc, lambda_entropy):
+    cfg = _base_cfg()
+    cfg["stage1_model"]["router"] = {
+        "enabled": True,
+        "type": "energy",
+        "n_experts": 3,
+        "temperature": 0.5,
+        "learn_centers": True,
+        "lambda_balance": lambda_balance,
+        "lambda_proc": lambda_proc,
+        "lambda_entropy": lambda_entropy,
+    }
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
+    return trainers["stage1"]
+
+
+def test_router_aux_losses_skipped_when_lambda_zero_but_run_when_positive():
+    """Gitea #31: FlowDDPMStageTrainer._compute must not call
+    router.balance_loss/classify_loss/entropy_loss when the corresponding
+    lambda is 0 (the default) -- those calls do their own router.gate(...)
+    forward pass that is wasted once the term is masked out of the total
+    loss anyway. Checked both ways: zero lambdas must skip all three calls,
+    positive lambdas must still make them (the guard must not accidentally
+    suppress the real path)."""
+    batch = _fake_batches(1, 4)[0]
+    device = torch.device("cpu")
+
+    trainer_zero = _routed_stage1_trainer(0.0, 0.0, 0.0)
+    router_zero = trainer_zero.router
+    router_zero.balance_loss = MagicMock(wraps=router_zero.balance_loss)
+    router_zero.classify_loss = MagicMock(wraps=router_zero.classify_loss)
+    router_zero.entropy_loss = MagicMock(wraps=router_zero.entropy_loss)
+    stats_zero = trainer_zero.step(batch, device, global_step=1)
+    assert router_zero.balance_loss.call_count == 0
+    assert router_zero.classify_loss.call_count == 0
+    assert router_zero.entropy_loss.call_count == 0
+    assert stats_zero["loss_balance"] == 0.0
+    assert stats_zero["loss_proc"] == 0.0
+    assert stats_zero["loss_entropy"] == 0.0
+
+    trainer_pos = _routed_stage1_trainer(0.1, 0.1, 0.01)
+    router_pos = trainer_pos.router
+    router_pos.balance_loss = MagicMock(wraps=router_pos.balance_loss)
+    router_pos.classify_loss = MagicMock(wraps=router_pos.classify_loss)
+    router_pos.entropy_loss = MagicMock(wraps=router_pos.entropy_loss)
+    trainer_pos.step(batch, device, global_step=1)
+    assert router_pos.balance_loss.call_count == 1
+    assert router_pos.classify_loss.call_count == 1
+    assert router_pos.entropy_loss.call_count == 1
+
+
 # --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
 
 
 @pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"])
 @pytest.mark.parametrize("history", ["markov", "attention"])
 @pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
-def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
-    teacher_forcing, history, stage2_generator
-):
+def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(teacher_forcing, history, stage2_generator):
     """v0.3.0 step 7: history='attention' and teacher_forcing in
     {'scheduled', 'never'} must actually train — a stage-2 AR trainer.step()
     must run and produce a finite loss, for every {history} x
@@ -629,9 +667,7 @@ def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(
     model_config = _model_config(cfg)
     models = build_models(model_config)
     critics = build_critics(model_config)
-    trainers = build_stage_trainers(
-        cfg, models, critics, torch.device("cpu"), total_train_batches=4
-    )
+    trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
     trainer = trainers["stage2"]
     batch = _fake_batches(1, 4)[0]
     stats = trainer.step(batch, torch.device("cpu"), global_step=1)
@@ -665,9 +701,7 @@ def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
         with open(out_dir / "metrics.csv", newline="") as f:
             rows = list(csv.DictReader(f))
         assert len(rows) == cfg["train"]["epochs"]
-        loss_col = (
-            "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
-        )
+        loss_col = "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
         assert all(math.isfinite(float(r[loss_col])) for r in rows)
 
 
@@ -706,7 +740,7 @@ def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
 
 
 def test_wgan_physical_omits_grad_norm_slice_columns():
-    cfg = _base_cfg()  # default stage2_model has no particle_type -> "physical"
+    cfg = _base_cfg()  # _base_cfg's stage2_model.particle_type.target is "physical"
     with tempfile.TemporaryDirectory() as tmp:
         out_dir = Path(tmp) / "run"
         _run_train(cfg, out_dir)

tests/test_transforms.py

@@ -162,9 +162,7 @@ def test_local_frame_rotation_normalizes_non_unit_pre_dir():
     post_dir = rng.standard_normal((N, 3)).astype(np.float32)
     post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
 
-    pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(
-        np.float32
-    )
+    pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(np.float32)
     expected = local_frame_rotation(pre_dir_unit, post_dir)
     result = local_frame_rotation(pre_dir_scaled, post_dir)
     np.testing.assert_allclose(result, expected, atol=1e-4)
@@ -200,14 +198,10 @@ def test_reconstruct_post_pos_straight_line():
     step_length = rng.uniform(0.1, 5.0, size=N).astype(np.float32)
     post_pos = pre_pos + step_length[:, None] * pre_dir
 
-    travel_dir_local = local_frame_rotation(
-        pre_dir, travel_direction(pre_pos, post_pos)
-    )
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
     np.testing.assert_allclose(travel_dir_local, np.tile([0, 0, 1], (N, 1)), atol=1e-4)
 
-    reconstructed = reconstruct_post_pos(
-        pre_pos, pre_dir, step_length, travel_dir_local
-    )
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
     np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
 
 
@@ -221,12 +215,8 @@ def test_reconstruct_post_pos_general_roundtrip():
     post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
     step_length = np.linalg.norm(post_pos - pre_pos, axis=1).astype(np.float32)
 
-    travel_dir_local = local_frame_rotation(
-        pre_dir, travel_direction(pre_pos, post_pos)
-    )
-    reconstructed = reconstruct_post_pos(
-        pre_pos, pre_dir, step_length, travel_dir_local
-    )
+    travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
+    reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
     np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
 
 
@@ -356,9 +346,7 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
 
 
 def test_build_features_proc_idx_zero_without_proc_map():
-    data = _minimal_step_data(
-        3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
-    )
+    data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
     *_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map)
@@ -367,9 +355,7 @@ def test_build_features_proc_idx_zero_without_proc_map():
 
 
 def test_build_features_proc_idx_looks_up_proc_map():
-    data = _minimal_step_data(
-        3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
-    )
+    data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
     proc_map = {"compt": 0, "phot": 1, "eIoni": 2}
 
@@ -396,9 +382,7 @@ def test_build_features_require_secondaries_ok_when_no_secondaries():
     data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
     pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
 
-    _, _, _, _, sec_cont, *_ = build_features(
-        data, pdg_map, mat_map, require_secondaries=True
-    )
+    _, _, _, _, sec_cont, *_ = build_features(data, pdg_map, mat_map, require_secondaries=True)
 
     assert not sec_cont.any()
 
@@ -413,11 +397,7 @@ def fake_material_props(monkeypatch):
     in by the user (see giant.materials.MaterialPropertiesNotFilledError)."""
     import giant.materials as gm
 
-    fake = {
-        "PbWO4": gm.MaterialProperties(
-            z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
-        )
-    }
+    fake = {"PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7)}
     monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
     return fake
 
@@ -455,9 +435,7 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props):
     assert cond_cont.shape[1] == COND_DIM
     mass, charge = particle_mass_charge(11)
     expected_log_mass = log_transform(np.array([mass]))[0]
-    np.testing.assert_allclose(
-        cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5
-    )
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5)
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], charge)
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 2], 75.6)  # z_eff
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 3], 205.3)  # a_eff
@@ -500,9 +478,7 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
         material_conditioning="physical",
     )
 
-    np.testing.assert_allclose(
-        cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
-    )
+    np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])))
     np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0])
 
 
@@ -714,9 +690,7 @@ def test_welford_accumulator_matches_naive_running_mean_reference():
     naive_M2 = np.zeros(F)
     naive_n = 0
     for chunk in chunks:
-        naive_mean, naive_M2, naive_n = naive_update(
-            naive_mean, naive_M2, naive_n, chunk
-        )
+        naive_mean, naive_M2, naive_n = naive_update(naive_mean, naive_M2, naive_n, chunk)
 
     acc = _WelfordAccumulator(F)
     for chunk in chunks:

tests/test_validate.py

@@ -2,6 +2,7 @@ import numpy as np
 import torch
 
 from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+from giant.data.dataset import StepBatch
 from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim
 from giant.validate import validate_marginals
 
@@ -46,8 +47,7 @@ def _tiny_models(particle_type_cfg: dict | None = None):
 
 
 def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
-    """A val_loader matching StreamingStepsDataset's 7-tuple batch shape:
-    (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx)."""
+    """A val_loader matching StreamingStepsDataset's StepBatch shape."""
     batches = []
     for _ in range(n_batches):
         cond_cont = torch.randn(B, COND_DIM)
@@ -57,9 +57,7 @@ def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int
         sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
         proc_idx = torch.zeros(B, dtype=torch.long)
         sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
-        batches.append(
-            (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)
-        )
+        batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
     return batches
 
 
@@ -71,9 +69,7 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch
     s1, s2 = _tiny_models()
     loader = _loader(n_sec_value=0)
 
-    def _fake_resolve_n_sec(
-        stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred
-    ):
+    def _fake_resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred):
         return torch.zeros(cond_cont.size(0), dtype=torch.long)
 
     monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)

tests/test_wgan.py

@@ -152,9 +152,7 @@ def test_sample_secondaries_wgan_shape():
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.randint(0, K_MAX, (B,))
-    sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(
-        model, cond_cont, cond_cat, stage1_out, n_sec_pred
-    )
+    sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(model, cond_cont, cond_cat, stage1_out, n_sec_pred)
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_phys.shape == (B, K_MAX, 2)
     assert sec_valid.shape == (B, K_MAX)
@@ -183,9 +181,7 @@ def test_gradient_penalty_masked():
     mask = _mask(B, n_sec)
     real = torch.randn(B, SEC_DIM) * mask
     fake = torch.randn(B, SEC_DIM) * mask
-    gp = gradient_penalty(
-        lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask
-    )
+    gp = gradient_penalty(lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask)
     assert gp.item() >= 0.0
 
 
@@ -195,9 +191,7 @@ def test_critic_loss_scalar_and_grad():
     cond_cont, cond_cat = _cond(B)
     real = torch.randn(B, X_DIM)
     fake = torch.randn(B, X_DIM)
-    loss = critic_loss(
-        lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0
-    )
+    loss = critic_loss(lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0)
     assert loss.shape == ()
     loss.backward()
     assert any(p.grad is not None for p in critic.parameters())

uv.lock

@@ -357,6 +357,105 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ae/8c/469afb6465b853afff216f9528ffda78a915ff880ed58813ba4faf4ba0b6/contourpy-1.3.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b7448cb5a725bb1e35ce88771b86fba35ef418952474492cf7c764059933ff8b", size = 203831, upload-time = "2025-07-26T12:02:51.449Z" },
 ]
 
+[[package]]
+name = "coverage"
+version = "7.15.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/be/c3/4f2195f512fb172aa425a8803a874b2baa9ba7f80ff7b6080998761fc701/coverage-7.15.4.tar.gz", hash = "sha256:0548198fff07ccf4faf469520bce1c2eceb1ce3e62891921138dec10907f9d00", size = 936952, upload-time = "2026-08-06T13:50:24.442Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/1d/48/bc8d4ba7b37551a767bd863f15b3f80182b271c2f55975356f5f7dbe94c2/coverage-7.15.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d4fedd1f7f428f9fe83b1ead5e7cc87a43427be31aadafbac3ac0636dc7abb22", size = 222543, upload-time = "2026-08-06T13:47:37.562Z" },
+    { url = "https://files.pythonhosted.org/packages/20/dd/88d6f83f1fffc974a3691a34a97951c5b12df7512a6782c5963883cbc058/coverage-7.15.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:37e2f0cdf58e2e1fed4e4d5a8f8786ae2f7eb80b478016876667dc4a01d60a97", size = 222905, upload-time = "2026-08-06T13:47:38.927Z" },
+    { url = "https://files.pythonhosted.org/packages/bd/5c/54ee0d4748585bb0acab9891cd8d92f2d3593165b4e59fc9de113bfb3140/coverage-7.15.4-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:fb55d0e70bb15f2e81477613627286581414693d74ac7963c93a790dd453ca9d", size = 254407, upload-time = "2026-08-06T13:47:40.488Z" },
+    { url = "https://files.pythonhosted.org/packages/8c/3f/f0642a372f494bd0d7dad3b497083b910194a5f1c88be2c94fef707c3b59/coverage-7.15.4-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:899b9da30f3c6c336566e3707495bb23e8302d39d862f01fa78c48b99b9437e2", size = 257145, upload-time = "2026-08-06T13:47:41.931Z" },
+    { url = "https://files.pythonhosted.org/packages/71/17/8b46d0ed68251016002ec972c8fc0119961a765d0984cafb8bf317c43758/coverage-7.15.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d15715e8c46552827e5e4f30a35575a2dbcad14454cf3284c54483946bd16931", size = 258257, upload-time = "2026-08-06T13:47:43.527Z" },
+    { url = "https://files.pythonhosted.org/packages/30/b8/8498a0e72d0adbe15477dd07463d2b3bb2c9f6a4815e8589e50939e2c3ae/coverage-7.15.4-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:002a438859f7b430bc99afeaf01a6d187dad1d0dc907b64cdeffc632a5db8fd8", size = 260517, upload-time = "2026-08-06T13:47:45.121Z" },
+    { url = "https://files.pythonhosted.org/packages/41/e1/7dce19c3bdb1e3dd63e769508216500edad81bd5f69a26d724e32aceaf78/coverage-7.15.4-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e4193a04b518f7968f3099755f5509ee7cccc6dc2b92a6b14841934d22e222c9", size = 254785, upload-time = "2026-08-06T13:47:46.541Z" },
+    { url = "https://files.pythonhosted.org/packages/dd/b1/e1494703c675a2561723cd9b89f45c9168782c31280c611b1f767851e57c/coverage-7.15.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e98dcc55d572b38e69d117da7e8e8efb8500f1f5eaf81ecd460a63220790b839", size = 256176, upload-time = "2026-08-06T13:47:48.155Z" },
+    { url = "https://files.pythonhosted.org/packages/73/76/a5629d270fb638a43a4b10466f51e2f49d532c1aa4da2913cbbb150bbe0a/coverage-7.15.4-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:af6c538498ce66c10d3fd541c2a8d5b03da5850355add34e6cba564210cb9e72", size = 254321, upload-time = "2026-08-06T13:47:49.757Z" },
+    { url = "https://files.pythonhosted.org/packages/ff/4f/9c44447218435d5766b911534f9d798144a5560f85e9a54ebe5f3f5d19f9/coverage-7.15.4-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:1d10025d96ea89fc2f73714dbc4cbd433fe012c1ac9e23f895d7728b238b6e52", size = 258390, upload-time = "2026-08-06T13:47:51.248Z" },
+    { url = "https://files.pythonhosted.org/packages/de/36/c1e127616fb3fa18a9ff71e76c417f2fd7424332a4870015ac224ef4c039/coverage-7.15.4-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:d802e1947603162ded419bff83ac7489820355d2b856dfb09206574e3a37ac0c", size = 253894, upload-time = "2026-08-06T13:47:52.816Z" },
+    { url = "https://files.pythonhosted.org/packages/e9/b9/fdb92c8ae7a8bb9b850cc253b7b3b9c8526f68130002048b5671cd510d09/coverage-7.15.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c2de40895718f91951b86712b4c5b694acaf9a0a49be13874896f599a1eed3f4", size = 255763, upload-time = "2026-08-06T13:47:54.296Z" },
+    { url = "https://files.pythonhosted.org/packages/6f/c0/a7d51b2587c7bdb76e71b0896d2565bf7d60436b5122fc83e511adb1f7cd/coverage-7.15.4-cp312-cp312-win32.whl", hash = "sha256:5c3431b2161279b7db5c2a1aa58ae02e5cb8c3c42d93a5094be3f5537bd5b11b", size = 224597, upload-time = "2026-08-06T13:47:56.074Z" },
+    { url = "https://files.pythonhosted.org/packages/49/b9/5c5f80cc55f5acaaca6dee677626bfcec8c87204a7809b438b08e84f4571/coverage-7.15.4-cp312-cp312-win_amd64.whl", hash = "sha256:6befeab5fb2b51c958ca4ac6c5d141a1e8240f4f76e46350f1911963deda49cd", size = 225135, upload-time = "2026-08-06T13:47:57.52Z" },
+    { url = "https://files.pythonhosted.org/packages/47/e4/2a4561f89ff6bf7c925c287d0f2cce8bdf139c3a33735c87e3203401cf94/coverage-7.15.4-cp312-cp312-win_arm64.whl", hash = "sha256:67bc345491ab55b837277d76f5775d057e8c7f1ac44d890d8c2c82adde258c6f", size = 224515, upload-time = "2026-08-06T13:47:58.977Z" },
+    { url = "https://files.pythonhosted.org/packages/f1/84/651a9310859673aaa3b3203f1aa1641ca60fcf2494683e1c9474c7172780/coverage-7.15.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c705b28feb2775dc82a25f1d473a370bc37ff93f5177f4e29ce2425f560f6921", size = 222565, upload-time = "2026-08-06T13:48:00.796Z" },
+    { url = "https://files.pythonhosted.org/packages/82/f9/4dcf700137e8af550670f4d74d1b63828ce93e1e2b05e5f10710eb2ea987/coverage-7.15.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3ff205ab5e3ecc670f6a4dd19d9cbf12ede53dd41cfc1e15716ec961ea6d314e", size = 222936, upload-time = "2026-08-06T13:48:02.391Z" },
+    { url = "https://files.pythonhosted.org/packages/07/4a/612ff1e780b3fbfd637486f542f84adc5503873d8b5d279dec1ffeef9414/coverage-7.15.4-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5172326e861a38b48b48befca15e0f477a26b283337a33a739c8fed229934e36", size = 253926, upload-time = "2026-08-06T13:48:04.382Z" },
+    { url = "https://files.pythonhosted.org/packages/b0/04/d1cff1c2ead4708a6a79c01d3736b6a25bd38a36678398f72a8dd33dfad9/coverage-7.15.4-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:12b59c90084e3234fb11184886bf4a40f4f16a8c8f867be2e087b81f8e8868d4", size = 256523, upload-time = "2026-08-06T13:48:05.996Z" },
+    { url = "https://files.pythonhosted.org/packages/b9/80/d34e13fb4b293cbdb9665838cf5522077b8ad14ef947550631a4bced36a5/coverage-7.15.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:349062d66f00b40fa2c1c222438bad25fabf755631b5d82937fe985c8008615c", size = 257759, upload-time = "2026-08-06T13:48:08.036Z" },
+    { url = "https://files.pythonhosted.org/packages/0f/e7/2c5fe7636fdb0732fe0f09f308a5b066864078b7fc61f6678e8478554f2e/coverage-7.15.4-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4256ced708e598e05209bc1a8ab4074e04a51dba4c62fb45926a229af675ace7", size = 259890, upload-time = "2026-08-06T13:48:09.834Z" },
+    { url = "https://files.pythonhosted.org/packages/92/28/9689f0858dfff59c2ea688938ab9fa2925631235df67126a42b6c5c70ae1/coverage-7.15.4-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d80f974b20782d9612c8b4c9beeca867074c7cf4079d1419843fa25a26428b25", size = 254121, upload-time = "2026-08-06T13:48:11.459Z" },
+    { url = "https://files.pythonhosted.org/packages/f9/e2/785077c230c157243eb5aa9a26c3be260ecd02001bead54a3cada3df8e03/coverage-7.15.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2e179f19bfe1d31f8eeeaa12990194d761c4f62f0759661000bca6cd8729f40b", size = 255891, upload-time = "2026-08-06T13:48:13.209Z" },
+    { url = "https://files.pythonhosted.org/packages/d4/90/e20371b17b40f912f21305c2db2f30efa3de306f7320fc916804872c85a4/coverage-7.15.4-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:8bc16bb47b7679670eceff71d78bfb7d6e5b143f6c2cd117487ec7c75e0d4b78", size = 253859, upload-time = "2026-08-06T13:48:14.736Z" },
+    { url = "https://files.pythonhosted.org/packages/05/49/25371987ee459a5f67c0427fb75c74f9358e65f2c71fe75bf41c1b6c5fcb/coverage-7.15.4-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:1cd685005cd2c4200adfc14cf39a603b9320efab3f18a8f7f156d20c9cc3345f", size = 258011, upload-time = "2026-08-06T13:48:16.464Z" },
+    { url = "https://files.pythonhosted.org/packages/30/6e/32e67467f6154bf4f1c4f63b05acc5097cba4237d45bbeeea446b52e8ac1/coverage-7.15.4-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:337399ad2c93b3acd2a937627dae8b3e86b66707cd3d3e856347999aadf1ef8d", size = 253676, upload-time = "2026-08-06T13:48:18.493Z" },
+    { url = "https://files.pythonhosted.org/packages/03/c1/8b24192e89286399765155251f99ee9f070a9d637109018ac23d99b99f6f/coverage-7.15.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:96e257121228ec5cd2bb919276e94ac11074471bc37d68dbae0e8308cce15fff", size = 255453, upload-time = "2026-08-06T13:48:20.057Z" },
+    { url = "https://files.pythonhosted.org/packages/16/6f/8b41ebdf67c87854e17c035336a90f1cfbad0c14c2a584301be6ff148718/coverage-7.15.4-cp313-cp313-win32.whl", hash = "sha256:c65a9e0dfc6143491879da4e13b5e30f8be192055de508d737fb14601edbd22c", size = 224605, upload-time = "2026-08-06T13:48:21.655Z" },
+    { url = "https://files.pythonhosted.org/packages/e0/e2/2946c7f0b42b152ecb21ff1bdad72e3d301e790c0c487e4a86e8c9f69347/coverage-7.15.4-cp313-cp313-win_amd64.whl", hash = "sha256:2ff8f5e9b8f7a94f0c11c45631eee103dbcb7d63274edd12c56efe1be690b3b4", size = 225148, upload-time = "2026-08-06T13:48:23.376Z" },
+    { url = "https://files.pythonhosted.org/packages/9e/83/3f4a69957f48ae7a0aba76c34743f88963d607b19e03f3f8e66f91cae0f9/coverage-7.15.4-cp313-cp313-win_arm64.whl", hash = "sha256:6e0a8a5083b096487d6cfced94cdd514d8f5db6f113610fb36c0620edb1028cf", size = 224536, upload-time = "2026-08-06T13:48:25.117Z" },
+    { url = "https://files.pythonhosted.org/packages/ea/ac/748cf29eeb2d6be34a3176ce26a4f49e38085ee08e8935f05f6f26ed7e0f/coverage-7.15.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:770e9325ab5ea6d56f77e59b29ecfe0ac20b57a82a601876f90494a4dda0386f", size = 222608, upload-time = "2026-08-06T13:48:26.806Z" },
+    { url = "https://files.pythonhosted.org/packages/0b/02/1abbf5c984677b0aa439cdacaccbf38d248939d8ef8fe1cc7a50d73edb77/coverage-7.15.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d12b33a3a50a1676b7784dc8d00a0c6d66a9f2add4b85a041c19b6a7e53ef23c", size = 222940, upload-time = "2026-08-06T13:48:28.432Z" },
+    { url = "https://files.pythonhosted.org/packages/eb/e1/ff8f9f53d9fcf586125b55d0b1f04ec1c14955fee41e83d5814bee141bb5/coverage-7.15.4-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5669c8378ebde86f5def7a25d29586631b58acc27ffde04399f678f3dfc6e082", size = 253985, upload-time = "2026-08-06T13:48:29.995Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/26/595759762e514e81be1d7d01ed03444303bcd152226a6529998d253f9201/coverage-7.15.4-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:ff97a14362eef486483ed44042ca2027ea257df6ff768e62358ee0c9776925ac", size = 256492, upload-time = "2026-08-06T13:48:31.634Z" },
+    { url = "https://files.pythonhosted.org/packages/24/68/b79aabac54d482be23b5fcdd4f4662bff24a78edc4ee29201726929936d5/coverage-7.15.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5a325e815318638aed1655d9c06e6d7c2d3d46c09231ce988070428a8762d734", size = 257837, upload-time = "2026-08-06T13:48:33.186Z" },
+    { url = "https://files.pythonhosted.org/packages/09/0f/bf7f297885a5bf6fd71e5782404e0ff059ca09e8711ceb3a08544abde45a/coverage-7.15.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:474223409d88eb20d2d6a0d37ea60e8647a65a90cc008dc1f0410af5f64f1e0d", size = 260152, upload-time = "2026-08-06T13:48:34.75Z" },
+    { url = "https://files.pythonhosted.org/packages/fd/f1/296744e854ff8368542343457414380465e9ceefb9192342feb9d3bc461d/coverage-7.15.4-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7f2f62ae3cd189dd2e13aece758c57b3eecbd27be070dbd4cbd10936049e5dbf", size = 253978, upload-time = "2026-08-06T13:48:36.434Z" },
+    { url = "https://files.pythonhosted.org/packages/55/b0/bbdb2e9057493e66220a2e149ca2d301ba0e3a58a83bd6b90de9826d16f3/coverage-7.15.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:39ece820e29e0a2ba34b3ecb3be83c27e997eed8926f2ba6fe7ce7a0bda5843b", size = 255846, upload-time = "2026-08-06T13:48:38.317Z" },
+    { url = "https://files.pythonhosted.org/packages/96/e4/38015b2b6d21258713bd17e76b59d033b191efb5703589cffd037dfbca20/coverage-7.15.4-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:f21b56dcace11dfe013014201f577dcd592b2a9b72182d930361b47cf6f73f25", size = 253808, upload-time = "2026-08-06T13:48:39.993Z" },
+    { url = "https://files.pythonhosted.org/packages/0b/64/0d515c1e60ee6fbfd1a0e79c07cd87d388a233b7adc37758735677203808/coverage-7.15.4-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:93a3a0b662abcc10c73a47cbc72cd60f63618d6989fb2d1286e50eacd974f303", size = 258081, upload-time = "2026-08-06T13:48:41.971Z" },
+    { url = "https://files.pythonhosted.org/packages/91/71/04d9e7a3642146c6351338aef4ef85ab11dbbb54744c13245caba1aad1c0/coverage-7.15.4-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:141fae2cabf5569b782c10afc4c850ce10f618c13f8db54765cba99cc839da1f", size = 253624, upload-time = "2026-08-06T13:48:43.731Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/a7/6c28b74c81ebff66987b0e2522ba5cffa3e90b0c33cb6a2eb264d4ee8cf1/coverage-7.15.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:81294c7e6ab30c5f74c0353b11b2fd6320e72d9bee6ac73b357caa8b916323a5", size = 255280, upload-time = "2026-08-06T13:48:45.58Z" },
+    { url = "https://files.pythonhosted.org/packages/52/af/bc19996a7014b98d7bbb0f0939453c67074af65784a3aa16a789a07381fa/coverage-7.15.4-cp314-cp314-win32.whl", hash = "sha256:7bbd7d6418e0dab31a206af5203bd43ae36edb8e7fba1940b055d3e9249290d7", size = 224768, upload-time = "2026-08-06T13:48:47.525Z" },
+    { url = "https://files.pythonhosted.org/packages/ee/90/219484e476d6e101ba0a444852579e05f5b75c37c611a42ed1190f73ef62/coverage-7.15.4-cp314-cp314-win_amd64.whl", hash = "sha256:f0204ed122758782970526057093f448051a39db9d810d4e344bb87a3546f425", size = 225259, upload-time = "2026-08-06T13:48:49.513Z" },
+    { url = "https://files.pythonhosted.org/packages/b7/66/fa77daf4e383e5f776dac62c2409b6af81910ae6fe326bd5170dba74cc63/coverage-7.15.4-cp314-cp314-win_arm64.whl", hash = "sha256:9e71e7bc71c686a123347ae47a0de33a175e797a85bb57b791492adf4eec8ed8", size = 224684, upload-time = "2026-08-06T13:48:51.235Z" },
+    { url = "https://files.pythonhosted.org/packages/58/5b/f03bf0ce362bbf3f785fa5219620d00778d4ac6fc9e407734828e9c672f6/coverage-7.15.4-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7c922735321eef3f87c280a3d39afff6b646723a2880b862cda4ac7a093b8aa8", size = 223338, upload-time = "2026-08-06T13:48:52.896Z" },
+    { url = "https://files.pythonhosted.org/packages/0f/76/e77d0ae22501831cc9f92193e8a957a5caa1dd177f90a6d1d9b106242d92/coverage-7.15.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f41c17c4668a655ce96d090d8d5ffdc24ef64b5a02f9753884d08483e8a4a41a", size = 223609, upload-time = "2026-08-06T13:48:54.688Z" },
+    { url = "https://files.pythonhosted.org/packages/82/1a/b1f089da8d38ac612fa2dd6dc7f4a1a7657d12f3e261d2996edd3a838d0b/coverage-7.15.4-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:46822e9b6ff1c6a72b518c162c44a8f45a61a1d609c51084bf5b16c023c5037b", size = 264970, upload-time = "2026-08-06T13:48:56.403Z" },
+    { url = "https://files.pythonhosted.org/packages/bf/31/e66d98d6e9c7fcc88470f1e234eaf6b1950dc0dfbf797f7282c1c861da24/coverage-7.15.4-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3d6f4955b73b5445271379a59e3792b0d978f42d4a01e0cf7a67d9c33a3bb0a5", size = 267088, upload-time = "2026-08-06T13:48:58.41Z" },
+    { url = "https://files.pythonhosted.org/packages/59/a1/ae94eb2c541add426378408379f233591e069040b1e2cdb33df9498a0682/coverage-7.15.4-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3fc9e047706fb4a9abb54f719d3aa643e80e5bb3818182c40aee01ac0f0247ba", size = 269508, upload-time = "2026-08-06T13:49:00.42Z" },
+    { url = "https://files.pythonhosted.org/packages/9c/c7/88a10694a1c6a213569766aba9f25847b28155d4ac731b13226db216356d/coverage-7.15.4-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:05e491d4f3165d62d4f5c8fd48dfeabf2ae8f42cbbd484319af33ea851b78982", size = 270629, upload-time = "2026-08-06T13:49:02.234Z" },
+    { url = "https://files.pythonhosted.org/packages/b3/34/d8b8232e5e55169933b59aabcef2fedfa4b9d8897361bb80fcbda146505f/coverage-7.15.4-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:226c66e80ec0598d3b9b4874123df167ccca342aca8714f77cac6829688ee09c", size = 264043, upload-time = "2026-08-06T13:49:04.102Z" },
+    { url = "https://files.pythonhosted.org/packages/7e/35/58b009dbf8c471c7224716478b9fed4a7e1af15320e1ed41660978504663/coverage-7.15.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ac41cc14bebda0dbfb0628036b7f75706935c95bcc07fefe9a0f93614aa60a57", size = 266963, upload-time = "2026-08-06T13:49:05.821Z" },
+    { url = "https://files.pythonhosted.org/packages/62/aa/57fbda1b42c892968273c56b6ee9dc0f1310850859230a507bc7873b1f65/coverage-7.15.4-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:8af623e5cd92080acddd02b38f2f406a2c3a0893c38950b211890361448fbf26", size = 264569, upload-time = "2026-08-06T13:49:07.706Z" },
+    { url = "https://files.pythonhosted.org/packages/98/8a/360e6e7f24d477b7e889703af0afa878d15b6d4d8d2a822b2835c169a879/coverage-7.15.4-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:07545711d4f0f32852a18f18ad11f76f0109909d09e78b9008b4cfc67e829429", size = 268299, upload-time = "2026-08-06T13:49:09.587Z" },
+    { url = "https://files.pythonhosted.org/packages/4e/89/6f701261aee21b6b5fa8f7872229406dc917e125069448292223bf213606/coverage-7.15.4-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:a0865421cfdc53654b342d515e5a233187590882d20b95752150e53f65460017", size = 263413, upload-time = "2026-08-06T13:49:11.604Z" },
+    { url = "https://files.pythonhosted.org/packages/3f/0f/6f04036edc260ed425af83e834f627fad48941ce97b50bfe6edd8b6fa623/coverage-7.15.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:460115e32ee40566476db5048f9bec1e842c127ad8e6f8be745aad3ac9cbc839", size = 265725, upload-time = "2026-08-06T13:49:13.38Z" },
+    { url = "https://files.pythonhosted.org/packages/c4/ce/d19b5d4d5c49a7bfb925fd74310fee7d28bc99520ac3367ccbc54e662518/coverage-7.15.4-cp314-cp314t-win32.whl", hash = "sha256:cbde877ef9dd7baf272b9bfef2b8a25edd45d9170fc326951dd20eb480335e85", size = 225079, upload-time = "2026-08-06T13:49:15.265Z" },
+    { url = "https://files.pythonhosted.org/packages/26/bb/7aa1b3b173faee0679037ca950bbbe1247273656697994d8d13f80f8d4b4/coverage-7.15.4-cp314-cp314t-win_amd64.whl", hash = "sha256:3da9e92d1c551fd7563833e9ade686efb0c4b7363ab7681a94283958c950bf5e", size = 225911, upload-time = "2026-08-06T13:49:17.279Z" },
+    { url = "https://files.pythonhosted.org/packages/81/1c/4ea9e47426d80038d9222db3c4534cb6021a74b237d3ff97ffd33b6600dd/coverage-7.15.4-cp314-cp314t-win_arm64.whl", hash = "sha256:3a54f5a0d85050c73a38f6793090ee83974531e67fe5e57a1da9bee11398aa5e", size = 225219, upload-time = "2026-08-06T13:49:19.293Z" },
+    { url = "https://files.pythonhosted.org/packages/2b/c4/dc5d2ac8f9142e7ec7de66e7bf0591db29d78955a040bd915870d9c0e657/coverage-7.15.4-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:2c9872e4d9dc5d3cf616bf4b382f5a00359305a5be666a3dd0b5cdb4e49597f9", size = 222604, upload-time = "2026-08-06T13:49:21.279Z" },
+    { url = "https://files.pythonhosted.org/packages/70/39/33e63df81fe2ee100897451841c821467635923e58e37c6bd4b46dd8106c/coverage-7.15.4-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:e101dbb4b9b72f0cddd8cdc8c9c5b47f456766f5e0ac82dbfb75e5c55409b78a", size = 222944, upload-time = "2026-08-06T13:49:23.187Z" },
+    { url = "https://files.pythonhosted.org/packages/99/1f/ef3ffb5557febc75a0d97aa459d0266d7d741110265121cc6d8539343d44/coverage-7.15.4-cp315-cp315-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:7d1abebdb047729e852b9c77a00497dfbeb11eb3a117e037d7dbc3ac8e5f5c54", size = 254050, upload-time = "2026-08-06T13:49:25.008Z" },
+    { url = "https://files.pythonhosted.org/packages/6f/f5/1f0f6f77698c3601ca0ae7431e34b24c62ca2f06fecb23b73ed1f651d2be/coverage-7.15.4-cp315-cp315-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:d28a4a899354d0ea6214cc59b4fa19eefbce1b9ff1688ab579acf49e894bd3fb", size = 256967, upload-time = "2026-08-06T13:49:26.896Z" },
+    { url = "https://files.pythonhosted.org/packages/03/7a/2ed9bed79925f4367c83c77f66a89e5ca7229c288d2d19ad5f36d1ca0070/coverage-7.15.4-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ffb3c2aacea411cc7e1d27712490c11108e2de1d39019ae32915493a59a8b9ed", size = 258587, upload-time = "2026-08-06T13:49:28.692Z" },
+    { url = "https://files.pythonhosted.org/packages/45/8c/fa34044f71b7cc4ecb6da9c2408770959b0591fa9b5fb6fb6bca38f94298/coverage-7.15.4-cp315-cp315-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a9447978a92f405d301123cfd39ff49895490efb769a758fe2734c7f631bf8ce", size = 260785, upload-time = "2026-08-06T13:49:30.472Z" },
+    { url = "https://files.pythonhosted.org/packages/4f/54/d5727ce36b4524a7394ab9f5f1df378e1f23affcdab01037dc8655185cc7/coverage-7.15.4-cp315-cp315-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:050467a7983b8e2fe7dd41a78bb30c3e7f8c0b8cafda14b1c46f8b5e3cf2dd3c", size = 254545, upload-time = "2026-08-06T13:49:32.271Z" },
+    { url = "https://files.pythonhosted.org/packages/dc/e6/6e3783e576719590194bdffb6dd6d85490801785b7c331e35a245d8cb8b5/coverage-7.15.4-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:d003b7a5708ddad5c206c79607a6b92abb6fc13c57d99d8a4468cc03a2941ced", size = 256682, upload-time = "2026-08-06T13:49:34.089Z" },
+    { url = "https://files.pythonhosted.org/packages/dc/f2/bacdbde18b69ed2de424fcf64d9fb0a4913753d4f0eca8bae9daad69f4bd/coverage-7.15.4-cp315-cp315-musllinux_1_2_i686.whl", hash = "sha256:c38efe30fd74e5c19e9433f11fb1f5dc9c6522770971b7c6145bbaa413dc8800", size = 254560, upload-time = "2026-08-06T13:49:36.052Z" },
+    { url = "https://files.pythonhosted.org/packages/6c/a3/1fb927196e3477c1b48831169ab58ba08f451ba87ae311ff1de68b26a616/coverage-7.15.4-cp315-cp315-musllinux_1_2_ppc64le.whl", hash = "sha256:1f4f826d70f772ab8b0c052329580d7fe8b8abd191e4ce0c8f81aec6614665d3", size = 258792, upload-time = "2026-08-06T13:49:38.01Z" },
+    { url = "https://files.pythonhosted.org/packages/41/58/30d4c149c69053de0edfe325614c1d28d508f62b1783e0e4a234d2e49136/coverage-7.15.4-cp315-cp315-musllinux_1_2_riscv64.whl", hash = "sha256:4a4bf917c9953f57c957be31c1cd504e3bd2f34d4a352b9d391a3025336f6768", size = 253968, upload-time = "2026-08-06T13:49:39.934Z" },
+    { url = "https://files.pythonhosted.org/packages/89/e4/77f639371b918aad30dda4051f95404b43578f7f2e2f87ba73e02ed1ff37/coverage-7.15.4-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:1c9bf40ebef178a45192c75c4964760bb261b0e6ad725da5fc4c93f674f19753", size = 255893, upload-time = "2026-08-06T13:49:41.825Z" },
+    { url = "https://files.pythonhosted.org/packages/5c/62/13be29b3ddab35f14c87967a4820a05106d2a3eccb4fa4ff550bf30b75e0/coverage-7.15.4-cp315-cp315-win32.whl", hash = "sha256:43619d04c3671792d2c4706ae8bf45e265dc87bbd4078189ef8b847ea1e74be2", size = 224768, upload-time = "2026-08-06T13:49:44.08Z" },
+    { url = "https://files.pythonhosted.org/packages/a1/70/af0c6be0f964af6954f6b74bc109b0dbca02824696d2520fb17fe1ab06e3/coverage-7.15.4-cp315-cp315-win_amd64.whl", hash = "sha256:be619439dbcd31a2eab10b32de9fff62c26ed4bab69dc32b8363fdaaa0882809", size = 225242, upload-time = "2026-08-06T13:49:45.899Z" },
+    { url = "https://files.pythonhosted.org/packages/4f/2d/f3bd3aab899fc9efc18b53133ee68f5f98574ef480649b23e12962226387/coverage-7.15.4-cp315-cp315-win_arm64.whl", hash = "sha256:def597967dafc2e8d97c9097ea453c464e0bb8ed38f193a43070f10dc623bb6d", size = 224674, upload-time = "2026-08-06T13:49:48.322Z" },
+    { url = "https://files.pythonhosted.org/packages/f5/ca/f69251cd63eabc6438321aea22148754cce758a26bde07dd490e3fe7cfc5/coverage-7.15.4-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:c7dbc748ac8a1e3e59a2b28bea47675e6e778081dbbf081bde0d75def2fcbe1d", size = 223333, upload-time = "2026-08-06T13:49:50.293Z" },
+    { url = "https://files.pythonhosted.org/packages/a7/a7/037b53b2885b0d8447064432491a4d5a1014cd9f97a594d53acd0c04541a/coverage-7.15.4-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:2413074a5ecbb61a01a7888fc72db0ca324d13588c5b38bc0dd8564cdcdfea26", size = 223630, upload-time = "2026-08-06T13:49:52.637Z" },
+    { url = "https://files.pythonhosted.org/packages/80/4f/152b8a4779ae90da11bb24f7467df8a59f0be48a5c52acb856325ca48289/coverage-7.15.4-cp315-cp315t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4e6f6f632b7b2f714bf7a1346e8f97b650ee71f3c298aaad42a2ab60f0f07645", size = 264489, upload-time = "2026-08-06T13:49:54.52Z" },
+    { url = "https://files.pythonhosted.org/packages/10/2d/84b4b9e0e1dd6528a51920ff7031f35b789382e467a28ec6a5a578cb8812/coverage-7.15.4-cp315-cp315t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8df457da2249d3c75ca2e5e835d59c725abfe92d27fdff6cd99eed85b51d5e9a", size = 267567, upload-time = "2026-08-06T13:49:56.721Z" },
+    { url = "https://files.pythonhosted.org/packages/53/fc/ba01cc25299f9f8a2c8b02d3b28c53f3543d9fbfbe4e74fa2760b48f163e/coverage-7.15.4-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:050f66a08805acb5b8a23c6d4a517b1ecf82c08e81ed0e4bd727df065e5c6624", size = 270123, upload-time = "2026-08-06T13:49:58.736Z" },
+    { url = "https://files.pythonhosted.org/packages/cf/d0/db2647cbf40b14f8c308f94ff7bf89c06d564e59f396906edf50086ec788/coverage-7.15.4-cp315-cp315t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1587fb771d1ccceef708fdde1e5af8c7ed24b486b61d13a321acb7d8145390aa", size = 271107, upload-time = "2026-08-06T13:50:00.811Z" },
+    { url = "https://files.pythonhosted.org/packages/70/ff/4d2d17924552c458bb4f77dd631f0e3bc92fbbdf2d2d916cd4b33bbfd5b1/coverage-7.15.4-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:8b4f1c3a69ca580f3fbd6b2046915f536d7f586874f25c1bb23add2a3c88d50f", size = 264955, upload-time = "2026-08-06T13:50:03.023Z" },
+    { url = "https://files.pythonhosted.org/packages/ee/de/dc010c7a3691f396d93bbc26bfcafa1c2a3a351cd520470f15faf5795bd5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:ffb58d7eff5b7f6ecc6fa21d6288ab7f968a212cb67d682c269c09b9eba3b66f", size = 267949, upload-time = "2026-08-06T13:50:05.557Z" },
+    { url = "https://files.pythonhosted.org/packages/78/ea/dc96a11375e83c045c2f7c61fb6918277cfe9401db7c0f7b1d111a84b2e5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_i686.whl", hash = "sha256:d9df165544774574ee004b953023d1bebada1894a80b1052a43d798b0f676e67", size = 264421, upload-time = "2026-08-06T13:50:07.612Z" },
+    { url = "https://files.pythonhosted.org/packages/c8/86/b77131a0f9503ce461cd577076147d7a9040f0c5dda772686f729e2cc9cb/coverage-7.15.4-cp315-cp315t-musllinux_1_2_ppc64le.whl", hash = "sha256:f9de0a24a4079b53e523b5c5e2c5945ec251ab486652659955187cf255a259bc", size = 269121, upload-time = "2026-08-06T13:50:09.58Z" },
+    { url = "https://files.pythonhosted.org/packages/24/24/944bc35007862955e7ebf05754e645419dcf5d7526c52735cfa2715e8ebf/coverage-7.15.4-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:150089274bdc9f940628552cb92844e0223c987f1902ab8efe9f45a2ec758d88", size = 264565, upload-time = "2026-08-06T13:50:11.722Z" },
+    { url = "https://files.pythonhosted.org/packages/c7/cc/a3bb9f93e7e740659163e2ea584f8196ddcd2c456a5dbe15f6c50105fec1/coverage-7.15.4-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:a58a94fed5da6997d258e8f7668c1e195fbd04a691d781b7558f1e468f9e68bc", size = 266522, upload-time = "2026-08-06T13:50:13.786Z" },
+    { url = "https://files.pythonhosted.org/packages/49/dd/e0e40f3560d878d888c580698ff5ad1179f5e1c3ac949684ef66b41a3817/coverage-7.15.4-cp315-cp315t-win32.whl", hash = "sha256:ebd5a6d8466ff30836572f3ba2cae8a5e8f85029b1c6d5e2ed338dc472a5166a", size = 225068, upload-time = "2026-08-06T13:50:15.825Z" },
+    { url = "https://files.pythonhosted.org/packages/c6/7e/37732ea80eebc30e976e4cdab15c190bc42d96959a42e38ddf6f8c60468f/coverage-7.15.4-cp315-cp315t-win_amd64.whl", hash = "sha256:288bde2a2d7ab6b6c2d7252fcde8b524387f2d970bdba9658fc6f8bbcaef0f9b", size = 225895, upload-time = "2026-08-06T13:50:17.928Z" },
+    { url = "https://files.pythonhosted.org/packages/c6/08/1e00f7923eaaba45fb3d51dd794125fc766304b1df264f3a9c6557bfb30e/coverage-7.15.4-cp315-cp315t-win_arm64.whl", hash = "sha256:68be5e1de60ff13c9095bbec0e5a7fa45b33b101752215b91345ea1f61c4a278", size = 225213, upload-time = "2026-08-06T13:50:19.981Z" },
+    { url = "https://files.pythonhosted.org/packages/b4/d9/e70c286c979378f061d8266e279b686ab0b0b688e1fe0af864684f23a77d/coverage-7.15.4-py3-none-any.whl", hash = "sha256:964730a1e9de9c0cf11be6a1a3c79ce419c34882842abd256086ba4698705e84", size = 214332, upload-time = "2026-08-06T13:50:22.192Z" },
+]
+
 [[package]]
 name = "cramjam"
 version = "2.11.0"
@@ -572,6 +671,7 @@ dev = [
     { name = "plotstyle" },
     { name = "polars" },
     { name = "pytest" },
+    { name = "pytest-cov" },
     { name = "ruff" },
     { name = "scikit-learn" },
     { name = "ty" },
@@ -599,6 +699,7 @@ requires-dist = [
     { name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" },
     { name = "pyarrow", specifier = ">=16,<25" },
     { name = "pytest", marker = "extra == 'dev'", specifier = ">=8,<10" },
+    { name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5,<8" },
     { name = "pyyaml", specifier = ">=6,<7" },
     { name = "ruff", marker = "extra == 'dev'", specifier = ">=0.15,<1" },
     { name = "scikit-learn", marker = "extra == 'geometry'", specifier = ">=1.4,<2" },
@@ -1713,6 +1814,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/8b/5a/ba30a81239b909821b3153e303e7def45178bf353da4f72380e6c5e8793b/pytest-9.1.0-py3-none-any.whl", hash = "sha256:8ebb0e7888bdf2bdfc602ec51f8f62d50200af37356c74e503c79a94f5c81f32", size = 386453, upload-time = "2026-06-13T18:52:44.045Z" },
 ]
 
+[[package]]
+name = "pytest-cov"
+version = "7.1.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "coverage" },
+    { name = "pluggy" },
+    { name = "pytest" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b1/51/a849f96e117386044471c8ec2bd6cfebacda285da9525c9106aeb28da671/pytest_cov-7.1.0.tar.gz", hash = "sha256:30674f2b5f6351aa09702a9c8c364f6a01c27aae0c1366ae8016160d1efc56b2", size = 55592, upload-time = "2026-03-21T20:11:16.284Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" },
+]
+
 [[package]]
 name = "python-dateutil"
 version = "2.9.0.post0"