chore: update changelog for v0.3.9 [skip ci]

Reviewed-on: #80

.bumpversion.toml

@@ -0,0 +1,17 @@
+[tool.bumpversion]
+current_version = "0.3.9"
+parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
+serialize = ["{major}.{minor}.{patch}"]
+search = "{current_version}"
+replace = "{new_version}"
+regex = false
+allow_dirty = false
+commit = true
+tag = false
+message = "chore: bump version {current_version} -> {new_version} [skip ci]"
+pre_commit_hooks = ["uv lock", "git add uv.lock"]
+
+[[tool.bumpversion.files]]
+filename = "pyproject.toml"
+search = "version = \"{current_version}\""
+replace = "version = \"{new_version}\""

.gitea/workflows/ci.yml

@@ -88,6 +88,92 @@ jobs:
           name: coverage-report
           path: coverage.xml
 
+  bump-version:
+    name: Bump version, tag, and update changelog on merge to master
+    needs: [ruff-check, ruff-format, type-check, test]
+    if: github.ref == 'refs/heads/master' && github.event_name == 'push'
+    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
+    steps:
+      # CI_TOKEN needs write:repository scope (not just read) — this job
+      # pushes commits and tags to master, unlike ruff-check/ruff-format/
+      # type-check/test above, which only need to check out the repo.
+      - uses: actions/checkout@v4
+        with:
+          token: ${{ secrets.CI_TOKEN }}
+          fetch-depth: 0
+      - name: Check whether this push is a merge commit
+        id: merge_check
+        run: |
+          PARENTS=$(git rev-parse HEAD^@ | wc -l)
+          echo "HEAD has $PARENTS parent(s)"
+          if [ "$PARENTS" -ge 2 ]; then
+            echo "is_merge=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "is_merge=false" >> "$GITHUB_OUTPUT"
+          fi
+      - uses: astral-sh/setup-uv@v5
+        if: steps.merge_check.outputs.is_merge == 'true'
+        with:
+          enable-cache: false
+      - run: |
+          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+        if: steps.merge_check.outputs.is_merge == 'true'
+      - run: uv sync --extra cpu --extra dev
+        if: steps.merge_check.outputs.is_merge == 'true'
+      - name: Configure git identity
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          git config user.name "gitea-actions"
+          git config user.email "actions@git.larsbogner.de"
+      - name: Bump patch version if this merge didn't already bump it
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          OLD_VERSION=$(git show "${{ github.event.before }}:pyproject.toml" 2>/dev/null | grep -m1 '^version = ' | sed -E 's/version = "(.*)"/\1/')
+          CURRENT_VERSION=$(uv version --short)
+          if [ -z "$OLD_VERSION" ]; then
+            echo "Could not read pyproject.toml at github.event.before; falling back to HEAD^1"
+            OLD_VERSION=$(git show "HEAD^1:pyproject.toml" | grep -m1 '^version = ' | sed -E 's/version = "(.*)"/\1/')
+          fi
+          if [ "$OLD_VERSION" = "$CURRENT_VERSION" ]; then
+            echo "Version unchanged by this merge ($CURRENT_VERSION); bumping patch"
+            uv run bump-my-version bump patch --current-version "$CURRENT_VERSION"
+          else
+            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); skipping auto-bump"
+          fi
+      - name: Update changelog for the current version if not already tagged
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          VERSION=$(uv version --short)
+          TAG="v$VERSION"
+          if git rev-parse "$TAG" >/dev/null 2>&1; then
+            echo "Tag $TAG already exists; skipping changelog update"
+          else
+            uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
+            git add CHANGELOG.md
+            if ! git diff --cached --quiet -- CHANGELOG.md; then
+              git commit -m "chore: update changelog for $TAG [skip ci]"
+            else
+              git restore --staged CHANGELOG.md
+            fi
+          fi
+      - name: Push commits and tag the current version
+        if: steps.merge_check.outputs.is_merge == 'true'
+        run: |
+          git push origin HEAD:master
+          VERSION=$(uv version --short)
+          TAG="v$VERSION"
+          if git rev-parse "$TAG" >/dev/null 2>&1; then
+            echo "Tag $TAG already exists"
+          else
+            git tag -a "$TAG" -m "$TAG"
+            git push origin "refs/tags/$TAG"
+          fi
+
   sync-version-on-tag:
     name: Sync project version with tag
     if: startsWith(github.ref, 'refs/tags/')

CHANGELOG.md

@@ -0,0 +1,68 @@
+# Changelog
+
+## [0.3.9] - 2026-08-24
+
+### Added
+
+- Add multi-rollout support to giant analyze [gitea #77](https://git.larsbogner.de/lars/giant/issues/77)
+
+
+### Changed
+
+- Escape LaTeX-special characters in plot titles/xlabels [gitea #81](https://git.larsbogner.de/lars/giant/issues/81)
+
+## [0.3.8] - 2026-08-24
+
+### Added
+
+- Add giant analyze metrics plots for training progress [gitea #75](https://git.larsbogner.de/lars/giant/issues/75)
+
+
+### Fixed
+
+- Fix LaTeX-unavailable skip check in analyze metrics smoke test
+
+## [0.3.7] - 2026-08-24
+
+### Added
+
+- Add rollout-quality distance, confusion, containment and router plots [gitea #76](https://git.larsbogner.de/lars/giant/issues/76)
+
+## [0.3.6] - 2026-08-24
+
+### Changed
+
+- Give CriticModel a registry-built trunk and StageModel base [gitea #57](https://git.larsbogner.de/lars/giant/issues/57)
+
+## [0.3.5] - 2026-08-24
+
+### Added
+
+- Add "none" variants for router, history, and trunk [gitea #45](https://git.larsbogner.de/lars/giant/issues/45)
+
+## [0.3.4] - 2026-08-23
+
+### Added
+
+- Add giant model summary command [gitea #46](https://git.larsbogner.de/lars/giant/issues/46)
+
+- Add per-stage init_from/freeze [gitea #42](https://git.larsbogner.de/lars/giant/issues/42)
+
+- Add bf16 autocast to the training loop [gitea #47](https://git.larsbogner.de/lars/giant/issues/47)
+
+- Add class-balanced secondary particle-type loss [gitea #44](https://git.larsbogner.de/lars/giant/issues/44)
+
+
+### Changed
+
+- Implement stage2_model.stage1_context = "sampled" [gitea #41](https://git.larsbogner.de/lars/giant/issues/41)
+
+- Bump patch version to 0.3.3
+
+- Offset event_id across multi-shard reference reads in giant analyze [gitea #22](https://git.larsbogner.de/lars/giant/issues/22)
+
+- Auto-bump patch version, tag, and update changelog on merge to master [gitea #50](https://git.larsbogner.de/lars/giant/issues/50)
+
+- Document CI_TOKEN's write:repository scope requirement [gitea #50](https://git.larsbogner.de/lars/giant/issues/50)
+
+# Changelog

CLAUDE.md

@@ -73,7 +73,7 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
 
 **Validation** (`giant/validate.py`): step-level marginal comparisons.
 
-**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_<id>/`) holding `shared.json`, `run_meta.json`, `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit rollout.yaml --chunks N` runs `prep` (recording the run's chunk count `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` (`catalog.py`) splits into a `compute_partial`/`finalize` pair so a plot's chunks can be summed/concatenated back together correctly (`chunkable=False` specs — the router diagnostics, already bounded/subsampled — always run as a single chunk regardless of `N`). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`), then turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
+**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive `giant rollout` runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + `RolloutSpec`/`RolloutSide` — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json` over the union of the reference and every rollout, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries; `Bundle.rollouts` is a name-keyed dict of `RolloutSide`, and every `compute_partial`/`finalize` builds a `Reduced.payload["series"]` dict keyed the same way, with `payload["reference"]` as the one distinguished non-rollout entry), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`; each rollout gets a stable `ps.get_color(i)` slot by its position in `series`, the reference always draws in one fixed dashed-ink style). The two heatmap-shaped specs (`marginal_distance_summary`, `n_sec_confusion`) and the router/type-embedding diagnostics (`router_gating.py`, `type_embedding_distance.py`) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout. **Input is one or more `giant rollout` YAML sidecars** (`condor.py:load_rollout_yamls`, wrapping the single-YAML `load_rollout_yaml`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML — matching pre-multi-rollout output exactly). `prep` derives its own **run directory** next to the *first* rollout's parquet (`<...>/analysis_<tag(s)>/`) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit a.yaml [b.yaml ...] --chunks N` runs `prep` (recording the run's chunk count `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` (`catalog.py`) splits into a `compute_partial`/`finalize` pair so a plot's chunks can be summed/concatenated back together correctly per rollout (`chunkable=False` specs — the router diagnostics, already bounded/subsampled — always run as a single chunk regardless of `N`). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`), then turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
 
 **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
 

README.md

@@ -139,6 +139,7 @@ Useful flags on `giant train`:
 - `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
 - `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
 - `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
+- `--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
 
 Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
 
@@ -151,10 +152,11 @@ Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_mod
 
 ```bash
 giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot (compute only)
+giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan  # N rollouts vs one shared reference
 giant analyze render <run_dir> --gallery                   # local: styled PDFs + HTML gallery (needs LaTeX)
 ```
 
-`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
+`<run_dir>` is derived next to the first rollout's parquet (`analyze prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
 
 ## Development
 

cliff.toml

@@ -0,0 +1,52 @@
+# git-cliff configuration — see https://git-cliff.org/docs/configuration
+#
+# Commit messages in this repo aren't Conventional Commits; they're plain
+# imperative summaries like "Add class-balanced secondary particle-type loss
+# (gitea #44)". Parsing here is tuned to that convention rather than to
+# feat:/fix:-style prefixes.
+
+[changelog]
+header = "# Changelog\n\n"
+body = """
+{% if version %}\
+## [{{ version | trim_start_matches(pat="v") }}] - {{ timestamp | date(format="%Y-%m-%d") }}
+{% else %}\
+## [Unreleased]
+{% endif %}\
+{% for group, commits in commits | group_by(attribute="group") %}
+### {{ group | striptags | trim | upper_first }}
+{% for commit in commits %}
+- {{ commit.message | upper_first }}
+{% endfor %}
+{% endfor %}
+"""
+trim = true
+render_always = true
+postprocessors = []
+
+[git]
+conventional_commits = false
+filter_unconventional = false
+require_conventional = false
+split_commits = false
+# Keep only the commit subject (first line), then linkify "(gitea #N)".
+commit_preprocessors = [
+    { pattern = "(?s)\n.*", replace = "" },
+    { pattern = "\\(gitea #(\\d+)\\)", replace = "[gitea #${1}](https://git.larsbogner.de/lars/giant/issues/${1})" },
+]
+protect_breaking_commits = false
+commit_parsers = [
+    { message = "^Merge ", skip = true },
+    { message = "\\[skip ci\\]", skip = true },
+    { message = "^Add", group = "<!-- 0 -->Added" },
+    { message = "^(Fix|Clamp|Clip)", group = "<!-- 1 -->Fixed" },
+    { message = "^(Remove|Drop|Deprecate)", group = "<!-- 2 -->Removed" },
+    { message = ".*", group = "<!-- 3 -->Changed" },
+]
+filter_commits = false
+link_parsers = []
+use_branch_tags = false
+topo_order = false
+topo_order_commits = true
+sort_commits = "oldest"
+recurse_submodules = false

giant/analysis/__init__.py

@@ -1,9 +1,10 @@
 """Rollout-vs-reference analysis: streaming compute + plotstyle rendering.
 
-Compares one autoregressive ``giant rollout`` against a held-out miniCaloSim
-reference file, producing publication-styled comparison plots generated in
-parallel on HTCondor (one job per plot x data chunk, compute/merge/render
-split).
+Compares one or more autoregressive ``giant rollout`` runs against a single
+held-out miniCaloSim reference file shared by all of them, producing
+publication-styled comparison plots (one colored series per rollout, one
+reference line) generated in parallel on HTCondor (one job per plot x data
+chunk, compute/merge/render split).
 
 Only ``render`` (and the ``render`` CLI path) imports plotstyle/LaTeX; everything
 re-exported here is plotstyle-free so it runs on a compute worker. Import
@@ -12,12 +13,14 @@ re-exported here is plotstyle-free so it runs on a compute worker. Import
 
 from giant.analysis.catalog import build_catalog, catalog_ids, get_spec
 from giant.analysis.condor import (
+    LoadedRollout,
     RunMeta,
     SubmitConfig,
     compute_one,
     compute_reduced,
     derive_run_dir,
     load_rollout_yaml,
+    load_rollout_yamls,
     merge_all,
     merge_one,
     prep,
@@ -26,18 +29,20 @@ from giant.analysis.condor import (
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial, Reduced
 from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
-from giant.analysis.sources import Side
+from giant.analysis.sources import RolloutSpec, Side
 
 __all__ = [
     "build_catalog",
     "catalog_ids",
     "get_spec",
+    "LoadedRollout",
     "RunMeta",
     "SubmitConfig",
     "compute_one",
     "compute_reduced",
     "derive_run_dir",
     "load_rollout_yaml",
+    "load_rollout_yamls",
     "merge_all",
     "merge_one",
     "prep",
@@ -46,6 +51,7 @@ __all__ = [
     "build_context",
     "Partial",
     "Reduced",
+    "RolloutSpec",
     "Side",
     "RUNTIME_SAFETY_MARGIN",
     "estimate_runtime_s",

giant/analysis/catalog.py

@@ -4,7 +4,7 @@ Each spec knows its stable ``id`` (used for the reduced-data filename, the PDF
 stem and the condor queue item), its gallery ``family`` (subdirectory), and a
 ``compute_partial(bundle) -> dict`` / ``finalize(parts, ctx) -> Reduced`` pair
 that together run the streaming reduction. ``compute_partial`` runs once per
-``(plot, chunk)`` condor job against a ``Bundle`` whose four LazyFrames are
+``(plot, chunk)`` condor job against a ``Bundle`` whose LazyFrames are
 already filtered to that chunk (see ``Bundle.open``'s ``chunk`` argument); it
 returns a small JSON-safe partial artifact — either a raw sum-mergeable count
 dict (histograms/species sums against fixed edges) or a raw per-event/
@@ -16,6 +16,19 @@ exactly what a single unchunked pass would produce. Specs marked
 ``chunkable=False`` (the router ones) always run as a single chunk regardless
 of the configured chunk count.
 
+Every ``compute_partial`` here returns ``{"r": {rollout_name: <shape>}, "t":
+<shape>}`` — one entry per rollout in ``Bundle.rollouts`` (insertion order,
+which is the order rollouts were given on the CLI) plus the single reference.
+``finalize`` merges each rollout's chunks independently and assembles a
+``Reduced.payload`` keyed the same way: ``"series": {name: ...}`` for the
+rollouts, ``"reference": ...`` as one distinguished entry (omitted on
+rollout-only plots like ``leakage_fraction``). The two heatmap-shaped specs
+(``marginal_distance_summary``, ``n_sec_confusion``) and the router
+diagnostics are inherently one-matrix/one-checkpoint per rollout, so their
+``"series"`` entries are whole per-rollout artifacts (a matrix, a gating
+dict) rather than a single number/array — ``render.py`` draws those as one
+panel per rollout instead of one line/bar per rollout.
+
 Rendering lives in ``render.py`` and dispatches on ``Reduced.kind`` — the
 catalog itself never imports plotstyle, so ``compute-one`` jobs stay LaTeX-free.
 
@@ -47,6 +60,7 @@ from giant.analysis.reduce import (
     leakage_fraction,
     profile_finalize,
     profile_partial,
+    sec_count_by_event,
     species_share,
     sum_merge,
     transverse_expr,
@@ -56,8 +70,9 @@ from giant.analysis.router_gating import (
     compute_router_gating,
     compute_router_share_by_pdg,
     compute_router_share_by_process,
+    compute_router_specialization,
 )
-from giant.analysis.sources import Side, open_side, physical_steps, secondaries
+from giant.analysis.sources import RolloutSide, RolloutSpec, Side, open_side, physical_steps, secondaries
 from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
 from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
 
@@ -67,51 +82,44 @@ class Bundle:
     """Everything a compute runs against — built once per ``compute-one`` job."""
 
     ctx: Context
-    r_all: pl.LazyFrame  # rollout, all rows (incl. synthetic termination rows)
+    rollouts: dict[str, RolloutSide]  # name -> frames, insertion order = CLI order
     t_all: pl.LazyFrame  # reference, all rows
-    r_phys: pl.LazyFrame  # rollout, physical steps only
     t_phys: pl.LazyFrame  # reference, physical steps only
-    checkpoint: str | None = None  # from the rollout YAML; router_gating only
-    # Diagnostic pre-aggregated at rollout time (giant.rollout.
-    # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
-    # only. Unlike checkpoint/router_gating, this needs no live model: it's
-    # already a finished histogram, just passed through.
-    type_embedding_l1_dist: dict | None = None
 
     @classmethod
     def open(
         cls,
-        rollout,
+        rollouts: list[RolloutSpec],
         reference,
         ctx: Context,
-        checkpoint=None,
         chunk: tuple[int, int] | None = None,
-        type_embedding_l1_dist: dict | None = None,
     ) -> "Bundle":
-        """Open both sides, optionally restricted to one event-disjoint chunk.
+        """Open the reference + every rollout, optionally restricted to one event-disjoint chunk.
 
-        ``chunk = (chunk_index, n_chunks)`` filters both sides to
+        ``chunk = (chunk_index, n_chunks)`` filters every side to
         ``event_id % n_chunks == chunk_index`` *before* deriving the physical/
         secondary views, so every downstream reduction (which is either
         row-local or a ``group_by("event_id")``) sees a self-contained,
         event-disjoint slice — no cross-chunk lookups are ever needed.
         """
-        r_all = open_side(rollout, Side.rollout)
         t_all = open_side(reference, Side.reference)
+        pred = None
         if chunk is not None:
             idx, n = chunk
             pred = pl.col("event_id") % n == idx
-            r_all = r_all.filter(pred)
             t_all = t_all.filter(pred)
-        return cls(
-            ctx=ctx,
-            r_all=r_all,
-            t_all=t_all,
-            r_phys=physical_steps(r_all, Side.rollout),
-            t_phys=physical_steps(t_all, Side.reference),
-            checkpoint=checkpoint,
-            type_embedding_l1_dist=type_embedding_l1_dist,
-        )
+        sides: dict[str, RolloutSide] = {}
+        for rs in rollouts:
+            r_all = open_side(rs.source, Side.rollout)
+            if pred is not None:
+                r_all = r_all.filter(pred)
+            sides[rs.name] = RolloutSide(
+                all=r_all,
+                phys=physical_steps(r_all, Side.rollout),
+                checkpoint=rs.checkpoint,
+                type_embedding_l1_dist=rs.type_embedding_l1_dist,
+            )
+        return cls(ctx=ctx, rollouts=sides, t_all=t_all, t_phys=physical_steps(t_all, Side.reference))
 
 
 @dataclass
@@ -144,8 +152,10 @@ def _unchunkable(
 # small numpy/hist helpers
 # ---------------------------------------------------------------------------
 
-_ROLL = "rollout"
-_REF = "reference"
+
+def _per_rollout(b: Bundle, fn: Callable[[RolloutSide], object]) -> dict[str, object]:
+    """``{name: fn(rollout_side)}`` over every rollout, preserving CLI order."""
+    return {name: fn(rs) for name, rs in b.rollouts.items()}
 
 
 def _counts(h: dict, key, nbins: int) -> list[int]:
@@ -166,14 +176,89 @@ 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]:
-    """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])
+def _np_hist_shared_edges(arrays: list[np.ndarray], nbins: int) -> tuple[np.ndarray, list[np.ndarray]]:
+    """Shared-edge histogram of several small per-event arrays (robust range).
+
+    The edges are sized from the union of every array (reference + all
+    rollouts), so every series in the resulting overlay is directly
+    comparable on one axis.
+    """
+    non_empty = [a for a in arrays if len(a)]
+    both = np.concatenate(non_empty) if non_empty else np.array([0.0, 1.0])
     lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
     if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
         lo, hi = lo - 0.5, hi + 0.5
     edges = np.linspace(lo, hi, nbins + 1)
-    return edges, np.histogram(r, edges)[0], np.histogram(t, edges)[0]
+    return edges, [np.histogram(a, edges)[0] for a in arrays]
+
+
+def _ks_statistic(r_counts, t_counts) -> float:
+    """KS statistic (max |CDF diff|) between two same-edge binned histograms.
+
+    ``nan`` when neither side has any mass (nothing to compare); 1.0 (maximal
+    mismatch) when exactly one side is entirely empty and the other isn't —
+    correctly the worst score rather than an undefined one.
+    """
+    r_counts = np.asarray(r_counts, dtype=np.float64)
+    t_counts = np.asarray(t_counts, dtype=np.float64)
+    r_tot, t_tot = r_counts.sum(), t_counts.sum()
+    if r_tot == 0 and t_tot == 0:
+        return float("nan")
+    if r_tot == 0 or t_tot == 0:
+        return 1.0
+    r_cdf = np.cumsum(r_counts) / r_tot
+    t_cdf = np.cumsum(t_counts) / t_tot
+    return float(np.max(np.abs(r_cdf - t_cdf)))
+
+
+def _integer_confusion(
+    t: np.ndarray, r: np.ndarray, max_bins: int = 21, cap: int | None = None
+) -> tuple[list[str], np.ndarray]:
+    """Confusion matrix of two paired small-integer arrays (e.g. secondary counts).
+
+    Bins are consecutive integers ``0..cap``, with the last bin an overflow
+    ``"cap+"`` bucket, so an occasional pathological count doesn't blow up the
+    heatmap. Returns ``(labels, matrix)`` with ``matrix[i, j]`` counting pairs
+    with ``t == i`` and ``r == j`` (both clipped into ``[0, cap]``).
+
+    ``cap``, if given, is used as-is instead of being derived from ``t``/``r``
+    — lets a multi-rollout caller fix one shared cap (and so one shared label
+    set) across every rollout's matrix rather than each panel picking its own.
+    """
+    if cap is None:
+        cap = min(max(int(t.max()) if len(t) else 0, int(r.max()) if len(r) else 0, 1), max_bins - 1)
+    t_c = np.clip(t.astype(np.int64), 0, cap)
+    r_c = np.clip(r.astype(np.int64), 0, cap)
+    n = cap + 1
+    mat = np.zeros((n, n), dtype=np.int64)
+    np.add.at(mat, (t_c, r_c), 1)
+    labels = [str(i) for i in range(cap)] + [f"{cap}+"]
+    return labels, mat
+
+
+def _containment_depths(mat: np.ndarray, edges: np.ndarray, quantile: float) -> np.ndarray:
+    """Per-event depth containing ``quantile`` of that event's deposited energy.
+
+    ``mat`` is a ``(n_events, n_bins)`` edep-per-depth-bin sum matrix (see
+    ``reduce.profile_partial``); bins are ordered by increasing depth (matching
+    ``edges``, monotonic). Zero-energy events are dropped — containment depth is
+    undefined for them.
+    """
+    totals = mat.sum(axis=1)
+    valid = totals > 0
+    mat, totals = mat[valid], totals[valid]
+    cum = np.cumsum(mat, axis=1) / totals[:, None]
+    idx = (cum >= quantile).argmax(axis=1)  # first bin whose cumulative fraction reaches quantile
+    return edges[1:][idx]
+
+
+def _group_keys(ctx: Context, axis: str) -> list:
+    """The group keys ``_marginal_grouped_finalize`` iterates for ``axis``."""
+    if axis == "pdg":
+        return list(ctx.top_pdgs)
+    if axis == "material":
+        return list(ctx.materials)
+    return list(range(len(ctx.energy_edges) - 1))  # energy
 
 
 # Human-readable figure titles per marginal variable (the axis labels carry units;
@@ -210,7 +295,7 @@ def _marginal_overall_partial(b: Bundle, var: str) -> dict:
     _, expr = _var(var)
     edges = _marginal_edges(b.ctx, var)
     return {
-        "r": _partial_hist(b.r_phys, expr, edges),
+        "r": _per_rollout(b, lambda rs: _partial_hist(rs.phys, expr, edges)),
         "t": _partial_hist(b.t_phys, expr, edges),
     }
 
@@ -219,7 +304,8 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
     label, _ = _var(var)
     edges = _marginal_edges(ctx, var)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
+    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
     t = sum_merge([p["t"] for p in parts])
     return Reduced(
         id=f"marginal_{var}",
@@ -229,8 +315,8 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
         xlabel=label,
         payload={
             "edges": edges.tolist(),
-            _ROLL: _finalize_counts(r, 0, nb),
-            _REF: _finalize_counts(t, 0, nb),
+            "series": series,
+            "reference": _finalize_counts(t, 0, nb),
             "log_y": True,
         },
     )
@@ -241,23 +327,24 @@ def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
     return pl.col("event_id").replace_strict(ids, bins, default=-1, return_dtype=pl.Int64)
 
 
+def _grouped_hist_dict(lf: pl.LazyFrame, expr: pl.Expr, edges: np.ndarray, axis: str, ctx: Context, nb: int) -> dict:
+    if axis == "pdg":
+        h = hist1d(lf, expr, edges, group=pl.col("pdg"))
+    elif axis == "material":
+        h = hist1d(lf, expr, edges, group=pl.col("material"))
+    else:  # energy
+        e_edges = np.asarray(ctx.energy_edges)
+        h = hist1d(lf, expr, edges, group=_energy_group_expr(lf, e_edges))
+    return {str(k): _counts(h, k, nb) for k in h}
+
+
 def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
     _, expr = _var(var)
     edges = _marginal_edges(b.ctx, var)
-    if axis == "pdg":
-        r = hist1d(b.r_phys, expr, edges, group=pl.col("pdg"))
-        t = hist1d(b.t_phys, expr, edges, group=pl.col("pdg"))
-    elif axis == "material":
-        r = hist1d(b.r_phys, expr, edges, group=pl.col("material"))
-        t = hist1d(b.t_phys, expr, edges, group=pl.col("material"))
-    else:  # energy
-        e_edges = np.asarray(b.ctx.energy_edges)
-        r = hist1d(b.r_phys, expr, edges, group=_energy_group_expr(b.r_phys, e_edges))
-        t = hist1d(b.t_phys, expr, edges, group=_energy_group_expr(b.t_phys, e_edges))
     nb = len(edges) - 1
     return {
-        "r": {str(k): _counts(r, k, nb) for k in r},
-        "t": {str(k): _counts(t, k, nb) for k in t},
+        "r": _per_rollout(b, lambda rs: _grouped_hist_dict(rs.phys, expr, edges, axis, b.ctx, nb)),
+        "t": _grouped_hist_dict(b.t_phys, expr, edges, axis, b.ctx, nb),
     }
 
 
@@ -265,29 +352,24 @@ def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis:
     label, _ = _var(var)
     edges = _marginal_edges(ctx, var)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
-    t = sum_merge([p["t"] for p in parts])
-    groups: dict[str, dict] = {}
+    names = list(parts[0]["r"])
+    r_merged = {name: sum_merge([p["r"][name] for p in parts]) for name in names}
+    t_merged = sum_merge([p["t"] for p in parts])
 
     if axis == "pdg":
-        for k in ctx.top_pdgs:
-            groups[pdg_label(k)] = {
-                _ROLL: _finalize_counts(r, k, nb),
-                _REF: _finalize_counts(t, k, nb),
-            }
+        keys, labels = ctx.top_pdgs, [pdg_label(k) for k in ctx.top_pdgs]
     elif axis == "material":
-        for m in ctx.materials:
-            groups[material_label(m)] = {
-                _ROLL: _finalize_counts(r, m, nb),
-                _REF: _finalize_counts(t, m, nb),
-            }
+        keys, labels = ctx.materials, [material_label(m) for m in ctx.materials]
     else:  # energy
         e_edges = np.asarray(ctx.energy_edges)
-        for bi, lbl in enumerate(energy_bin_labels(e_edges)):
-            groups[lbl] = {
-                _ROLL: _finalize_counts(r, bi, nb),
-                _REF: _finalize_counts(t, bi, nb),
-            }
+        keys, labels = list(range(len(e_edges) - 1)), energy_bin_labels(e_edges)
+
+    groups: dict[str, dict] = {}
+    for k, lbl in zip(keys, labels):
+        groups[lbl] = {
+            "series": {name: _finalize_counts(r_merged[name], k, nb) for name in names},
+            "reference": _finalize_counts(t_merged, k, nb),
+        }
 
     return Reduced(
         id=f"marginal_{var}_by_{axis}",
@@ -299,22 +381,96 @@ def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis:
     )
 
 
+# ---------------------------------------------------------------------------
+# distance summary: a var x group-axis scorecard per rollout, reusing the marginal hists
+# ---------------------------------------------------------------------------
+
+
+def _distance_summary_partial(b: Bundle) -> dict:
+    out: dict[str, dict] = {}
+    for var in MARGINAL_VARS:
+        out[var] = {"overall": _marginal_overall_partial(b, var)}
+        for axis in GROUPING_AXES:
+            out[var][axis] = _marginal_grouped_partial(b, var, axis)
+    return out
+
+
+def _distance_summary_finalize(parts: list[dict], ctx: Context) -> Reduced:
+    col_labels = ["overall", *GROUPING_AXES]
+    names = list(parts[0][MARGINAL_VARS[0]]["overall"]["r"])
+    matrices: dict[str, list[list[float]]] = {name: [] for name in names}
+
+    for var in MARGINAL_VARS:
+        edges = _marginal_edges(ctx, var)
+        nb = len(edges) - 1
+
+        t_overall = sum_merge([p[var]["overall"]["t"] for p in parts])
+        r_overall = {name: sum_merge([p[var]["overall"]["r"][name] for p in parts]) for name in names}
+        row: dict[str, list[float]] = {name: [] for name in names}
+        for name in names:
+            row[name].append(
+                _ks_statistic(_finalize_counts(r_overall[name], 0, nb), _finalize_counts(t_overall, 0, nb))
+            )
+
+        for axis in GROUPING_AXES:
+            t_grp = sum_merge([p[var][axis]["t"] for p in parts])
+            r_grp = {name: sum_merge([p[var][axis]["r"][name] for p in parts]) for name in names}
+            for name in names:
+                dists, weights = [], []
+                for k in _group_keys(ctx, axis):
+                    rc, tc = _finalize_counts(r_grp[name], k, nb), _finalize_counts(t_grp, k, nb)
+                    w = sum(rc) + sum(tc)
+                    if w == 0:
+                        continue
+                    dists.append(_ks_statistic(rc, tc))
+                    weights.append(w)
+                row[name].append(float(np.average(dists, weights=weights)) if dists else float("nan"))
+
+        for name in names:
+            matrices[name].append(row[name])
+
+    return Reduced(
+        id="marginal_distance_summary",
+        family="quality",
+        kind="heatmap",
+        title="Marginal distance summary (KS statistic, rollout vs reference)",
+        xlabel="grouping axis",
+        payload={
+            "series": matrices,
+            "row_labels": [_TITLE_NAMES[v] for v in MARGINAL_VARS],
+            "col_labels": col_labels,
+            "ylabel": "marginal variable",
+            "cbar_label": "KS statistic (0 = identical, 1 = maximal mismatch)",
+            "vmin": 0.0,
+            "vmax": 1.0,
+        },
+    )
+
+
 # ---------------------------------------------------------------------------
 # per-event scalar observables
 # ---------------------------------------------------------------------------
 
 
 def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
-    r_lf, t_lf = (b.r_all, b.t_all) if use_all else (b.r_phys, b.t_phys)
-    r = event_scalars(r_lf)[col].to_numpy()
-    t = event_scalars(t_lf)[col].to_numpy()
-    return {"r": r.tolist(), "t": t.tolist()}
+    t_lf = b.t_all if use_all else b.t_phys
+
+    def _vals(rs: RolloutSide) -> list[float]:
+        lf = rs.all if use_all else rs.phys
+        return event_scalars(lf)[col].to_numpy().tolist()
+
+    return {
+        "r": _per_rollout(b, _vals),
+        "t": event_scalars(t_lf)[col].to_numpy().tolist(),
+    }
 
 
 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])
+    names = list(parts[0]["r"])
+    r_arrays = {name: np.concatenate([np.asarray(p["r"][name], dtype=float) for p in parts]) for name in names}
     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)
+    edges, counts = _np_hist_shared_edges([t, *(r_arrays[n] for n in names)], ctx.n_marginal_bins)
+    t_counts, *r_counts = counts
     return Reduced(
         id=spec_id,
         family="event",
@@ -323,40 +479,44 @@ def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title:
         xlabel=xlabel,
         payload={
             "edges": edges.tolist(),
-            _ROLL: rc.astype(np.int64).tolist(),
-            _REF: tc.astype(np.int64).tolist(),
+            "series": {name: c.astype(np.int64).tolist() for name, c in zip(names, r_counts)},
+            "reference": t_counts.astype(np.int64).tolist(),
             "log_y": False,
         },
     )
 
 
 def _event_total_edep_by_energy_partial(b: Bundle) -> dict:
-    r = event_scalars(b.r_all)
     t = event_scalars(b.t_all)
+
+    def _vals(rs: RolloutSide) -> dict:
+        r = event_scalars(rs.all)
+        return {"incident": r["incident_E"].to_list(), "edep": r["total_edep"].to_list()}
+
     return {
-        "r_incident": r["incident_E"].to_list(),
-        "r_edep": r["total_edep"].to_list(),
-        "t_incident": t["incident_E"].to_list(),
-        "t_edep": t["total_edep"].to_list(),
+        "r": _per_rollout(b, _vals),
+        "t": {"incident": t["incident_E"].to_list(), "edep": t["total_edep"].to_list()},
     }
 
 
 def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
     e_edges = np.asarray(ctx.energy_edges)
-    r_inc = np.concatenate([np.asarray(p["r_incident"], dtype=float) for p in parts])
-    r_val = np.concatenate([np.asarray(p["r_edep"], dtype=float) for p in parts])
-    t_inc = np.concatenate([np.asarray(p["t_incident"], dtype=float) for p in parts])
-    t_val = np.concatenate([np.asarray(p["t_edep"], dtype=float) for p in parts])
-    r_bin = np.clip(np.digitize(r_inc, e_edges[1:-1]), 0, len(e_edges) - 2)
+    names = list(parts[0]["r"])
+    t_inc = np.concatenate([np.asarray(p["t"]["incident"], dtype=float) for p in parts])
+    t_val = np.concatenate([np.asarray(p["t"]["edep"], dtype=float) for p in parts])
+    r_inc = {n: np.concatenate([np.asarray(p["r"][n]["incident"], dtype=float) for p in parts]) for n in names}
+    r_val = {n: np.concatenate([np.asarray(p["r"][n]["edep"], dtype=float) for p in parts]) for n in names}
+
+    edges, _ = _np_hist_shared_edges([t_val, *(r_val[n] for n in names)], ctx.n_marginal_bins)
     t_bin = np.clip(np.digitize(t_inc, e_edges[1:-1]), 0, len(e_edges) - 2)
-    edges, _, _ = _np_hist_pair(r_val, t_val, ctx.n_marginal_bins)
+    r_bin = {n: np.clip(np.digitize(r_inc[n], e_edges[1:-1]), 0, len(e_edges) - 2) for n in names}
+
     groups: dict[str, dict] = {}
     for bi, lbl in enumerate(energy_bin_labels(e_edges)):
-        rc = np.histogram(r_val[r_bin == bi], edges)[0]
         tc = np.histogram(t_val[t_bin == bi], edges)[0]
         groups[lbl] = {
-            _ROLL: rc.astype(np.int64).tolist(),
-            _REF: tc.astype(np.int64).tolist(),
+            "series": {n: np.histogram(r_val[n][r_bin[n] == bi], edges)[0].astype(np.int64).tolist() for n in names},
+            "reference": tc.astype(np.int64).tolist(),
         }
     return Reduced(
         id="event_total_edep_by_energy",
@@ -375,15 +535,15 @@ def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Red
 
 def _profile_partial(b: Bundle, coord_fn, edges_key: str) -> dict:
     edges = np.asarray(getattr(b.ctx, edges_key))
-    r_lf = attach_entry_axis(b.r_all, entry_axis(b.r_all))
-    t_lf = attach_entry_axis(b.t_all, entry_axis(b.t_all))
-    r_ids, r_mat = profile_partial(r_lf, coord_fn(), edges, pl.col("edep"))
-    t_ids, t_mat = profile_partial(t_lf, coord_fn(), edges, pl.col("edep"))
+
+    def _mat(lf: pl.LazyFrame) -> dict:
+        lf2 = attach_entry_axis(lf, entry_axis(lf))
+        ids, mat = profile_partial(lf2, coord_fn(), edges, pl.col("edep"))
+        return {"ids": ids.tolist(), "mat": mat.tolist()}
+
     return {
-        "r_ids": r_ids.tolist(),
-        "r_mat": r_mat.tolist(),
-        "t_ids": t_ids.tolist(),
-        "t_mat": t_mat.tolist(),
+        "r": _per_rollout(b, lambda rs: _mat(rs.all)),
+        "t": _mat(b.t_all),
     }
 
 
@@ -415,12 +575,19 @@ def _profile_finalize(
 ) -> Reduced:
     edges = np.asarray(getattr(ctx, edges_key))
     nb = len(edges) - 1
-    _assert_event_disjoint([p["r_ids"] for p in parts], spec_id, "rollout")
-    _assert_event_disjoint([p["t_ids"] for p in parts], spec_id, "reference")
-    r_mats = [np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts]
-    t_mats = [np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts]
-    r_mean, r_std = profile_finalize(r_mats)
+    names = list(parts[0]["r"])
+
+    _assert_event_disjoint([p["t"]["ids"] for p in parts], spec_id, "reference")
+    t_mats = [np.asarray(p["t"]["mat"], dtype=float).reshape(-1, nb) for p in parts]
     t_mean, t_std = profile_finalize(t_mats)
+
+    series: dict[str, dict] = {}
+    for name in names:
+        _assert_event_disjoint([p["r"][name]["ids"] for p in parts], spec_id, name)
+        mats = [np.asarray(p["r"][name]["mat"], dtype=float).reshape(-1, nb) for p in parts]
+        mean, std = profile_finalize(mats)
+        series[name] = {"mean": mean.tolist(), "std": std.tolist()}
+
     return Reduced(
         id=spec_id,
         family="shower",
@@ -429,35 +596,85 @@ def _profile_finalize(
         xlabel=xlabel,
         payload={
             "edges": edges.tolist(),
-            "rollout_mean": r_mean.tolist(),
-            "rollout_std": r_std.tolist(),
-            "reference_mean": t_mean.tolist(),
-            "reference_std": t_std.tolist(),
+            "series": series,
+            "reference": {"mean": t_mean.tolist(), "std": t_std.tolist()},
             "ylabel": "mean deposited energy per event [MeV]",
         },
     )
 
 
+# ---------------------------------------------------------------------------
+# shower containment depth (reuses the longitudinal profile's per-event matrix)
+# ---------------------------------------------------------------------------
+
+_CONTAINMENT_QUANTILES: list[tuple[float, str]] = [
+    (0.90, "shower_containment_depth_90"),
+    (0.95, "shower_containment_depth_95"),
+]
+
+
+def _containment_finalize(parts: list[dict], ctx: Context, spec_id: str, quantile: float) -> Reduced:
+    edges = np.asarray(ctx.depth_edges)
+    nb = len(edges) - 1
+    names = list(parts[0]["r"])
+
+    _assert_event_disjoint([p["t"]["ids"] for p in parts], spec_id, "reference")
+    t_full = np.concatenate([np.asarray(p["t"]["mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
+    t_depth = _containment_depths(t_full, edges, quantile)
+
+    r_depths: dict[str, np.ndarray] = {}
+    for name in names:
+        _assert_event_disjoint([p["r"][name]["ids"] for p in parts], spec_id, name)
+        full = np.concatenate([np.asarray(p["r"][name]["mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
+        r_depths[name] = _containment_depths(full, edges, quantile)
+
+    hedges, counts = _np_hist_shared_edges([t_depth, *(r_depths[n] for n in names)], ctx.n_marginal_bins)
+    t_counts, *r_counts = counts
+    return Reduced(
+        id=spec_id,
+        family="shower",
+        kind="overlay_hist",
+        title=f"Shower containment depth ({quantile:.0%} of deposited energy)",
+        xlabel=f"depth containing {quantile:.0%} of deposited energy [mm]",
+        payload={
+            "edges": hedges.tolist(),
+            "series": {name: c.astype(np.int64).tolist() for name, c in zip(names, r_counts)},
+            "reference": t_counts.astype(np.int64).tolist(),
+            "log_y": False,
+        },
+    )
+
+
 # ---------------------------------------------------------------------------
 # species share + leakage
 # ---------------------------------------------------------------------------
 
 
 def _species_share_partial(b: Bundle) -> dict:
-    r = species_share(b.r_all)
     t = species_share(b.t_all)
+
+    def _map(rs: RolloutSide) -> dict[str, float]:
+        r = species_share(rs.all)
+        return {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())}
+
     return {
-        "r": {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())},
+        "r": _per_rollout(b, _map),
         "t": {str(k): v for k, v in zip(t["pdg"].to_list(), t["total_edep"].to_list())},
     }
 
 
 def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    r_map = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
+    r_maps = {n: sum_merge([p["r"][n] for p in parts]) for n in names}
     t_map = sum_merge([p["t"] for p in parts])
-    r_tot = sum(r_map.values()) or 1.0
     t_tot = sum(t_map.values()) or 1.0
     labels = [pdg_label(k) for k in ctx.top_pdgs]
+
+    series: dict[str, list[float]] = {}
+    for n in names:
+        r_tot = sum(r_maps[n].values()) or 1.0
+        series[n] = [r_maps[n].get(str(k), 0.0) / r_tot for k in ctx.top_pdgs]
+
     return Reduced(
         id="species_edep_share",
         family="species",
@@ -466,22 +683,23 @@ def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="species",
         payload={
             "labels": labels,
-            _ROLL: [r_map.get(str(k), 0.0) / r_tot for k in ctx.top_pdgs],
-            _REF: [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs],
+            "series": series,
+            "reference": [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs],
             "ylabel": "fraction of total deposited energy",
         },
     )
 
 
 def _leakage_partial(b: Bundle) -> dict:
-    frac = leakage_fraction(b.r_all)
-    return {"frac": frac.tolist()}
+    return {"r": _per_rollout(b, lambda rs: leakage_fraction(rs.all).tolist())}
 
 
 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)
-    counts = np.histogram(frac, edges)[0]
+    names = list(parts[0]["r"])
+    arrays = {n: np.concatenate([np.asarray(p["r"][n], dtype=float) for p in parts]) for n in names}
+    max_val = max((float(a.max()) for a in arrays.values() if len(a)), default=1e-3)
+    edges = np.linspace(0.0, max(max_val, 1e-3), ctx.n_marginal_bins + 1)
+    series = {n: np.histogram(arrays[n], edges)[0].astype(np.int64).tolist() for n in names}
     return Reduced(
         id="leakage_fraction",
         family="species",
@@ -490,7 +708,7 @@ def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="escaped energy fraction",
         payload={
             "edges": edges.tolist(),
-            _ROLL: counts.astype(np.int64).tolist(),
+            "series": series,
             "log_y": True,
             "note": "rollout only; the reference has no detector-escape concept",
         },
@@ -502,24 +720,36 @@ def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
 # ---------------------------------------------------------------------------
 
 
-def _sec_frames(b: Bundle):
-    return (
-        secondaries(b.r_phys, Side.rollout),
-        secondaries(b.t_all, Side.reference),
-    )
+def _t_sec(b: Bundle) -> pl.LazyFrame:
+    return secondaries(b.t_all, Side.reference)
+
+
+def _r_sec(rs: RolloutSide) -> pl.LazyFrame:
+    return secondaries(rs.phys, Side.rollout)
 
 
 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()
-    return {"r": r.tolist(), "t": t.tolist()}
+    t = _t_sec(b).group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
+
+    def _r(rs: RolloutSide) -> list[float]:
+        return (
+            _r_sec(rs)
+            .group_by("event_id")
+            .agg(pl.len().alias("n"))
+            .collect(engine="streaming")["n"]
+            .to_numpy()
+            .tolist()
+        )
+
+    return {"r": _per_rollout(b, _r), "t": t.tolist()}
 
 
 def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
+    names = list(parts[0]["r"])
     t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
-    edges, rc, tc = _np_hist_pair(r, t, min(ctx.n_marginal_bins, 40))
+    r = {n: np.concatenate([np.asarray(p["r"][n], dtype=float) for p in parts]) for n in names}
+    edges, counts = _np_hist_shared_edges([t, *(r[n] for n in names)], min(ctx.n_marginal_bins, 40))
+    t_c, *r_cs = counts
     return Reduced(
         id="sec_count_per_event",
         family="secondaries",
@@ -528,8 +758,8 @@ def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="secondaries per event",
         payload={
             "edges": edges.tolist(),
-            _ROLL: rc.astype(np.int64).tolist(),
-            _REF: tc.astype(np.int64).tolist(),
+            "series": {n: c.astype(np.int64).tolist() for n, c in zip(names, r_cs)},
+            "reference": t_c.astype(np.int64).tolist(),
             "log_y": False,
         },
     )
@@ -541,14 +771,22 @@ def _counts_by_pdg(sec_lf: pl.LazyFrame) -> dict[str, int]:
 
 
 def _sec_count_per_species_partial(b: Bundle) -> dict:
-    r_sec, t_sec = _sec_frames(b)
-    return {"r": _counts_by_pdg(r_sec), "t": _counts_by_pdg(t_sec)}
+    return {"r": _per_rollout(b, lambda rs: _counts_by_pdg(_r_sec(rs))), "t": _counts_by_pdg(_t_sec(b))}
 
 
 def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
+    r_maps = {n: sum_merge([p["r"][n] for p in parts]) for n in names}
     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)]
+
+    all_keys = set(t)
+    for m in r_maps.values():
+        all_keys |= set(m)
+
+    def _total(k: str) -> float:
+        return t.get(k, 0) + sum(m.get(k, 0) for m in r_maps.values())
+
+    keys = sorted(all_keys, key=lambda k: -_total(k))[: len(ctx.top_pdgs)]
     return Reduced(
         id="sec_count_per_species",
         family="secondaries",
@@ -557,39 +795,34 @@ def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="species",
         payload={
             "labels": [pdg_label(int(k)) for k in keys],
-            _ROLL: [float(r.get(k, 0)) for k in keys],
-            _REF: [float(t.get(k, 0)) for k in keys],
+            "series": {n: [float(r_maps[n].get(k, 0)) for k in keys] for n in names},
+            "reference": [float(t.get(k, 0)) for k in keys],
             "ylabel": "secondary count",
         },
     )
 
 
 def _sec_energy_partial(b: Bundle) -> dict:
-    r_sec, t_sec = _sec_frames(b)
     edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1)
     return {
-        "r": _partial_hist(r_sec, pl.col("energy"), edges),
-        "t": _partial_hist(t_sec, pl.col("energy"), edges),
+        "r": _per_rollout(b, lambda rs: _partial_hist(_r_sec(rs), pl.col("energy"), edges)),
+        "t": _partial_hist(_t_sec(b), pl.col("energy"), edges),
     }
 
 
 def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
     edges = np.linspace(*ctx.sec_energy_range, ctx.n_sec_bins + 1)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
     t = sum_merge([p["t"] for p in parts])
+    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
     return Reduced(
         id="sec_energy",
         family="secondaries",
         kind="overlay_hist",
         title="Secondary birth energy",
         xlabel="secondary energy [MeV]",
-        payload={
-            "edges": edges.tolist(),
-            _ROLL: _finalize_counts(r, 0, nb),
-            _REF: _finalize_counts(t, 0, nb),
-            "log_y": True,
-        },
+        payload={"edges": edges.tolist(), "series": series, "reference": _finalize_counts(t, 0, nb), "log_y": True},
     )
 
 
@@ -603,26 +836,80 @@ def _sec_cos_angle_partial(b: Bundle) -> dict:
         ea = entry_axis(steps_lf)
         return _partial_hist(attach_entry_axis(sec_lf, ea), cos, edges)
 
-    r_sec, t_sec = _sec_frames(b)
-    return {"r": _side(r_sec, b.r_phys), "t": _side(t_sec, b.t_all)}
+    return {
+        "r": _per_rollout(b, lambda rs: _side(_r_sec(rs), rs.phys)),
+        "t": _side(_t_sec(b), b.t_all),
+    }
 
 
 def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
     edges = np.linspace(-1.0, 1.0, ctx.n_sec_bins + 1)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
     t = sum_merge([p["t"] for p in parts])
+    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
     return Reduced(
         id="sec_cos_angle",
         family="secondaries",
         kind="overlay_hist",
         title="Secondary emission angle relative to the shower axis",
         xlabel="cos of emission angle",
+        payload={"edges": edges.tolist(), "series": series, "reference": _finalize_counts(t, 0, nb), "log_y": False},
+    )
+
+
+def _n_sec_confusion_partial(b: Bundle) -> dict:
+    t_ids, t_n = sec_count_by_event(b.t_all, _t_sec(b))
+
+    def _r(rs: RolloutSide) -> dict:
+        ids, n = sec_count_by_event(rs.phys, _r_sec(rs))
+        return {"ids": ids.tolist(), "n": n.tolist()}
+
+    return {"r": _per_rollout(b, _r), "t": {"ids": t_ids.tolist(), "n": t_n.tolist()}}
+
+
+def _n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
+    names = list(parts[0]["r"])
+    # event-disjoint chunking (see Bundle.open) means each event_id appears in
+    # exactly one part on each side, so a plain dict build is a safe merge.
+    t_ids = np.concatenate([np.asarray(p["t"]["ids"], dtype=np.int64) for p in parts])
+    t_n = np.concatenate([np.asarray(p["t"]["n"], dtype=np.int64) for p in parts])
+    t_map = dict(zip(t_ids.tolist(), t_n.tolist()))
+
+    pairs: dict[str, tuple[np.ndarray, np.ndarray]] = {}
+    max_val = 0
+    for name in names:
+        r_ids = np.concatenate([np.asarray(p["r"][name]["ids"], dtype=np.int64) for p in parts])
+        r_n = np.concatenate([np.asarray(p["r"][name]["n"], dtype=np.int64) for p in parts])
+        r_map = dict(zip(r_ids.tolist(), r_n.tolist()))
+        common = sorted(set(r_map) & set(t_map))
+        true_n = np.array([t_map[e] for e in common], dtype=np.int64)
+        pred_n = np.array([r_map[e] for e in common], dtype=np.int64)
+        pairs[name] = (true_n, pred_n)
+        if len(true_n):
+            max_val = max(max_val, int(true_n.max()), int(pred_n.max()))
+
+    cap = min(max(max_val, 1), 20)
+    matrices: dict[str, list[list[int]]] = {}
+    labels: list[str] = []
+    for name in names:
+        true_n, pred_n = pairs[name]
+        labels, mat = _integer_confusion(true_n, pred_n, cap=cap)
+        matrices[name] = mat.tolist()
+
+    return Reduced(
+        id="n_sec_confusion",
+        family="secondaries",
+        kind="heatmap",
+        title="Predicted vs true secondary count per event",
+        xlabel="predicted secondaries (rollout)",
         payload={
-            "edges": edges.tolist(),
-            _ROLL: _finalize_counts(r, 0, nb),
-            _REF: _finalize_counts(t, 0, nb),
-            "log_y": False,
+            "series": matrices,
+            "row_labels": labels,
+            "col_labels": labels,
+            "ylabel": "true secondaries (reference)",
+            "cbar_label": "event count",
+            "vmin": 0.0,
         },
     )
 
@@ -631,17 +918,18 @@ def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
 # router diagnostics (not chunked — already bounded/subsampled)
 # ---------------------------------------------------------------------------
 
-_router_gating_partial, _router_gating_finalize = _unchunkable(
-    lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys)
-)
+_router_gating_partial, _router_gating_finalize = _unchunkable(lambda b: compute_router_gating(b.rollouts, 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.rollouts, 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)
+    lambda b: compute_router_share_by_process(b.rollouts, b.t_phys)
+)
+_router_specialization_partial, _router_specialization_finalize = _unchunkable(
+    lambda b: compute_router_specialization(b.rollouts, 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)
+    lambda b: compute_type_embedding_l1_distance(b.rollouts)
 )
 
 
@@ -676,6 +964,15 @@ def build_catalog() -> list[PlotSpec]:
                 )
             )
 
+    specs.append(
+        PlotSpec(
+            "marginal_distance_summary",
+            "quality",
+            compute_partial=_distance_summary_partial,
+            finalize=_distance_summary_finalize,
+        )
+    )
+
     specs += [
         PlotSpec(
             "event_total_edep",
@@ -745,6 +1042,17 @@ def build_catalog() -> list[PlotSpec]:
                 "transverse_edges",
             ),
         ),
+    ]
+    for quantile, spec_id in _CONTAINMENT_QUANTILES:
+        specs.append(
+            PlotSpec(
+                spec_id,
+                "shower",
+                compute_partial=lambda b: _profile_partial(b, depth_expr, "depth_edges"),
+                finalize=lambda parts, ctx, q=quantile, sid=spec_id: _containment_finalize(parts, ctx, sid, q),
+            )
+        )
+    specs += [
         PlotSpec(
             "species_edep_share",
             "species",
@@ -781,6 +1089,12 @@ def build_catalog() -> list[PlotSpec]:
             compute_partial=_sec_cos_angle_partial,
             finalize=_sec_cos_angle_finalize,
         ),
+        PlotSpec(
+            "n_sec_confusion",
+            "secondaries",
+            compute_partial=_n_sec_confusion_partial,
+            finalize=_n_sec_confusion_finalize,
+        ),
         PlotSpec(
             "router_gating",
             "model",
@@ -802,6 +1116,13 @@ def build_catalog() -> list[PlotSpec]:
             finalize=_router_share_process_finalize,
             chunkable=False,
         ),
+        PlotSpec(
+            "router_specialization",
+            "model",
+            compute_partial=_router_specialization_partial,
+            finalize=_router_specialization_finalize,
+            chunkable=False,
+        ),
         PlotSpec(
             "type_embedding_l1_distance",
             "model",

giant/analysis/condor.py

@@ -1,4 +1,4 @@
-"""HTCondor orchestration driven by a ``giant rollout`` YAML sidecar.
+"""HTCondor orchestration driven by one or more ``giant rollout`` YAML sidecars.
 
 A rollout writes a YAML sidecar (``giant/cli.py:_write_prediction_ref`` +
 rollout extras) that already names both files we need and carries the run's
@@ -10,8 +10,11 @@ provenance:
 * ``checkpoint``, ``geometry_oracle``, ``energy_cutoff``, ``steps``, ... —
   metadata that flows straight into every plot's gallery ``metadata.yaml``.
 
-So the analysis takes that one YAML as input, derives its own **run directory**
-next to the rollout parquet, and lays everything out under it:
+The analysis takes N such YAMLs — one series per rollout, all required to
+share the same ``dataset`` (the premise is "N candidates vs one ground
+truth") — resolves each one's series name (``load_rollout_yamls``), derives
+its own **run directory** next to the first rollout's parquet, and lays
+everything out under it:
 
     <run_dir>/shared.json     fixed bin edges / group sets (prep)
     <run_dir>/run_meta.json   resolved rollout/reference paths + plot metadata
@@ -44,6 +47,7 @@ from __future__ import annotations
 import json
 import shutil
 import sys
+from collections.abc import Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -54,7 +58,7 @@ from giant.analysis.catalog import Bundle, catalog_ids, get_spec
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial
 from giant.analysis.runtime_estimate import estimate_runtime_s
-from giant.analysis.sources import Side, open_side
+from giant.analysis.sources import RolloutSpec, Side, open_side
 
 # Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
 _PLOT_META_KEYS = (
@@ -111,8 +115,60 @@ def load_rollout_yaml(path: str | Path) -> dict:
     return d
 
 
+@dataclass
+class LoadedRollout:
+    """One rollout YAML plus its resolved series ``name`` (see ``load_rollout_yamls``)."""
+
+    name: str
+    yaml: dict
+
+
+def load_rollout_yamls(
+    paths: Sequence[str | Path], labels: Sequence[str] | None = None
+) -> tuple[list[LoadedRollout], str]:
+    """Load every rollout YAML, resolve each one's series name, and verify they
+    all share one reference (``dataset``) file — the premise is "N candidates
+    vs one ground truth", not N independent comparisons.
+
+    Names: an explicit ``labels[i]`` if given (``labels`` must be empty or
+    exactly ``len(paths)`` long); otherwise the YAML's stem for N>1, or
+    ``"rollout"`` for the single-YAML case — matching today's one-series
+    legend/payload key, so a single-rollout run renders identically to
+    before this feature existed. Raises ``ValueError`` if two rollouts
+    resolve to the same name, or if the YAMLs don't all name the same
+    ``dataset``.
+    """
+    if labels and len(labels) != len(paths):
+        raise ValueError(f"--label given {len(labels)} time(s) but {len(paths)} rollout YAML(s) were passed")
+    yamls = [load_rollout_yaml(p) for p in paths]
+    if labels:
+        names = list(labels)
+    elif len(paths) == 1:
+        names = ["rollout"]
+    else:
+        names = [Path(p).stem for p in paths]
+    if len(set(names)) != len(names):
+        dupes = sorted({n for n in names if names.count(n) > 1})
+        raise ValueError(f"rollout series names collide: {dupes} — pass --label to disambiguate")
+
+    references = {str(y["dataset"]) for y in yamls}
+    if len(references) > 1:
+        detail = "\n".join(f"  {p}: dataset={y['dataset']!r}" for p, y in zip(paths, yamls))
+        raise ValueError(
+            "all rollout YAMLs must be seeded from the same reference (dataset) "
+            f"file — got {len(references)} distinct ones:\n{detail}"
+        )
+
+    return [LoadedRollout(name=n, yaml=y) for n, y in zip(names, yamls)], yamls[0]["dataset"]
+
+
+def _run_tag(y: dict) -> str:
+    rollout = Path(y["output"])
+    return str(y.get("prediction_id") or rollout.stem)[:8]
+
+
 def derive_run_dir(
-    rollout_yaml: dict,
+    rollout_yamls: list[dict],
     run_dir: str | Path | None = None,
     default_base: str | Path | None = None,
 ) -> Path:
@@ -122,14 +178,24 @@ def derive_run_dir(
     ``default_base / analysis_<tag>`` if ``default_base`` is given (the CLI
     passes the repo's gitignored ``analysis_runs/``, so run directories don't
     pile up on ``/ceph`` next to the rollout parquet). Falls back to next to
-    the rollout parquet — the original convention — for callers that don't
-    care where the run directory lives.
+    the *first* rollout's parquet — the original convention — for callers
+    that don't care where the run directory lives.
+
+    ``tag`` is a single rollout's ``prediction_id``/output stem (matching
+    today's single-rollout convention exactly) when there's only one; for
+    N>1 it joins up to three tags with ``-``, then ``-plus<K>`` for any
+    beyond that, so a many-rollout run still gets a short, stable directory
+    name.
     """
     if run_dir is not None:
         return Path(run_dir)
-    rollout = Path(rollout_yaml["output"])
-    tag = str(rollout_yaml.get("prediction_id") or rollout.stem)[:8]
-    base = Path(default_base) if default_base is not None else rollout.parent
+    tags = [_run_tag(y) for y in rollout_yamls]
+    if len(tags) == 1:
+        tag = tags[0]
+    else:
+        shown, rest = tags[:3], tags[3:]
+        tag = "-".join(shown) + (f"-plus{len(rest)}" if rest else "")
+    base = Path(default_base) if default_base is not None else Path(rollout_yamls[0]["output"]).parent
     return base / f"analysis_{tag}"
 
 
@@ -139,17 +205,21 @@ def _plot_meta(rollout_yaml: dict) -> dict:
 
 @dataclass
 class RunMeta:
-    """Resolved paths + plot metadata for one analysis run (``run_meta.json``)."""
+    """Resolved paths + plot metadata for one analysis run (``run_meta.json``).
 
-    rollout: str
+    ``rollouts`` is ``[{"name", "path", "plot_meta"}, ...]``, insertion order
+    = the order rollouts were given on the CLI (and so the order every
+    ``Reduced.payload["series"]`` dict is built in — see ``catalog.py``).
+    """
+
+    rollouts: list[dict]
     reference: str
     run_dir: str
     title: str
-    plot_meta: dict
     n_chunks: int = 1
-    # rollout+reference row count of each event_id-disjoint chunk, and the
-    # dataset total — inputs to `runtime_estimate.estimate_runtime_s`. Empty/0
-    # on run directories written before this field existed.
+    # combined rollout+reference row count of each event_id-disjoint chunk,
+    # and the dataset total — inputs to `runtime_estimate.estimate_runtime_s`.
+    # Empty/0 on run directories written before this field existed.
     rows_per_chunk: list[int] = field(default_factory=list)
     total_rows: int = 0
 
@@ -161,8 +231,8 @@ 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]:
-    """Rollout+reference row count of each ``event_id % n_chunks`` chunk.
+def _rows_per_chunk(rollouts: list[str | Path], reference: str | Path, n_chunks: int) -> list[int]:
+    """Combined rollout+reference row count of each ``event_id % n_chunks`` chunk.
 
     One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
     the sizing input every job's estimated walltime
@@ -178,7 +248,8 @@ def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -
         )
 
     out = [0] * n_chunks
-    for lf in (open_side(rollout, Side.rollout), open_side(reference, Side.reference)):
+    sides = [open_side(reference, Side.reference)] + [open_side(r, Side.rollout) for r in rollouts]
+    for lf in sides:
         df = counts(lf)
         for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
             out[c] += n
@@ -186,20 +257,22 @@ def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -
 
 
 def prep(
-    rollout_yaml: str | Path,
+    rollout_yamls: Sequence[str | Path],
     run_dir: str | Path | None = None,
     n_chunks: int = 1,
     default_base: str | Path | None = None,
+    labels: Sequence[str] | None = None,
     **ctx_kwargs,
 ) -> Path:
-    """Read the rollout YAML, build the shared context, and lay out the run dir.
+    """Read the rollout YAML(s), build the shared context, and lay out the run dir.
 
     Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
     ``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one``
     job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is
     resolved once here rather than re-passed (and risking disagreement) at every
     later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
-    resolve the actual directory.
+    resolve the actual directory, and ``load_rollout_yamls`` for how
+    ``labels``/YAML stems resolve each rollout's series name.
 
     Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
     this same ``run_dir``: partial files carry no record of what context
@@ -208,8 +281,8 @@ def prep(
     rollout/reference files changed) would otherwise let ``merge_one`` silently
     merge stale partials against the new ``shared.json``.
     """
-    y = load_rollout_yaml(rollout_yaml)
-    run_path = derive_run_dir(y, run_dir, default_base=default_base)
+    loaded, reference = load_rollout_yamls(list(rollout_yamls), labels)
+    run_path = derive_run_dir([lr.yaml for lr in loaded], run_dir, default_base=default_base)
     run_path.mkdir(parents=True, exist_ok=True)
 
     for stale in ("reduced_partial", "reduced"):
@@ -217,19 +290,22 @@ def prep(
         if stale_dir.exists():
             shutil.rmtree(stale_dir)
 
-    rollout, reference = y["output"], y["dataset"]
-    ctx = build_context(rollout, reference, **ctx_kwargs)
+    rollout_specs = [RolloutSpec(name=lr.name, source=lr.yaml["output"]) for lr in loaded]
+    ctx = build_context(rollout_specs, reference, **ctx_kwargs)
     ctx.save(run_path / "shared.json")
 
-    rows_per_chunk = _rows_per_chunk(rollout, reference, n_chunks)
+    rows_per_chunk = _rows_per_chunk([lr.yaml["output"] for lr in loaded], reference, n_chunks)
+
+    rollouts_meta = [
+        {"name": lr.name, "path": str(lr.yaml["output"]), "plot_meta": _plot_meta(lr.yaml)} for lr in loaded
+    ]
+    ckpts = ", ".join(Path(lr.yaml.get("checkpoint", "")).name or "rollout" for lr in loaded)
 
-    ckpt = Path(y.get("checkpoint", "")).name or "rollout"
     RunMeta(
-        rollout=str(rollout),
+        rollouts=rollouts_meta,
         reference=str(reference),
         run_dir=str(run_path),
-        title=f"GIANT rollout analysis — {ckpt}",
-        plot_meta=_plot_meta(y),
+        title=f"GIANT rollout analysis — {ckpts}",
         n_chunks=n_chunks,
         rows_per_chunk=rows_per_chunk,
         total_rows=sum(rows_per_chunk),
@@ -244,17 +320,19 @@ def prep(
 
 def compute_reduced(
     spec_id: str,
-    rollout: str | Path,
+    rollouts: list[dict],
     reference: str | Path,
     shared: str | Path,
     out: str | Path,
-    checkpoint: str | None = None,
     chunk_index: int = 0,
     n_chunks: int = 1,
-    type_embedding_l1_dist: dict | None = None,
 ) -> Path:
     """Core: run one (plot, chunk)'s partial reduction against explicit paths.
 
+    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?},
+    ...]``, one per rollout series (insertion order preserved through to every
+    plot's ``Reduced.payload["series"]``).
+
     Writes a ``Partial`` JSON — the raw, not-yet-merged output of
     ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one``
     is what combines every chunk's ``Partial`` for a plot into the final
@@ -268,14 +346,16 @@ def compute_reduced(
         raise ValueError(
             f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})"
         )
-    bundle = Bundle.open(
-        rollout,
-        reference,
-        ctx,
-        checkpoint=checkpoint,
-        chunk=(chunk_index, effective_n),
-        type_embedding_l1_dist=type_embedding_l1_dist,
-    )
+    rollout_specs = [
+        RolloutSpec(
+            name=r["name"],
+            source=r["path"],
+            checkpoint=r.get("checkpoint"),
+            type_embedding_l1_dist=r.get("type_embedding_l1_dist"),
+        )
+        for r in rollouts
+    ]
+    bundle = Bundle.open(rollout_specs, reference, ctx, chunk=(chunk_index, effective_n))
     partial = Partial(
         id=spec_id,
         family=spec.family,
@@ -291,16 +371,23 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
     """Run one (plot, chunk)'s partial reduction from a prepped run directory."""
     run_path = Path(run_dir)
     meta = RunMeta.load(run_path / "run_meta.json")
+    rollouts = [
+        {
+            "name": ro["name"],
+            "path": ro["path"],
+            "checkpoint": ro["plot_meta"].get("checkpoint"),
+            "type_embedding_l1_dist": ro["plot_meta"].get("type_embedding_l1_dist"),
+        }
+        for ro in meta.rollouts
+    ]
     return compute_reduced(
         spec_id,
-        meta.rollout,
+        rollouts,
         meta.reference,
         run_path / "shared.json",
         run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json",
-        checkpoint=meta.plot_meta.get("checkpoint"),
         chunk_index=chunk_index,
         n_chunks=meta.n_chunks,
-        type_embedding_l1_dist=meta.plot_meta.get("type_embedding_l1_dist"),
     )
 
 

giant/analysis/context.py

@@ -26,7 +26,7 @@ from giant.analysis.reduce import (
     entry_axis,
     transverse_expr,
 )
-from giant.analysis.sources import Side, open_side, physical_steps, secondaries
+from giant.analysis.sources import RolloutSpec, Side, open_side, physical_steps, secondaries
 from giant.analysis.variables import RANGED_VARS
 
 
@@ -74,9 +74,9 @@ 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]:
-    """Robust (lo_q, hi_q) range over the union of two value samples."""
-    both = np.concatenate([r_vals, t_vals])
+def _combined_quantiles(vals: list[np.ndarray], lo_q: float, hi_q: float) -> tuple[float, float]:
+    """Robust (lo_q, hi_q) range over the union of several value samples."""
+    both = np.concatenate(vals)
     lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
     if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
         lo, hi = lo - 0.5, hi + 0.5
@@ -84,7 +84,7 @@ def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_
 
 
 def build_context(
-    rollout: str | Path | pl.LazyFrame,
+    rollouts: list[RolloutSpec],
     reference: str | Path | pl.LazyFrame,
     *,
     n_energy_bins: int = 4,
@@ -94,41 +94,51 @@ def build_context(
     sample_rows: int = 1_000_000,
     seed: int = 0,
 ) -> Context:
-    """Resolve the shared context from the two files (the ``prep`` step)."""
-    r_all = open_side(rollout, Side.rollout)
+    """Resolve the shared context from the reference + every rollout (the ``prep`` step).
+
+    Every range/quantile below is the union of the reference and *all*
+    rollouts, so a single set of fixed bin edges/group sets is valid for
+    every series a compute job streams over.
+    """
     t_all = open_side(reference, Side.reference)
-    r_lf = physical_steps(r_all, Side.rollout)
     t_lf = physical_steps(t_all, Side.reference)
+    r_lfs = {rs.name: physical_steps(open_side(rs.source, Side.rollout), Side.rollout) for rs in rollouts}
 
     # 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 = {
+        name: _row_subsample(lf, sample_rows, seed).select(exprs).collect(engine="streaming")
+        for name, lf in r_lfs.items()
+    }
     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([t_s[name].to_numpy(), *(df[name].to_numpy() for df in r_s.values())], _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()
 
-    r_inc, t_inc = _incident(r_lf), _incident(t_lf)
-    energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
+    t_inc = _incident(t_lf)
+    r_inc = {name: _incident(lf) for name, lf in r_lfs.items()}
+    energy_edges = energy_bin_edges(np.concatenate([t_inc, *r_inc.values()]), n_energy_bins)
 
     # Top PDG species and material list (cheap single-column group_bys).
     def _counts(lf: pl.LazyFrame, col: str) -> pl.DataFrame:
         return lf.group_by(col).agg(pl.len().alias("n")).collect(engine="streaming")
 
     pdg_counts = (
-        pl.concat([_counts(r_lf, "pdg"), _counts(t_lf, "pdg")])
+        pl.concat([_counts(t_lf, "pdg"), *(_counts(lf, "pdg") for lf in r_lfs.values())])
         .group_by("pdg")
         .agg(pl.col("n").sum())
         .sort("n", descending=True)
     )
     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())
-    )
+    material_set: set[str] = set(_counts(t_lf, "material")["material"].to_list())
+    for lf in r_lfs.values():
+        material_set |= set(_counts(lf, "material")["material"].to_list())
+    materials = sorted(material_set)
 
     # Shower depth / transverse ranges from a subsampled proxy.
     def _proxy(lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
@@ -140,19 +150,20 @@ def build_context(
         )
         return sub["d"].to_numpy(), sub["t"].to_numpy()
 
-    r_d, r_t = _proxy(r_lf)
     t_d, t_t = _proxy(t_lf)
-    d_lo, d_hi = _combined_quantiles(r_d, t_d, _LO_Q, _HI_Q)
+    r_proxy = {name: _proxy(lf) for name, lf in r_lfs.items()}
+    d_lo, d_hi = _combined_quantiles([t_d, *(p[0] for p in r_proxy.values())], _LO_Q, _HI_Q)
     depth_edges = np.linspace(d_lo, d_hi, n_marginal_bins + 1)
-    t_hi = max(float(np.quantile(np.concatenate([r_t, t_t]), _HI_Q)), 1e-6)
+    t_hi = max(float(np.quantile(np.concatenate([t_t, *(p[1] for p in r_proxy.values())]), _HI_Q)), 1e-6)
     transverse_edges = np.linspace(0.0, t_hi, n_marginal_bins + 1)
 
     # Secondary energy range.
-    r_se = secondaries(r_lf, Side.rollout).select("energy")
-    t_se = secondaries(t_all, Side.reference).select("energy")
-    r_se = _row_sample_col(r_se, sample_rows, seed)
-    t_se = _row_sample_col(t_se, sample_rows, seed)
-    sec_energy_range = _combined_quantiles(r_se, t_se, _LO_Q, _HI_Q)
+    t_se = _row_sample_col(secondaries(t_all, Side.reference).select("energy"), sample_rows, seed)
+    r_se = {
+        name: _row_sample_col(secondaries(lf, Side.rollout).select("energy"), sample_rows, seed)
+        for name, lf in r_lfs.items()
+    }
+    sec_energy_range = _combined_quantiles([t_se, *r_se.values()], _LO_Q, _HI_Q)
 
     return Context(
         n_marginal_bins=n_marginal_bins,
@@ -165,8 +176,8 @@ def build_context(
         sec_energy_range=sec_energy_range,
         n_sec_bins=n_sec_bins,
         n_events={
-            "rollout": len(r_inc),
             "reference": len(t_inc),
+            **{name: len(arr) for name, arr in r_inc.items()},
         },
     )
 

giant/analysis/reduce.py

@@ -271,3 +271,20 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
     escaped = per_event["escaped"].fill_null(0.0).to_numpy()
     total = deposited + escaped
     return np.where(total > 0, escaped / total, 0.0)
+
+
+def sec_count_by_event(lf_all: pl.LazyFrame, sec_lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
+    """Per-event secondary count, zero-filled for events that produced none.
+
+    Two bounded per-event ``group_by``s — the full event set (from ``lf_all``)
+    and the secondary counts (from ``sec_lf``, see ``sources.secondaries``) —
+    merged in Python via a dict. Both results are event-granularity (not
+    per-row), so this stays in the same bounded-memory budget as
+    ``event_scalars``; a plain ``group_by`` on ``sec_lf`` alone would silently
+    drop zero-secondary events instead of zero-filling them.
+    """
+    ev = lf_all.select("event_id").unique().collect(engine="streaming")["event_id"].to_numpy()
+    cnt_df = sec_lf.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")
+    cnt = dict(zip(cnt_df["event_id"].to_list(), cnt_df["n"].to_list()))
+    counts = np.array([cnt.get(int(e), 0) for e in ev], dtype=np.int64)
+    return ev, counts

giant/analysis/reduced.py

@@ -11,15 +11,24 @@ import json
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
 
-# Reduced.kind values:
-#   "overlay_hist"   rollout vs reference density histogram over shared edges
+# Reduced.kind values (payload keys a rollout series by name under
+# payload["series"], with the reference — where one exists — kept as one
+# distinguished payload["reference"] entry; see catalog.py's module
+# docstring for the full per-kind payload shape):
+#   "overlay_hist"   N-rollout-series vs reference density histogram over shared edges
 #   "grouped_hist"   one panel per group (energy/pdg/material), each an overlay
-#   "profile"        edep-weighted mean +/- event-RMS vs depth/radius, two series
-#   "bar"            per-category rollout vs reference bars (share / counts)
-#   "single_hist"    one series only (e.g. rollout leakage; reference has none)
-#   "router_gating"  stacked mean MoE gate weight vs energy, rollout + reference
-#   "router_share"   stacked bar of MoE top-1 dispatch share by category
-#   "unavailable"    plot not applicable to this run (e.g. non-MoE checkpoint)
+#   "profile"        edep-weighted mean +/- event-RMS vs depth/radius, N series + reference
+#   "bar"            per-category N-rollout-series vs reference bars (share / counts)
+#   "single_hist"    rollout-only series (e.g. leakage; reference has none)
+#   "router_gating"  stacked mean MoE gate weight vs energy, one rollout+reference
+#                    panel-pair per rollout with an enabled MoE router
+#   "router_share"   stacked bar of MoE top-1 dispatch share by category, one
+#                    panel per rollout with an enabled MoE router
+#   "router_specialization"  max gate weight vs energy (one scalar trend line
+#                    summarizing "router_gating"), per rollout with an enabled router
+#   "heatmap"        row x col matrix + colorbar, one panel per rollout (a
+#                    distance scorecard or a predicted-vs-true confusion matrix)
+#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint)
 
 
 @dataclass

giant/analysis/render.py

@@ -9,10 +9,19 @@ streaming compute.
 For each reduced artifact it writes ``<out>/<family>/<id>.pdf`` plus a sibling
 ``<id>.yaml`` (per-plot gallery metadata) and a per-family ``metadata.yaml``.
 Optionally runs ``gallery generate`` to build the static HTML site.
+
+Every rollout series gets a stable color via ``ps.get_color(i)``, ``i`` being
+its position in ``payload["series"]`` — that position is fixed by the run's
+YAML/``--label`` order (threaded unchanged from ``condor.RunMeta.rollouts``
+through every ``PlotSpec``), so a given rollout keeps the same color across
+every plot in a run. The reference, where a plot has one, always draws in one
+fixed, distinct style (dark ink, dashed) instead of taking a slot in that
+cycle.
 """
 
 from __future__ import annotations
 
+import dataclasses
 import subprocess
 from pathlib import Path
 
@@ -22,7 +31,32 @@ import yaml
 
 from giant.analysis.reduced import Reduced
 
-_SERIES_LABELS = {"rollout": "rollout", "reference": "reference (Geant4)"}
+_REFERENCE_LABEL = "reference (Geant4)"
+
+_TEX_ESCAPE_MAP = {
+    "\\": r"\textbackslash{}",
+    "%": r"\%",
+    "&": r"\&",
+    "#": r"\#",
+    "$": r"\$",
+    "_": r"\_",
+    "{": r"\{",
+    "}": r"\}",
+}
+
+
+def _tex_escape(text: str) -> str:
+    """Escape characters LaTeX treats specially in catalog-authored title/xlabel
+    text (e.g. a literal ``%`` in a "90% of deposited energy" title, which
+    ``usetex`` otherwise reads as a comment marker and aborts the whole figure —
+    see gitea #81). A single pass over the *original* characters, so the
+    backslashes an escape itself introduces (e.g. ``\textbackslash{}``) are
+    never re-escaped."""
+    return "".join(_TEX_ESCAPE_MAP.get(c, c) for c in text)
+
+
+def _ref_color() -> str:
+    return ps.colors.INK["primary"]
 
 
 def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
@@ -33,10 +67,13 @@ def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
     return counts / (total * (edges[1] - edges[0]))
 
 
-def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> None:
-    for key in ("reference", "rollout"):
-        if key in series:
-            ax.stairs(_density(series[key], edges), edges, label=_SERIES_LABELS[key])
+def _overlay(ax, edges: np.ndarray, payload: dict, log_y: bool) -> None:
+    if "reference" in payload:
+        ax.stairs(
+            _density(payload["reference"], edges), edges, label=_REFERENCE_LABEL, color=_ref_color(), linestyle="--"
+        )
+    for i, (name, counts) in enumerate(payload.get("series", {}).items()):
+        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
     if log_y:
         ax.set_yscale("log")
 
@@ -47,8 +84,8 @@ def _router_summary(router_cfg: dict) -> str:
     return f"{router_cfg.get('type', '?')}×{router_cfg.get('n_experts', '?')}"
 
 
-def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
-    """`_figure_params` for a new-shape (nested) `model_config` — has a
+def _figure_params_v2(mc: dict, meta: dict) -> dict:
+    """`_figure_params_single` for a new-shape (nested) `model_config` — has a
     `stage1_model` key. Reports stage 1's architecture (the headline
     generator); stage 2's generator is only added (`mode_s2`) when it
     differs from stage 1's, since a mixed run (the `stage1=flow` +
@@ -71,23 +108,24 @@ def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
     if particle_type is not None:
         params["conditioning"] = particle_type
     params["router"] = _router_summary(s1.get("router") or {})
-    if run_meta.get("training_epoch") is not None:
-        params["epoch"] = run_meta["training_epoch"]
-    if run_meta.get("best_val_loss") is not None:
-        params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
+    if meta.get("training_epoch") is not None:
+        params["epoch"] = meta["training_epoch"]
+    if meta.get("best_val_loss") is not None:
+        params["best_val_loss"] = round(meta["best_val_loss"], 4)
     if mode == "wgan":
         noise_dim = (s1.get("wgan") or {}).get("noise_dim")
         if noise_dim is not None:
             params["noise_dim"] = noise_dim
-    elif run_meta.get("steps") is not None:
-        params["steps"] = run_meta["steps"]
+    elif meta.get("steps") is not None:
+        params["steps"] = meta["steps"]
     return params
 
 
-def _figure_params(run_meta: dict) -> dict:
-    """Curated run identity for the figure subtitle (``new_figure(params=...)``).
+def _figure_params_single(meta: dict) -> dict:
+    """Curated run identity for the figure subtitle (``new_figure(params=...)``),
+    for exactly one rollout's ``plot_meta``.
 
-    ``run_meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
+    ``meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
     carry every threaded model/training/rollout/dataset parameter for
     after-the-fact lookup — this picks only the handful that matter for
     telling figures apart at a glance while flipping through a gallery, since
@@ -99,9 +137,9 @@ def _figure_params(run_meta: dict) -> dict:
     Handles both a v0.2 checkpoint's flat ``model_config`` and a v0.3.0
     nested one (has a ``stage1_model`` key — see ``_figure_params_v2``).
     """
-    mc = run_meta.get("model_config") or {}
+    mc = meta.get("model_config") or {}
     if "stage1_model" in mc:
-        return _figure_params_v2(mc, run_meta)
+        return _figure_params_v2(mc, meta)
 
     mode = mc.get("mode")
     params: dict = {}
@@ -114,18 +152,35 @@ def _figure_params(run_meta: dict) -> dict:
     if mc.get("conditioning") is not None:
         params["conditioning"] = mc["conditioning"]
     params["router"] = _router_summary(mc.get("router") or {})
-    if run_meta.get("training_epoch") is not None:
-        params["epoch"] = run_meta["training_epoch"]
-    if run_meta.get("best_val_loss") is not None:
-        params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
+    if meta.get("training_epoch") is not None:
+        params["epoch"] = meta["training_epoch"]
+    if meta.get("best_val_loss") is not None:
+        params["best_val_loss"] = round(meta["best_val_loss"], 4)
     if mode == "wgan":
         if mc.get("noise_dim") is not None:
             params["noise_dim"] = mc["noise_dim"]
-    elif run_meta.get("steps") is not None:
-        params["steps"] = run_meta["steps"]
+    elif meta.get("steps") is not None:
+        params["steps"] = meta["steps"]
     return params
 
 
+def _figure_params(run_meta: dict) -> dict:
+    """Curated run identity for the figure subtitle.
+
+    A single-rollout run reuses that rollout's ``plot_meta`` (same curated
+    model/training/rollout subset as always — see ``_figure_params_single``);
+    a multi-rollout run instead names the series being compared, since no
+    single ``model_config`` applies to the figure as a whole (each plot's own
+    gallery YAML still carries every rollout's full ``plot_meta`` for
+    after-the-fact lookup, via ``_plot_metadata``).
+    """
+    rollouts = run_meta.get("rollouts") or {}
+    if len(rollouts) == 1:
+        ((_, meta),) = rollouts.items()
+        return _figure_params_single(meta)
+    return {"rollouts": ", ".join(rollouts)} if rollouts else {}
+
+
 def _render_overlay(r: Reduced, params: dict):
     edges = np.asarray(r.payload["edges"])
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
@@ -139,7 +194,8 @@ 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"])
+    for i, (name, counts) in enumerate(r.payload.get("series", {}).items()):
+        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
     if r.payload.get("log_y"):
         ax.set_yscale("log")
     if r.payload.get("log_x"):
@@ -181,11 +237,17 @@ def _render_profile(r: Reduced, params: dict):
     edges = np.asarray(r.payload["edges"])
     centers = 0.5 * (edges[:-1] + edges[1:])
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    for key in ("reference", "rollout"):
-        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())
+    if "reference" in r.payload:
+        ref = r.payload["reference"]
+        mean, std = np.asarray(ref["mean"]), np.asarray(ref["std"])
+        color = _ref_color()
+        ax.plot(centers, mean, label=_REFERENCE_LABEL, color=color, linestyle="--")
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
+    for i, (name, side) in enumerate(r.payload.get("series", {}).items()):
+        mean, std = np.asarray(side["mean"]), np.asarray(side["std"])
+        color = ps.get_color(i)
+        ax.plot(centers, mean, label=name, color=color)
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
     ax.set_xlabel(r.xlabel)
     ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
     ps.style_legend(ax, title="source")
@@ -195,10 +257,19 @@ def _render_profile(r: Reduced, params: dict):
 def _render_bar(r: Reduced, params: dict):
     labels = r.payload["labels"]
     x = np.arange(len(labels))
-    width = 0.4
+    series = r.payload.get("series", {})
+    has_ref = "reference" in r.payload
+    n_bars = len(series) + (1 if has_ref else 0)
+    width = 0.8 / max(n_bars, 1)
+    offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, n_bars)
     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["rollout"], width, label=_SERIES_LABELS["rollout"])
+    idx = 0
+    if has_ref:
+        ax.bar(x + offsets[idx], r.payload["reference"], width, label=_REFERENCE_LABEL, color=_ref_color())
+        idx += 1
+    for i, (name, vals) in enumerate(series.items()):
+        ax.bar(x + offsets[idx], vals, width, label=name, color=ps.get_color(i))
+        idx += 1
     ax.set_xticks(x)
     ax.set_xticklabels(labels, rotation=45, ha="right")
     ax.set_ylabel(r.payload.get("ylabel", "value"))
@@ -207,56 +278,137 @@ def _render_bar(r: Reduced, params: dict):
 
 
 def _render_router_gating(r: Reduced, params: dict):
-    n_experts = r.payload["n_experts"]
+    series = r.payload.get("series", {})
+    names = list(series)
     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)
-    flat = axes.ravel()
-    for ax, key in zip(flat, ("rollout", "reference")):
-        side = r.payload.get(key, {})
-        centers = np.asarray(side.get("centers", []))
-        means = np.asarray(side.get("means", []))
-        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}")
-                cum = cum + means[:, i]
-        if log_x:
-            ax.set_xscale("log")
-        ax.set_ylim(0, 1)
-        ax.set_title(_SERIES_LABELS[key], fontsize=8)
-        ax.set_xlabel(r.xlabel)
-    flat[0].set_ylabel("mean gate weight")
-    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=2, squeeze=False)
+    for row, name in enumerate(names):
+        entry = series[name]
+        n_experts = entry["n_experts"]
+        for col, key in enumerate(("rollout", "reference")):
+            ax = axes[row, col]
+            side = entry.get(key, {})
+            centers = np.asarray(side.get("centers", []))
+            means = np.asarray(side.get("means", []))
+            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}")
+                    cum = cum + means[:, i]
+            if log_x:
+                ax.set_xscale("log")
+            ax.set_ylim(0, 1)
+            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
+            ax.set_title(f"{name} — {panel_label}", fontsize=8)
+            if row == len(names) - 1:
+                ax.set_xlabel(r.xlabel)
+        axes[row, 0].set_ylabel("mean gate weight")
+    if names:
+        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
     return fig
 
 
 def _render_router_share(r: Reduced, params: dict):
-    categories = r.payload["categories"]
-    n_experts = r.payload["n_experts"]
-    x = np.arange(len(categories))
-    present = [k for k in ("rollout", "reference") if k in r.payload]
+    series = r.payload.get("series", {})
+    names = list(series)
+    present: tuple[str, ...] = ("rollout", "reference")
+    if names:
+        present = tuple(k for k in ("rollout", "reference") if k in series[names[0]])
+    ncols = max(len(present), 1)
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=ncols, squeeze=False)
+    for row, name in enumerate(names):
+        entry = series[name]
+        n_experts = entry["n_experts"]
+        cats = entry["categories"]
+        x = np.arange(len(cats))
+        for col, key in enumerate(present):
+            ax = axes[row, col]
+            side = entry.get(key)
+            if side is not None:
+                shares = np.array([side[c] for c in cats])  # (n_cat, n_experts)
+                bottom = np.zeros(len(cats))
+                for i in range(n_experts):
+                    ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
+                    bottom += shares[:, i]
+            ax.set_xticks(x)
+            ax.set_xticklabels(cats, rotation=45, ha="right")
+            ax.set_ylim(0, 1)
+            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
+            ax.set_title(f"{name} — {panel_label}", fontsize=8)
+        axes[row, 0].set_ylabel("share of rows dispatched to expert")
+    if names:
+        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
+    return fig
+
+
+def _render_router_specialization(r: Reduced, params: dict):
+    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
+    series = r.payload.get("series", {})
+    chance_levels: set[float] = set()
+    for i, (name, entry) in enumerate(series.items()):
+        color = ps.get_color(i)
+        if entry.get("chance_level") is not None:
+            chance_levels.add(entry["chance_level"])
+        for key, linestyle, label in (
+            ("rollout", "-", name),
+            ("reference", "--", f"{name} ({_REFERENCE_LABEL})"),
+        ):
+            side = entry.get(key)
+            if side and side["centers"]:
+                ax.plot(
+                    side["centers"],
+                    side["score"],
+                    label=label,
+                    color=color,
+                    linestyle=linestyle,
+                    marker="o",
+                    markersize=3,
+                )
+    for lvl in sorted(chance_levels):
+        ax.axhline(lvl, linestyle=":", color="gray")
+    if r.payload.get("log_x"):
+        ax.set_xscale("log")
+    ax.set_ylim(0, 1)
+    ax.set_xlabel(r.xlabel)
+    ax.set_ylabel("max gate weight")
+    ps.style_legend(ax, title="router")
+    return fig
+
+
+def _render_heatmap(r: Reduced, params: dict):
+    series = r.payload["series"]
+    row_labels = r.payload["row_labels"]
+    col_labels = r.payload["col_labels"]
+    names = list(series)
     fig, axes = ps.new_figure(
-        "slide-16x9",
+        "slide-16x9" if len(names) > 1 else "thesis-single",
         title=r.title,
         params=params,
         nrows=1,
-        ncols=len(present),
+        ncols=len(names),
         squeeze=False,
     )
     flat = axes.ravel()
-    for ax, key in zip(flat, present):
-        side = r.payload[key]
-        shares = np.array([side[c] for c in categories])  # (n_cat, n_experts)
-        bottom = np.zeros(len(categories))
-        for i in range(n_experts):
-            ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
-            bottom += shares[:, i]
-        ax.set_xticks(x)
-        ax.set_xticklabels(categories, rotation=45, ha="right")
-        ax.set_ylim(0, 1)
-        ax.set_title(_SERIES_LABELS[key], fontsize=8)
-    flat[0].set_ylabel("share of rows dispatched to expert")
-    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    im = None
+    for ax, name in zip(flat, names):
+        mat = np.asarray(series[name], dtype=float)
+        im = ax.imshow(
+            mat,
+            origin="upper",
+            aspect="auto",
+            cmap=r.payload.get("cmap", "viridis"),
+            vmin=r.payload.get("vmin"),
+            vmax=r.payload.get("vmax"),
+        )
+        ax.set_xticks(range(len(col_labels)))
+        ax.set_xticklabels(col_labels, rotation=45, ha="right")
+        ax.set_yticks(range(len(row_labels)))
+        ax.set_yticklabels(row_labels)
+        ax.set_xlabel(r.xlabel)
+        if len(names) > 1:
+            ax.set_title(name, fontsize=8)
+    flat[0].set_ylabel(r.payload.get("ylabel", ""))
+    fig.colorbar(im, ax=list(flat), label=r.payload.get("cbar_label", "value"))
     return fig
 
 
@@ -284,13 +436,21 @@ _RENDERERS = {
     "bar": _render_bar,
     "router_gating": _render_router_gating,
     "router_share": _render_router_share,
+    "router_specialization": _render_router_specialization,
+    "heatmap": _render_heatmap,
     "unavailable": _render_unavailable,
 }
 
 
 def render(r: Reduced, run_meta: dict | None = None):
-    """Build the matplotlib figure for one reduced artifact (dispatch on kind)."""
-    return _RENDERERS[r.kind](r, _figure_params(run_meta or {}))
+    """Build the matplotlib figure for one reduced artifact (dispatch on kind).
+
+    ``title``/``xlabel`` are LaTeX-escaped here, at the one point every kind's
+    renderer draws them from — ``_plot_metadata`` deliberately keeps using the
+    unescaped ``r`` for the gallery YAML, which isn't LaTeX.
+    """
+    escaped = dataclasses.replace(r, title=_tex_escape(r.title), xlabel=_tex_escape(r.xlabel))
+    return _RENDERERS[r.kind](escaped, _figure_params(run_meta or {}))
 
 
 def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
@@ -349,7 +509,7 @@ def render_all(
         yaml.safe_dump(
             {
                 "title": run_meta.get("title", "GIANT rollout analysis"),
-                "description": "Autoregressive rollout compared against held-out Geant4 reference steps.",
+                "description": "Autoregressive rollout(s) compared against a held-out Geant4 reference steps file.",
                 "experiment": "GIANT",
                 "parameters": {k: v for k, v in run_meta.items() if k != "title"},
             },
@@ -382,8 +542,7 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
     meta = RunMeta.load(run_dir / "run_meta.json")
     run_meta = {
         "title": meta.title,
-        "rollout": meta.rollout,
         "reference": meta.reference,
-        **meta.plot_meta,
+        "rollouts": {ro["name"]: ro["plot_meta"] for ro in meta.rollouts},
     }
     return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)

giant/analysis/router_gating.py

@@ -35,6 +35,7 @@ from giant.analysis.reduced import Reduced
 if TYPE_CHECKING:
     import torch
 
+    from giant.analysis.sources import RolloutSide
     from giant.data.transforms import Normalizer
 
 _SAMPLE_ROWS = 200_000
@@ -203,6 +204,7 @@ _TITLES = {
     "router_gating": "Router gating (mixture-of-experts decision boundaries)",
     "router_share_by_pdg": "Router expert share by particle species",
     "router_share_by_process": "Router expert share by physics process",
+    "router_specialization": "Router specialization score vs energy (max gate weight)",
 }
 
 
@@ -217,51 +219,103 @@ def _unavailable(spec_id: str) -> Reduced:
     )
 
 
-def compute_router_gating(
-    checkpoint: str | Path | None,
-    r_phys: pl.LazyFrame,
-    t_phys: pl.LazyFrame,
-    seed: int = 0,
-) -> Reduced:
-    """`Reduced` for the router-gating figure, or an explanatory note if n/a."""
+def _gating_entry(checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, seed: int) -> dict | None:
+    """One rollout's ``router_gating`` panel data, or ``None`` if not a MoE checkpoint."""
     handle = load_router(checkpoint) if checkpoint else None
     if handle is None:
-        return _unavailable("router_gating")
-
+        return None
     sides: dict[str, dict] = {}
     for name, lf in (("rollout", r_phys), ("reference", t_phys)):
         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": []}
+    return {"router_type": handle.router_type, "n_experts": handle.router.n_experts, **sides}
 
+
+def compute_router_gating(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+    """`Reduced` for the router-gating figure: one panel-pair per rollout with
+    an enabled MoE router, or an explanatory note if none of them have one."""
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _gating_entry(rs.checkpoint, rs.phys, t_phys, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_gating")
     return Reduced(
         id="router_gating",
         family="model",
         kind="router_gating",
         title=_TITLES["router_gating"],
         xlabel="pre-step energy [MeV]",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "log_x": True,
-            **sides,
-        },
+        payload={"series": series, "log_x": True},
     )
 
 
-def compute_router_share_by_pdg(
-    checkpoint: str | Path | None,
-    r_phys: pl.LazyFrame,
-    t_phys: pl.LazyFrame,
-    top_pdgs: list[int],
-    seed: int = 0,
-) -> Reduced:
-    """Stacked-bar share of each particle species dispatched to each expert."""
+def _specialization_entry(
+    checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, seed: int
+) -> dict | None:
+    """One rollout's ``router_specialization`` curve data, or ``None`` if not a MoE checkpoint.
+
+    Scalar specialization trend: max gate weight vs energy, per side.
+    Summarizes `router_gating`'s full per-expert stacked area into one curve —
+    the routing plan's own "how sharp is the boundary here" number (1/n_experts
+    = uniform/no specialization, 1.0 = one expert fully owns that energy). Same
+    quantile energy bins as `router_gating` (`_quantile_bins`), so this is
+    directly comparable to that plot's ceiling described in the roadmap's MoE
+    writeup.
+    """
     handle = load_router(checkpoint) if checkpoint else None
     if handle is None:
-        return _unavailable("router_share_by_pdg")
+        return None
+    sides: dict[str, dict] = {}
+    for name, lf in (("rollout", r_phys), ("reference", t_phys)):
+        df = _subsample(lf, _SAMPLE_ROWS, seed)
+        df, gate = _gate_for_df(handle, df)
+        x = df["pre_E"].to_numpy()
+        if len(x):
+            binned = _quantile_bins(x, gate, _N_BINS)
+            means = np.asarray(binned["means"])
+            score = means.max(axis=1).tolist() if means.size else []
+            sides[name] = {"centers": binned["centers"], "score": score}
+        else:
+            sides[name] = {"centers": [], "score": []}
+    return {
+        "router_type": handle.router_type,
+        "n_experts": handle.router.n_experts,
+        "chance_level": 1.0 / handle.router.n_experts,
+        **sides,
+    }
 
+
+def compute_router_specialization(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+    """`Reduced` for the router-specialization figure, one curve per rollout with
+    an enabled MoE router (see `_specialization_entry`)."""
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _specialization_entry(rs.checkpoint, rs.phys, t_phys, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_specialization")
+    return Reduced(
+        id="router_specialization",
+        family="model",
+        kind="router_specialization",
+        title=_TITLES["router_specialization"],
+        xlabel="pre-step energy [MeV]",
+        payload={"series": series, "log_x": True},
+    )
+
+
+def _share_by_pdg_entry(
+    checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int
+) -> dict | None:
+    """One rollout's ``router_share_by_pdg`` panel-pair data, or ``None`` if not a MoE checkpoint."""
+    handle = load_router(checkpoint) if checkpoint else None
+    if handle is None:
+        return None
     labels = [pdg_label(p) for p in top_pdgs]
     sides: dict[str, dict] = {}
     for name, lf in (("rollout", r_phys), ("reference", t_phys)):
@@ -273,41 +327,36 @@ def compute_router_share_by_pdg(
         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)}
+    return {"router_type": handle.router_type, "n_experts": handle.router.n_experts, "categories": labels, **sides}
 
+
+def compute_router_share_by_pdg(
+    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int = 0
+) -> Reduced:
+    """`Reduced` for the router expert-share-by-species figure, one panel-pair
+    per rollout with an enabled MoE router."""
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _share_by_pdg_entry(rs.checkpoint, rs.phys, t_phys, top_pdgs, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_share_by_pdg")
     return Reduced(
         id="router_share_by_pdg",
         family="model",
         kind="router_share",
         title=_TITLES["router_share_by_pdg"],
         xlabel="particle species",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "categories": labels,
-            **sides,
-        },
+        payload={"series": series},
     )
 
 
-def compute_router_share_by_process(
-    checkpoint: str | Path | None,
-    t_phys: pl.LazyFrame,
-    seed: int = 0,
-    top_k: int = _TOP_K_PROCESS,
-) -> Reduced:
-    """Stacked-bar share of each physics process dispatched to each expert.
-
-    Reference-only: ``process`` is the true post-step physics process — a
-    label the rollout side has no equivalent of (see
-    `giant.model.network.ProcessRouter`, which predicts it from pre-step
-    conditioning alone, never observes it at eval time). This plot instead
-    checks *after the fact*, on real data, how well the router's conditioning
-    -based dispatch lines up with the true process.
-    """
+def _share_by_process_entry(checkpoint: str | Path | None, t_phys: pl.LazyFrame, seed: int, top_k: int) -> dict | None:
+    """One rollout checkpoint's ``router_share_by_process`` panel data (reference-only), or ``None`` if not MoE."""
     handle = load_router(checkpoint) if checkpoint else None
     if handle is None:
-        return _unavailable("router_share_by_process")
-
+        return None
     df = _subsample(t_phys, _SAMPLE_ROWS, seed, extra_cols=("process",))
     df, gate = _gate_for_df(handle, df)
     if len(df):
@@ -317,17 +366,39 @@ def compute_router_share_by_process(
         shares = _top1_shares(df["process"].to_numpy(), idx, order, handle.router.n_experts)
     else:
         order, shares = [], {}
+    return {
+        "router_type": handle.router_type,
+        "n_experts": handle.router.n_experts,
+        "categories": order,
+        "reference": {p: shares[p] for p in order},
+    }
 
+
+def compute_router_share_by_process(
+    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0, top_k: int = _TOP_K_PROCESS
+) -> Reduced:
+    """Stacked-bar share of each physics process dispatched to each expert, one
+    panel per rollout checkpoint with an enabled MoE router.
+
+    Reference-only: ``process`` is the true post-step physics process — a
+    label the rollout side has no equivalent of (see
+    `giant.model.network.ProcessRouter`, which predicts it from pre-step
+    conditioning alone, never observes it at eval time). This plot instead
+    checks *after the fact*, on real data, how well each checkpoint's router
+    -based dispatch lines up with the true process.
+    """
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _share_by_process_entry(rs.checkpoint, t_phys, seed, top_k)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_share_by_process")
     return Reduced(
         id="router_share_by_process",
         family="model",
         kind="router_share",
         title=_TITLES["router_share_by_process"],
         xlabel="physics process",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "categories": order,
-            "reference": {p: shares[p] for p in order},
-        },
+        payload={"series": series},
     )

giant/analysis/runtime_estimate.py

@@ -6,8 +6,8 @@ streaming `group_by` pass(es) over the chunk (see `catalog.py`/`reduce.py`).
 `_COST_MODEL` below is ``spec_id -> (intercept_s, seconds_per_row)``.
 ``n_rows`` is the combined rollout+reference row count of the job's input:
 the chunk's row count for `chunkable=True` specs, the whole dataset's for the
-three `chunkable=False` router specs (they always run as a single job
-regardless of chunk count).
+`chunkable=False` router specs in `_ROUTER_IDS` (they always run as a single
+job regardless of chunk count).
 
 Calibrated 2026-07-27 from real HTCondor timings (`condor_history`
 ``RemoteWallClockTime``) of a production run: prediction ``563f5ee3``
@@ -54,7 +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", "router_specialization"})
 
 # 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

@@ -1,7 +1,12 @@
-"""Canonical world-frame LazyFrame builders for the two sides of a comparison.
+"""Canonical world-frame LazyFrame builders for the two kinds of comparison input.
 
-The analysis compares one autoregressive ``giant rollout`` (the *generated* side)
-against a raw miniCaloSim steps file (the *reference* / real side). Both carry a
+The analysis compares one or more autoregressive ``giant rollout`` runs (the
+*generated* side — one named series each, see ``RolloutSpec``) against a single
+raw miniCaloSim steps file shared by all of them (the *reference* / real side).
+Every rollout is the same *kind* of file regardless of how many there are, so
+``Side`` stays binary: it describes a file's schema (rollout column layout +
+synthetic-termination rows + per-track secondary view, vs. reference
+``sec_*_list`` columns), not series identity. Both kinds carry a
 **shared world-frame physical column subset** under identical names, so no
 renaming or coordinate decode is needed — everything is already in world-frame
 mm / MeV:
@@ -26,6 +31,7 @@ HTCondor workers that have no LaTeX toolchain.
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 
@@ -40,6 +46,11 @@ from giant.constants import (
     TERM_MAX_STEPS,
     TERM_UNKNOWN_PDG,
 )
+from giant.data.loader import event_id_offset, find_parquet_files
+
+# Helper column name for the per-shard offset join in open_side; dropped before
+# the LazyFrame is returned, so it never leaks into a caller's schema.
+_SOURCE_PATH_COL = "__source_path"
 
 # The world-frame physical columns both sides share under identical names.
 PHYS_COLS: tuple[str, ...] = (
@@ -77,12 +88,44 @@ SYNTHETIC_TERMINATION_REASONS: frozenset[str] = frozenset(
 
 
 class Side(str, Enum):
-    """Which of the two comparison inputs a file is."""
+    """Which of the two comparison-input *kinds* a file is."""
 
     rollout = "rollout"
     reference = "reference"
 
 
+@dataclass
+class RolloutSpec:
+    """One named rollout input, as fed to ``build_context``/``Bundle.open``.
+
+    ``name`` is the series' identity throughout the rest of the pipeline (a
+    plot's ``payload["series"]`` key, a figure's legend label, its color) —
+    resolved once in ``condor.load_rollout_yamls`` from ``--label`` or the
+    YAML stem, then threaded through unchanged. ``checkpoint`` /
+    ``type_embedding_l1_dist`` are only used by the router/type-embedding
+    diagnostics (``catalog.py``'s ``chunkable=False`` specs).
+    """
+
+    name: str
+    source: str | Path | pl.LazyFrame
+    checkpoint: str | None = None
+    type_embedding_l1_dist: dict | None = None
+
+
+@dataclass
+class RolloutSide:
+    """One rollout's opened frames + per-checkpoint diagnostic inputs (``catalog.Bundle.rollouts`` value)."""
+
+    all: pl.LazyFrame  # rollout, all rows (incl. synthetic termination rows)
+    phys: pl.LazyFrame  # rollout, physical steps only
+    checkpoint: str | None = None  # from the rollout YAML; router_gating only
+    # Diagnostic pre-aggregated at rollout time (giant.rollout.
+    # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
+    # only. Unlike checkpoint, this needs no live model: it's already a
+    # finished histogram, just passed through.
+    type_embedding_l1_dist: dict | None = None
+
+
 def _check_rollout_metadata(path: Path) -> None:
     """Raise if ``path`` carries coord metadata that isn't the rollout tag.
 
@@ -108,6 +151,22 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
     reference file's upstream ROOT→parquet conversion don't agree on integer
     width, and an uncast mismatch only surfaces later as a ``pl.concat``
     ``SchemaError`` (e.g. in ``build_context``'s pdg-count merge).
+
+    The reference (a rollout's seed ``dataset``) may be a directory of parquet
+    shards, or a ``.manifest`` naming a subset, rather than a single file — each
+    such shard is a separate Geant4 job whose own ``event_id`` numbering
+    restarts from 0, so a multi-shard load offsets every shard's ids by
+    ``giant.data.loader.event_id_offset(file_index)`` to keep them globally
+    unique, exactly as the training/rollout data pipeline already does
+    (``giant/data/loader.py``). ``file_index`` comes from
+    ``find_parquet_files``'s deterministic ordering — the same list and
+    ordering ``giant rollout`` used (via ``_seed_from_data``) to offset the
+    rollout side's own ``event_id``s, so both sides agree on what an
+    ``event_id`` means. There is no overflow guard here (unlike
+    ``loader._offset_event_id``): checking it would cost an eager
+    ``event_id``-column read per shard in every condor compute job, and
+    ``giant rollout`` already ran that check over this exact file list when it
+    produced the seed.
     """
     if isinstance(source, pl.LazyFrame):
         return source.with_columns(pl.col("pdg").cast(pl.Int64))
@@ -116,9 +175,18 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
         _check_rollout_metadata(path)
         lf = pl.scan_parquet(path)
     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)
+        files = find_parquet_files(path)
+        if len(files) == 1:
+            lf = pl.scan_parquet(files[0])
+        else:
+            offsets = {str(p): event_id_offset(i) for i, p in enumerate(files)}
+            lf = (
+                pl.scan_parquet(files, include_file_paths=_SOURCE_PATH_COL)
+                .with_columns(
+                    pl.col("event_id") + pl.col(_SOURCE_PATH_COL).replace_strict(offsets, return_dtype=pl.Int64)
+                )
+                .drop(_SOURCE_PATH_COL)
+            )
     return lf.with_columns(pl.col("pdg").cast(pl.Int64))
 
 

giant/analysis/type_embedding_distance.py

@@ -20,26 +20,37 @@ redesign exists to fix.
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from giant.analysis.reduced import Reduced
 
+if TYPE_CHECKING:
+    from giant.analysis.sources import RolloutSide
+
 _NOTE_NOT_APPLICABLE = (
-    "not applicable: this rollout's checkpoint doesn't use "
+    "not applicable: none of these rollouts' checkpoints use "
     "stage2_model.particle_type.target='embedding' (or generated no "
     "secondaries), so giant rollout recorded no type_embedding_l1_dist "
-    "diagnostic in its YAML sidecar"
+    "diagnostic in their YAML sidecar"
 )
 
 
-def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
-    """`Reduced` for the type-embedding-distance figure, or an explanatory
-    note if this checkpoint never populated the diagnostic.
+def compute_type_embedding_l1_distance(rollouts: dict[str, "RolloutSide"]) -> Reduced:
+    """`Reduced` for the type-embedding-distance figure: one series per rollout
+    whose checkpoint populated the diagnostic, or an explanatory note if none did.
 
-    `l1_dist`: `giant.rollout.L1DistCollector.summary()`'s dict, as recorded
-    in the rollout YAML's `type_embedding_l1_dist` key (`Bundle.
-    type_embedding_l1_dist`) — `{"n", "mean", "std", "min", "max",
-    "hist_edges", "hist_counts"}`.
+    Each rollout's `RolloutSide.type_embedding_l1_dist` is
+    `giant.rollout.L1DistCollector.summary()`'s dict, as recorded in that
+    rollout's YAML `type_embedding_l1_dist` key — `{"n", "mean", "std",
+    "min", "max", "hist_edges", "hist_counts"}`. Every collector uses the
+    same fixed log-spaced edges (`L1DistCollector.__init__`'s defaults, never
+    overridden — see `giant/cli.py`'s rollout command), so it's safe to plot
+    every rollout's counts against the first one's edges.
     """
-    if l1_dist is None:
+    entries = {
+        name: rs.type_embedding_l1_dist for name, rs in rollouts.items() if rs.type_embedding_l1_dist is not None
+    }
+    if not entries:
         return Reduced(
             id="type_embedding_l1_distance",
             family="model",
@@ -49,6 +60,11 @@ def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
             payload={"note": _NOTE_NOT_APPLICABLE},
         )
 
+    edges = next(iter(entries.values()))["hist_edges"]
+    notes = [
+        f"{name}: n={d['n']:,} mean={d['mean']:.4g} std={d['std']:.4g} min={d['min']:.4g} max={d['max']:.4g}"
+        for name, d in entries.items()
+    ]
     return Reduced(
         id="type_embedding_l1_distance",
         family="model",
@@ -56,15 +72,10 @@ def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
         title="Secondary-type embedding L1 distance (predicted vector -> nearest PDG row)",
         xlabel="L1 distance",
         payload={
-            "edges": l1_dist["hist_edges"],
-            "rollout": l1_dist["hist_counts"],
+            "edges": edges,
+            "series": {name: d["hist_counts"] for name, d in entries.items()},
             "log_y": True,
             "log_x": True,
-            "note": (
-                f"n={l1_dist['n']:,}  mean={l1_dist['mean']:.4g}  "
-                f"std={l1_dist['std']:.4g}  min={l1_dist['min']:.4g}  "
-                f"max={l1_dist['max']:.4g}; rollout only, no reference "
-                "concept for a raw pre-decode vector"
-            ),
+            "note": "; ".join(notes) + "; rollout only, no reference concept for a raw pre-decode vector",
         },
     )

giant/cli.py

@@ -41,6 +41,7 @@ from giant.data.transforms import (
 )
 from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
 from giant.geometry import GeometryOracle
+from giant.materials import MATERIAL_PROPERTIES
 from giant.pipeline import run_train_job
 from giant.rollout import (
     L1DistCollector,
@@ -487,6 +488,36 @@ def train(
             help="WGAN-GP (--mode wgan only): critic depth for stage 2 (default: same as generator's n_res_blocks)",
         ),
     ] = None,
+    stage1_init_from: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--stage1-init-from",
+            help="Checkpoint .pt to load stage 1's weights from before training starts "
+            "(gitea #42) — combine with --stage1-freeze to retrain stage 2 alone "
+            "against a fixed, known-good stage 1",
+        ),
+    ] = None,
+    stage1_freeze: Annotated[
+        Optional[bool],
+        typer.Option(
+            "--stage1-freeze/--no-stage1-freeze",
+            help="Never update stage 1's weights (requires --stage1-init-from, or --resume)",
+        ),
+    ] = None,
+    stage2_init_from: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--stage2-init-from",
+            help="Checkpoint .pt to load stage 2's weights from before training starts (gitea #42)",
+        ),
+    ] = None,
+    stage2_freeze: Annotated[
+        Optional[bool],
+        typer.Option(
+            "--stage2-freeze/--no-stage2-freeze",
+            help="Never update stage 2's weights (requires --stage2-init-from, or --resume)",
+        ),
+    ] = None,
     val_fraction: Annotated[Optional[float], typer.Option("--val-fraction", "-f")] = None,
     seed: Annotated[
         Optional[int],
@@ -581,6 +612,14 @@ def train(
             "steps (default: 50); per-epoch metrics always log in full",
         ),
     ] = None,
+    precision: Annotated[
+        Optional[str],
+        typer.Option(
+            "--precision",
+            help="Training-step autocast precision: 'fp32' (default) or "
+            "'bf16'. No 'fp16' — see giant.training.amp.resolve_autocast",
+        ),
+    ] = None,
 ) -> None:
     """Train the GIANT surrogate model."""
     batch_size_auto = False
@@ -616,6 +655,7 @@ def train(
         "wandb_project": wandb_project,
         "wandb_run_name": wandb_run_name,
         "wandb_log_every": wandb_log_every,
+        "precision": precision,
         "hidden_dim": hidden_dim,
         "n_blocks": n_blocks,
         "dropout": dropout,
@@ -651,11 +691,15 @@ def train(
         "stage1_critic_n_res_blocks": stage1_critic_n_res_blocks,
         "stage2_critic_hidden_dim": stage2_critic_hidden_dim,
         "stage2_critic_n_res_blocks": stage2_critic_n_res_blocks,
+        "stage1_init_from": str(stage1_init_from) if stage1_init_from is not None else None,
+        "stage1_freeze": stage1_freeze,
+        "stage2_init_from": str(stage2_init_from) if stage2_init_from is not None else None,
+        "stage2_freeze": stage2_freeze,
     }
     overrides = gconfig.overrides_from_flags(flag_values)
 
     cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides)
-    gconfig.validate_config(cfg)
+    gconfig.validate_config(cfg, resume=resume is not None)
     t = cfg["train"]
 
     _device = torch.device(device) if device else gconfig.auto_device()
@@ -690,6 +734,7 @@ def train(
 
     typer.echo(f"device:  {_device}")
     typer.echo(f"out_dir: {out_dir}")
+    typer.echo(f"precision: {t['precision']}")
 
     run_train_job(
         data=data,
@@ -735,6 +780,10 @@ def new_run(
     stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None,
     stage2_context_dim: Annotated[Optional[int], typer.Option("--stage2-context-dim")] = None,
     stage2_stage1_context: Annotated[Optional[Stage1Context], typer.Option("--stage2-stage1-context")] = None,
+    stage1_init_from: Annotated[Optional[Path], typer.Option("--stage1-init-from")] = None,
+    stage1_freeze: Annotated[Optional[bool], typer.Option("--stage1-freeze/--no-stage1-freeze")] = None,
+    stage2_init_from: Annotated[Optional[Path], typer.Option("--stage2-init-from")] = None,
+    stage2_freeze: Annotated[Optional[bool], typer.Option("--stage2-freeze/--no-stage2-freeze")] = None,
     conditioning: Annotated[Optional[Conditioning], typer.Option("--conditioning")] = None,
     router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
     router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
@@ -795,6 +844,10 @@ def new_run(
         "stage2_k_max": stage2_k_max,
         "stage2_context_dim": stage2_context_dim,
         "stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None,
+        "stage1_init_from": str(stage1_init_from) if stage1_init_from is not None else None,
+        "stage1_freeze": stage1_freeze,
+        "stage2_init_from": str(stage2_init_from) if stage2_init_from is not None else None,
+        "stage2_freeze": stage2_freeze,
         "mode": mode.value if mode is not None else None,
         "stage1_generator": stage1_generator.value if stage1_generator is not None else None,
         "stage2_generator": stage2_generator.value if stage2_generator is not None else None,
@@ -854,6 +907,52 @@ def new_run(
     typer.echo(f"  giant train {data_arg} --config {config_path} --out {run_dir}")
 
 
+model_app = typer.Typer(
+    no_args_is_help=True,
+    help="Inspect a resolved model architecture without training.",
+)
+app.add_typer(model_app, name="model")
+
+
+@model_app.command("summary")
+def model_summary(
+    config: Annotated[
+        Optional[Path],
+        typer.Option("--config", "-c", help="TOML config file (default: built-in defaults)"),
+    ] = None,
+    pdg_vocab: Annotated[
+        int,
+        typer.Option(
+            "--pdg-vocab",
+            help="Placeholder PDG vocab size for conditioning.particle.type='embedding' "
+            "or a pdg/process router (no dataset attached to derive the real training vocab)",
+        ),
+    ] = 300,
+    mat_vocab: Annotated[
+        int,
+        typer.Option(
+            "--mat-vocab",
+            help="Placeholder material vocab size for conditioning.material.type='embedding' "
+            "or a process router (default: the number of known materials in giant.materials)",
+        ),
+    ] = len(MATERIAL_PROPERTIES),
+) -> None:
+    """Build the resolved model graph from a config with no dataset attached, and
+    print per-module parameter counts, trunk widths, which heads exist, and which
+    conditioning/stage1_model/stage2_model config keys actually shaped it."""
+    from giant.model.summary import render_summary, summarize_model
+
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, {})
+    try:
+        gconfig.validate_config(cfg)
+    except ValueError as exc:
+        typer.echo(f"error: {exc}", err=True)
+        raise typer.Exit(1)
+
+    summary = summarize_model(cfg, pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    typer.echo(render_summary(summary))
+
+
 @app.command()
 def predict(
     data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
@@ -1457,10 +1556,23 @@ app.add_typer(analyze_app, name="analyze")
 
 @analyze_app.command("prep")
 def analyze_prep(
-    rollout_yaml: Annotated[
-        Path,
-        typer.Argument(help="giant rollout YAML sidecar (names the rollout + reference files)"),
+    rollout_yamls: Annotated[
+        list[Path],
+        typer.Argument(
+            help="giant rollout YAML sidecar(s) (names the rollout + reference files). "
+            "Multiple compare N rollouts against one shared reference — every YAML must "
+            "name the same `dataset`."
+        ),
     ],
+    label: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--label",
+            help="Series name for a rollout YAML, positionally matched to it — give none, "
+            'or exactly one per YAML. Defaults to the YAML stem (or "rollout" for a '
+            "single YAML).",
+        ),
+    ] = None,
     run_dir: Annotated[
         Optional[Path],
         typer.Option(
@@ -1477,14 +1589,15 @@ def analyze_prep(
         typer.Option("--chunks", help="Split each plot's data into this many event_id chunks"),
     ] = 1,
 ) -> None:
-    """Read the rollout YAML → shared.json + run_meta.json in the run directory."""
+    """Read the rollout YAML(s) → shared.json + run_meta.json in the run directory."""
     from giant.analysis import prep
 
     path = prep(
-        rollout_yaml,
+        rollout_yamls,
         run_dir,
         n_chunks=chunks,
         default_base=Path.cwd() / "analysis_runs",
+        labels=label,
         n_energy_bins=n_energy_bins,
         n_marginal_bins=n_marginal_bins,
         top_k_pdg=top_k_pdg,
@@ -1545,10 +1658,44 @@ def analyze_render(
     typer.echo(f"rendered {len(pdfs)} plots → {Path(run_dir) / 'plots'}")
 
 
+@analyze_app.command("metrics")
+def analyze_metrics(
+    run_dir: Annotated[Path, typer.Argument(help="Run directory containing metrics.csv (from `giant train`)")],
+    out_dir: Annotated[
+        Optional[Path],
+        typer.Option(
+            "--out",
+            "-o",
+            help="Override the output directory (default: <cwd>/analysis_runs/metrics_<run_dir name>)",
+        ),
+    ] = None,
+) -> None:
+    """Render training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) from <run_dir>/metrics.csv."""
+    from giant.training.plots import render_metrics
+
+    paths = render_metrics(run_dir, out_dir, default_base=Path.cwd() / "analysis_runs")
+    typer.echo(f"rendered {len(paths)} plots -> {paths[0].parent if paths else '(nothing to render)'}")
+
+
 @analyze_app.command("submit")
 def analyze_submit(
-    rollout_yaml: Annotated[Path, typer.Argument(help="giant rollout YAML sidecar")],
+    rollout_yamls: Annotated[
+        list[Path],
+        typer.Argument(
+            help="giant rollout YAML sidecar(s). Multiple compare N rollouts against one "
+            "shared reference — every YAML must name the same `dataset`."
+        ),
+    ],
     accounting_group: Annotated[str, typer.Option("--accounting-group")],
+    label: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--label",
+            help="Series name for a rollout YAML, positionally matched to it — give none, "
+            'or exactly one per YAML. Defaults to the YAML stem (or "rollout" for a '
+            "single YAML).",
+        ),
+    ] = None,
     run_dir: Annotated[
         Optional[Path],
         typer.Option(
@@ -1581,10 +1728,11 @@ def analyze_submit(
     from giant.analysis import SubmitConfig, prep, write_submit
 
     path = prep(
-        rollout_yaml,
+        rollout_yamls,
         run_dir,
         n_chunks=chunks,
         default_base=Path.cwd() / "analysis_runs",
+        labels=label,
         n_energy_bins=n_energy_bins,
         n_marginal_bins=n_marginal_bins,
         top_k_pdg=top_k_pdg,

giant/config.py

@@ -468,6 +468,18 @@ class ParticleTypeConfig:
     # silently the same number). 0 = inherit conditioning.particle.emb_dim,
     # preserving pre-#29 behavior.
     n_classes: int = 0
+    # Class-balances the target = "onehot" cross-entropy loss against the
+    # secondary-species long tail (gitea #44: the failure mode motivating the
+    # v0.3.0 pivot was specifically a species collapse — zero photon
+    # secondaries, hallucinated antineutrinos). "none": plain CE (pre-#44
+    # behavior). "inverse_freq": CE weighted by 1/count per class,
+    # normalized to mean 1 across classes so lambda_weight doesn't need
+    # retuning when this is switched on. validate_config requires target =
+    # "onehot" and stage2_model.generator != "wgan" whenever this isn't
+    # "none" — "embedding"/"physical" have no class CE to weight, and the
+    # WGAN stage-2 path feeds its type slice to the critic via a
+    # straight-through Gumbel relaxation instead of a CE loss.
+    class_weighting: str = "none"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
@@ -477,6 +489,7 @@ class ParticleTypeConfig:
             lambda_weight=d.get("lambda", 1.0),
             other_policy=d.get("other_policy", "sample"),
             n_classes=d.get("n_classes", 0),
+            class_weighting=d.get("class_weighting", "none"),
         )
 
     def to_dict(self) -> dict:
@@ -485,6 +498,7 @@ class ParticleTypeConfig:
             "lambda": self.lambda_weight,
             "other_policy": self.other_policy,
             "n_classes": self.n_classes,
+            "class_weighting": self.class_weighting,
         }
 
 
@@ -594,6 +608,19 @@ class Stage1ModelConfig:
     # false skips building/training stage 1 entirely. The resulting
     # checkpoint holds only stage 2 and cannot be rolled out.
     active: bool = True
+    # Checkpoint .pt to load this stage's weights from before training starts
+    # (its own "model"/"sec_decoder" key, not this run's own resume state) —
+    # "" means start from a fresh init. See `freeze` below for the partial-
+    # retrain use case this exists for (gitea #42).
+    init_from: str = ""
+    # true keeps this stage's weights exactly as loaded from `init_from` —
+    # forward/backward still run every batch (so its loss/grad_norm metrics
+    # stay meaningful, and a WGAN stage's critic still gets a real signal to
+    # report), but its optimizer never steps. Lets a rollout-capable
+    # checkpoint retrain only the *other* stage against a fixed, known-good
+    # one (gitea #42) — `validate_config` requires `init_from` to be set
+    # whenever this is true, unless the run is a `--resume`.
+    freeze: bool = False
     # "flow": conditional flow matching (~10 ODE steps at inference).
     # "ddpm": cosine-schedule diffusion baseline.
     # "wgan": WGAN-GP, single forward pass at inference.
@@ -620,6 +647,8 @@ class Stage1ModelConfig:
         d = d or {}
         return cls(
             active=d.get("active", True),
+            init_from=d.get("init_from", ""),
+            freeze=d.get("freeze", False),
             generator=d.get("generator", "flow"),
             hidden_dim=d.get("hidden_dim", 256),
             n_res_blocks=d.get("n_res_blocks", 6),
@@ -636,6 +665,8 @@ class Stage1ModelConfig:
     def to_dict(self) -> dict:
         return {
             "active": self.active,
+            "init_from": self.init_from,
+            "freeze": self.freeze,
             "generator": self.generator,
             "hidden_dim": self.hidden_dim,
             "n_res_blocks": self.n_res_blocks,
@@ -655,6 +686,9 @@ class Stage2ModelConfig:
     # false trains stage 1 alone. giant rollout must then refuse the
     # checkpoint; giant predict still works.
     active: bool = True
+    # See Stage1ModelConfig.init_from/.freeze — same semantics, this stage.
+    init_from: str = ""
+    freeze: bool = False
     # "one_shot": predict all k_max slots simultaneously with padded slots
     # masked from the loss (v0.2 behaviour).
     # "autoregressive": emit one secondary at a time in descending-energy
@@ -677,6 +711,13 @@ class Stage2ModelConfig:
     # output, closing the train/inference gap at the cost of a sampling pass
     # per batch and a moving target early in training.
     stage1_context: str = "truth"
+    # Ramp for "sampled": P(condition on the ground-truth stage-1 outcome
+    # rather than a fresh sample), linearly interpolated from ctx_p_start
+    # (epoch 0) to ctx_p_end (the final epoch) — the same scheduled-sampling
+    # shape as autoregressive.tf_p_start/tf_p_end, so stage 2 doesn't chase a
+    # wildly moving stage-1 target in early epochs. Unread under "truth".
+    ctx_p_start: float = 1.0
+    ctx_p_end: float = 0.0
     n_sec: NSecConfig = field(default_factory=NSecConfig)
     particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
     autoregressive: AutoregressiveConfig = field(default_factory=AutoregressiveConfig)
@@ -692,6 +733,8 @@ class Stage2ModelConfig:
         d = d or {}
         return cls(
             active=d.get("active", True),
+            init_from=d.get("init_from", ""),
+            freeze=d.get("freeze", False),
             decoder=d.get("decoder", "autoregressive"),
             generator=d.get("generator", "wgan"),
             hidden_dim=d.get("hidden_dim", 256),
@@ -701,6 +744,8 @@ class Stage2ModelConfig:
             k_max=d.get("k_max", 15),
             context_dim=d.get("context_dim", 64),
             stage1_context=d.get("stage1_context", "truth"),
+            ctx_p_start=d.get("ctx_p_start", 1.0),
+            ctx_p_end=d.get("ctx_p_end", 0.0),
             n_sec=NSecConfig.from_dict(d.get("n_sec")),
             particle_type=ParticleTypeConfig.from_dict(d.get("particle_type")),
             autoregressive=AutoregressiveConfig.from_dict(d.get("autoregressive")),
@@ -715,6 +760,8 @@ class Stage2ModelConfig:
     def to_dict(self) -> dict:
         return {
             "active": self.active,
+            "init_from": self.init_from,
+            "freeze": self.freeze,
             "decoder": self.decoder,
             "generator": self.generator,
             "hidden_dim": self.hidden_dim,
@@ -724,6 +771,8 @@ class Stage2ModelConfig:
             "k_max": self.k_max,
             "context_dim": self.context_dim,
             "stage1_context": self.stage1_context,
+            "ctx_p_start": self.ctx_p_start,
+            "ctx_p_end": self.ctx_p_end,
             "n_sec": self.n_sec.to_dict(),
             "particle_type": self.particle_type.to_dict(),
             "autoregressive": self.autoregressive.to_dict(),
@@ -766,6 +815,12 @@ class TrainConfig:
     # thousands of steps. Per-epoch metrics (the metrics.csv row) always log
     # in full.
     wandb_log_every: int = 50
+    # Training-step autocast dtype: "fp32" (default, no autocast) or "bf16".
+    # No "fp16" — GradScaler and the double-backward in
+    # giant.model.wgan.gradient_penalty don't mix well, and bf16 alone covers
+    # every training GPU in the fleet (Ampere and newer). See
+    # giant.training.amp.resolve_autocast (gitea #47).
+    precision: str = "fp32"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "TrainConfig":
@@ -787,6 +842,7 @@ class TrainConfig:
             wandb_project=d.get("wandb_project", "giant"),
             wandb_run_name=d.get("wandb_run_name", ""),
             wandb_log_every=d.get("wandb_log_every", 50),
+            precision=d.get("precision", "fp32"),
         )
 
     def to_dict(self) -> dict:
@@ -807,6 +863,7 @@ class TrainConfig:
             "wandb_project": self.wandb_project,
             "wandb_run_name": self.wandb_run_name,
             "wandb_log_every": self.wandb_log_every,
+            "precision": self.precision,
         }
 
 
@@ -843,6 +900,25 @@ class GiantConfig:
 DEFAULT_CONFIG: dict = GiantConfig().to_dict()
 
 
+def leaf_paths(node: dict, prefix: str = "") -> list[str]:
+    """Every dotted leaf path in a DEFAULT_CONFIG-shaped dict, e.g.
+    "stage1_model.router.n_experts". `[meta]` (run provenance, no schema
+    counterpart) is skipped at the top level, matching `validate_config_keys`.
+    Shared by `tests/test_config_consumed_keys.py` (the static per-identifier
+    audit) and `giant.model.summary` (the runtime per-config audit, gitea
+    #46) so both walk the exact same tree."""
+    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 git_hash() -> str:
     try:
         return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
@@ -1060,6 +1136,7 @@ FLAG_SPECS: tuple[FlagSpec, ...] = (
     FlagSpec("wandb_project", ("train.wandb_project",)),
     FlagSpec("wandb_run_name", ("train.wandb_run_name",)),
     FlagSpec("wandb_log_every", ("train.wandb_log_every",)),
+    FlagSpec("precision", ("train.precision",)),
     # --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat
     # shorthands (they predate stage2_model having its own flags);
     # --stage1-* wins when both are given.
@@ -1112,6 +1189,13 @@ FLAG_SPECS: tuple[FlagSpec, ...] = (
     FlagSpec("stage1_critic_n_res_blocks", ("stage1_model.wgan.critic_n_res_blocks",)),
     FlagSpec("stage2_critic_hidden_dim", ("stage2_model.wgan.critic_hidden_dim",)),
     FlagSpec("stage2_critic_n_res_blocks", ("stage2_model.wgan.critic_n_res_blocks",)),
+    # Partial-retrain (gitea #42): stage-scoped only, no shared alias — a
+    # shared "freeze both stages from the same file" flag has no sensible
+    # meaning (a checkpoint has one set of weights per stage).
+    FlagSpec("stage1_init_from", ("stage1_model.init_from",)),
+    FlagSpec("stage1_freeze", ("stage1_model.freeze",)),
+    FlagSpec("stage2_init_from", ("stage2_model.init_from",)),
+    FlagSpec("stage2_freeze", ("stage2_model.freeze",)),
 )
 
 
@@ -1344,7 +1428,7 @@ def merge_cli_overrides(
     return cfg
 
 
-def validate_config(cfg: dict) -> None:
+def validate_config(cfg: dict, *, resume: bool = False) -> None:
     """Cross-block validation the per-block schema can't express on its own.
 
     Raises ValueError with a clear message on the first violation found. Call
@@ -1352,6 +1436,10 @@ def validate_config(cfg: dict) -> None:
     these checks need to see across blocks, so they don't belong in
     `migrate_config` (which only ever sees one dict's own keys) or in any
     single block's defaults.
+
+    `resume=True` (only `giant train --resume` passes this) relaxes the
+    `stage{1,2}_model.freeze` -> `.init_from` requirement below: a resumed
+    frozen stage's weights come from the resume checkpoint, not `init_from`.
     """
     particle_type = _get_path(cfg, "conditioning.particle.type")
 
@@ -1364,7 +1452,34 @@ def validate_config(cfg: dict) -> None:
             f"{particle_type!r})"
         )
 
+    class_weighting = _get_path(cfg, "stage2_model.particle_type.class_weighting")
+    if class_weighting not in ("none", "inverse_freq"):
+        raise ValueError(
+            f"stage2_model.particle_type.class_weighting = {class_weighting!r} — must be 'none' or 'inverse_freq'"
+        )
+    if class_weighting != "none" and pt_target != "onehot":
+        raise ValueError(
+            "stage2_model.particle_type.class_weighting != 'none' requires "
+            f"stage2_model.particle_type.target = 'onehot' (there is no class "
+            f"cross-entropy to weight under target = {pt_target!r})"
+        )
+    if class_weighting != "none" and _get_path(cfg, "stage2_model.generator") == "wgan":
+        raise ValueError(
+            "stage2_model.particle_type.class_weighting != 'none' is "
+            "incompatible with stage2_model.generator = 'wgan' — that path "
+            "feeds the type slice to the critic via a straight-through "
+            "Gumbel relaxation instead of a class cross-entropy, so there is "
+            "nothing to weight"
+        )
+
     for stage_name in ("stage1_model", "stage2_model"):
+        if _get_path(cfg, f"{stage_name}.freeze") and not _get_path(cfg, f"{stage_name}.init_from") and not resume:
+            raise ValueError(
+                f"{stage_name}.freeze = true requires {stage_name}.init_from "
+                "to be set (or --resume) — freezing a randomly-initialized "
+                "model is almost certainly a mistake"
+            )
+
         router = _get_path(cfg, f"{stage_name}.router") or {}
         if router.get("enabled") and router.get("type") in ("pdg", "process") and particle_type == "physical":
             raise ValueError(
@@ -1401,14 +1516,36 @@ def validate_config(cfg: dict) -> None:
     if stop_sampling not in ("greedy", "sample"):
         raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
 
-    if _get_path(cfg, "stage2_model.stage1_context") == "sampled":
+    precision = _get_path(cfg, "train.precision")
+    if precision not in ("fp32", "bf16"):
         raise ValueError(
-            "stage2_model.stage1_context = 'sampled' is accepted by the schema "
-            "but not implemented — trainers.py always trains stage 2 against "
-            "the ground-truth stage-1 output; use 'truth' (default) instead "
-            "(see issues.md Issue 16 for the planned implementation)"
+            f"train.precision = {precision!r} — must be 'fp32' or 'bf16' "
+            "('fp16' is not supported: see giant.training.amp.resolve_autocast)"
         )
 
+    stage1_context = _get_path(cfg, "stage2_model.stage1_context")
+    if stage1_context not in ("truth", "sampled"):
+        raise ValueError(f"stage2_model.stage1_context = {stage1_context!r} — must be 'truth' or 'sampled'")
+    if stage1_context == "sampled":
+        if not (_get_path(cfg, "stage1_model.active") and _get_path(cfg, "stage2_model.active")):
+            raise ValueError(
+                "stage2_model.stage1_context = 'sampled' requires both "
+                "stage1_model.active and stage2_model.active = true — there is "
+                "no stage-1 model to sample from in a stage-2-only run"
+            )
+        ctx_p_start = _get_path(cfg, "stage2_model.ctx_p_start")
+        ctx_p_end = _get_path(cfg, "stage2_model.ctx_p_end")
+        for name, value in (("ctx_p_start", ctx_p_start), ("ctx_p_end", ctx_p_end)):
+            if not (0.0 <= value <= 1.0):
+                raise ValueError(f"stage2_model.{name} = {value} — must be in [0, 1]")
+        if ctx_p_start == 1.0 and ctx_p_end == 1.0:
+            raise ValueError(
+                "stage2_model.stage1_context = 'sampled' with ctx_p_start = "
+                "ctx_p_end = 1.0 always conditions on the ground truth — "
+                "identical to 'truth' but silently so; use 'truth' instead or "
+                "lower ctx_p_end"
+            )
+
     if (
         _get_path(cfg, "stage2_model.n_sec.mode") == "truth"
         and _get_path(cfg, "stage1_model.active")
@@ -1569,6 +1706,8 @@ _OUT_DIR_NAME_CANDIDATES = [
     ),
     ("particle_conditioning", _conditioning_candidate("particle", "c")),
     ("material_conditioning", _conditioning_candidate("material", "m")),
+    ("stage1_freeze", _path_candidate("stage1_model.freeze", "s1frozen", formatter=lambda _: "")),
+    ("stage2_freeze", _path_candidate("stage2_model.freeze", "s2frozen", formatter=lambda _: "")),
     ("stage1_hidden_dim", _path_candidate("stage1_model.hidden_dim", "h")),
     ("stage2_hidden_dim", _path_candidate("stage2_model.hidden_dim", "s2h")),
     ("stage1_n_res_blocks", _path_candidate("stage1_model.n_res_blocks", "b")),

giant/data/loader.py

@@ -1,4 +1,4 @@
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from pathlib import Path
 from typing import Iterator
 
@@ -256,24 +256,32 @@ def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
     return counts
 
 
-def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]:
+def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]:
     """Frequency-capped value->index map: the `n_classes - 1` most frequent
     keys get their own index; every rarer key is bucketed into a shared
     "other" index (`n_classes - 1`).
 
-    Returns `(class_map, other_members)` — `other_members` is `{key: count}`
-    for every key bucketed into "other" (the empirical within-bucket
-    distribution, for `other_policy = "sample"` at rollout).
+    Returns `(class_map, other_members, class_counts)` — `other_members` is
+    `{key: count}` for every key bucketed into "other" (the empirical
+    within-bucket distribution, for `other_policy = "sample"` at rollout);
+    `class_counts` is `{index: total_count}` for every resulting class index
+    (0-indexed; the "other" index's count is the sum of `other_members`),
+    the per-class frequencies `stage2_model.particle_type.class_weighting`
+    (gitea #44) needs and that would otherwise be dropped once `counts` is
+    collapsed into `class_map`.
     """
     ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
     keep = ranked[: max(n_classes - 1, 0)]
     class_map = {k: i for i, k in enumerate(keep)}
+    class_counts = {i: counts[k] for i, k in enumerate(keep)}
     other_idx = n_classes - 1
     other_members: dict = {}
     for k in ranked[len(keep) :]:
         class_map[k] = other_idx
         other_members[k] = counts[k]
-    return class_map, other_members
+    if other_members:
+        class_counts[other_idx] = sum(other_members.values())
+    return class_map, other_members, class_counts
 
 
 def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
@@ -287,7 +295,7 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
     fixed-width n_sec_head classifier.
     """
     counts = _rank_by_frequency_from_files(files, "process", str)
-    class_map, _ = _topn_plus_other_map(counts, n_experts)
+    class_map, _, _ = _topn_plus_other_map(counts, n_experts)
     return class_map
 
 
@@ -299,6 +307,11 @@ class TopNMap:
 
     class_map: dict
     other_members: dict
+    # {class_index: total_count} — see _topn_plus_other_map. Empty for a
+    # TopNMap decoded from a checkpoint/sidecar predating gitea #44; only
+    # stage2_model.particle_type.class_weighting reads it, and it raises
+    # loudly if it needs counts that aren't there (giant/training/trainers.py).
+    class_counts: dict = field(default_factory=dict)
 
 
 def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
@@ -315,8 +328,8 @@ def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, ca
     free during this same scan.
     """
     counts = _rank_by_frequency_from_files(files, column, cast)
-    class_map, other_members = _topn_plus_other_map(counts, n_classes)
-    return TopNMap(class_map=class_map, other_members=other_members)
+    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
 
 
 def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
@@ -347,5 +360,5 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
         if has_sec:
             exploded = df["sec_pdg_list"].explode().dropna()
             _accumulate_value_counts(counts, exploded, int)
-    class_map, other_members = _topn_plus_other_map(counts, n_classes)
-    return TopNMap(class_map=class_map, other_members=other_members)
+    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)

giant/data/setup_cache.py

@@ -35,7 +35,10 @@ from giant.data.transforms import Normalizer, sorted_membership
 # v3: NormalizerEntry.energy_reservoir_sample (100k raw values) replaced by
 # energy_quantiles (a fixed ENERGY_QUANTILE_LEVELS-point quantile grid) — a
 # v2 sidecar has no such grid to fall back on, so it must be recomputed.
-_CACHE_FORMAT_VERSION = 3
+# v4: TopNMap gained class_counts (gitea #44, stage2_model.particle_type.
+# class_weighting) — a v3 sidecar's cached topn_maps have no counts, so they
+# must be rebuilt rather than silently cached with class_counts={}.
+_CACHE_FORMAT_VERSION = 4
 
 _DIMS = {
     "COND_DIM": COND_DIM,
@@ -131,6 +134,7 @@ def topnmap_to_json(m: TopNMap) -> dict:
     return {
         "class_map": {str(k): v for k, v in m.class_map.items()},
         "other_members": {str(k): v for k, v in m.other_members.items()},
+        "class_counts": {str(k): v for k, v in m.class_counts.items()},
     }
 
 
@@ -139,6 +143,11 @@ def topnmap_from_json(d: dict, axis: str) -> TopNMap:
     return TopNMap(
         class_map={cast(k): v for k, v in d["class_map"].items()},
         other_members={cast(k): v for k, v in d["other_members"].items()},
+        # Missing for a checkpoint's topn maps predating gitea #44 — {} is
+        # the correct decode there (inference never reads class_counts; only
+        # stage2_model.particle_type.class_weighting does, at train time, and
+        # it raises loudly if it needs counts a checkpoint doesn't have).
+        class_counts={int(k): v for k, v in d.get("class_counts", {}).items()},
     )
 
 

giant/model/builders.py

@@ -202,6 +202,8 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
             cond_out_dim=cond_out_dim,
             dropout=s1_spec.dropout,
             stage="stage1",
+            trunk_type=s1_spec.trunk.type,
+            block_conditioning=s1_spec.trunk.block_conditioning,
         )
 
     if s2_spec.active and build_objective(s2_spec.generator).is_adversarial:
@@ -225,6 +227,8 @@ def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
             dropout=s2_spec.dropout,
             stage="stage2",
             context_dim=s2_spec.context_dim,
+            trunk_type=s2_spec.trunk.type,
+            block_conditioning=s2_spec.trunk.block_conditioning,
         )
 
     return result

giant/model/history.py

@@ -63,6 +63,23 @@ def build_history(name: str, in_dim: int, out_dim: int, **kwargs) -> HistoryEnco
     return cls(in_dim, out_dim, **filtered)
 
 
+@register_history("none")
+class NoHistory(HistoryEncoder):
+    """No history signal at all — ignores feat/has_prev entirely and always
+    returns zeros. Ablates whether the AR decoder's history conditioning is
+    earning its parameters. `init_cache`/`step` use the base class's O(1)
+    defaults unmodified (this encoder's own `forward` is already O(1) per
+    call regardless of prefix length)."""
+
+    def __init__(self, in_dim: int, out_dim: int) -> None:
+        super().__init__()
+        self.out_dim = out_dim
+
+    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
+        B, K, _ = feat.shape
+        return torch.zeros(B, K, self.out_dim, device=feat.device, dtype=feat.dtype)
+
+
 @register_history("markov")
 class MarkovHistory(HistoryEncoder):
     """Summarizes the previous secondary's own `(energy_fraction, direction,

giant/model/models.py

@@ -8,7 +8,7 @@ from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
 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 HistoryEncoder, build_history
-from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding, build_mlp_head
+from giant.model.layers import ContextAdapter, SinusoidalEmbedding, build_mlp_head
 from giant.model.objectives import build_objective
 from giant.model.routers import Router
 from giant.model.trunks import build_trunk
@@ -188,6 +188,26 @@ class StageModel(nn.Module):
             hidden = max(1, round(hidden_dim * head_cfg.hidden_ratio))
             self.stop_head = build_mlp_head(cond_out_dim, 1, hidden, head_cfg.depth)
 
+    def _build_context_fusion(self, x_dim: int, context_dim: int, cond_out_dim: int) -> None:
+        """Builds `self.context_adapter`/`self.fuse` — the stage-2-style
+        context-fusion pattern (project the previous stage's outcome down to
+        `context_dim` via `ContextAdapter`, concat onto the base conditioning,
+        project back to `cond_out_dim`) shared by `Stage2OneShot` and a
+        `stage="stage2"` `CriticModel` (gitea #57). Call from a subclass's
+        `__init__` before using `_cond_embed`."""
+        self.context_adapter = ContextAdapter(x_dim, context_dim)
+        self.fuse = nn.Sequential(
+            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            nn.SiLU(),
+        )
+
+    def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
+        """Fuses base conditioning with the previous stage's outcome — pairs
+        with `_build_context_fusion`."""
+        base = self.cond_enc(cond_cont, cond_cat)
+        ctx = self.context_adapter(stage1_out)
+        return self.fuse(torch.cat([base, ctx], dim=-1))
+
     def _require_n_sec_head(self) -> None:
         if self.n_sec_head is None:
             raise RuntimeError(
@@ -363,11 +383,7 @@ class Stage2OneShot(StageModel):
             particle_type_cfg=particle_type_cfg,
             cond_enc=cond_enc,
         )
-        self.context_adapter = ContextAdapter(x_dim, context_dim)
-        self.fuse = nn.Sequential(
-            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
-            nn.SiLU(),
-        )
+        self._build_context_fusion(x_dim, context_dim, cond_out_dim)
         target = self.particle_type_cfg.target
         type_head_out_dim = None if target == "physical" else k_max * self.type_dim
         self._build_trunk_and_heads(
@@ -386,11 +402,6 @@ class Stage2OneShot(StageModel):
             type_head_cfg=type_head_cfg,
         )
 
-    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,
@@ -702,11 +713,25 @@ class Stage2Autoregressive(StageModel):
         return self.stop_head(c_emb.reshape(B * K, -1)).view(B, K)
 
 
-class CriticModel(nn.Module):
+class CriticModel(StageModel):
     """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"`."""
+    `Stage2OneShot`, via `StageModel._build_context_fusion`/`_cond_embed`).
+    Used only when that stage's `generator == "wgan"`.
+
+    Subclasses `StageModel` for the `cond_enc` construction and (stage 2)
+    context-fusion scaffolding only — its trunk is built directly via
+    `build_trunk` (output width 1) rather than through
+    `_build_trunk_and_heads`, since that helper is shaped around a
+    generator's `Objective`/time-embedding/flow-matching concerns
+    (`forward`'s `(x_t, cond) -> vector` shape) that don't apply to a critic's
+    `(x, cond) -> scalar` (gitea #57). `generator="wgan"` is passed to the
+    base purely because that's factually when a critic exists; nothing here
+    ever calls `_build_trunk_and_heads`, so no head/time-embedding machinery
+    is built from it. Never routed (MoE) — that's a separate, unrequested
+    axis of scope; see gitea #57's proposal, which covers only the trunk/
+    block registries."""
 
     def __init__(
         self,
@@ -722,22 +747,26 @@ class CriticModel(nn.Module):
         stage: str = "stage1",
         context_dim: int = 64,
         context_in_dim: int = X_DIM,
+        trunk_type: str = "resmlp",
+        block_conditioning: str = "add",
     ) -> None:
-        super().__init__()
+        super().__init__(
+            pdg_vocab,
+            mat_vocab,
+            particle_cfg,
+            material_cfg,
+            cond_out_dim=cond_out_dim,
+            generator="wgan",
+            noise_dim=0,
+        )
         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)
+            self._build_context_fusion(context_in_dim, context_dim, cond_out_dim)
+        self.trunk = build_trunk(
+            None, trunk_type, in_dim, 1, hidden_dim, n_res_blocks, cond_out_dim, dropout, block_conditioning
+        )
 
     def forward(
         self,
@@ -746,13 +775,9 @@ class CriticModel(nn.Module):
         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))
+            assert stage1_out is not None, "stage='stage2' CriticModel requires stage1_out"
+            cond = self._cond_embed(cond_cont, cond_cat, stage1_out)
         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)
+            cond = self.cond_enc(cond_cont, cond_cat)
+        return self.trunk(x, cond, cond_cont, cond_cat).squeeze(-1)

giant/model/network.py

@@ -15,6 +15,7 @@ from giant.model.history import (
     AttentionHistory,
     HistoryEncoder,
     MarkovHistory,
+    NoHistory,
     _CausalAttnBlock,
     build_history,
     register_history,
@@ -54,6 +55,7 @@ from giant.model.routers import (
     ROUTER_REGISTRY,
     ComposedRouter,
     EnergyRouter,
+    NoneRouter,
     PdgRouter,
     ProcessRouter,
     Router,
@@ -67,6 +69,7 @@ from giant.model.routers import (
 from giant.model.trunks import (
     TRUNK_REGISTRY,
     ExpertTrunk,
+    LinearTrunk,
     RoutedTrunk,
     Trunk,
     _route_forward,
@@ -90,7 +93,10 @@ __all__ = [
     "FlowObjective",
     "HISTORY_REGISTRY",
     "HistoryEncoder",
+    "LinearTrunk",
     "MarkovHistory",
+    "NoHistory",
+    "NoneRouter",
     "OBJECTIVE_REGISTRY",
     "Objective",
     "PdgRouter",

giant/model/routers.py

@@ -45,21 +45,36 @@ class Router(nn.Module):
         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.
+
+        Forced fp32 (`torch.autocast(..., enabled=False)`) regardless of the
+        caller's ambient `train.precision` autocast region: `clamp_min(1e-8)`
+        below sits under bf16's precision but *above* fp16's ~6e-8 subnormal
+        floor, so `log_probs` degrading here is exactly the kind of quiet
+        drift that cost a whole rollout benchmark before (see the MoE section
+        of CLAUDE.md's Roadmap) — cheap to rule out (gitea #47).
         """
-        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)
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            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)
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            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
+        """Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017).
+
+        Forced fp32 — `importance` sums `gate()` over the whole batch (a
+        large-magnitude accumulation in reduced precision), then takes a
+        `std/mean` ratio: a classic catastrophic-cancellation shape (gitea
+        #47)."""
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            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.
@@ -76,12 +91,17 @@ class Router(nn.Module):
 
     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
+        the full explanation, unchanged in v0.3.0.
+
+        Forced fp32, same rationale as `balance_loss`/`combine_weights`: the
+        `+ 1e-8` epsilon here is `entropy_loss`'s training-loss path too, not
+        just a diagnostic (gitea #47)."""
+        with torch.autocast(cond_cont.device.type, enabled=False):
+            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]] = {}
@@ -124,6 +144,24 @@ def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
     return math.log(p / (1 - p))
 
 
+@register_router("none")
+class NoneRouter(Router):
+    """Uniform 1/n_experts gate — no learned routing signal at all.
+
+    Still builds n_experts expert trunks via RoutedTrunk (same parameter
+    budget as a real router), but every row gets an identical weight
+    regardless of conditioning. Ablates whether the *learned routing
+    signal* — as opposed to simply having multiple experts — is earning
+    its parameters. `top1()` (the base class default) always dispatches to
+    expert 0 (argmax of a uniform vector), which still exercises
+    RoutedTrunk's real per-expert grouped-dispatch code path at eval time.
+    """
+
+    def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
+        B = cond_cont.shape[0]
+        return torch.full((B, self.n_experts), 1.0 / self.n_experts, device=cond_cont.device)
+
+
 @register_router("energy")
 class EnergyRouter(Router):
     """Soft turn-on gate over normalized pre-step log-energy.

giant/model/summary.py

@@ -0,0 +1,318 @@
+"""Build-only model introspection (gitea #46): construct the resolved
+Stage1/Stage2/critic graph from a config with no dataset attached, and report
+per-module parameter counts, trunk widths, which heads exist, and — via
+differential probing — which `conditioning`/`stage1_model`/`stage2_model`
+config keys actually shape the built model. This is the runtime counterpart
+to `tests/test_config_consumed_keys.py`'s static per-identifier audit: that
+test asks "does any code reference this key's name at all", this module asks
+"given *this* resolved config, does the key change what `build_models`/
+`build_critics` (`giant/model/builders.py`) actually produces".
+
+Differential probing, not identifier matching: build the model once from the
+resolved config and take a structural fingerprint (`_fingerprint` — which
+submodules exist, every parameter's/buffer's shape+dtype, every plain scalar
+attribute stored on any module). Then, for each in-scope leaf key, perturb
+just that one value (`_perturb`), rebuild, and re-fingerprint. A changed
+fingerprint — or a rebuild that raises — means the key was consumed; an
+identical fingerprint means construction never looked at it under this
+particular config. A key can be genuinely inert under one config and live
+under another (e.g. any `stage1_model.router.*` key when `router.enabled =
+false`) — that config-dependence is exactly the "silently degenerate
+combination" issue #46 is after, so it is reported per-run rather than
+baked into a static table.
+
+Keys legitimately owned by the trainer/sampler/rollout rather than by
+`build_models`/`build_critics` (loss weights, WGAN-GP training
+hyperparameters, teacher-forcing and stage1-context schedules, ...) are
+cataloged in `_NOT_BUILD_TIME` below so the report doesn't flag them as
+suspicious. One leaf is inert under every config today —
+`stage2_model.autoregressive.order` — matching
+`tests/test_config_consumed_keys.py`'s own `_KNOWN_UNUSED` entry; it is
+deliberately *not* in `_NOT_BUILD_TIME`, since "always inert" is itself the
+finding those two tests independently converge on.
+"""
+
+import copy
+from dataclasses import dataclass, field
+
+import torch.nn as nn
+
+from giant.config import _get_path, _set_path, leaf_paths
+from giant.model.builders import build_critics, build_models
+from giant.model.trunks import RoutedTrunk
+
+_IN_SCOPE_ROOTS = ("conditioning", "stage1_model", "stage2_model")
+
+_PROBE_STR = "__giant_model_summary_probe__"
+
+# A handful of string leaves branch on equality against one specific literal
+# (e.g. `builders.py`: `stop_token = s2_spec.n_sec.mode == "stop_token"`),
+# where every value other than that literal behaves identically. A single
+# generic sentinel probe would then falsely read as inert whenever the
+# config's *current* value is already one of those identically-behaving
+# "other" values (e.g. mode="head") — it never crosses the one boundary that
+# actually matters. Named here so probing tries the real alternative(s) too;
+# every other string leaf is registry-validated (raises on garbage, still
+# correctly detected as consumed) or genuinely value-independent, so doesn't
+# need an entry.
+_STRING_ALTERNATIVES: dict[str, tuple[str, ...]] = {
+    "stage2_model.n_sec.owner": ("stage1", "stage2"),
+    "stage2_model.n_sec.mode": ("stop_token", "head", "truth"),
+    "stage2_model.particle_type.target": ("physical", "onehot", "embedding"),
+}
+
+# Verified by reading giant/training/trainers.py, giant/training/stage2_inputs.py
+# and giant/rollout.py while implementing gitea #46 — not auto-derived, so a
+# future reader touching these fields should re-check this table still holds.
+_NOT_BUILD_TIME: dict[str, str] = {
+    "stage1_model.init_from": "training/checkpoint.py's init_stages_from_checkpoints, run before build_stage_trainers (gitea #42)",
+    "stage1_model.freeze": "trainers.py: StageSpec.freeze, gates StageTrainer._step_optimizer (gitea #42)",
+    "stage2_model.init_from": "training/checkpoint.py's init_stages_from_checkpoints, run before build_stage_trainers (gitea #42)",
+    "stage2_model.freeze": "trainers.py: StageSpec.freeze, gates StageTrainer._step_optimizer (gitea #42)",
+    "stage1_model.lambda": "trainers.py: StageSpec.lambda_weight, the total-loss mix weight",
+    "stage2_model.lambda": "trainers.py: StageSpec.lambda_weight, the total-loss mix weight",
+    "stage2_model.n_sec.lambda": "trainers.py: StageSpec.n_sec_lambda, the n_sec-head loss weight",
+    "stage2_model.particle_type.lambda": "trainers.py: Stage2Trainer.particle_type_lambda, the type-head loss weight",
+    "stage2_model.particle_type.other_policy": "giant/rollout.py: resolves an 'other'-bucket secondary's PDG code at inference",
+    "stage2_model.particle_type.class_weighting": "trainers.py: FlowDDPMStageTrainer.type_class_weights, shapes the type-head loss, not the built graph (gitea #44)",
+    "stage2_model.autoregressive.teacher_forcing": "giant/training/stage2_inputs.py's training-time input assembly",
+    "stage2_model.autoregressive.tf_p_start": "trainers.py's teacher-forcing schedule",
+    "stage2_model.autoregressive.tf_p_end": "trainers.py's teacher-forcing schedule",
+    "stage2_model.stage1_context": "trainers.py's stage1/stage2 boundary — StageTrainer._stage1_context",
+    "stage2_model.ctx_p_start": "trainers.py's stage1-context sampling schedule",
+    "stage2_model.ctx_p_end": "trainers.py's stage1-context sampling schedule",
+    "stage1_model.router.lambda_balance": "trainers.py's load-balancing auxiliary loss weight",
+    "stage1_model.router.lambda_entropy": "trainers.py's entropy-regularization auxiliary loss weight",
+    "stage1_model.router.lambda_proc": "trainers.py's supervised process-classification auxiliary loss weight",
+    "stage1_model.router.gumbel_tau_start": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage1_model.router.gumbel_tau_end": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage2_model.router.lambda_balance": "trainers.py's load-balancing auxiliary loss weight",
+    "stage2_model.router.lambda_entropy": "trainers.py's entropy-regularization auxiliary loss weight",
+    "stage2_model.router.lambda_proc": "trainers.py's supervised process-classification auxiliary loss weight",
+    "stage2_model.router.gumbel_tau_start": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage2_model.router.gumbel_tau_end": "trainers.py's expert-combination Gumbel-softmax temperature anneal",
+    "stage1_model.wgan.n_critic": "trainers.py's WGAN-GP critic-update cadence",
+    "stage1_model.wgan.gp_weight": "trainers.py's WGAN-GP gradient-penalty coefficient",
+    "stage1_model.wgan.critic_lr": "trainers.py's critic optimizer learning rate",
+    "stage2_model.wgan.n_critic": "trainers.py's WGAN-GP critic-update cadence",
+    "stage2_model.wgan.gp_weight": "trainers.py's WGAN-GP gradient-penalty coefficient",
+    "stage2_model.wgan.critic_lr": "trainers.py's critic optimizer learning rate",
+    "stage2_model.wgan.gumbel_tau_start": "trainers.py's type-slice Gumbel-softmax temperature anneal (type_gumbel_tau_start)",
+    "stage2_model.wgan.gumbel_tau_end": "trainers.py's type-slice Gumbel-softmax temperature anneal (type_gumbel_tau_end)",
+}
+
+
+@dataclass
+class ModelSummary:
+    modules: dict[str, nn.Module]
+    consumed: list[str]
+    inert: list[str]
+    elsewhere: list[str]
+    pdg_vocab: int
+    mat_vocab: int
+    vocab_caveats: list[str] = field(default_factory=list)
+
+
+def _build_model_config(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict:
+    return {
+        "pdg_vocab": pdg_vocab,
+        "mat_vocab": mat_vocab,
+        "conditioning": cfg["conditioning"],
+        "stage1_model": cfg["stage1_model"],
+        "stage2_model": cfg["stage2_model"],
+    }
+
+
+def _built_modules(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict[str, nn.Module]:
+    model_config = _build_model_config(cfg, pdg_vocab, mat_vocab)
+    modules: dict[str, nn.Module] = {}
+    for name, m in build_models(model_config).items():
+        if m is not None:
+            modules[name] = m
+    for name, m in build_critics(model_config).items():
+        if m is not None:
+            modules[f"{name}_critic"] = m
+    return modules
+
+
+def _fingerprint(modules: dict[str, nn.Module]) -> list:
+    """A config-shape fingerprint of the built graph: which submodules
+    exist, every parameter's/buffer's shape+dtype (never values — those are
+    randomly initialized and irrelevant to *structure*), and every plain
+    scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
+    .stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
+    non-parametric key's effect on construction observable."""
+    sig = []
+    for stage_name, module in modules.items():
+        for mod_name, m in module.named_modules():
+            full = f"{stage_name}.{mod_name}" if mod_name else stage_name
+            for k, v in vars(m).items():
+                if k.startswith("_"):
+                    continue
+                if v is None or isinstance(v, (bool, int, float, str)):
+                    sig.append((full, k, v))
+        for pname, p in module.named_parameters():
+            sig.append((stage_name, "param", pname, tuple(p.shape), str(p.dtype)))
+        for bname, b in module.named_buffers():
+            sig.append((stage_name, "buffer", bname, tuple(b.shape), str(b.dtype)))
+    return sorted(sig, key=repr)
+
+
+def _perturb_candidates(path: str, value) -> list:
+    """Values to try perturbing `path`'s current `value` to, in order —
+    probing stops at the first one that changes the fingerprint or raises.
+    Almost always a single candidate; see `_STRING_ALTERNATIVES`."""
+    if isinstance(value, bool):
+        return [not value]
+    if isinstance(value, int):
+        return [value + 1]
+    if isinstance(value, float):
+        return [value + 1.0]
+    if isinstance(value, str):
+        alternatives = [v for v in _STRING_ALTERNATIVES.get(path, ()) if v != value]
+        return [*alternatives, _PROBE_STR]
+    raise TypeError(f"gitea #46 probing: unsupported leaf value type {type(value)!r} ({value!r})")
+
+
+def _vocab_caveats(cfg: dict) -> list[str]:
+    caveats = []
+    if _get_path(cfg, "conditioning.particle.type") == "embedding":
+        caveats.append(
+            "conditioning.particle.type = 'embedding' -- pdg_vocab below is a "
+            "placeholder (no dataset attached to derive the real training vocab size)"
+        )
+    if _get_path(cfg, "conditioning.material.type") == "embedding":
+        caveats.append(
+            "conditioning.material.type = 'embedding' -- mat_vocab below is a "
+            "placeholder (no dataset attached to derive the real training vocab size)"
+        )
+    for stage in ("stage1_model", "stage2_model"):
+        router_type = _get_path(cfg, f"{stage}.router.type")
+        if _get_path(cfg, f"{stage}.router.enabled") and router_type in ("pdg", "process"):
+            caveats.append(
+                f"{stage}.router.type = {router_type!r} builds its own pdg_vocab-sized "
+                "embedding -- the count above is a placeholder"
+            )
+    return caveats
+
+
+def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
+    """Build `cfg`'s model with no dataset attached and report its resolved
+    graph, plus which `conditioning`/`stage1_model`/`stage2_model` config
+    keys actually shaped it (differential probing — see module docstring).
+    `cfg` must already be a fully-merged v0.3 config (`merge_cli_overrides`
+    output) — this does not migrate or validate it."""
+    modules = _built_modules(cfg, pdg_vocab, mat_vocab)
+    baseline_fp = _fingerprint(modules)
+
+    in_scope = [p for p in leaf_paths(cfg) if p.split(".", 1)[0] in _IN_SCOPE_ROOTS]
+    consumed: list[str] = []
+    inert: list[str] = []
+    elsewhere: list[str] = []
+    for path in in_scope:
+        original = _get_path(cfg, path)
+        changed = False
+        for candidate in _perturb_candidates(path, original):
+            probe_cfg = copy.deepcopy(
+                {
+                    "conditioning": cfg["conditioning"],
+                    "stage1_model": cfg["stage1_model"],
+                    "stage2_model": cfg["stage2_model"],
+                }
+            )
+            _set_path(probe_cfg, path, candidate)
+            try:
+                changed = _fingerprint(_built_modules(probe_cfg, pdg_vocab, mat_vocab)) != baseline_fp
+            except Exception:
+                changed = True
+            if changed:
+                break
+        if changed:
+            consumed.append(path)
+        elif path in _NOT_BUILD_TIME:
+            elsewhere.append(path)
+        else:
+            inert.append(path)
+
+    return ModelSummary(
+        modules=modules,
+        consumed=sorted(consumed),
+        inert=sorted(inert),
+        elsewhere=sorted(elsewhere),
+        pdg_vocab=pdg_vocab,
+        mat_vocab=mat_vocab,
+        vocab_caveats=_vocab_caveats(cfg),
+    )
+
+
+def _tree_lines(module: nn.Module, name: str, indent: int = 0) -> list[str]:
+    total = sum(p.numel() for p in module.parameters())
+    in_dim = getattr(module, "in_dim", None)
+    out_dim = getattr(module, "out_dim", None)
+    widths = f"  [in={in_dim}, out={out_dim}]" if in_dim is not None and out_dim is not None else ""
+    lines = [f"{'  ' * indent}{name} ({type(module).__name__}): {total:,}{widths}"]
+    for child_name, child in module.named_children():
+        lines.extend(_tree_lines(child, child_name, indent + 1))
+    return lines
+
+
+_HEAD_NAMES = ("n_sec_head", "type_head", "stop_head")
+
+
+def _stage_header(name: str, module: nn.Module) -> list[str]:
+    total = sum(p.numel() for p in module.parameters())
+    lines = [f"{name}: {type(module).__name__} -- {total:,} parameters"]
+    generator = getattr(module, "generator_kind", None)
+    if generator is not None:
+        lines.append(f"  generator: {generator}")
+    trunk = getattr(module, "trunk", None)
+    if trunk is not None:
+        in_dim = getattr(trunk, "in_dim", "?")
+        out_dim = getattr(trunk, "out_dim", "?")
+        if isinstance(trunk, RoutedTrunk):
+            detail = f"routed, n_experts={trunk.router.n_experts}, expert type={type(trunk.experts[0]).__name__}"
+        else:
+            detail = f"unrouted, {type(trunk).__name__}"
+        lines.append(f"  trunk: {detail}, in={in_dim}, out={out_dim}")
+    history_kind = getattr(module, "history_kind", None)
+    if history_kind is not None:
+        lines.append(f"  autoregressive history: {history_kind}")
+    present = [h for h in _HEAD_NAMES if getattr(module, h, None) is not None]
+    absent = [h for h in _HEAD_NAMES if hasattr(module, h) and getattr(module, h) is None]
+    if present or absent:
+        lines.append(f"  heads present: {', '.join(present) if present else 'none'}")
+        if absent:
+            lines.append(f"  heads absent: {', '.join(absent)}")
+    return lines
+
+
+def render_summary(summary: ModelSummary) -> str:
+    lines: list[str] = []
+    for name, module in summary.modules.items():
+        lines.extend(_stage_header(name, module))
+        lines.extend(_tree_lines(module, name, indent=1))
+        lines.append("")
+
+    lines.append(
+        f"config keys read during construction: {len(summary.consumed)} / "
+        f"read elsewhere (trainer/sampler/rollout): {len(summary.elsewhere)} / "
+        f"inert under this config: {len(summary.inert)}"
+    )
+    if summary.elsewhere:
+        lines.append("read elsewhere, not by construction:")
+        for path in summary.elsewhere:
+            lines.append(f"  {path}  ({_NOT_BUILD_TIME[path]})")
+    lines.append("inert under this config (declared, parsed, but doing nothing here):")
+    if summary.inert:
+        for path in summary.inert:
+            lines.append(f"  {path}")
+    else:
+        lines.append("  (none)")
+
+    if summary.vocab_caveats:
+        lines.append("")
+        lines.append("vocab placeholder caveats:")
+        for caveat in summary.vocab_caveats:
+            lines.append(f"  {caveat}")
+
+    return "\n".join(lines)

giant/model/trunks.py

@@ -73,6 +73,7 @@ class ExpertTrunk(nn.Module):
         block_conditioning: str = "add",
     ) -> None:
         super().__init__()
+        self.in_dim = in_dim
         self.out_dim = out_dim
         self.input_proj = nn.Linear(in_dim, hidden_dim)
         self.blocks = nn.ModuleList(
@@ -93,6 +94,49 @@ class ExpertTrunk(nn.Module):
         return self.out_proj(x)
 
 
+@register_trunk("linear")
+class LinearTrunk(nn.Module):
+    """`nn.Linear(in_dim + cond_dim, out_dim)` over `concat([x, cond])` —
+    the trivial trunk body: no hidden layer, no ResBlock stack, no
+    nonlinearity. Ablates whether trunk depth/nonlinearity is earning its
+    parameters, holding everything else (heads, ConditionEncoder,
+    generator, ...) fixed. Composes for free with `router.enabled = true`
+    (gitea #33): a RoutedTrunk of n_experts linear bodies is "mixture of
+    trivial linear experts". `hidden_dim`/`n_blocks`/`dropout`/
+    `block_conditioning` are accepted and ignored, matching
+    `build_expert_body`'s shared factory signature.
+
+    `x` — the trunk's own input (e.g. the noised primary vector for flow
+    matching) — does not already carry conditioning; that's fused in
+    per-body via `cond`. So this concatenates `x` and `cond` itself to
+    remain a valid, conditioning-dependent model.
+    """
+
+    def __init__(
+        self,
+        in_dim: int,
+        out_dim: int,
+        hidden_dim: int,
+        n_blocks: int,
+        cond_dim: int,
+        dropout: float = 0.0,
+        block_conditioning: str = "add",
+    ) -> None:
+        super().__init__()
+        self.in_dim = in_dim
+        self.out_dim = out_dim
+        self.linear = nn.Linear(in_dim + cond_dim, out_dim)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        cond: torch.Tensor,
+        cond_cont: torch.Tensor | None = None,
+        cond_cat: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.linear(torch.cat([x, cond], dim=-1))
+
+
 def _route_forward(
     experts: nn.ModuleList,
     router: Router,
@@ -108,21 +152,50 @@ def _route_forward(
     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.
+
+    The accumulator's dtype is deferred to the first expert call rather than
+    fixed at fp32: under autocast (`train.precision = "bf16"`, gitea #47) an
+    expert's `ResBlock` stack returns bf16, and an fp32-fixed accumulator
+    would silently upcast every mixture term (train mode) or downcast every
+    dispatched row via `index_put_` (eval mode) — making a `RoutedTrunk`
+    return a different dtype than the unrouted `ExpertTrunk` it's a drop-in
+    replacement for, purely because `router.enabled` was set.
+
+    `router.combine_weights` is deliberately fp32 internally (it forces its
+    own autocast-disabled region — see `Router.combine_weights`'s docstring),
+    so `weights` itself is always fp32 regardless of the ambient precision.
+    Left as-is, `weights[:, i:i+1] * expert(x, cond)` would type-promote the
+    whole mixture back to fp32 by ordinary PyTorch promotion rules — the same
+    dtype-mismatch bug this function exists to avoid, just moved one line
+    over. `weights` is cast down to each expert's own output dtype right
+    before combining: the softmax stays numerically stable at fp32, but its
+    *result* (values in [0, 1], not precision-sensitive to represent) loses
+    nothing meaningful by then being used at bf16.
     """
     if training:
-        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts)
-        out = torch.zeros(x.shape[0], experts[0].out_dim, device=x.device)
+        weights = router.combine_weights(cond_cont, cond_cat)  # (B, n_experts), fp32
+        out = None
         for i, expert in enumerate(experts):
-            out = out + weights[:, i : i + 1] * expert(x, cond)
+            expert_out = expert(x, cond)
+            term = weights[:, i : i + 1].to(expert_out.dtype) * expert_out
+            out = term if out is None else out + term
+        assert out is not None, "RoutedTrunk built with zero experts"
         return out
 
     idx = router.top1(cond_cont, cond_cat)  # (B,)
-    out_dim = experts[0].out_dim
-    out = torch.zeros(x.shape[0], out_dim, device=x.device)
+    out = None
     for i, expert in enumerate(experts):
         mask = idx == i
         if mask.any():
-            out[mask] = expert(x[mask], cond[mask])
+            expert_out = expert(x[mask], cond[mask])
+            if out is None:
+                out = torch.zeros(x.shape[0], expert_out.shape[-1], device=x.device, dtype=expert_out.dtype)
+            out[mask] = expert_out
+    if out is None:
+        # No row was ever dispatched (only reachable with an empty batch,
+        # x.shape[0] == 0) — nothing to infer a dtype from, so fall back to
+        # x's own, matching this function's pre-autocast behavior.
+        out = torch.zeros(x.shape[0], experts[0].out_dim, device=x.device, dtype=x.dtype)
     return out
 
 
@@ -130,7 +203,10 @@ class Trunk(nn.Module):
     """Interface implemented by a standalone trunk body (any `TRUNK_REGISTRY`
     entry, e.g. `ExpertTrunk`) and by `RoutedTrunk`: everything downstream of
     the fused conditioning vector, i.e. the actual generative trunk of a
-    stage."""
+    stage. Implementations are expected to expose `in_dim`/`out_dim`
+    attributes (as `ExpertTrunk`/`RoutedTrunk` do) — `giant.model.summary`
+    (gitea #46) reads them to report trunk widths without needing to know the
+    body architecture."""
 
     def forward(
         self,
@@ -157,6 +233,8 @@ class RoutedTrunk(Trunk):
     ) -> None:
         super().__init__()
         self.router = router
+        self.in_dim = in_dim
+        self.out_dim = out_dim
         self.experts = nn.ModuleList(
             [
                 build_expert_body(

giant/model/wgan.py

@@ -22,21 +22,31 @@ def gradient_penalty(
     norm to 1 — `x_hat`/`grad` are forced to all-zero for such a row, which
     would otherwise contribute a constant `(||0|| - 1)^2 == 1` bias to the
     mean regardless of critic behavior — so they're excluded from the mean.
+
+    Deliberately kept fp32 (`torch.autocast(..., enabled=False)`) regardless
+    of the caller's ambient `train.precision` autocast region: this is a
+    `create_graph=True` double-backward, and `grad.norm(2, dim=1)` sums
+    squares over the critic's full input width (hundreds of dims for stage
+    2), which overflows bf16's range at gradient magnitudes well within
+    normal early-WGAN-GP territory. Disclosed cost: the critic forward
+    inside this function always runs fp32, even when the rest of the WGAN
+    stage's step is bf16 (gitea #47).
     """
-    eps = torch.rand(real.size(0), 1, device=real.device)
-    x_hat = eps * real + (1 - eps) * fake
-    if mask is not None:
-        x_hat = x_hat * mask
-    x_hat = x_hat.requires_grad_(True)
-    scores = critic_fn(x_hat)
-    grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0]
-    if mask is not None:
-        grad = grad * mask
-    penalty = (grad.norm(2, dim=1) - 1) ** 2
-    if mask is not None:
-        valid = (mask.sum(dim=1) > 0).float()
-        return (penalty * valid).sum() / valid.sum().clamp_min(1.0)
-    return penalty.mean()
+    with torch.autocast(real.device.type, enabled=False):
+        eps = torch.rand(real.size(0), 1, device=real.device)
+        x_hat = eps * real.float() + (1 - eps) * fake.float()
+        if mask is not None:
+            x_hat = x_hat * mask
+        x_hat = x_hat.requires_grad_(True)
+        scores = critic_fn(x_hat)
+        grad = torch.autograd.grad(outputs=scores.sum(), inputs=x_hat, create_graph=True)[0]
+        if mask is not None:
+            grad = grad * mask
+        penalty = (grad.norm(2, dim=1) - 1) ** 2
+        if mask is not None:
+            valid = (mask.sum(dim=1) > 0).float()
+            return (penalty * valid).sum() / valid.sum().clamp_min(1.0)
+        return penalty.mean()
 
 
 def critic_loss(

giant/pipeline.py

@@ -359,7 +359,7 @@ def run_train_job(
             "section)"
         )
 
-    config.validate_config(cfg)
+    config.validate_config(cfg, resume=resume is not None)
     particle_conditioning = cfg["conditioning"]["particle"]["type"]
     material_conditioning = cfg["conditioning"]["material"]["type"]
     k_max = cfg["stage2_model"]["k_max"]

giant/tools/profile_analysis_costs.py

@@ -28,6 +28,7 @@ import polars as pl
 from giant.analysis.catalog import catalog_ids, get_spec
 from giant.analysis.condor import compute_reduced
 from giant.analysis.context import build_context
+from giant.analysis.sources import RolloutSpec
 
 # Row counts (per side) to benchmark at. Kept in local memory/CPU range so the
 # whole sweep finishes in about a minute; the fit is linear so it extrapolates
@@ -165,11 +166,10 @@ def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path)
     t0 = time.perf_counter()
     compute_reduced(
         spec_id,
-        rollout,
+        [{"name": "rollout", "path": str(rollout)}],
         reference,
         shared,
         out,
-        checkpoint=None,
         chunk_index=0,
         n_chunks=1,
     )
@@ -191,7 +191,7 @@ def main() -> None:
 
             shared = tmp_path / f"shared_{n_side}.json"
             ctx = build_context(
-                rollout,
+                [RolloutSpec(name="rollout", source=rollout)],
                 reference,
                 n_energy_bins=4,
                 n_marginal_bins=50,

giant/training/__init__.py

@@ -5,7 +5,7 @@ Split out of the former single-module `giant/train.py`. The public surface is
 that tests and tooling construct directly.
 """
 
-from giant.training.checkpoint import build_checkpoint, load_checkpoint
+from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
 from giant.training.metrics import MetricsCollector, MetricSpec
 from giant.training.loop import train
 from giant.training.trainers import (
@@ -25,6 +25,7 @@ __all__ = [
     "WGANStageTrainer",
     "build_checkpoint",
     "build_stage_trainers",
+    "init_stages_from_checkpoints",
     "load_checkpoint",
     "train",
 ]

giant/training/amp.py

@@ -0,0 +1,47 @@
+"""Mixed-precision training support (`train.precision`, gitea #47).
+
+Only `"fp32"` (no autocast) and `"bf16"` are supported — no `"fp16"`/
+`GradScaler`. bf16 needs no gradient scaler and covers every training GPU in
+the fleet (Ampere and newer: A100, L40S, H200, RTX 4070); fp16 would need a
+scaler *and* fixes to two fragile spots that stay correct under bf16 but break
+under fp16's narrower range — `giant.model.routers`' `1e-8` epsilons (below
+fp16's ~6e-8 subnormal floor) and `giant.model.wgan.gradient_penalty`'s
+sum-of-squares gradient norm (overflows fp16 above ~65504). Revisit if a
+pre-Ampere (V100) training target ever shows up.
+"""
+
+import torch
+
+_SUPPORTED_DEVICE_TYPES = ("cuda", "cpu")
+
+
+def resolve_autocast(precision: str, device: torch.device) -> tuple[str, torch.dtype, bool]:
+    """Resolves `train.precision` + a target device into the
+    `(device_type, dtype, enabled)` triple `torch.autocast` takes as kwargs —
+    computed once per `StageTrainer` rather than re-derived every step.
+
+    Raises `ValueError` rather than silently falling back to fp32: a training
+    run that's quietly not using the mixed precision it was configured for is
+    a wasted GPU-week, not a warning.
+    """
+    if precision == "fp32":
+        return device.type, torch.float32, False
+    if precision != "bf16":
+        raise ValueError(f"unknown precision {precision!r}; must be 'fp32' or 'bf16'")
+
+    if device.type == "cuda":
+        if not torch.cuda.is_bf16_supported():
+            cap = torch.cuda.get_device_capability(device)
+            raise ValueError(
+                f"train.precision = 'bf16' but {torch.cuda.get_device_name(device)} "
+                f"(compute capability {cap[0]}.{cap[1]}) has no native bf16 support "
+                "(needs Ampere/sm_80 or newer) — use train.precision = 'fp32' instead"
+            )
+        return "cuda", torch.bfloat16, True
+    if device.type == "cpu":
+        # torch 2.3's CPU autocast supports bf16 unconditionally — this is
+        # also what lets the bf16 training path be tested without a GPU.
+        return "cpu", torch.bfloat16, True
+    raise ValueError(
+        f"train.precision = 'bf16' is not supported on device type {device.type!r} (only {_SUPPORTED_DEVICE_TYPES} are)"
+    )

giant/training/checkpoint.py

@@ -8,6 +8,8 @@ and per-stage `optimizer_<stage>` / `optimizer_d_<stage>` / `lr_sched_<stage>`
 entries.
 """
 
+import torch
+
 from giant.training.trainers import StageTrainer
 
 #: Stage name -> the checkpoint key its weights live under. Historical: stage
@@ -46,6 +48,37 @@ def build_checkpoint(
     return ckpt
 
 
+def init_stages_from_checkpoints(trainers: dict[str, StageTrainer]) -> list[str]:
+    """Load each trainer's `spec.init_from` checkpoint (gitea #42) into its
+    model, before training starts — the partial-retrain counterpart to
+    `load_checkpoint`'s full-run `--resume`. Only weights move: unlike
+    `load_checkpoint`, this never touches optimizer/lr_sched/epoch state, so
+    it composes cleanly with `--resume` (call this first; a resume's own
+    `load_checkpoint` then overwrites whatever this loaded with the resumed
+    run's own weights).
+
+    A stage with no `init_from` set (`""`, the default) is left alone. The
+    EMA companion (`<key>_ema`) is loaded too when both the source checkpoint
+    and this trainer have one, so `--weights ema` at inference still sees the
+    source's EMA shadow rather than a copy of its raw weights. Returns one
+    description string per stage actually initialized, for the caller to
+    echo.
+    """
+    loaded = []
+    for name, trainer in trainers.items():
+        init_from = trainer.spec.init_from
+        if not init_from:
+            continue
+        key = _STAGE_KEY[name]
+        ckpt = torch.load(init_from, map_location="cpu", weights_only=False)
+        trainer.model.load_state_dict(ckpt[key])
+        ema_key = f"{key}_ema"
+        if trainer.ema_model is not None and ema_key in ckpt:
+            trainer.ema_model.load_state_dict(ckpt[ema_key])
+        loaded.append(f"{name}: loaded from {init_from}" + (" (frozen)" if trainer.frozen else ""))
+    return loaded
+
+
 def load_checkpoint(trainers: dict[str, StageTrainer], ckpt: dict, lr: float) -> None:
     """Restore every active stage, then hand `lr`'s authority back to the
     config — `load_state_dict` would otherwise leave the checkpoint's own

giant/training/loop.py

@@ -20,7 +20,7 @@ from tqdm import tqdm
 
 from giant.data.loader import TopNMap
 from giant.data.setup_cache import topnmap_to_json
-from giant.training.checkpoint import build_checkpoint, load_checkpoint
+from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
 from giant.training.metrics import MetricsCollector
 from giant.training.trainers import (
     FlowDDPMStageTrainer,
@@ -132,9 +132,12 @@ def train(
     validate_steps = t.get("validate_steps", 10)
     max_val_batches = t.get("max_val_batches", 0)
 
-    trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches)
+    sec_type_class_counts = sec_type_topn_map.class_counts if sec_type_topn_map is not None else None
+    trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches, sec_type_class_counts)
     if not trainers:
         raise ValueError("no active stage — stage1_model.active and stage2_model.active are both false")
+    for line in init_stages_from_checkpoints(trainers):
+        print(line)
     has_adversarial = any(not tr.supports_val_loss for tr in trainers.values())
 
     checkpoint_extras = {

giant/training/plots.py

@@ -0,0 +1,356 @@
+"""Training-progress plots from `<run_dir>/metrics.csv` (gitea #75).
+
+`MetricsCollector` (`giant.training.metrics`) writes one row per epoch with a
+column set that varies by run — flow/ddpm vs wgan, routed vs not (see the
+`MetricSpec` declarations in `giant.training.trainers`). This module reads
+that header dynamically rather than hardcoding a column list, buckets columns
+by the fixed naming convention `MetricsCollector` itself documents
+(`<stage>/train/<key>`, `<stage>/val/<key>`, `<stage>/router/<key>`,
+`<stage>/<key>` for point-in-time values, and an unprefixed run-level tail —
+see `giant.training.metrics`'s module docstring), and renders one PDF per
+applicable figure with the same `plotstyle` conventions
+`giant.analysis.render` uses, for visual consistency with the
+rollout-vs-reference plots.
+
+Unlike `giant.analysis`, there is no reduce/chunk/condor split here — the CSV
+is tiny and this always runs as one local pass — but the CLI entry point
+still lives under `giant analyze` (`analyze metrics`) as the shared home for
+plotstyle-rendered diagnostics, and shares its `analysis_runs/` output
+convention (see `derive_metrics_dir`) so training-progress plots don't get
+written into the training run directory itself.
+"""
+
+from __future__ import annotations
+
+import csv
+import math
+from dataclasses import dataclass
+from pathlib import Path
+
+# Stage names are always exactly these two — hardcoded in
+# `giant.training.trainers.build_stage_trainers` — so a column belongs to a
+# stage iff it's prefixed by one of these, and everything else (bar `epoch`)
+# is run-level. This is what makes dynamic header parsing tractable without
+# needing to know the per-run metric keys themselves.
+_STAGE_NAMES = ("stage1", "stage2")
+
+_ACC_KEYS = {"nsec_acc", "stop_acc", "type_acc"}
+_WGAN_BALANCE_KEYS = {"d_loss", "g_loss", "wasserstein", "gp_loss"}
+_ROUTER_KEYS = ("entropy", "util_min", "util_max", "util_std")
+
+
+@dataclass
+class MetricsTable:
+    """`<run_dir>/metrics.csv`, parsed with no hardcoded column list."""
+
+    epochs: list[int]
+    columns: dict[str, list[float]]
+
+    @classmethod
+    def load(cls, path: str | Path) -> "MetricsTable":
+        with open(path, newline="") as f:
+            rows = list(csv.DictReader(f))
+        epochs = [int(float(r["epoch"])) for r in rows]
+        fieldnames = rows[0].keys() if rows else []
+        columns = {name: [float(r[name]) for r in rows] for name in fieldnames if name != "epoch"}
+        return cls(epochs=epochs, columns=columns)
+
+    def best_epochs(self) -> list[int]:
+        is_best = self.columns.get("is_best")
+        if not is_best:
+            return []
+        return [epoch for epoch, flag in zip(self.epochs, is_best) if flag]
+
+
+# --- column classification --------------------------------------------------
+
+
+def _stages(columns: dict) -> list[str]:
+    return [s for s in _STAGE_NAMES if any(name.startswith(f"{s}/") for name in columns)]
+
+
+def _split(columns: dict, stage: str, split: str) -> dict[str, str]:
+    prefix = f"{stage}/{split}/"
+    return {name[len(prefix) :]: name for name in columns if name.startswith(prefix)}
+
+
+def _point_in_time(columns: dict, stage: str) -> dict[str, str]:
+    prefix = f"{stage}/"
+    out = {}
+    for name in columns:
+        if not name.startswith(prefix):
+            continue
+        rest = name[len(prefix) :]
+        head = rest.split("/", 1)[0]
+        if head not in ("train", "val", "router"):
+            out[rest] = name
+    return out
+
+
+def _router(columns: dict, stage: str) -> dict[str, str]:
+    prefix = f"{stage}/router/"
+    return {name[len(prefix) :]: name for name in columns if name.startswith(prefix)}
+
+
+def _run_level(columns: dict) -> dict[str, str]:
+    known_prefixes = tuple(f"{s}/" for s in _STAGE_NAMES)
+    return {name: name for name in columns if not name.startswith(known_prefixes)}
+
+
+def _loss_keys(train: dict[str, str], val: dict[str, str]) -> list[str]:
+    keys = {k for k in train if k not in _ACC_KEYS and k not in _WGAN_BALANCE_KEYS and k != "grad_norm"}
+    keys |= {k for k in val if k not in _ACC_KEYS and k not in _WGAN_BALANCE_KEYS and k != "grad_norm"}
+    return sorted(keys)
+
+
+# --- output location ---------------------------------------------------------
+
+
+def derive_metrics_dir(
+    run_dir: str | Path,
+    out_dir: str | Path | None = None,
+    default_base: str | Path | None = None,
+) -> Path:
+    """Plots output directory.
+
+    Precedence: an explicit `out_dir` always wins. Otherwise
+    `default_base / f"metrics_{run_dir.name}"` (the CLI passes the repo's
+    gitignored `analysis_runs/`, matching `giant.analysis.condor.derive_run_dir`'s
+    convention) — training-progress plots live alongside rollout-vs-reference
+    analysis runs, not inside the training run directory itself.
+    """
+    if out_dir is not None:
+        return Path(out_dir)
+    base = Path(default_base) if default_base is not None else Path.cwd() / "analysis_runs"
+    return base / f"metrics_{Path(run_dir).name}"
+
+
+# --- figures ------------------------------------------------------------------
+
+
+def _mark_best(ax, table: MetricsTable) -> None:
+    for epoch in table.best_epochs():
+        ax.axvline(epoch, color="grey", linestyle="--", linewidth=0.8, alpha=0.7)
+
+
+def _overview_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    run_level = _run_level(table.columns)
+    if "val/loss" not in run_level:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title="training overview")
+    ax.plot(table.epochs, table.columns["val/loss"], label="val/loss")
+    if "val/marginal_kl" in run_level:
+        kl = table.columns["val/marginal_kl"]
+        if any(math.isfinite(v) for v in kl):
+            ax.plot(table.epochs, kl, label="val/marginal_kl")
+    _mark_best(ax, table)
+    best = table.best_epochs()
+    if best:
+        idx = table.epochs.index(best[-1])
+        ax.annotate(
+            f"best: epoch {best[-1]}\nval/loss={table.columns['val/loss'][idx]:.4g}",
+            xy=(best[-1], table.columns["val/loss"][idx]),
+            xytext=(0.98, 0.95),
+            textcoords="axes fraction",
+            ha="right",
+            va="top",
+            fontsize=8,
+        )
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("loss")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _loss_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    train = _split(table.columns, stage, "train")
+    val = _split(table.columns, stage, "val")
+    keys = _loss_keys(train, val)
+    if not keys:
+        return None
+    n = len(keys)
+    ncols = min(3, n)
+    nrows = (n + ncols - 1) // ncols
+    fig, axes = ps.new_figure(
+        "slide-16x9",
+        title=f"{stage} loss",
+        nrows=nrows,
+        ncols=ncols,
+        squeeze=False,
+    )
+    flat = axes.ravel()
+    for ax, key in zip(flat, keys):
+        if key in train:
+            ax.plot(table.epochs, table.columns[train[key]], label="train")
+        if key in val:
+            ax.plot(table.epochs, table.columns[val[key]], label="val")
+        ax.set_yscale("log")
+        ax.set_title(key, fontsize=8)
+        ax.set_xlabel("epoch")
+    for j in range(n, len(flat)):
+        flat[j].set_visible(False)
+    ps.style_legend(flat[0], title="series")
+    return fig
+
+
+def _lr_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    series: dict[str, str] = {}
+    for stage in _stages(table.columns):
+        for key, col in _point_in_time(table.columns, stage).items():
+            series[f"{stage}/{key}"] = col
+    if not series:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title="learning rate schedule")
+    for label, col in series.items():
+        ax.plot(table.epochs, table.columns[col], label=label)
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("learning rate")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _accuracy_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    train = _split(table.columns, stage, "train")
+    val = _split(table.columns, stage, "val")
+    keys = sorted((set(train) | set(val)) & _ACC_KEYS)
+    if not keys:
+        return None
+    n = len(keys)
+    fig, axes = ps.new_figure("slide-16x9", title=f"{stage} accuracy", nrows=1, ncols=n, squeeze=False)
+    flat = axes.ravel()
+    for ax, key in zip(flat, keys):
+        if key in train:
+            ax.plot(table.epochs, table.columns[train[key]], label="train")
+        if key in val:
+            ax.plot(table.epochs, table.columns[val[key]], label="val")
+        ax.set_title(key, fontsize=8)
+        ax.set_xlabel("epoch")
+        ax.set_ylim(0, 1)
+    ps.style_legend(flat[0], title="series")
+    return fig
+
+
+def _grad_norm_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    run_level = _run_level(table.columns)
+    if "grad_norm" not in run_level:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title="gradient norm")
+    ax.plot(table.epochs, table.columns["grad_norm"], label="grad_norm")
+    for stage in _stages(table.columns):
+        train = _split(table.columns, stage, "train")
+        for key in ("grad_norm_d", "grad_norm_g", "grad_norm_type_slice", "grad_norm_cont_slice"):
+            if key in train:
+                ax.plot(table.epochs, table.columns[train[key]], label=f"{stage}/{key}")
+    ax.set_yscale("log")
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("grad norm")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _router_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    router = _router(table.columns, stage)
+    if "entropy" not in router:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title=f"{stage} router health")
+    ax.plot(table.epochs, table.columns[router["entropy"]], label="entropy", color="black")
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("entropy [bits]")
+    ax2 = ax.twinx()
+    for key in ("util_min", "util_max", "util_std"):
+        if key in router:
+            ax2.plot(table.epochs, table.columns[router[key]], label=key, linestyle="--")
+    ax2.set_ylabel("expert utilization")
+    ax2.set_ylim(0, 1)
+    lines1, labels1 = ax.get_legend_handles_labels()
+    lines2, labels2 = ax2.get_legend_handles_labels()
+    ax.legend(lines1 + lines2, labels1 + labels2, loc="upper right", frameon=False, fontsize=7)
+    return fig
+
+
+def _wgan_balance_figure(table: MetricsTable, stage: str):
+    import plotstyle as ps
+
+    train = _split(table.columns, stage, "train")
+    keys = [k for k in _WGAN_BALANCE_KEYS if k in train]
+    if not keys:
+        return None
+    fig, ax = ps.new_figure("thesis-single", title=f"{stage} WGAN critic/generator balance")
+    for key in sorted(keys):
+        ax.plot(table.epochs, table.columns[train[key]], label=key)
+    ax.set_xlabel("epoch")
+    ax.set_ylabel("value")
+    ps.style_legend(ax, title="series")
+    return fig
+
+
+def _throughput_figure(table: MetricsTable):
+    import plotstyle as ps
+
+    run_level = _run_level(table.columns)
+    keys = [k for k in ("samples_per_sec", "gpu_mem_mb", "epoch_time_s") if k in run_level]
+    if not keys:
+        return None
+    fig, axes = ps.new_figure("slide-16x9", title="throughput / resources", nrows=1, ncols=len(keys), squeeze=False)
+    flat = axes.ravel()
+    for ax, key in zip(flat, keys):
+        ax.plot(table.epochs, table.columns[key])
+        _mark_best(ax, table)
+        ax.set_title(key, fontsize=8)
+        ax.set_xlabel("epoch")
+    return fig
+
+
+# --- entry point ---------------------------------------------------------
+
+
+def render_metrics(
+    run_dir: str | Path,
+    out_dir: str | Path | None = None,
+    default_base: str | Path | None = None,
+) -> list[Path]:
+    """`<run_dir>/metrics.csv` -> `<plots dir>/<name>.pdf`.
+
+    See `derive_metrics_dir` for how the plots directory is resolved.
+    """
+    import matplotlib.pyplot as plt
+    import plotstyle as ps
+
+    ps.use()
+    table = MetricsTable.load(Path(run_dir) / "metrics.csv")
+    plots_dir = derive_metrics_dir(run_dir, out_dir, default_base)
+    plots_dir.mkdir(parents=True, exist_ok=True)
+
+    figures = [("overview", _overview_figure(table))]
+    for stage in _stages(table.columns):
+        figures.append((f"{stage}_loss", _loss_figure(table, stage)))
+    figures.append(("lr", _lr_figure(table)))
+    for stage in _stages(table.columns):
+        figures.append((f"{stage}_accuracy", _accuracy_figure(table, stage)))
+    figures.append(("grad_norm", _grad_norm_figure(table)))
+    for stage in _stages(table.columns):
+        figures.append((f"{stage}_router", _router_figure(table, stage)))
+        figures.append((f"{stage}_wgan_balance", _wgan_balance_figure(table, stage)))
+    figures.append(("throughput", _throughput_figure(table)))
+
+    paths: list[Path] = []
+    for name, fig in figures:
+        if fig is None:
+            continue
+        path = plots_dir / name
+        ps.savefig(fig, str(path), formats=("pdf",))
+        plt.close(fig)
+        paths.append(path.with_suffix(".pdf"))
+    return paths

giant/training/stage2_inputs.py

@@ -120,9 +120,17 @@ def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor:
     slot i: `1.0` at `i=0`, `prod_{j<i}(1-fraction_j)` for `i>=1`
     ("no re-derivation needed": the existing
     stick-breaking encoding is already scale-free, so this is derivable from
-    the batch's ground-truth stick logits alone, no `e_sec` required)."""
-    cumprod = torch.cumprod(1.0 - fraction, dim=1)
-    return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1)
+    the batch's ground-truth stick logits alone, no `e_sec` required).
+
+    Forced fp32 regardless of the caller's ambient `train.precision` autocast
+    region: a `cumprod` over `K_MAX` slots in bf16 underflows to zero within a
+    handful of slots, killing `remaining_frac` as a conditioning signal — the
+    numpy encoder (`giant.data.transforms.encode_secondaries`'s stick-breaking
+    twin) already promotes to float64 for exactly this reason (gitea #47)."""
+    with torch.autocast(fraction.device.type, enabled=False):
+        fraction = fraction.float()
+        cumprod = torch.cumprod(1.0 - fraction, dim=1)
+        return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1)
 
 
 def _shift_prev(x: torch.Tensor) -> torch.Tensor:
@@ -202,21 +210,37 @@ def _assemble_stage2_ar_inputs(
     return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)}
 
 
+def _linear_schedule(p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
+    """Linear interpolation from `p_start` (epoch 0) to `p_end` (the final
+    epoch) — standard scheduled sampling (Bengio et al. 2015), shared by
+    every train-time schedule keyed on epoch."""
+    frac = epoch / max(total_epochs - 1, 1)
+    frac = min(max(frac, 0.0), 1.0)
+    return p_start + (p_end - p_start) * frac
+
+
 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`).
     `"always"`/`"never"` are the two degenerate constants; `"scheduled"`
-    linearly interpolates
-    `p_start` (epoch 0) to `p_end` (the final epoch) — standard scheduled
-    sampling (Bengio et al. 2015)."""
+    linearly interpolates `p_start` to `p_end` via `_linear_schedule`."""
     if mode == "always":
         return 1.0
     if mode == "never":
         return 0.0
-    frac = epoch / max(total_epochs - 1, 1)
-    frac = min(max(frac, 0.0), 1.0)
-    return p_start + (p_end - p_start) * frac
+    return _linear_schedule(p_start, p_end, epoch, total_epochs)
+
+
+def _ctx_truth_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
+    """P(condition stage 2 on the TRUE stage-1 outcome rather than a fresh
+    stage-1 sample), for the current epoch (`stage2_model.stage1_context`).
+    `"truth"` is the degenerate constant 1.0; `"sampled"` linearly
+    interpolates `ctx_p_start` to `ctx_p_end` via `_linear_schedule` — the
+    stage-boundary counterpart of `_stage2_tf_prob`."""
+    if mode == "truth":
+        return 1.0
+    return _linear_schedule(p_start, p_end, epoch, total_epochs)
 
 
 def _history_repr_from_ar_sample(

giant/training/trainers.py

@@ -27,10 +27,13 @@ from giant.constants import CONT_SLOT_DIM
 from giant.data.dataset import StepBatch
 from giant.model.network import Router, build_objective, resolve_type_n_classes, stage2_type_dim
 from giant.model.wgan import generator_loss, gradient_penalty
+from giant.sample import sample_stage1
+from giant.training.amp import resolve_autocast
 from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric
 from giant.training.stage2_inputs import (
     _assemble_stage2_ar_inputs_scheduled,
     _assemble_stage2_ar_target,
+    _ctx_truth_prob,
     _gumbel_tau,
     _relax_onehot_type_slice,
     _stage2_tf_prob,
@@ -73,6 +76,41 @@ def _batch_to_device(batch: StepBatch, device: torch.device) -> StepBatch:
     return type(batch)(*(t.to(device) for t in batch))
 
 
+def _type_class_weight_vector(class_counts: dict[int, int], n_classes: int, scheme: str) -> list[float] | None:
+    """Per-class `F.cross_entropy(weight=...)` vector for the stage-2 type
+    head's `class_weighting` (gitea #44), or `None` under `"none"` (the
+    pre-#44 unweighted-CE behavior — the caller must pass that through as
+    `weight=None`, not a vector of ones, so old runs stay bit-identical).
+
+    `"inverse_freq"`: `1 / count` per class, normalized to mean 1 over
+    `n_classes` so switching this on doesn't rescale the type loss against
+    `particle_type.lambda` / the generator loss it's summed with. A class
+    with zero training examples (fewer distinct species than `n_classes - 1`
+    slots) clamps its count to 1 — its weight is otherwise undefined, and
+    since it never appears in a batch's labels the value is inert anyway.
+
+    Raises if `scheme != "none"` and `class_counts` is empty: that means the
+    `TopNMap` behind this run predates gitea #44 (a stale checkpoint's decode
+    map, or a not-yet-rebuilt setup-cache sidecar) and truly has no
+    frequency information to weight by — silently falling back to uniform
+    weights would look like the feature is active when it isn't.
+    """
+    if scheme == "none":
+        return None
+    if not class_counts:
+        raise ValueError(
+            f"stage2_model.particle_type.class_weighting = {scheme!r} requires "
+            "per-class counts, but this run's sec_type_topn_map has none "
+            "(class_counts={}) — it was built before gitea #44 or loaded "
+            "from a stale setup-cache sidecar/checkpoint; rebuild the setup "
+            "cache (giant train --rebuild-setup-cache) or retrain."
+        )
+    counts = [max(class_counts.get(i, 0), 1) for i in range(n_classes)]
+    inv = [1.0 / c for c in counts]
+    mean_inv = sum(inv) / len(inv)
+    return [w / mean_inv for w in inv]
+
+
 @dataclass(frozen=True)
 class StageSpec:
     """One stage's resolved training configuration.
@@ -87,6 +125,10 @@ class StageSpec:
     generator: str
     decoder: str = "one_shot"
 
+    # partial-retrain (gitea #42)
+    init_from: str = ""
+    freeze: bool = False
+
     # loss weights
     lambda_weight: float = 1.0
     n_sec_lambda: float = 0.1
@@ -95,6 +137,10 @@ class StageSpec:
     # particle-type target (stage 2 only)
     particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
     particle_type_n_classes: int = 16
+    # Resolved by from_config from sec_type_class_counts (dataset-derived,
+    # not itself a cfg value — see _type_class_weight_vector) crossed with
+    # particle_type.class_weighting (gitea #44). None under "none".
+    type_class_weights: list[float] | None = None
 
     # optimization
     lr: float = 3e-4
@@ -103,6 +149,7 @@ class StageSpec:
     warmup_epochs: int = 0
     epochs: int = 1
     steps_per_epoch: int = 1
+    precision: str = "fp32"
 
     # routing auxiliaries
     lambda_balance: float = 0.0
@@ -117,6 +164,11 @@ class StageSpec:
     tf_p_end: float = 1.0
     ar_sample_steps: int = 10
 
+    # stage-1/stage-2 boundary (stage 2 only)
+    stage1_context: str = "truth"
+    ctx_p_start: float = 1.0
+    ctx_p_end: float = 0.0
+
     # generator-specific
     ddpm_n_steps: int = 1000
     n_critic: int = 5
@@ -126,7 +178,18 @@ 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":
+    def from_config(
+        cls,
+        cfg: dict,
+        name: str,
+        is_stage2: bool,
+        steps_per_epoch: int,
+        sec_type_class_counts: dict[int, int] | None = None,
+    ) -> "StageSpec":
+        """`sec_type_class_counts` is dataset-derived (`sec_type_topn_map.class_counts`,
+        gitea #44), not a `cfg` value — it's the one input to `StageSpec` that
+        doesn't come from `cfg`, kept separate from the "only place that
+        reads `cfg`" invariant below on purpose."""
         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
@@ -139,17 +202,23 @@ class StageSpec:
         # 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"])
+        particle_type_n_classes = resolve_type_n_classes(
+            s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"]
+        )
         return cls(
             name=name,
             is_stage2=is_stage2,
             generator=stage_spec.generator,
             decoder=s2_spec.decoder if is_stage2 else "one_shot",
+            init_from=stage_spec.init_from,
+            freeze=stage_spec.freeze,
             lambda_weight=stage_spec.lambda_weight,
             n_sec_lambda=s2_spec.n_sec.lambda_weight,
             n_sec_mode=s2_spec.n_sec.mode,
             particle_type=s2_spec.particle_type,
-            particle_type_n_classes=resolve_type_n_classes(
-                s2_spec.particle_type, cfg["conditioning"]["particle"]["emb_dim"]
+            particle_type_n_classes=particle_type_n_classes,
+            type_class_weights=_type_class_weight_vector(
+                sec_type_class_counts or {}, particle_type_n_classes, s2_spec.particle_type.class_weighting
             ),
             # train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
             # (giant/config.py), so TrainConfig.from_dict never has to fall
@@ -161,6 +230,7 @@ class StageSpec:
             warmup_epochs=t.warmup_epochs,
             epochs=t.epochs,
             steps_per_epoch=max(steps_per_epoch, 1),
+            precision=t.precision,
             lambda_balance=stage_spec.router.lambda_balance,
             lambda_proc=stage_spec.router.lambda_proc,
             lambda_entropy=stage_spec.router.lambda_entropy,
@@ -169,6 +239,9 @@ class StageSpec:
             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,
+            stage1_context=s2_spec.stage1_context if is_stage2 else cls.stage1_context,
+            ctx_p_start=s2_spec.ctx_p_start if is_stage2 else cls.ctx_p_start,
+            ctx_p_end=s2_spec.ctx_p_end if is_stage2 else cls.ctx_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
@@ -186,13 +259,16 @@ class StageSpec:
 class StageTrainer:
     """One active stage's optimizer(s), EMA, and per-batch step.
 
-    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.
+    Reads only the shared `StepBatch` (`giant.data.dataset`) by default — stage
+    2 conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context =
+    "truth"`, stage-level teacher forcing), so "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. `stage2_model.stage1_context = "sampled"` is the one
+    exception — `build_stage_trainers` wires the stage-2 trainer to the
+    stage-1 one via `attach_stage1` so it can draw a real stage-1 sample
+    (`giant.sample.sample_stage1`) instead, scheduled by `ctx_p_start`/
+    `ctx_p_end` (see `_stage1_context`).
 
     Grad-norm clipping is per-stage here — v0.2's single shared optimizer
     clipped both stages' gradients jointly; splitting per stage is a small,
@@ -231,11 +307,18 @@ class StageTrainer:
         self.is_stage2 = spec.is_stage2
         self.generator = spec.generator
         self.decoder = spec.decoder
+        self.frozen = spec.freeze
         self.device = device
         self.model = model.to(device)
         self.router = _stage_router(self.model)
         self._modules = (self.model, *extra_modules)
 
+        # Resolved once (not re-derived every step) — see
+        # giant.training.amp.resolve_autocast (gitea #47).
+        self._autocast_device_type, self._autocast_dtype, self._autocast_enabled = resolve_autocast(
+            spec.precision, device
+        )
+
         self.particle_type_cfg = spec.particle_type
         self.particle_type_n_classes = spec.particle_type_n_classes
         self.ema_decay = spec.ema_decay
@@ -246,6 +329,18 @@ class StageTrainer:
             for p in self.ema_model.parameters():
                 p.requires_grad_(False)
 
+        #: Set by `attach_stage1` when `stage2_model.stage1_context =
+        #: "sampled"` — the stage-1 `StageTrainer` this (stage-2) trainer
+        #: draws its context sample from. `None` for stage 1 itself, and for
+        #: stage 2 under "truth".
+        self.stage1_source: "StageTrainer | None" = None
+
+    def attach_stage1(self, stage1_trainer: "StageTrainer") -> None:
+        """Wires this (stage-2) trainer to the stage-1 trainer it should
+        sample from under `stage2_model.stage1_context = "sampled"` — see
+        `build_stage_trainers`."""
+        self.stage1_source = stage1_trainer
+
     # --- schedule -------------------------------------------------------
 
     def _init_lr_schedule(self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int) -> None:
@@ -289,6 +384,58 @@ class StageTrainer:
         for module in self._modules:
             module.eval()
 
+    # --- stage-1/stage-2 boundary (shared by both trainer subclasses) ---
+
+    def _stage1_context(
+        self,
+        x1_s1: torch.Tensor,
+        cond_cont: torch.Tensor,
+        cond_cat: torch.Tensor,
+        epoch: int | None,
+    ) -> torch.Tensor:
+        """The stage-1 outcome stage 2 conditions on this batch.
+
+        `epoch=None` means "always ground truth" regardless of
+        `spec.stage1_context` — the same val-loss convention `_ar_inputs`
+        uses, so validation stays a stable, non-stochastic comparison.
+        Otherwise, under `stage1_context = "sampled"`, each example
+        independently uses the ground truth with probability `p_truth`
+        (`_ctx_truth_prob`, ramped by `ctx_p_start`/`ctx_p_end`) and a fresh
+        `giant.sample.sample_stage1` draw from `stage1_source.sampling_model()`
+        otherwise — a real sampling pass, not a cheap proxy, matching
+        `_assemble_stage2_ar_inputs_scheduled`'s precedent for the equivalent
+        in-stage-2 self-sample. Mixed per example (not per-dimension): a row
+        is either the real ground-truth 9D vector or a real sample, never an
+        elementwise blend of the two.
+        """
+        x1_s1 = x1_s1.detach()
+        if self.stage1_source is None or epoch is None:
+            return x1_s1
+        p_truth = _ctx_truth_prob(
+            self.spec.stage1_context,
+            self.spec.ctx_p_start,
+            self.spec.ctx_p_end,
+            epoch,
+            self.spec.epochs,
+        )
+        if p_truth >= 1.0:
+            return x1_s1
+
+        stage1_model = self.stage1_source.sampling_model()
+        was_training = stage1_model.training
+        sampled, _ = sample_stage1(
+            stage1_model,
+            cond_cont,
+            cond_cat,
+            steps=self.spec.ar_sample_steps,
+            ddpm_steps=self.stage1_source.spec.ddpm_n_steps,
+        )
+        if was_training:
+            stage1_model.train()
+
+        use_truth = torch.rand(x1_s1.size(0), 1, device=x1_s1.device) < p_truth
+        return torch.where(use_truth, x1_s1, sampled).detach()
+
     # --- stage-2 secondary assembly (shared by both trainer subclasses) ---
 
     def _ar_inputs(
@@ -434,14 +581,36 @@ class StageTrainer:
         stop_acc = (((logits >= 0).float() == target).float() * mask_f).sum() / denom
         return l_stop, stop_acc
 
-    @staticmethod
-    def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
+    def _autocast(self) -> torch.autocast:
+        """The training-step autocast region (`train.precision`, gitea #47).
+
+        Only wraps forward/loss computation — `backward()`/`optimizer.step()`
+        stay outside, and `val_loss` never calls this at all, so validation
+        (and the best-checkpoint selection it drives) stays precision-
+        independent and comparable against every fp32-only run recorded so
+        far. `enabled=False` under `precision = "fp32"` (the default) makes
+        this a true no-op, so callers never need to branch on precision
+        themselves."""
+        return torch.autocast(
+            self._autocast_device_type,
+            dtype=self._autocast_dtype,
+            enabled=self._autocast_enabled,
+        )
+
+    def _step_optimizer(self, 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."""
+        the pre-clip grad norm. The one place the grad-clip constant lives.
+
+        `self.frozen` (`stage{1,2}_model.freeze`, gitea #42) skips only the
+        final `optimizer.step()` — backward/clip still run so loss/grad_norm
+        stay meaningful to watch, but the stage's weights (and, for a WGAN
+        stage, its critic's — this same method is both trainers' single
+        optimizer-step choke point) never move."""
         optimizer.zero_grad()
         loss.backward()
         grad_norm = torch.nn.utils.clip_grad_norm_(params, 1.0)
-        optimizer.step()
+        if not self.frozen:
+            optimizer.step()
         return grad_norm.item()
 
     def _extra_state(self) -> dict:
@@ -504,6 +673,12 @@ class FlowDDPMStageTrainer(StageTrainer):
         # "onehot"/"embedding" pull it out into model.type_head instead (0
         # here).
         self._flow_type_dim = None if self.particle_type_cfg.target == "physical" else 0
+        # gitea #44: None under class_weighting = "none" (the default),
+        # matching F.cross_entropy's own unweighted default — a real tensor
+        # only materializes when the config asked for one.
+        self.type_class_weights = (
+            None if spec.type_class_weights is None else torch.tensor(spec.type_class_weights, device=device)
+        )
 
         self.params = list(self.model.parameters())
         self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
@@ -597,8 +772,12 @@ class FlowDDPMStageTrainer(StageTrainer):
         mask = sec_mask.float()
         denom = mask.sum().clamp(min=1)
         if self.particle_type_cfg.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, weight=self.type_class_weights, reduction="none"
+            )
             l_type = (ce * mask).sum() / denom
+            # Unweighted, deliberately — type_acc is a diagnostic of raw
+            # per-slot correctness, not the (possibly class-weighted) loss.
             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()
@@ -608,9 +787,10 @@ class FlowDDPMStageTrainer(StageTrainer):
 
     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
-        the training `step` path schedules `p_tf` by epoch."""
+        forcing (`p_tf=1.0`) and the ground-truth stage-1 context, regardless
+        of `spec.teacher_forcing`/`spec.stage1_context` — validation should
+        stay a stable, non-stochastic ground-truth comparison; only the
+        training `step` path schedules `p_tf`/`p_truth` by epoch."""
         (
             cond_cont,
             cond_cat,
@@ -621,7 +801,7 @@ class FlowDDPMStageTrainer(StageTrainer):
             sec_type_idx,
         ) = _batch_to_device(batch, device)
         sec_mask = self._sec_mask(n_sec, sec_cont.size(1), device)
-        stage1_ctx = x1_s1.detach()
+        stage1_ctx = self._stage1_context(x1_s1, cond_cont, cond_cat, epoch)
 
         x1_s2 = None
         ar_inputs = None
@@ -689,10 +869,12 @@ class FlowDDPMStageTrainer(StageTrainer):
                 self.spec.gumbel_tau_end,
             )
         epoch = global_step // self.spec.steps_per_epoch
-        out = self._compute(batch, device, epoch=epoch)
+        with self._autocast():
+            out = self._compute(batch, device, epoch=epoch)
         grad_norm = self._step_optimizer(self.optimizer, out["loss"], self.params)
-        self.lr_sched.step()
-        if self.ema_model is not None:
+        if not self.frozen:
+            self.lr_sched.step()
+        if self.ema_model is not None and not self.frozen:
             _update_ema(self.ema_model, self.model, self.ema_decay)
         stats = {key: value.item() for key, value in out.items()}
         stats["grad_norm"] = grad_norm
@@ -846,53 +1028,58 @@ class WGANStageTrainer(StageTrainer):
             sec_type_idx,
         ) = _batch_to_device(batch, device)
         B = cond_cont.size(0)
-        stage1_ctx = x1_s1.detach()
+        epoch = global_step // self.spec.steps_per_epoch
+        stage1_ctx = self._stage1_context(x1_s1, cond_cont, cond_cat, epoch)
         grad_probe: dict[str, float] = {}
 
         ar_inputs = None
-        if not self.is_stage2:
-            real = x1_s1
+        with self._autocast():
+            if not self.is_stage2:
+                real = x1_s1
 
-            def critic_fn(x):
-                return self.critic(x, cond_cont, cond_cat)
+                def critic_fn(x):
+                    return self.critic(x, cond_cont, cond_cat)
 
-            z = torch.randn(B, self.model.noise_dim, device=device)
-            fake = self.model(z, cond_cont, cond_cat)
-            mask = None
-        else:
-            real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake(
-                _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
-                stage1_ctx,
-                global_step,
-                device,
-            )
-            if self.particle_type_cfg.target == "onehot":
-                # Straight-through Gumbel-softmax relaxation of the type
-                # slice only — the critic must see a hard one-hot forward
-                # (matching what "real" data looks like) while gradient
-                # still flows smoothly to the generator. grad_probe captures
-                # the gradient-magnitude instrumentation — see
-                # _relax_onehot_type_slice's docstring.
-                tau = _gumbel_tau(
+                z = torch.randn(B, self.model.noise_dim, device=device)
+                fake = self.model(z, cond_cont, cond_cat)
+                mask = None
+            else:
+                real, fake_raw, mask, critic_fn, ar_inputs = self._stage2_real_and_fake(
+                    _Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
+                    stage1_ctx,
                     global_step,
-                    self.total_steps,
-                    self.spec.type_gumbel_tau_start,
-                    self.spec.type_gumbel_tau_end,
+                    device,
                 )
-                fake_raw = _relax_onehot_type_slice(
-                    fake_raw,
-                    sec_cont.size(1),
-                    CONT_SLOT_DIM,
-                    stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes),
-                    tau,
-                    grad_probe=grad_probe,
-                )
-            fake = fake_raw * mask
+                if self.particle_type_cfg.target == "onehot":
+                    # Straight-through Gumbel-softmax relaxation of the type
+                    # slice only — the critic must see a hard one-hot forward
+                    # (matching what "real" data looks like) while gradient
+                    # still flows smoothly to the generator. grad_probe captures
+                    # the gradient-magnitude instrumentation — see
+                    # _relax_onehot_type_slice's docstring.
+                    tau = _gumbel_tau(
+                        global_step,
+                        self.total_steps,
+                        self.spec.type_gumbel_tau_start,
+                        self.spec.type_gumbel_tau_end,
+                    )
+                    fake_raw = _relax_onehot_type_slice(
+                        fake_raw,
+                        sec_cont.size(1),
+                        CONT_SLOT_DIM,
+                        stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes),
+                        tau,
+                        grad_probe=grad_probe,
+                    )
+                fake = fake_raw * mask
 
-        # --- critic step (every batch) ---
-        fake_detached = fake.detach()
-        real_score = critic_fn(real)
-        fake_score = critic_fn(fake_detached)
+            # --- critic step (every batch) ---
+            fake_detached = fake.detach()
+            real_score = critic_fn(real)
+            fake_score = critic_fn(fake_detached)
+
+        # gradient_penalty forces its own fp32 region internally (see its
+        # docstring) regardless of the ambient autocast above.
         gp = gradient_penalty(critic_fn, real, fake_detached, mask=mask)
         d_loss = fake_score.mean() - real_score.mean() + self.gp_weight * gp
         wasserstein = (real_score.mean() - fake_score.mean()).detach()
@@ -901,8 +1088,9 @@ 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_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)
+        with self._autocast():
+            l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
+            l_stop, stop_acc = self._stop_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device, ar_inputs)
 
         # On a non-generator-step batch with no n_sec_head/stop_head on this
         # stage (n_sec now defaults to stage 2), there's nothing for
@@ -910,7 +1098,8 @@ class WGANStageTrainer(StageTrainer):
         # be a graph-less zero tensor, which .backward() rejects outright.
         skip_g_step = not did_g_step and self.model.n_sec_head is None and self.model.stop_head is None
         if did_g_step:
-            g_loss_adv = generator_loss(critic_fn, fake)
+            with self._autocast():
+                g_loss_adv = generator_loss(critic_fn, fake)
             g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * (l_nsec + l_stop)
         else:
             g_loss_adv = torch.zeros((), device=device)
@@ -921,8 +1110,9 @@ class WGANStageTrainer(StageTrainer):
             grad_norm_g = self._step_optimizer(self.optimizer, g_loss, self.g_params)
 
         if did_g_step:
-            self.lr_sched.step()
-            if self.ema_model is not None:
+            if not self.frozen:
+                self.lr_sched.step()
+            if self.ema_model is not None and not self.frozen:
                 _update_ema(self.ema_model, self.model, self.ema_decay)
 
         return {
@@ -994,15 +1184,27 @@ def build_stage_trainers(
     critics: dict[str, torch.nn.Module | None],
     device: torch.device,
     total_train_batches: int,
+    sec_type_class_counts: dict[int, int] | None = None,
 ) -> dict[str, StageTrainer]:
     """One trainer per active stage — `models[name] is None` means that stage
-    is `active = false` and is simply never constructed."""
+    is `active = false` and is simply never constructed.
+
+    `stage2_model.stage1_context = "sampled"` additionally wires the
+    stage-2 trainer to the stage-1 one (`StageTrainer.attach_stage1`) so it
+    can draw a real stage-1 sample instead of only ever seeing the
+    ground-truth stage-1 outcome — `validate_config` already guarantees both
+    stages are active whenever that config value is set.
+
+    `sec_type_class_counts` (`sec_type_topn_map.class_counts`, gitea #44) is
+    the one dataset-derived input `StageSpec.from_config` needs beyond `cfg`
+    — `None`/absent whenever `stage2_model.particle_type.class_weighting =
+    "none"` (the default), which never reads it."""
     trainers: dict[str, StageTrainer] = {}
     for name, is_stage2 in (("stage1", False), ("stage2", True)):
         model = models.get(name)
         if model is None:
             continue
-        spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1))
+        spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1), sec_type_class_counts)
         if build_objective(spec.generator).is_adversarial:
             critic = critics.get(name)
             assert critic is not None, (
@@ -1011,4 +1213,9 @@ def build_stage_trainers(
             trainers[name] = WGANStageTrainer(spec, model, critic, device)
         else:
             trainers[name] = FlowDDPMStageTrainer(spec, model, device)
+
+    stage2 = trainers.get("stage2")
+    stage1 = trainers.get("stage1")
+    if stage2 is not None and stage1 is not None and stage2.spec.stage1_context == "sampled":
+        stage2.attach_stage1(stage1)
     return trainers

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.2"
+version = "0.3.9"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
@@ -26,6 +26,8 @@ dev = [
     "pytest-cov>=5,<8",
     "ruff>=0.15,<1",
     "ty>=0.0.50,<0.1",
+    "bump-my-version>=1.2,<2",
+    "git-cliff>=2,<3",
     "giant[convert,analysis,geometry,wandb]",
 ]
 geometry = [

tests/test_amp.py

@@ -0,0 +1,136 @@
+"""Tests for giant/training/amp.py (gitea #47)."""
+
+import tempfile
+from pathlib import Path
+
+import pytest
+import torch
+
+from giant.model.routers import EnergyRouter
+from giant.model.wgan import gradient_penalty
+from giant.training.amp import resolve_autocast
+from giant.training.stage2_inputs import _remaining_energy_fraction
+from test_train import _base_cfg, _run_train
+
+# ---------------------------------------------------------------------------
+# resolve_autocast
+# ---------------------------------------------------------------------------
+
+
+def test_resolve_autocast_fp32_is_disabled():
+    device_type, dtype, enabled = resolve_autocast("fp32", torch.device("cpu"))
+    assert device_type == "cpu"
+    assert dtype is torch.float32
+    assert enabled is False
+
+
+def test_resolve_autocast_bf16_on_cpu_is_enabled():
+    """CPU bf16 autocast is what lets the mixed-precision path be tested
+    without a GPU (torch 2.3 supports it)."""
+    device_type, dtype, enabled = resolve_autocast("bf16", torch.device("cpu"))
+    assert device_type == "cpu"
+    assert dtype is torch.bfloat16
+    assert enabled is True
+
+
+def test_resolve_autocast_bf16_on_unsupported_cuda_raises(monkeypatch):
+    monkeypatch.setattr(torch.cuda, "is_bf16_supported", lambda: False)
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda device=None: (7, 0))
+    monkeypatch.setattr(torch.cuda, "get_device_name", lambda device=None: "Tesla V100")
+    with pytest.raises(ValueError, match="bf16"):
+        resolve_autocast("bf16", torch.device("cuda"))
+
+
+def test_resolve_autocast_bf16_on_mps_raises():
+    with pytest.raises(ValueError, match="bf16"):
+        resolve_autocast("bf16", torch.device("mps"))
+
+
+def test_resolve_autocast_unknown_precision_raises():
+    with pytest.raises(ValueError, match="fp32.*bf16"):
+        resolve_autocast("fp16", torch.device("cpu"))
+
+
+# ---------------------------------------------------------------------------
+# End-to-end: train() under bf16 on CPU
+# ---------------------------------------------------------------------------
+
+
+def test_train_end_to_end_bf16_cpu_completes_and_stores_fp32_params():
+    """Reuses tests/test_train.py's synthetic-batch harness — train() itself
+    is device-agnostic, and CPU bf16 autocast is real (not mocked) in torch
+    2.3, so this is a genuine exercise of the autocast region added to
+    FlowDDPMStageTrainer.step/WGANStageTrainer.step, not just a config
+    passthrough check.
+
+    Also asserts the checkpoint's stored parameters are fp32: autocast only
+    changes the dtype of intermediate activations, never the model's own
+    stored weights — a regression here would mean something accidentally
+    cast the model itself (e.g. `model.to(dtype=torch.bfloat16)`) rather than
+    using autocast."""
+    cfg = _base_cfg()
+    cfg["train"]["precision"] = "bf16"
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        assert (out_dir / "metrics.csv").exists()
+        ckpt = torch.load(out_dir / "last.pt", weights_only=False)
+        for stage_key in ("model", "sec_decoder"):
+            if stage_key not in ckpt:
+                continue
+            for name, tensor in ckpt[stage_key].items():
+                if tensor.is_floating_point():
+                    assert tensor.dtype == torch.float32, f"{stage_key}.{name} is {tensor.dtype}, expected fp32"
+
+
+@pytest.mark.parametrize("generator", ["wgan", "flow"])
+def test_train_end_to_end_bf16_cpu_stage2_generators(generator):
+    """bf16 covers both trainer subclasses (FlowDDPMStageTrainer and
+    WGANStageTrainer) — the wgan default in _base_cfg exercises the
+    generator-forward/critic-scoring autocast region added to
+    WGANStageTrainer.step, and flow exercises the plain _compute wrap."""
+    cfg = _base_cfg()
+    cfg["train"]["precision"] = "bf16"
+    cfg["stage2_model"]["generator"] = generator
+    with tempfile.TemporaryDirectory() as tmp:
+        _run_train(cfg, Path(tmp) / "run")
+
+
+# ---------------------------------------------------------------------------
+# fp32 guards: correct in fp32, quietly degrade in bf16 — stay fp32 even
+# under an active bf16 autocast region.
+# ---------------------------------------------------------------------------
+
+
+def test_remaining_energy_fraction_stays_fp32_under_bf16_autocast():
+    fraction = torch.rand(4, 5).to(torch.bfloat16)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        out = _remaining_energy_fraction(fraction)
+    assert out.dtype == torch.float32
+
+
+def test_gradient_penalty_stays_fp32_under_bf16_autocast():
+    critic = torch.nn.Linear(6, 1)
+
+    def critic_fn(x):
+        return critic(x)
+
+    real = torch.randn(4, 6)
+    fake = torch.randn(4, 6)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        gp = gradient_penalty(critic_fn, real, fake)
+    assert gp.dtype == torch.float32
+
+
+def test_router_balance_and_entropy_loss_stay_fp32_under_bf16_autocast():
+    router = EnergyRouter(n_experts=3)
+    cond_cont = torch.randn(8, 15)
+    cond_cat = torch.zeros(8, 2, dtype=torch.long)
+    with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+        balance = router.balance_loss(cond_cont, cond_cat)
+        entropy = router.entropy_loss(cond_cont, cond_cat)
+        weights = router.combine_weights(cond_cont, cond_cat)
+    assert balance.dtype == torch.float32
+    assert entropy.dtype == torch.float32
+    assert weights.dtype == torch.float32

tests/test_analysis_reduce.py

@@ -10,9 +10,11 @@ from giant.analysis import reduce as R
 from giant.analysis.sources import (
     SYNTHETIC_TERMINATION_REASONS,
     Side,
+    open_side,
     physical_steps,
     secondaries,
 )
+from giant.data.loader import EVENT_ID_FILE_STRIDE
 
 
 def _rollout_frame() -> pl.LazyFrame:
@@ -157,6 +159,20 @@ def test_secondaries_rollout_vs_reference_align():
     assert t["pdg"].to_list() == [22, 22]
 
 
+def test_sec_count_by_event_zero_fills_events_with_no_secondaries():
+    r_phys = physical_steps(_rollout_frame(), Side.rollout)
+    r_sec = secondaries(_rollout_frame(), Side.rollout)
+    ev, n = R.sec_count_by_event(r_phys, r_sec)
+    # event 1 has one secondary track; event 2 has none and must still appear (as 0),
+    # not silently drop out of a plain group_by on the secondaries frame alone.
+    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}
+
+    t_all = _reference_frame()
+    t_sec = secondaries(t_all, Side.reference)
+    ev, n = R.sec_count_by_event(t_all, t_sec)
+    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1}
+
+
 def test_leakage_fraction():
     frac = R.leakage_fraction(_rollout_frame())
     # event 1: escaped pre_E=30, deposited=90 -> 30/120 = 0.25; event 2: 0
@@ -194,3 +210,35 @@ def test_pdg_and_material_labels():
     assert G.pdg_label(22) == "gamma"
     assert G.pdg_label(999999) == "999999"
     assert G.material_label("G4_PbWO4") == "PbWO4"
+
+
+def _write_shard(path, event_ids, edeps):
+    pl.DataFrame({"event_id": event_ids, "pdg": [11] * len(event_ids), "edep": edeps}).write_parquet(path)
+
+
+def test_open_side_reference_offsets_event_ids_across_shards(tmp_path):
+    # Each shard is a separate Geant4 job whose own event_id numbering restarts
+    # from 0 — a naive multi-shard scan collides on event_id across shards.
+    _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
+    _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
+    df = open_side(tmp_path, Side.reference).sort("event_id").collect()
+    assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
+    assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]
+    assert "__source_path" not in df.columns
+
+
+def test_open_side_reference_single_file_unchanged(tmp_path):
+    _write_shard(tmp_path / "only.parquet", [0, 1], [1.0, 2.0])
+    df = open_side(tmp_path / "only.parquet", Side.reference).sort("event_id").collect()
+    assert df["event_id"].to_list() == [0, 1]
+    assert "__source_path" not in df.columns
+
+
+def test_open_side_reference_manifest(tmp_path):
+    _write_shard(tmp_path / "a.parquet", [0, 1], [1.0, 2.0])
+    _write_shard(tmp_path / "b.parquet", [0, 1], [3.0, 4.0])
+    manifest = tmp_path / "shards.manifest"
+    manifest.write_text("a.parquet\nb.parquet\n")
+    df = open_side(manifest, Side.reference).sort("event_id").collect()
+    assert df["event_id"].to_list() == [0, 1, EVENT_ID_FILE_STRIDE, EVENT_ID_FILE_STRIDE + 1]
+    assert df["edep"].to_list() == [1.0, 2.0, 3.0, 4.0]

tests/test_catalog.py

@@ -6,14 +6,29 @@ import numpy as np
 import pytest
 
 from giant.analysis import build_catalog, catalog_ids, get_spec
-from giant.analysis.catalog import Bundle, PlotSpec
+from giant.analysis.catalog import (
+    Bundle,
+    PlotSpec,
+    _containment_depths,
+    _integer_confusion,
+    _ks_statistic,
+)
 from giant.analysis.context import Context, build_context
+from giant.analysis.sources import RolloutSpec
 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(
+        [RolloutSpec("rollout", r)], t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
+    )
+
+
+def _two_rollout_specs() -> list[RolloutSpec]:
+    # Two distinct rollout sources so multi-series merging/finalize code is
+    # exercised even though the underlying frame is the same fixture.
+    return [RolloutSpec("flow", _rollout_frame()), RolloutSpec("wgan", _rollout_frame())]
 
 
 @pytest.fixture(scope="module")
@@ -21,9 +36,20 @@ def ctx() -> Context:
     return _build_ctx()
 
 
+@pytest.fixture(scope="module")
+def two_ctx() -> Context:
+    t = _reference_frame()
+    return build_context(_two_rollout_specs(), t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
+
+
 @pytest.fixture(scope="module")
 def bundle(ctx: Context) -> Bundle:
-    return Bundle.open(_rollout_frame(), _reference_frame(), ctx)
+    return Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)
+
+
+@pytest.fixture(scope="module")
+def two_bundle(two_ctx: Context) -> Bundle:
+    return Bundle.open(_two_rollout_specs(), _reference_frame(), two_ctx)
 
 
 def test_catalog_ids_unique_and_nonempty():
@@ -53,41 +79,76 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle):
             "single_hist",
             "router_gating",
             "router_share",
+            "router_specialization",
+            "heatmap",
             "unavailable",
         }
         assert r.title and r.xlabel
-        _validate_payload(r)
+        _validate_payload(r, ["rollout"])
 
 
-def _validate_payload(r) -> None:
+def test_every_spec_computes_valid_reduced_with_two_rollouts(two_bundle: Bundle):
+    for spec in build_catalog():
+        r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
+        assert r.id == spec.id
+        _validate_payload(r, ["flow", "wgan"])
+
+
+def _validate_payload(r, names: list[str]) -> None:
     p = r.payload
     if r.kind == "overlay_hist":
         n = len(p["edges"]) - 1
-        assert len(p["rollout"]) == n and len(p["reference"]) == n
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(v) == n
+        assert len(p["reference"]) == n
     elif r.kind == "single_hist":
-        assert len(p["rollout"]) == len(p["edges"]) - 1
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(v) == len(p["edges"]) - 1
     elif r.kind == "grouped_hist":
         n = len(p["edges"]) - 1
         assert p["groups"], "grouped hist must have at least one group"
         for g in p["groups"].values():
-            assert len(g["rollout"]) == n and len(g["reference"]) == n
+            assert list(g["series"]) == names
+            for v in g["series"].values():
+                assert len(v) == n
+            assert len(g["reference"]) == n
     elif r.kind == "profile":
         n = len(p["edges"]) - 1
-        for k in ("rollout_mean", "rollout_std", "reference_mean", "reference_std"):
-            assert len(p[k]) == n
+        assert list(p["series"]) == names
+        for side in p["series"].values():
+            assert len(side["mean"]) == n and len(side["std"]) == n
+        assert len(p["reference"]["mean"]) == n and len(p["reference"]["std"]) == n
     elif r.kind == "bar":
-        assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(p["labels"]) == len(v)
+        assert len(p["labels"]) == len(p["reference"])
     elif r.kind == "unavailable":
         assert p["note"]
     elif r.kind == "router_gating":
-        for side in ("rollout", "reference"):
-            if side in p:
-                assert len(p[side]["centers"]) == len(p[side]["means"])
-    elif r.kind == "router_share":
-        for cat in p["categories"]:
+        for entry in p["series"].values():
             for side in ("rollout", "reference"):
-                if side in p:
-                    assert cat in p[side]
+                if side in entry:
+                    assert len(entry[side]["centers"]) == len(entry[side]["means"])
+    elif r.kind == "router_share":
+        for entry in p["series"].values():
+            for cat in entry["categories"]:
+                for side in ("rollout", "reference"):
+                    if side in entry:
+                        assert cat in entry[side]
+    elif r.kind == "router_specialization":
+        for entry in p["series"].values():
+            for side in ("rollout", "reference"):
+                if side in entry:
+                    assert len(entry[side]["centers"]) == len(entry[side]["score"])
+    elif r.kind == "heatmap":
+        assert list(p["series"]) == names
+        for mat in p["series"].values():
+            assert len(mat) == len(p["row_labels"])
+            for row in mat:
+                assert len(row) == len(p["col_labels"])
 
 
 # ---------------------------------------------------------------------------
@@ -98,7 +159,10 @@ def _validate_payload(r) -> None:
 # sec_count_per_species via pdg-keyed sums), concat-then-finalize with
 # data-dependent edges (event_total_edep), concat-then-mean/std (shower_
 # longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
-# ratio (species_edep_share), and a chunkable=False passthrough (router_gating).
+# ratio (species_edep_share), a chunkable=False passthrough (router_gating),
+# nested sum-merge into a scorecard (marginal_distance_summary), concat-then-
+# event-id-join (n_sec_confusion), and concat-then-per-event-derived-quantity
+# (shower_containment_depth_90, reusing the profile matrix's own merge shape).
 _CHUNK_EQUIVALENCE_IDS = [
     "marginal_edep",
     "species_edep_share",
@@ -107,6 +171,9 @@ _CHUNK_EQUIVALENCE_IDS = [
     "leakage_fraction",
     "sec_count_per_species",
     "router_gating",
+    "marginal_distance_summary",
+    "n_sec_confusion",
+    "shower_containment_depth_90",
 ]
 
 
@@ -128,21 +195,81 @@ def _assert_payload_close(a, b, path: str = "payload") -> None:
 
 
 @pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS)
-def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
+def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
     """A plot computed over N event-disjoint chunks then merged must equal the
     same plot computed in one unchunked pass — the core chunking correctness
-    guarantee (see the analysis-rollout-plots chunking plan)."""
+    guarantee (see the analysis-rollout-plots chunking plan). Exercised with
+    two rollout series so the per-rollout merge path is covered too."""
     spec: PlotSpec = get_spec(spec_id)
-    r, t = _rollout_frame(), _reference_frame()
+    rollouts, t = _two_rollout_specs(), _reference_frame()
 
-    unchunked_bundle = Bundle.open(r, t, ctx)
-    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)
+    unchunked_bundle = Bundle.open(rollouts, t, two_ctx)
+    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], two_ctx)
 
     # 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)]
-    chunked = spec.finalize(parts, ctx)
+    parts = [spec.compute_partial(Bundle.open(rollouts, t, two_ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
+    chunked = spec.finalize(parts, two_ctx)
 
     assert chunked.id == unchunked.id
     assert chunked.kind == unchunked.kind
     _assert_payload_close(unchunked.payload, chunked.payload)
+
+
+# ---------------------------------------------------------------------------
+# new (gitea #76) reductions: KS distance, confusion matrix, containment depth
+# ---------------------------------------------------------------------------
+
+
+def test_ks_statistic():
+    assert _ks_statistic([10, 10], [10, 10]) == 0.0  # identical shape -> 0
+    assert _ks_statistic([10, 0], [0, 10]) == 1.0  # fully disjoint -> 1
+    assert _ks_statistic([0, 0], [0, 0]) != _ks_statistic([0, 0], [0, 0])  # nan (no data either side)
+    assert _ks_statistic([10, 0], [0, 0]) == 1.0  # one side empty, other isn't -> maximal mismatch
+
+
+def test_integer_confusion_matches_event_pairing():
+    # true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
+    labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
+    assert labels == ["0", "1+"]
+    assert mat.tolist() == [[0, 0], [1, 1]]  # row=true, col=pred
+
+
+def test_integer_confusion_caps_pathological_outliers():
+    labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
+    assert labels[-1] == "4+"
+    assert mat.shape == (5, 5)
+    assert mat.sum() == 2
+
+
+def test_integer_confusion_explicit_cap_overrides_local_range():
+    # Even though this pair's own max is 1, an explicit shared cap forces a
+    # wider (and so cross-rollout-consistent) label set.
+    labels, mat = _integer_confusion(np.array([1, 1]), np.array([0, 1]), cap=3)
+    assert labels == ["0", "1", "2", "3+"]
+    assert mat.shape == (4, 4)
+
+
+def test_containment_depths_simple_ramp():
+    # one event, edep concentrated in the first bin -> 90%/95% containment
+    # depth is the first bin's right edge; a zero-energy event is dropped.
+    mat = np.array([[9.0, 1.0, 0.0], [0.0, 0.0, 0.0]])
+    edges = np.array([0.0, 1.0, 2.0, 3.0])
+    depths = _containment_depths(mat, edges, 0.90)
+    assert depths.tolist() == [1.0]
+
+
+def test_n_sec_confusion_spec(bundle):
+    spec = get_spec("n_sec_confusion")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
+    assert r.payload["series"]["rollout"] == [[0, 0], [1, 1]]
+
+
+def test_n_sec_confusion_shares_one_cap_across_rollouts(two_bundle):
+    spec = get_spec("n_sec_confusion")
+    r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
+    assert list(r.payload["series"]) == ["flow", "wgan"]
+    # both rollouts share the same fixture data here, so their matrices (and
+    # the shared label set) must be identical.
+    assert r.payload["series"]["flow"] == r.payload["series"]["wgan"]

tests/test_cli_new_run.py

@@ -91,6 +91,31 @@ def test_dry_run_writes_nothing(tmp_path: Path):
     assert not out_dir.exists()
 
 
+def test_stage1_init_from_and_freeze_flags_scaffold_a_partial_retrain_config(tmp_path: Path):
+    """gitea #42."""
+    out_dir = tmp_path / "run5"
+    result = runner.invoke(
+        app,
+        [
+            "new-run",
+            "--out",
+            str(out_dir),
+            "--stage1-init-from",
+            "ckpt/stage1_good/best.pt",
+            "--stage1-freeze",
+        ],
+    )
+    assert result.exit_code == 0, result.output
+
+    with open(out_dir / "config.toml", "rb") as f:
+        cfg = tomllib.load(f)
+
+    assert cfg["stage1_model"]["init_from"] == "ckpt/stage1_good/best.pt"
+    assert cfg["stage1_model"]["freeze"] is True
+    assert cfg["stage2_model"]["init_from"] == ""
+    assert cfg["stage2_model"]["freeze"] is False
+
+
 def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
     out_dir = tmp_path / "run5"
     out_dir.mkdir()

tests/test_cli_train_overrides.py

@@ -97,6 +97,33 @@ def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
     assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
 
 
+def test_stage1_init_from_and_freeze_flags_land_in_cfg_and_dont_touch_stage2(monkeypatch, tmp_path):
+    """gitea #42: --stage{1,2}-init-from/--stage{1,2}-freeze are stage-scoped
+    only. --stage1-freeze alone would fail validate_config (freeze requires
+    init_from or --resume), so both flags are passed together here."""
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--stage1-init-from", "ckpt/stage1_good/best.pt", "--stage1-freeze"],
+    )
+    assert cfg["stage1_model"]["init_from"] == "ckpt/stage1_good/best.pt"
+    assert cfg["stage1_model"]["freeze"] is True
+    assert cfg["stage2_model"]["init_from"] == ""
+    assert cfg["stage2_model"]["freeze"] is False
+
+
+def test_stage2_init_from_and_freeze_flags_land_in_cfg_and_dont_touch_stage1(monkeypatch, tmp_path):
+    cfg = _invoke_and_capture_cfg(
+        monkeypatch,
+        tmp_path,
+        ["--stage2-init-from", "ckpt/stage2_good/best.pt", "--stage2-freeze"],
+    )
+    assert cfg["stage2_model"]["init_from"] == "ckpt/stage2_good/best.pt"
+    assert cfg["stage2_model"]["freeze"] is True
+    assert cfg["stage1_model"]["init_from"] == ""
+    assert cfg["stage1_model"]["freeze"] is False
+
+
 def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
     monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
     result = runner.invoke(

tests/test_condor.py

@@ -1,4 +1,4 @@
-"""Tests for the rollout-YAML → run-directory flow, compute, and submit."""
+"""Tests for the rollout-YAML(s) → run-directory flow, compute, and submit."""
 
 from __future__ import annotations
 
@@ -17,6 +17,7 @@ from giant.analysis import (
     compute_reduced,
     derive_run_dir,
     load_rollout_yaml,
+    load_rollout_yamls,
     merge_one,
     prep,
     write_submit,
@@ -28,13 +29,17 @@ from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 
+def _write_rollout(path: Path) -> None:
+    tbl = _rollout_frame().collect().to_arrow()
+    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
+    pq.write_table(tbl, path)
+
+
 def _write_inputs(tmp_path: Path) -> Path:
     """Materialize rollout+reference parquet and a rollout YAML; return the YAML path."""
     rollout = tmp_path / "rollout.parquet"
     reference = tmp_path / "reference.parquet"
-    tbl = _rollout_frame().collect().to_arrow()
-    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
-    pq.write_table(tbl, rollout)
+    _write_rollout(rollout)
     _reference_frame().collect().write_parquet(reference)
 
     yaml_path = tmp_path / "run.yaml"
@@ -54,6 +59,33 @@ def _write_inputs(tmp_path: Path) -> Path:
     return yaml_path
 
 
+def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
+    """Two rollout YAMLs (distinct output files) sharing one reference file."""
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+
+    paths = []
+    for tag, pred_id in (("a", "aaaa1111ef"), ("b", "bbbb2222ef")):
+        rollout = tmp_path / f"rollout_{tag}.parquet"
+        _write_rollout(rollout)
+        yaml_path = tmp_path / f"run_{tag}.yaml"
+        yaml_path.write_text(
+            yaml.safe_dump(
+                {
+                    "prediction_id": pred_id,
+                    "output": str(rollout),
+                    "dataset": str(reference),
+                    "checkpoint": f"/ckpt/{tag}.pt",
+                    "kind": "rollout",
+                    "energy_cutoff": 0.1,
+                    "steps": 10,
+                }
+            )
+        )
+        paths.append(yaml_path)
+    return paths[0], paths[1]
+
+
 def _fake_venv(repo_dir: Path) -> None:
     """Stand in for a `uv sync`'d venv: write_submit checks `.venv/bin/giant` exists."""
     giant = repo_dir / ".venv" / "bin" / "giant"
@@ -62,12 +94,13 @@ 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_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
     """``prep`` with small test-sized context bins/sampling."""
     return prep(
-        rollout_yaml,
+        rollout_yamls,
         run_dir,
         n_chunks=chunks,
+        labels=labels,
         n_energy_bins=2,
         n_marginal_bins=8,
         top_k_pdg=3,
@@ -82,39 +115,108 @@ def test_load_rollout_yaml_requires_paths(tmp_path: Path):
         load_rollout_yaml(bad)
 
 
+def test_load_rollout_yamls_single_defaults_to_rollout_name(tmp_path: Path):
+    yaml_path = _write_inputs(tmp_path)
+    loaded, reference = load_rollout_yamls([yaml_path])
+    assert [lr.name for lr in loaded] == ["rollout"]
+    assert reference.endswith("reference.parquet")
+
+
+def test_load_rollout_yamls_multi_defaults_to_stem(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    loaded, _ = load_rollout_yamls([a, b])
+    assert [lr.name for lr in loaded] == ["run_a", "run_b"]
+
+
+def test_load_rollout_yamls_explicit_labels(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    loaded, _ = load_rollout_yamls([a, b], labels=["flow", "wgan"])
+    assert [lr.name for lr in loaded] == ["flow", "wgan"]
+
+
+def test_load_rollout_yamls_label_count_mismatch(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    with pytest.raises(ValueError, match="--label"):
+        load_rollout_yamls([a, b], labels=["only-one"])
+
+
+def test_load_rollout_yamls_rejects_duplicate_names(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    with pytest.raises(ValueError, match="collide"):
+        load_rollout_yamls([a, b], labels=["same", "same"])
+
+
+def test_load_rollout_yamls_rejects_mismatched_reference(tmp_path: Path):
+    a, _ = _write_two_inputs(tmp_path)
+    other_ref = tmp_path / "other_reference.parquet"
+    _reference_frame().collect().write_parquet(other_ref)
+    c = tmp_path / "run_c.yaml"
+    c.write_text(
+        yaml.safe_dump(
+            {"prediction_id": "cccc3333ef", "output": str(tmp_path / "rollout_c.parquet"), "dataset": str(other_ref)}
+        )
+    )
+    _write_rollout(tmp_path / "rollout_c.parquet")
+    with pytest.raises(ValueError, match="same reference"):
+        load_rollout_yamls([a, c])
+
+
 def test_derive_run_dir_next_to_rollout():
     y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
-    assert derive_run_dir(y) == Path("/data/analysis_abcd1234")
-    assert derive_run_dir(y, "/somewhere") == Path("/somewhere")
+    assert derive_run_dir([y]) == Path("/data/analysis_abcd1234")
+    assert derive_run_dir([y], "/somewhere") == Path("/somewhere")
 
 
 def test_derive_run_dir_default_base():
     y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
-    assert derive_run_dir(y, default_base="/work/lbogner/giant2/analysis_runs") == Path(
+    assert derive_run_dir([y], default_base="/work/lbogner/giant2/analysis_runs") == Path(
         "/work/lbogner/giant2/analysis_runs/analysis_abcd1234"
     )
     # an explicit run_dir still wins over default_base
-    assert derive_run_dir(y, "/somewhere", default_base="/other") == Path("/somewhere")
+    assert derive_run_dir([y], "/somewhere", default_base="/other") == Path("/somewhere")
+
+
+def test_derive_run_dir_multi_rollout_joins_tags():
+    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(2)]
+    assert derive_run_dir(ys, default_base="/base") == Path("/base/analysis_tag0xxxx-tag1xxxx")
+
+
+def test_derive_run_dir_many_rollouts_truncates_with_plus_count():
+    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(5)]
+    run_dir = derive_run_dir(ys, default_base="/base")
+    assert run_dir == Path("/base/analysis_tag0xxxx-tag1xxxx-tag2xxxx-plus2")
 
 
 def test_prep_lays_out_run_dir(tmp_path: Path):
     yaml_path = _write_inputs(tmp_path)
-    run_dir = _prep(yaml_path)
+    run_dir = _prep([yaml_path])
     assert run_dir == tmp_path / "analysis_abcd1234"
     assert (run_dir / "shared.json").exists()
     ctx = Context.load(run_dir / "shared.json")
     assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
     meta = RunMeta.load(run_dir / "run_meta.json")
     assert meta.reference.endswith("reference.parquet")
-    assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt"
+    assert [ro["name"] for ro in meta.rollouts] == ["rollout"]
+    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/best.pt"
     assert "best.pt" in meta.title
     assert meta.n_chunks == 1
     assert meta.rows_per_chunk == [meta.total_rows]  # single chunk holds everything
     assert meta.total_rows == 8  # 5 rollout rows + 3 reference rows
 
 
+def test_prep_multi_rollout_lays_out_run_dir(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    run_dir = _prep([a, b], labels=["flow", "wgan"])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert [ro["name"] for ro in meta.rollouts] == ["flow", "wgan"]
+    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/a.pt"
+    assert meta.rollouts[1]["plot_meta"]["checkpoint"] == "/ckpt/b.pt"
+    # 5 rows from each rollout + 3 from the shared reference
+    assert meta.total_rows == 13
+
+
 def test_prep_splits_rows_per_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
     meta = RunMeta.load(run_dir / "run_meta.json")
     assert len(meta.rows_per_chunk) == 2
     assert sum(meta.rows_per_chunk) == meta.total_rows == 8
@@ -125,7 +227,7 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
     partials on disk for merge_one to silently merge against the new
     context (they'd be keyed/sized for the old n_chunks)."""
     yaml_path = _write_inputs(tmp_path)
-    run_dir = _prep(yaml_path, chunks=2)
+    run_dir = _prep([yaml_path], chunks=2)
     compute_one("marginal_edep", run_dir, chunk_index=0)
     compute_one("marginal_edep", run_dir, chunk_index=1)
     stale = run_dir / "reduced_partial" / "marginal_edep__0.json"
@@ -133,7 +235,7 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
     (run_dir / "reduced").mkdir(exist_ok=True)
     (run_dir / "reduced" / "marginal_edep.json").write_text("{}")
 
-    _prep(yaml_path, run_dir, chunks=1)
+    _prep([yaml_path], run_dir, chunks=1)
 
     assert not stale.exists()
     assert not (run_dir / "reduced" / "marginal_edep.json").exists()
@@ -141,20 +243,22 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
 
 
 def test_compute_one_from_run_dir(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
     out = compute_one("marginal_edep", run_dir)
     assert out == run_dir / "reduced_partial" / "marginal_edep__0.json"
     partial = Partial.load(out)
     assert partial.id == "marginal_edep" and partial.chunk == 0
     assert "r" in partial.data and "t" in partial.data
+    assert list(partial.data["r"]) == ["rollout"]
 
 
 def test_compute_reduced_explicit_paths(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
     meta = RunMeta.load(run_dir / "run_meta.json")
+    rollouts = [{"name": ro["name"], "path": ro["path"]} for ro in meta.rollouts]
     out = compute_reduced(
         "marginal_step_length",
-        meta.rollout,
+        rollouts,
         meta.reference,
         run_dir / "shared.json",
         tmp_path / "r.json",
@@ -163,17 +267,17 @@ def test_compute_reduced_explicit_paths(tmp_path: Path):
 
 
 def test_merge_one_produces_reduced(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
     compute_one("marginal_edep", run_dir)
     out = merge_one("marginal_edep", run_dir)
     assert out == run_dir / "reduced" / "marginal_edep.json"
     reduced = Reduced.load(out)
     assert reduced.id == "marginal_edep"
-    assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
+    assert len(reduced.payload["series"]["rollout"]) == len(reduced.payload["edges"]) - 1
 
 
 def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
     compute_one("marginal_edep", run_dir, chunk_index=0)  # chunk 1 never computed
     with pytest.raises(FileNotFoundError, match="missing chunk"):
         merge_one("marginal_edep", run_dir)
@@ -182,11 +286,11 @@ def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
 def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
     (tmp_path / "a").mkdir()
     (tmp_path / "b").mkdir()
-    unchunked_dir = _prep(_write_inputs(tmp_path / "a"))
+    unchunked_dir = _prep([_write_inputs(tmp_path / "a")])
     compute_one("marginal_step_length", unchunked_dir)
     unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir))
 
-    chunked_dir = _prep(_write_inputs(tmp_path / "b"), chunks=2)
+    chunked_dir = _prep([_write_inputs(tmp_path / "b")], chunks=2)
     for k in range(2):
         compute_one("marginal_step_length", chunked_dir, chunk_index=k)
     chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir))
@@ -194,14 +298,23 @@ def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
     assert chunked.payload == unchunked.payload
 
 
+def test_two_rollout_compute_and_merge_produces_both_series(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    run_dir = _prep([a, b], labels=["flow", "wgan"])
+    compute_one("marginal_edep", run_dir)
+    reduced = Reduced.load(merge_one("marginal_edep", run_dir))
+    assert list(reduced.payload["series"]) == ["flow", "wgan"]
+    assert "reference" in reduced.payload
+
+
 def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))  # n_chunks=1 (default)
+    run_dir = _prep([_write_inputs(tmp_path)])  # n_chunks=1 (default)
     with pytest.raises(ValueError, match="out of range"):
         compute_one("marginal_edep", run_dir, chunk_index=1)
 
 
 def test_write_submit_description(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_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)
     txt = write_submit(cfg).read_text()
@@ -223,7 +336,7 @@ def test_write_submit_description(tmp_path: Path):
 
 
 def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
-    run_dir = _prep(_write_inputs(tmp_path))
+    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.
@@ -233,7 +346,7 @@ def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.M
 
 
 def test_write_submit_remote_flag(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_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)
     txt = write_submit(cfg).read_text()
@@ -243,7 +356,7 @@ def test_write_submit_remote_flag(tmp_path: Path):
 
 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)
+    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)
     write_submit(cfg)
@@ -260,7 +373,7 @@ def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
     with — RunMeta.rows_per_chunk is sized to the prepped value, so a
     mismatch would otherwise surface as a confusing IndexError deep inside
     _job_walltimes instead of a clear error here."""
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
     _fake_venv(tmp_path)
     cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
     with pytest.raises(ValueError, match="n_chunks"):
@@ -282,7 +395,7 @@ def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
     """A chunked run's later job walltimes track that chunk's row count."""
     from giant.analysis.runtime_estimate import estimate_runtime_s
 
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    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)

tests/test_config.py

@@ -97,6 +97,19 @@ def test_trunk_config_defaults_block_conditioning_to_add_for_both_stages():
     assert gconfig.DEFAULT_CONFIG["stage2_model"]["trunk"]["block_conditioning"] == "add"
 
 
+def test_init_from_freeze_default_to_unset_for_both_stages():
+    """gitea #42: a pre-existing config with no init_from/freeze key must
+    reproduce today's from-scratch, always-training behaviour exactly."""
+    assert gconfig.Stage1ModelConfig().init_from == ""
+    assert gconfig.Stage1ModelConfig().freeze is False
+    assert gconfig.Stage2ModelConfig().init_from == ""
+    assert gconfig.Stage2ModelConfig().freeze is False
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["init_from"] == ""
+    assert gconfig.DEFAULT_CONFIG["stage1_model"]["freeze"] is False
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["init_from"] == ""
+    assert gconfig.DEFAULT_CONFIG["stage2_model"]["freeze"] is False
+
+
 def test_heads_config_defaults_reproduce_pre_gitea_36_hardcoded_shape():
     """gitea #36: a pre-existing config with no `heads` key must reproduce
     today's hardcoded `hidden_dim // 2`, one-hidden-layer architecture
@@ -119,7 +132,15 @@ def test_particle_type_config_n_classes_defaults_to_zero_and_round_trips():
     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_particle_type_config_class_weighting_defaults_to_none_and_round_trips():
+    """gitea #44: an existing config.toml with no
+    stage2_model.particle_type.class_weighting key must reproduce the
+    pre-#44 unweighted-CE behavior exactly."""
+    assert gconfig.ParticleTypeConfig().class_weighting == "none"
+    spec = gconfig.ParticleTypeConfig.from_dict({"class_weighting": "inverse_freq"})
+    assert spec.class_weighting == "inverse_freq"
 
 
 def test_router_config_extra_round_trips_composed_axis_keys():
@@ -691,6 +712,42 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning():
     gconfig.validate_config(cfg)  # must not raise
 
 
+def test_validate_config_bad_class_weighting_rejected():
+    cfg = _cfg_with(**{"stage2_model.particle_type.class_weighting": "effective_num"})
+    with pytest.raises(ValueError, match="class_weighting"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_requires_onehot_target():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.particle_type.class_weighting": "inverse_freq",
+            "stage2_model.particle_type.target": "physical",
+        }
+    )
+    with pytest.raises(ValueError, match="onehot"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_incompatible_with_wgan_generator():
+    # stage2_model.generator defaults to "wgan" and particle_type.target
+    # defaults to "onehot", so only class_weighting needs overriding here.
+    cfg = _cfg_with(**{"stage2_model.particle_type.class_weighting": "inverse_freq"})
+    with pytest.raises(ValueError, match="wgan"):
+        gconfig.validate_config(cfg)
+
+
+def test_validate_config_class_weighting_passes_with_onehot_and_flow():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.particle_type.class_weighting": "inverse_freq",
+            "stage2_model.particle_type.target": "onehot",
+            "stage2_model.generator": "flow",
+        }
+    )
+    gconfig.validate_config(cfg)  # must not raise
+
+
 def test_validate_config_mixed_particle_material_conditioning_is_valid():
     """The particle and material conditioning axes are configured
     independently and may mix freely — e.g. material
@@ -735,6 +792,29 @@ def test_validate_config_tie_to_stage1_requires_stage1_active():
         assert "tie_to_stage1" in str(e)
 
 
+@pytest.mark.parametrize("stage_name", ["stage1_model", "stage2_model"])
+def test_validate_config_freeze_without_init_from_or_resume_rejected(stage_name):
+    cfg = _cfg_with(**{f"{stage_name}.freeze": True})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "init_from" in str(e)
+        assert "--resume" in str(e)
+
+
+@pytest.mark.parametrize("stage_name", ["stage1_model", "stage2_model"])
+def test_validate_config_freeze_with_init_from_passes(stage_name):
+    cfg = _cfg_with(**{f"{stage_name}.freeze": True, f"{stage_name}.init_from": "ckpt/best.pt"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+@pytest.mark.parametrize("stage_name", ["stage1_model", "stage2_model"])
+def test_validate_config_freeze_without_init_from_passes_under_resume(stage_name):
+    cfg = _cfg_with(**{f"{stage_name}.freeze": True})
+    gconfig.validate_config(cfg, resume=True)  # must not raise
+
+
 def test_validate_config_stop_token_accepted_under_autoregressive():
     """DEFAULT_CONFIG's stage2_model.decoder is already "autoregressive"
     (see test_stage2_model_config_defaults_match_documented_v030_intent), so
@@ -780,13 +860,98 @@ def test_validate_config_bad_stop_sampling_rejected():
         assert "stop_sampling" in str(e)
 
 
-def test_validate_config_stage1_context_sampled_not_implemented():
-    cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"})
+def test_validate_config_default_precision_is_fp32():
+    assert gconfig.DEFAULT_CONFIG["train"]["precision"] == "fp32"
+
+
+def test_validate_config_bf16_precision_accepted():
+    cfg = _cfg_with(**{"train.precision": "bf16"})
+    gconfig.validate_config(cfg)  # no raise
+
+
+@pytest.mark.parametrize("bad", ["fp16", "bogus", ""])
+def test_validate_config_bad_precision_rejected(bad):
+    cfg = _cfg_with(**{"train.precision": bad})
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "sampled" in str(e)
+        assert "precision" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_accepted_with_both_stages_active():
+    """gitea #41: 'sampled' is now implemented, so DEFAULT_CONFIG's
+    stage1_model/stage2_model.active = true (both) must let it through."""
+    cfg = _cfg_with(**{"stage2_model.stage1_context": "sampled"})
+    gconfig.validate_config(cfg)  # must not raise
+
+
+def test_validate_config_bad_stage1_context_rejected():
+    cfg = _cfg_with(**{"stage2_model.stage1_context": "bogus"})
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "stage1_context" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_requires_stage1_active():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage1_model.active": False,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e) and "stage1_model.active" in str(e)
+
+
+def test_validate_config_stage1_context_sampled_requires_stage2_active():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage2_model.active": False,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "sampled" in str(e) and "stage2_model.active" in str(e)
+
+
+@pytest.mark.parametrize("key", ["ctx_p_start", "ctx_p_end"])
+@pytest.mark.parametrize("value", [-0.1, 1.1])
+def test_validate_config_ctx_p_out_of_range_rejected(key, value):
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            f"stage2_model.{key}": value,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert key in str(e)
+
+
+def test_validate_config_stage1_context_sampled_always_truth_rejected_as_noop():
+    cfg = _cfg_with(
+        **{
+            "stage2_model.stage1_context": "sampled",
+            "stage2_model.ctx_p_start": 1.0,
+            "stage2_model.ctx_p_end": 1.0,
+        }
+    )
+    try:
+        gconfig.validate_config(cfg)
+        assert False, "expected ValueError"
+    except ValueError as e:
+        assert "ctx_p_start" in str(e) and "ctx_p_end" in str(e)
 
 
 def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint():
@@ -1075,6 +1240,11 @@ def test_overrides_from_flags_train_block_passthrough():
     assert overrides == {"train": {"epochs": 5, "lr": 1e-3}}
 
 
+def test_overrides_from_flags_precision_passthrough():
+    overrides = gconfig.overrides_from_flags({"precision": "bf16"})
+    assert overrides == {"train": {"precision": "bf16"}}
+
+
 @pytest.mark.parametrize(
     ("shorthand", "explicit", "path_key"),
     [
@@ -1193,6 +1363,21 @@ def test_overrides_from_flags_critic_sizing_is_stage_scoped_only(stage_flag, sta
     assert overrides == {stage_model: {"wgan": {path_key: 32}}}
 
 
+@pytest.mark.parametrize(
+    ("init_from_flag", "freeze_flag", "stage_model"),
+    [
+        ("stage1_init_from", "stage1_freeze", "stage1_model"),
+        ("stage2_init_from", "stage2_freeze", "stage2_model"),
+    ],
+)
+def test_overrides_from_flags_init_from_freeze_is_stage_scoped_only(init_from_flag, freeze_flag, stage_model):
+    """gitea #42: no shared alias — a checkpoint has one set of weights per
+    stage, so "freeze both stages from the same file" has no sensible
+    meaning."""
+    overrides = gconfig.overrides_from_flags({init_from_flag: "ckpt/best.pt", freeze_flag: True})
+    assert overrides == {stage_model: {"init_from": "ckpt/best.pt", "freeze": True}}
+
+
 # ---------------------------------------------------------------------------
 # checkpoint config-mismatch warnings (unchanged surface, still exercised)
 # ---------------------------------------------------------------------------

tests/test_config_consumed_keys.py

@@ -28,6 +28,7 @@ import ast
 from pathlib import Path
 
 from giant.config import DEFAULT_CONFIG
+from giant.config import leaf_paths as _leaf_paths
 
 _REPO_ROOT = Path(__file__).resolve().parents[1]
 
@@ -53,13 +54,6 @@ _EXCLUDED_FILES = ("giant/model/_legacy.py",)
 # 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 "
@@ -72,19 +66,6 @@ _KNOWN_UNUSED = {
 _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)

tests/test_loader.py

@@ -195,6 +195,10 @@ def test_build_topn_map_from_files_keeps_most_frequent(tmp_path):
     assert m.class_map["G4_Fe"] == 2  # "other" (n_classes - 1)
     assert m.class_map["G4_Pb"] == 2
     assert m.other_members == {"G4_Fe": 2, "G4_Pb": 1}
+    # class_counts (gitea #44): per resulting index, "other" is the sum of
+    # everything folded into it (2 + 1 = 3), and the total equals row count.
+    assert m.class_counts == {0: 5, 1: 3, 2: 3}
+    assert sum(m.class_counts.values()) == len(materials)
 
 
 def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
@@ -205,6 +209,8 @@ def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
 
     assert m.class_map == {"G4_AIR": 0, "PbWO4": 1}
     assert m.other_members == {}
+    # No "other" bucket ever populated -> no entry for its index either.
+    assert m.class_counts == {0: 1, 1: 1}
 
 
 def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path):
@@ -224,6 +230,7 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path)
     # pooled: 11 -> 5, 22 -> 1 (primary) + 10 (secondary) = 11
     assert m.class_map[22] == 0
     assert m.class_map[11] == 1
+    assert m.class_counts == {0: 11, 1: 5}
 
 
 def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):

tests/test_model_summary.py

@@ -0,0 +1,137 @@
+"""Tests for `giant model summary` (gitea #46)."""
+
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from typer.testing import CliRunner
+
+from giant import config as gconfig
+from giant.cli import app
+from giant.materials import MATERIAL_PROPERTIES
+from giant.model.summary import _NOT_BUILD_TIME, _built_modules, _vocab_caveats, summarize_model
+
+runner = CliRunner()
+
+_PDG_VOCAB = 300
+_MAT_VOCAB = len(MATERIAL_PROPERTIES)
+
+
+def _cfg(overrides: dict | None = None) -> dict:
+    return gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, overrides or {})
+
+
+@pytest.fixture(scope="module")
+def default_summary():
+    return summarize_model(_cfg(), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB)
+
+
+def test_default_config_builds_both_stages_with_a_real_tree(default_summary):
+    assert set(default_summary.modules) >= {"stage1", "stage2"}
+    for module in default_summary.modules.values():
+        assert sum(p.numel() for p in module.parameters()) > 0
+    stage1 = default_summary.modules["stage1"]
+    assert hasattr(stage1, "cond_enc")
+    assert hasattr(stage1, "trunk")
+    assert {"input_proj", "blocks", "out_proj"} <= {n for n, _ in stage1.trunk.named_children()}
+
+
+def test_every_in_scope_leaf_is_classified(default_summary):
+    in_scope = {
+        p
+        for p in gconfig.leaf_paths(gconfig.DEFAULT_CONFIG)
+        if p.split(".", 1)[0] in ("conditioning", "stage1_model", "stage2_model")
+    }
+    classified = set(default_summary.consumed) | set(default_summary.inert) | set(default_summary.elsewhere)
+    assert classified == in_scope
+
+
+def test_not_build_time_allow_list_has_no_stale_entries():
+    in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
+    stale = set(_NOT_BUILD_TIME) - in_scope
+    assert not stale, f"_NOT_BUILD_TIME entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
+
+
+def test_router_disabled_by_default_so_its_fields_are_inert(default_summary):
+    assert "stage1_model.router.n_experts" in default_summary.inert
+    assert "stage1_model.router.temperature" in default_summary.inert
+
+
+def test_markov_history_leaves_attention_dims_inert_but_history_itself_consumed(default_summary):
+    assert "stage2_model.autoregressive.attn_n_heads" in default_summary.inert
+    assert "stage2_model.autoregressive.attn_n_layers" in default_summary.inert
+    assert "stage2_model.autoregressive.history" in default_summary.consumed
+
+
+def test_single_literal_branch_fields_are_correctly_seen_as_consumed(default_summary):
+    """Regression guard: n_sec.owner ("stage2"), n_sec.mode ("head") and
+    particle_type.target ("onehot") each branch as `== "one specific other
+    literal"` in giant/model/builders.py|models.py. A naive single generic
+    sentinel probe lands in the same "not that literal" bucket as the
+    current value and never crosses the boundary that actually matters --
+    this is exactly what _STRING_ALTERNATIVES exists to fix."""
+    assert "stage2_model.n_sec.owner" in default_summary.consumed
+    assert "stage2_model.n_sec.mode" in default_summary.consumed
+    assert "stage2_model.particle_type.target" in default_summary.consumed
+
+
+def test_stage1_wgan_generator_swaps_flow_time_dim_for_critic_dims():
+    summary = summarize_model(_cfg({"stage1_model": {"generator": "wgan"}}), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB)
+    assert "stage1_model.flow.time_dim" in summary.inert
+    assert "stage1_model.wgan.noise_dim" in summary.consumed
+    assert "stage1_model.wgan.critic_hidden_dim" in summary.consumed
+
+
+def test_stage2_one_shot_decoder_makes_autoregressive_block_inert():
+    summary = summarize_model(
+        _cfg({"stage2_model": {"decoder": "one_shot"}}), pdg_vocab=_PDG_VOCAB, mat_vocab=_MAT_VOCAB
+    )
+    assert "stage2_model.autoregressive.history" in summary.inert
+    assert "history_encoder" not in {n for n, _ in summary.modules["stage2"].named_children()}
+
+
+def test_energy_router_enabled_consumes_core_fields_but_not_process_only_fields():
+    summary = summarize_model(
+        _cfg({"stage1_model": {"router": {"enabled": True, "type": "energy", "n_experts": 4}}}),
+        pdg_vocab=_PDG_VOCAB,
+        mat_vocab=_MAT_VOCAB,
+    )
+    assert "stage1_model.router.n_experts" in summary.consumed
+    assert "stage1_model.router.temperature" in summary.consumed
+    # emb_dim/hidden_dim are pdg/process-router-only kwargs -- build_router's
+    # signature filter drops them for an energy router.
+    assert "stage1_model.router.hidden_dim" in summary.inert
+    assert "stage1_model.router.emb_dim" in summary.inert
+
+
+def test_vocab_caveat_text_for_embedding_particle_conditioning():
+    cfg = _cfg({"conditioning": {"particle": {"type": "embedding"}}})
+    caveats = _vocab_caveats(cfg)
+    assert any("pdg_vocab" in c and "embedding" in c for c in caveats)
+    assert not any("mat_vocab" in c for c in caveats)
+
+
+def test_pdg_vocab_flag_changes_embedding_table_size():
+    cfg = _cfg({"conditioning": {"particle": {"type": "embedding"}}})
+    small = _built_modules(cfg, pdg_vocab=10, mat_vocab=_MAT_VOCAB)
+    big = _built_modules(cfg, pdg_vocab=1000, mat_vocab=_MAT_VOCAB)
+    assert big["stage1"].cond_enc.pdg_emb.weight.numel() > small["stage1"].cond_enc.pdg_emb.weight.numel()
+
+
+def test_invalid_combo_exits_nonzero_with_validate_config_message(tmp_path: Path):
+    config_path = tmp_path / "bad.toml"
+    config_path.write_text('[meta]\nconfig_version = 3\n\n[stage2_model.particle_type]\ntarget = "embedding"\n')
+    result = runner.invoke(app, ["model", "summary", "--config", str(config_path)])
+    assert result.exit_code == 1
+    assert "requires conditioning.particle.type = 'embedding'" in result.output
+
+
+def test_cli_default_smoke():
+    result = runner.invoke(app, ["model", "summary"])
+    assert result.exit_code == 0, result.output
+    assert "stage1" in result.output
+    assert "stage2" in result.output
+    assert "parameters" in result.output
+    assert "trunk" in result.output
+    assert "inert under this config" in result.output

tests/test_network.py

@@ -8,8 +8,12 @@ from giant.model.network import (
     HISTORY_REGISTRY,
     AttentionHistory,
     ConditionEncoder,
+    CriticModel,
+    FilmResBlock,
     HistoryEncoder,
+    LinearTrunk,
     MarkovHistory,
+    NoHistory,
     SinusoidalEmbedding,
     Stage1Model,
     Stage2Autoregressive,
@@ -465,16 +469,52 @@ def test_attention_history_step_matches_forward():
     assert torch.allclose(stepped, expected, atol=1e-5)
 
 
+# --- NoHistory (gitea #45) ----------------------------------------------------
+
+
+def test_no_history_shape():
+    hist = NoHistory(in_dim=7, out_dim=12)
+    B, K = 3, 5
+    feat = torch.randn(B, K, 7)
+    has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    out = hist(feat, has_prev)
+    assert out.shape == (B, K, 12)
+
+
+def test_no_history_ignores_feat_and_has_prev():
+    hist = NoHistory(in_dim=4, out_dim=6)
+    B, K = 2, 3
+    has_prev_a = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
+    has_prev_b = torch.zeros(B, K, dtype=torch.bool)
+    feat_a = torch.randn(B, K, 4)
+    feat_b = torch.randn(B, K, 4) * 100
+    out_a = hist(feat_a, has_prev_a)
+    out_b = hist(feat_b, has_prev_b)
+    assert torch.equal(out_a, torch.zeros(B, K, 6))
+    assert torch.equal(out_a, out_b)
+
+
+def test_no_history_uses_base_class_o1_defaults():
+    hist = NoHistory(in_dim=4, out_dim=6)
+    assert hist.init_cache() is None
+    feat = torch.randn(2, 1, 4)
+    has_prev = torch.ones(2, 1, dtype=torch.bool)
+    out, cache = hist.step(feat, has_prev, "unused-cache")
+    assert torch.equal(out, torch.zeros(2, 1, 6))
+    assert cache == "unused-cache"
+
+
 # --- HISTORY_REGISTRY / build_history (gitea #35) ----------------------------
 
 
-def test_history_registry_has_exactly_the_two_known_histories():
-    assert set(HISTORY_REGISTRY) == {"markov", "attention"}
+def test_history_registry_has_exactly_the_known_histories():
+    assert set(HISTORY_REGISTRY) == {"markov", "attention", "none"}
 
 
 def test_build_history_returns_correct_concrete_type():
     assert isinstance(build_history("markov", 4, 6), MarkovHistory)
     assert isinstance(build_history("attention", 4, 8), AttentionHistory)
+    assert isinstance(build_history("none", 4, 6), NoHistory)
 
 
 def test_build_history_unknown_name_raises():
@@ -605,7 +645,7 @@ def test_stage2_autoregressive_n_sec_head_and_type_head_cfg_control_hidden_width
 
 @pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
 @pytest.mark.parametrize("generator", ["wgan", "flow"])
-@pytest.mark.parametrize("history", ["markov", "attention"])
+@pytest.mark.parametrize("history", ["markov", "attention", "none"])
 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)
@@ -907,7 +947,7 @@ def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
     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
+    assert wider_critic.trunk.input_proj.in_features > default_n_classes_critic.trunk.input_proj.in_features
 
 
 # ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
@@ -984,12 +1024,12 @@ def test_build_critics_omitted_particle_type_matches_default_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
+    onehot_in_dim = onehot_critic.trunk.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
+    physical_in_dim = physical_critic.trunk.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
@@ -1009,15 +1049,15 @@ def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_genera
 
     inherited = build_critics(cfg)["stage1"]
     assert inherited is not None
-    assert inherited.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.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
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
 
 
 def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
@@ -1028,15 +1068,15 @@ def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_genera
 
     inherited = build_critics(cfg)["stage2"]
     assert inherited is not None
-    assert inherited.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    assert inherited.trunk.input_proj.out_features == 8
+    assert len(inherited.trunk.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
+    assert overridden.trunk.input_proj.out_features == 16
+    assert len(overridden.trunk.blocks) == 3
 
 
 # ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive
@@ -1197,3 +1237,34 @@ def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator):
     assert model.generator_kind == generator
     assert model.noise_dim == 8
     assert (model.time_emb is not None) == build_objective(generator).needs_time
+
+
+# ── CriticModel uses the trunk/block registries + StageModel base (gitea #57) ─
+
+
+def test_critic_model_is_stagemodel_subclass():
+    assert issubclass(CriticModel, StageModel)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_trunk_type_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"type": "linear"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert isinstance(critic.trunk, LinearTrunk)
+
+
+@pytest.mark.parametrize("stage", ["stage1", "stage2"])
+def test_build_critics_threads_block_conditioning_from_generator_config(stage):
+    cfg = _minimal_model_config(share_stages=False)
+    cfg["stage1_model"]["generator"] = "wgan"
+    cfg["stage2_model"]["generator"] = "wgan"
+    cfg[f"{stage}_model"]["trunk"] = {"block_conditioning": "film"}
+
+    critic = build_critics(cfg)[stage]
+    assert critic is not None
+    assert all(isinstance(block, FilmResBlock) for block in critic.trunk.blocks)

tests/test_pipeline.py

@@ -231,7 +231,7 @@ def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, da
 
 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}
+    cfg["stage2_model"]["particle_type"].update({"target": "physical", "lambda": 1.0})
     echo = _run(data, tmp_path / "out", cfg=cfg)
     assert not any("top-N map" in m for m in echo)
 

tests/test_release_tooling.py

@@ -0,0 +1,77 @@
+"""Config-correctness tests for the CI version-bump/tag/changelog automation
+(gitea #50). The workflow YAML itself can only be exercised by a real push to
+master, so these check the two config files it drives (.bumpversion.toml,
+cliff.toml) against real repo content instead.
+"""
+
+import re
+import shutil
+import subprocess
+import tomllib
+from pathlib import Path
+
+import pytest
+
+_ROOT = Path(__file__).resolve().parents[1]
+
+
+def test_bumpversion_search_pattern_matches_pyproject():
+    bump_config = tomllib.loads((_ROOT / ".bumpversion.toml").read_text())["tool"]["bumpversion"]
+    current_version = bump_config["current_version"]
+    search = bump_config["files"][0]["search"].format(current_version=current_version)
+
+    pyproject = (_ROOT / "pyproject.toml").read_text()
+    assert search in pyproject, (
+        f"bumpversion search pattern {search!r} (rendered from .bumpversion.toml's "
+        f"current_version={current_version!r}) not found in pyproject.toml — "
+        "the bump would silently edit nothing"
+    )
+
+
+@pytest.mark.skipif(shutil.which("git-cliff") is None, reason="git-cliff binary not on PATH")
+def test_cliff_config_groups_and_links_commits(tmp_path):
+    repo = tmp_path / "repo"
+    repo.mkdir()
+    subprocess.run(["git", "init", "-q"], cwd=repo, check=True)
+    subprocess.run(["git", "config", "user.name", "test"], cwd=repo, check=True)
+    subprocess.run(["git", "config", "user.email", "test@example.com"], cwd=repo, check=True)
+
+    _commit(repo, "Add class-balanced secondary particle-type loss (gitea #44)")
+    _commit(repo, "Fix leaking secondary energy budget")
+    _commit(repo, "Merge pull request 'Add X' (#1) from fix/issue-1 into master")
+    _commit(repo, "chore: bump version 0.3.3 -> 0.3.4 [skip ci]")
+
+    result = subprocess.run(
+        [
+            "git-cliff",
+            "--config",
+            str(_ROOT / "cliff.toml"),
+            "--repository",
+            str(repo),
+            "--tag",
+            "v0.3.4",
+            "--unreleased",
+        ],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    changelog = result.stdout
+
+    assert "## [0.3.4]" in changelog
+    assert "### Added" in changelog
+    assert "### Fixed" in changelog
+    assert re.search(
+        r"\[gitea #44\]\(https://git\.larsbogner\.de/lars/giant/issues/44\)",
+        changelog,
+    )
+    assert "Add class-balanced secondary particle-type loss" in changelog
+    assert "Fix leaking secondary energy budget" in changelog
+    assert "Merge pull request" not in changelog
+    assert "skip ci" not in changelog
+
+
+def _commit(repo: Path, message: str) -> None:
+    (repo / "f.txt").write_text(message)
+    subprocess.run(["git", "add", "f.txt"], cwd=repo, check=True)
+    subprocess.run(["git", "commit", "-q", "-m", message], cwd=repo, check=True)

tests/test_render.py

@@ -25,41 +25,96 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
     reduced = [
         Reduced(
             "rg",
-            "router",
+            "model",
             "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]],
+                "series": {
+                    "flow": {
+                        "n_experts": 2,
+                        "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]],
+                        },
+                    },
+                    "wgan": {
+                        "n_experts": 2,
+                        "router_type": "energy",
+                        "rollout": {"centers": [1.0], "means": [[0.5, 0.5]]},
+                        "reference": {"centers": [1.0], "means": [[0.5, 0.5]]},
+                    },
                 },
             },
         ),
         Reduced(
             "rs",
-            "router",
+            "model",
             "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]},
+                "series": {
+                    "flow": {
+                        "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(
+            "rp",
+            "model",
+            "router_share",
+            "Router share by process (reference-only)",
+            "process",
+            {
+                "series": {
+                    "flow": {
+                        "categories": ["compt", "phot"],
+                        "n_experts": 2,
+                        "router_type": "energy",
+                        "reference": {"compt": [0.4, 0.6], "phot": [0.9, 0.1]},
+                    },
+                },
+            },
+        ),
+        Reduced(
+            "rz",
+            "model",
+            "router_specialization",
+            "Router specialization",
+            "pre-step energy [MeV]",
+            {
+                "log_x": True,
+                "series": {
+                    "flow": {
+                        "n_experts": 2,
+                        "chance_level": 0.5,
+                        "rollout": {"centers": [1.0, 10.0], "score": [0.6, 0.7]},
+                        "reference": {"centers": [1.0, 10.0], "score": [0.55, 0.65]},
+                    },
+                    "wgan": {
+                        "n_experts": 4,
+                        "chance_level": 0.25,
+                        "rollout": {"centers": [1.0, 10.0], "score": [0.3, 0.4]},
+                        "reference": {"centers": [], "score": []},
+                    },
+                },
             },
         ),
         Reduced(
             "ru",
-            "router",
+            "model",
             "unavailable",
             "Router unavailable",
             "x",
@@ -73,7 +128,10 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
             "x",
             {
                 "edges": [0, 1, 2],
-                "groups": {lbl: {"rollout": [1, 2], "reference": [2, 1]} for lbl in ("a", "b", "c", "d")},
+                "groups": {
+                    lbl: {"series": {"flow": [1, 2], "wgan": [2, 1]}, "reference": [2, 1]}
+                    for lbl in ("a", "b", "c", "d")
+                },
                 "log_y": True,
             },
         ),
@@ -83,7 +141,22 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
             "single_hist",
             "Single (log-x)",
             "x",
-            {"edges": [1, 10, 100], "rollout": [5, 1], "log_x": True, "log_y": True},
+            {"edges": [1, 10, 100], "series": {"flow": [5, 1], "wgan": [3, 2]}, "log_x": True, "log_y": True},
+        ),
+        Reduced(
+            "hm",
+            "quality",
+            "heatmap",
+            "Distance summary (2 rollouts)",
+            "grouping axis",
+            {
+                "series": {"flow": [[0.1, 0.2], [0.3, 0.4]], "wgan": [[0.5, 0.6], [0.7, 0.8]]},
+                "row_labels": ["step_length", "edep"],
+                "col_labels": ["overall", "energy"],
+                "cbar_label": "KS statistic",
+                "vmin": 0.0,
+                "vmax": 1.0,
+            },
         ),
     ]
     try:
@@ -117,7 +190,7 @@ def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
             "single_hist",
             "Single",
             "x",
-            {"edges": [0, 1, 2], "rollout": [5, 1]},
+            {"edges": [0, 1, 2], "series": {"rollout": [5, 1]}},
         )
     ]
     for r in reduced:
@@ -142,15 +215,14 @@ def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkey
     merge_calls = []
     monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
     meta = condor_mod.RunMeta(
-        rollout="rollout.parquet",
+        rollouts=[{"name": "rollout", "path": "rollout.parquet", "plot_meta": {"checkpoint": "ckpt/best.pt"}}],
         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(
+    Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "series": {"rollout": [1]}}).save(
         run_dir / "reduced" / "s.json"
     )
 
@@ -177,7 +249,7 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "x",
             {
                 "edges": [0, 1, 2, 3],
-                "rollout": [1, 2, 3],
+                "series": {"flow": [1, 2, 3], "wgan": [2, 2, 2]},
                 "reference": [3, 2, 1],
                 "log_y": False,
             },
@@ -190,7 +262,7 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "x",
             {
                 "edges": [0, 1, 2],
-                "groups": {"a": {"rollout": [1, 2], "reference": [2, 1]}},
+                "groups": {"a": {"series": {"flow": [1, 2]}, "reference": [2, 1]}},
                 "log_y": False,
             },
         ),
@@ -202,10 +274,8 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "depth",
             {
                 "edges": [0, 1, 2],
-                "rollout_mean": [1, 2],
-                "rollout_std": [0.1, 0.2],
-                "reference_mean": [1.1, 1.9],
-                "reference_std": [0.1, 0.1],
+                "series": {"flow": {"mean": [1, 2], "std": [0.1, 0.2]}},
+                "reference": {"mean": [1.1, 1.9], "std": [0.1, 0.1]},
                 "ylabel": "e",
             },
         ),
@@ -217,7 +287,7 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "species",
             {
                 "labels": ["e-", "gamma"],
-                "rollout": [0.6, 0.4],
+                "series": {"flow": [0.6, 0.4], "wgan": [0.55, 0.45]},
                 "reference": [0.5, 0.5],
                 "ylabel": "frac",
             },
@@ -228,7 +298,20 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "single_hist",
             "Single",
             "x",
-            {"edges": [0, 1, 2], "rollout": [5, 1], "log_y": True},
+            {"edges": [0, 1, 2], "series": {"flow": [5, 1]}, "log_y": True},
+        ),
+        Reduced(
+            "hm1",
+            "secondaries",
+            "heatmap",
+            "Confusion (single rollout)",
+            "predicted",
+            {
+                "series": {"flow": [[1, 0], [0, 1]]},
+                "row_labels": ["0", "1+"],
+                "col_labels": ["0", "1+"],
+                "cbar_label": "count",
+            },
         ),
     ]
     try:
@@ -276,8 +359,8 @@ def test_figure_params_v2_basics_and_router_and_epoch():
         },
         "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})
+    meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10, "model_config": mc}
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params == {
         "hidden_dim": 256,
         "n_res_blocks": 4,
@@ -297,8 +380,8 @@ def test_figure_params_v2_wgan_reports_noise_dim_not_steps():
             "wgan": {"noise_dim": 32},
         },
     }
-    run_meta = {"model_config": mc, "steps": 10}
-    params = render_mod._figure_params(run_meta)
+    meta = {"model_config": mc, "steps": 10}
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params["mode"] == "wgan"
     assert params["noise_dim"] == 32
     assert "steps" not in params
@@ -309,18 +392,18 @@ def test_figure_params_v2_reports_mode_s2_only_when_it_differs():
         "stage1_model": {"generator": "flow"},
         "stage2_model": {"generator": "flow"},
     }
-    assert "mode_s2" not in render_mod._figure_params({"model_config": same})
+    assert "mode_s2" not in render_mod._figure_params({"rollouts": {"rollout": {"model_config": same}}})
 
     mixed = {
         "stage1_model": {"generator": "flow"},
         "stage2_model": {"generator": "wgan"},
     }
-    params = render_mod._figure_params({"model_config": mixed})
+    params = render_mod._figure_params({"rollouts": {"rollout": {"model_config": mixed}}})
     assert params["mode_s2"] == "wgan"
 
 
 def test_figure_params_old_shape_basics():
-    run_meta = {
+    meta = {
         "model_config": {
             "hidden_dim": 128,
             "n_blocks": 3,
@@ -332,7 +415,7 @@ def test_figure_params_old_shape_basics():
         "best_val_loss": 0.5,
         "steps": 20,
     }
-    params = render_mod._figure_params(run_meta)
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params == {
         "hidden_dim": 128,
         "n_blocks": 3,
@@ -346,20 +429,30 @@ def test_figure_params_old_shape_basics():
 
 
 def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps():
-    run_meta = {
+    meta = {
         "model_config": {"mode": "wgan", "noise_dim": 16},
         "steps": 20,
     }
-    params = render_mod._figure_params(run_meta)
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params["noise_dim"] == 16
     assert "steps" not in params
 
 
+def test_figure_params_multi_rollout_names_the_series():
+    run_meta = {"rollouts": {"flow": {"model_config": {"mode": "flow"}}, "wgan": {"model_config": {"mode": "wgan"}}}}
+    assert render_mod._figure_params(run_meta) == {"rollouts": "flow, wgan"}
+
+
+def test_figure_params_empty_rollouts_is_empty():
+    assert render_mod._figure_params({}) == {}
+    assert render_mod._figure_params({"rollouts": {}}) == {}
+
+
 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"})
+    meta = render_mod._plot_metadata(r, {"title": "run-1", "reference": "ref.parquet", "rollouts": {"rollout": {}}})
     assert meta["note"] == "no router data"
-    assert meta["parameters"] == {"checkpoint": "ckpt.pt"}
+    assert meta["parameters"] == {"reference": "ref.parquet", "rollouts": {"rollout": {}}}
     assert "title" not in meta["parameters"]
 
 
@@ -368,3 +461,30 @@ def test_plot_metadata_omits_parameters_when_run_meta_empty():
     meta = render_mod._plot_metadata(r, {})
     assert "parameters" not in meta
     assert "note" not in meta
+
+
+def test_tex_escape_handles_percent_and_other_special_chars():
+    assert render_mod._tex_escape("90% of deposited energy") == r"90\% of deposited energy"
+    assert render_mod._tex_escape(r"a_b & c#d $e {f} \bar") == r"a\_b \& c\#d \$e \{f\} \textbackslash{}bar"
+
+
+def test_render_survives_title_and_xlabel_with_literal_percent(tmp_path: Path):
+    # Regression test for gitea #81: a literal "%" in a catalog title (e.g.
+    # "Shower containment depth (90% of deposited energy)") crashed the whole
+    # LaTeX render, since usetex treats an unescaped "%" as a comment marker.
+    reduced = [
+        Reduced(
+            "shower_containment_depth_90",
+            "shower",
+            "single_hist",
+            "Shower containment depth (90% of deposited energy)",
+            "depth containing 90% of deposited energy [mm]",
+            {"edges": [0, 1, 2], "series": {"flow": [5, 1]}},
+        ),
+    ]
+    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) == 1
+    assert pdfs[0].exists()

tests/test_router.py

@@ -13,6 +13,8 @@ from giant.model.network import (
     EnergyRouter,
     ExpertTrunk,
     FilmResBlock,
+    LinearTrunk,
+    NoneRouter,
     PdgRouter,
     ProcessRouter,
     ROUTER_REGISTRY,
@@ -73,6 +75,34 @@ def test_energy_router_registered():
     assert ROUTER_REGISTRY["energy"] is EnergyRouter
 
 
+# ── NoneRouter (gitea #45) ──────────────────────────────────────────────────
+
+
+def test_none_router_registered():
+    assert ROUTER_REGISTRY["none"] is NoneRouter
+
+
+def test_none_router_gate_is_uniform():
+    router = NoneRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    g = router.gate(cond_cont, cond_cat)
+    assert g.shape == (16, 4)
+    torch.testing.assert_close(g, torch.full((16, 4), 0.25))
+
+
+def test_none_router_gate_ignores_conditioning():
+    router = NoneRouter(n_experts=3)
+    cond_cont_a, cond_cat_a = _cond(8)
+    cond_cont_b, cond_cat_b = _cond(8)
+    torch.testing.assert_close(router.gate(cond_cont_a, cond_cat_a), router.gate(cond_cont_b, cond_cat_b))
+
+
+def test_none_router_top1_always_expert_zero():
+    router = NoneRouter(n_experts=4)
+    cond_cont, cond_cat = _cond(16)
+    assert torch.equal(router.top1(cond_cont, cond_cat), torch.zeros(16, dtype=torch.long))
+
+
 def test_energy_router_gate_partition_of_unity():
     router = EnergyRouter(n_experts=4)
     cond_cont, cond_cat = _cond(16)
@@ -168,6 +198,47 @@ def test_build_expert_body_unknown_type_raises():
     raise AssertionError("expected ValueError for unknown trunk type")
 
 
+# ── LinearTrunk (gitea #45) ──────────────────────────────────────────────────
+
+
+def test_trunk_registry_has_linear():
+    assert "linear" in TRUNK_REGISTRY
+    assert TRUNK_REGISTRY["linear"] is LinearTrunk
+
+
+def test_linear_trunk_forward_shape():
+    trunk = build_expert_body("linear", in_dim=9, out_dim=9, hidden_dim=64, n_blocks=6, cond_dim=12)
+    assert trunk.in_dim == 9
+    assert trunk.out_dim == 9
+    x = torch.randn(5, 9)
+    cond = torch.randn(5, 12)
+    out = trunk(x, cond)
+    assert out.shape == (5, 9)
+
+
+def test_linear_trunk_depends_on_x_and_cond():
+    trunk = build_expert_body("linear", in_dim=9, out_dim=9, hidden_dim=64, n_blocks=6, cond_dim=12)
+    x = torch.randn(5, 9)
+    cond_a = torch.randn(5, 12)
+    cond_b = torch.randn(5, 12)
+    assert not torch.allclose(trunk(x, cond_a), trunk(x, cond_b))
+
+
+def test_routed_linear_trunk_is_mixture_of_trivial_experts():
+    """trunk.type = 'linear' composes for free with router.enabled = true
+    (gitea #33's comment on this issue) — a RoutedTrunk of n_experts linear
+    bodies."""
+    router = build_router("energy", n_experts=3)
+    trunk = RoutedTrunk(router, "linear", in_dim=9, out_dim=9, hidden_dim=64, n_res_blocks=6, cond_dim=12)
+    assert len(trunk.experts) == 3
+    assert all(isinstance(e, LinearTrunk) for e in trunk.experts)
+    x = torch.randn(5, 9)
+    cond = torch.randn(5, 12)
+    cond_cont, cond_cat = _cond(5)
+    out = trunk(x, cond, cond_cont, cond_cat)
+    assert out.shape == (5, 9)
+
+
 # ── BLOCK_REGISTRY / build_block (gitea #34) ────────────────────────────────
 
 
@@ -1195,3 +1266,37 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow():
     )
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_valid.shape == (B, K_MAX)
+
+
+def test_routed_and_unrouted_trunk_agree_on_dtype_under_bf16_autocast():
+    """gitea #47 regression: `_route_forward`'s accumulator (giant/model/
+    trunks.py) used to be a hard-fp32 `torch.zeros`, so under autocast a
+    `RoutedTrunk` returned fp32 while an unrouted `ExpertTrunk` returned
+    bf16 — `router.enabled` alone silently changed the model's output dtype.
+    Checked in both train mode (the differentiable mixture sum) and eval
+    mode (the masked `out[mask] = expert(...)` dispatch) — the two branches
+    of `_route_forward` had independent copies of the bug."""
+    torch.manual_seed(0)
+    cond_cont, cond_cat = _cond(B=6)
+    x = torch.randn(6, X_DIM)
+    t = torch.rand(6)
+
+    unrouted = Stage1Model(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        hidden_dim=16,
+        n_res_blocks=2,
+    )
+    routed = _routed_stage1(n_experts=3)
+
+    for train_mode in (True, False):
+        unrouted.train(train_mode)
+        routed.train(train_mode)
+        with torch.autocast("cpu", dtype=torch.bfloat16, enabled=True):
+            out_unrouted = unrouted(x, cond_cont, cond_cat, t=t)
+            out_routed = routed(x, cond_cont, cond_cat, t=t)
+        assert out_unrouted.dtype == out_routed.dtype, (
+            f"train={train_mode}: unrouted returned {out_unrouted.dtype}, routed returned {out_routed.dtype}"
+        )

tests/test_router_gating.py

@@ -10,7 +10,9 @@ from giant.analysis.router_gating import (
     compute_router_gating,
     compute_router_share_by_pdg,
     compute_router_share_by_process,
+    compute_router_specialization,
 )
+from giant.analysis.sources import RolloutSide
 from giant.data.transforms import Normalizer
 from giant.model.network import build_models
 
@@ -34,7 +36,7 @@ def _model_cfg() -> dict:
     }
 
 
-def _write_checkpoint(tmp_path) -> str:
+def _write_checkpoint(tmp_path, name: str = "ckpt.pt") -> str:
     cfg = _model_cfg()
     stage1 = build_models(cfg)["stage1"]
     assert stage1 is not None
@@ -48,7 +50,7 @@ def _write_checkpoint(tmp_path) -> str:
         "mat_map": _MAT_MAP,
         "normalizer": {"cond": norm.to_dict()},
     }
-    path = tmp_path / "ckpt.pt"
+    path = tmp_path / name
     torch.save(ckpt, path)
     return str(path)
 
@@ -86,42 +88,80 @@ def _steps_frame(process: bool = False) -> pl.LazyFrame:
     return pl.DataFrame(data).lazy()
 
 
+def _side(checkpoint: str | None, lf: pl.LazyFrame) -> RolloutSide:
+    return RolloutSide(all=lf, phys=lf, checkpoint=checkpoint)
+
+
 def test_compute_router_gating_shapes(tmp_path):
     checkpoint = _write_checkpoint(tmp_path)
     lf = _steps_frame()
-    r = compute_router_gating(checkpoint, lf, lf)
+    r = compute_router_gating({"rollout": _side(checkpoint, lf)}, lf)
     assert r.kind == "router_gating"
-    assert r.payload["n_experts"] == 2
+    assert list(r.payload["series"]) == ["rollout"]
+    entry = r.payload["series"]["rollout"]
+    assert entry["n_experts"] == 2
     for side in ("rollout", "reference"):
-        means = r.payload[side]["means"]
+        means = entry[side]["means"]
         assert means, f"{side} produced no bins"
         assert all(abs(sum(row) - 1.0) < 1e-5 for row in means)
 
 
 def test_compute_router_gating_missing_checkpoint_is_unavailable():
     lf = _steps_frame()
-    r = compute_router_gating(None, lf, lf)
+    r = compute_router_gating({"rollout": _side(None, lf)}, lf)
     assert r.kind == "unavailable"
     assert "note" in r.payload
     assert r.title
 
 
+def test_compute_router_gating_two_rollouts_only_moe_ones_included(tmp_path):
+    lf = _steps_frame()
+    ckpt = _write_checkpoint(tmp_path)
+    rollouts = {"flow": _side(None, lf), "moe": _side(ckpt, lf)}
+    r = compute_router_gating(rollouts, lf)
+    assert list(r.payload["series"]) == ["moe"]
+
+
+def test_compute_router_specialization_two_rollouts(tmp_path):
+    lf = _steps_frame()
+    ckpt_a = _write_checkpoint(tmp_path, "a.pt")
+    ckpt_b = _write_checkpoint(tmp_path, "b.pt")
+    rollouts = {"a": _side(ckpt_a, lf), "b": _side(ckpt_b, lf)}
+    r = compute_router_specialization(rollouts, lf)
+    assert r.kind == "router_specialization"
+    assert list(r.payload["series"]) == ["a", "b"]
+    for entry in r.payload["series"].values():
+        assert entry["chance_level"] == 0.5
+        assert len(entry["rollout"]["centers"]) == len(entry["rollout"]["score"])
+
+
 def test_compute_router_share_by_pdg(tmp_path):
     checkpoint = _write_checkpoint(tmp_path)
     lf = _steps_frame()
-    r = compute_router_share_by_pdg(checkpoint, lf, lf, top_pdgs=[11, 22])
+    r = compute_router_share_by_pdg({"rollout": _side(checkpoint, lf)}, lf, top_pdgs=[11, 22])
     assert r.kind == "router_share"
+    entry = r.payload["series"]["rollout"]
     for side in ("rollout", "reference"):
-        assert set(r.payload[side]) == {"e-", "gamma"}
-        for shares in r.payload[side].values():
+        assert set(entry[side]) == {"e-", "gamma"}
+        for shares in entry[side].values():
             assert abs(sum(shares) - 1.0) < 1e-5
 
 
 def test_compute_router_share_by_process(tmp_path):
     checkpoint = _write_checkpoint(tmp_path)
     lf = _steps_frame(process=True)
-    r = compute_router_share_by_process(checkpoint, lf)
+    r = compute_router_share_by_process({"rollout": _side(checkpoint, lf)}, lf)
     assert r.kind == "router_share"
-    assert set(r.payload["categories"]) <= {"eIoni", "compt"}
-    for shares in r.payload["reference"].values():
+    entry = r.payload["series"]["rollout"]
+    assert set(entry["categories"]) <= {"eIoni", "compt"}
+    for shares in entry["reference"].values():
         assert abs(sum(shares) - 1.0) < 1e-5
+
+
+def test_no_moe_rollouts_are_unavailable(tmp_path):
+    lf = _steps_frame()
+    rollouts = {"flow": _side(None, lf), "wgan": _side(None, lf)}
+    assert compute_router_gating(rollouts, lf).kind == "unavailable"
+    assert compute_router_share_by_pdg(rollouts, lf, top_pdgs=[11, 22]).kind == "unavailable"
+    assert compute_router_share_by_process(rollouts, lf).kind == "unavailable"
+    assert compute_router_specialization(rollouts, lf).kind == "unavailable"

tests/test_setup_cache.py

@@ -99,7 +99,7 @@ def test_save_load_round_trip_topn_maps(tmp_path):
 
     cache = SetupCache.empty(files)
     cache.topn_maps[setup_cache.topn_key("pdg", 3)] = TopNMap(
-        class_map={22: 0, 11: 1, 2212: 2}, other_members={2212: 5}
+        class_map={22: 0, 11: 1, 2212: 2}, other_members={2212: 5}, class_counts={0: 100, 1: 50, 2: 5}
     )
     cache.topn_maps[setup_cache.topn_key("material", 2)] = TopNMap(
         class_map={"G4_AIR": 0, "PbWO4": 1}, other_members={}
@@ -114,9 +114,21 @@ def test_save_load_round_trip_topn_maps(tmp_path):
     assert pdg_m.other_members == {2212: 5}
     # key type is int (matches pdg_map's own key type), not str
     assert all(isinstance(k, int) for k in pdg_m.class_map)
+    # class_counts (gitea #44) round-trips too, keyed by class index (always
+    # int, independent of the pdg/material axis's own key type).
+    assert pdg_m.class_counts == {0: 100, 1: 50, 2: 5}
+    assert all(isinstance(k, int) for k in pdg_m.class_counts)
 
     mat_m = loaded.topn_maps[setup_cache.topn_key("material", 2)]
     assert mat_m.class_map == {"G4_AIR": 0, "PbWO4": 1}
+    assert mat_m.class_counts == {}
+
+
+def test_topnmap_from_json_missing_class_counts_defaults_empty():
+    """A checkpoint's topn map predating gitea #44 has no class_counts key at
+    all — must decode to {}, not raise, since inference never reads it."""
+    m = setup_cache.topnmap_from_json({"class_map": {"11": 0}, "other_members": {}}, axis="pdg")
+    assert m.class_counts == {}
 
 
 def test_topn_key_unknown_axis_raises():

tests/test_train.py

@@ -5,10 +5,12 @@ import csv
 import math
 import tempfile
 from pathlib import Path
-from unittest.mock import MagicMock
+from unittest.mock import MagicMock, patch
 
+import numpy as np
 import pytest
 import torch
+import torch.nn.functional as F
 
 from giant.config import ParticleTypeConfig
 from giant.constants import (
@@ -19,16 +21,22 @@ from giant.constants import (
     SEC_SLOT_DIM,
     X_DIM,
 )
+from giant.checkpoint_io import load_for_inference
 from giant.data.dataset import StepBatch
+from giant.data.transforms import Normalizer
 from giant.model.network import Stage2Autoregressive, build_critics, build_models
+from giant.sample import sample_stage1 as trainers_sample_stage1
 from giant.training import (
     FlowDDPMStageTrainer,
     StageSpec,
     WGANStageTrainer,
+    build_checkpoint,
     build_stage_trainers,
+    init_stages_from_checkpoints,
     train,
 )
 from giant.training.metrics import _wandb_run_config
+from giant.training.trainers import _type_class_weight_vector
 from giant.training.stage2_inputs import (
     _ar_has_prev,
     _assemble_stage2_ar_inputs,
@@ -351,7 +359,7 @@ def _model_config(cfg):
     }
 
 
-def _run_train(cfg, out_dir, resume_path=None):
+def _run_train(cfg, out_dir, resume_path=None, normalizer_dict=None):
     model_config = _model_config(cfg)
     models = build_models(model_config)
     critics = build_critics(model_config)
@@ -365,7 +373,7 @@ def _run_train(cfg, out_dir, resume_path=None):
         val_loader=val_loader,
         device=torch.device("cpu"),
         out_dir=out_dir,
-        normalizer_dict={"cond": {}, "target": {}, "sec_phys": {}},
+        normalizer_dict=normalizer_dict or {"cond": {}, "target": {}, "sec_phys": {}},
         pdg_map={"22": 0},
         mat_map={"G4_AIR": 0},
         proc_map=None,
@@ -593,6 +601,313 @@ def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
         FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
 
 
+# --- gitea #44: class-balanced secondary particle-type loss -----------------
+
+
+def test_type_class_weight_vector_none_scheme_returns_none():
+    assert _type_class_weight_vector({0: 100, 1: 5}, n_classes=2, scheme="none") is None
+
+
+def test_type_class_weight_vector_raises_without_counts():
+    with pytest.raises(ValueError, match="class_counts"):
+        _type_class_weight_vector({}, n_classes=4, scheme="inverse_freq")
+
+
+def test_type_class_weight_vector_inverse_freq_favors_rare_class_and_has_mean_one():
+    weights = _type_class_weight_vector({0: 1000, 1: 10, 2: 1, 3: 1}, n_classes=4, scheme="inverse_freq")
+    assert weights is not None
+    assert len(weights) == 4
+    assert weights[1] > weights[0]  # rarer class -> larger weight
+    assert math.isclose(sum(weights) / len(weights), 1.0, rel_tol=1e-9)
+
+
+def test_type_class_weight_vector_missing_index_clamps_to_count_one():
+    # n_classes=3 but only index 0 was ever observed (e.g. a tiny dataset) —
+    # indices 1/2 must not divide by zero.
+    weights = _type_class_weight_vector({0: 10}, n_classes=3, scheme="inverse_freq")
+    assert weights is not None
+    assert all(math.isfinite(w) for w in weights)
+
+
+def _onehot_flow_stage2_setup():
+    """A built stage-2 model + a batch, under target='onehot' + generator='flow'
+    (mirrors the 'stage2_onehot_target_flow' case in test_train_end_to_end)."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0}
+    model_config = _model_config(cfg)
+    model = build_models(model_config)["stage2"]
+    assert model is not None
+    batch = _fake_batches(1, 8)[0]
+    device = torch.device("cpu")
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+    # Mostly class 0 (common), a few slot 1's set to class 1 (rare) —
+    # PARTICLE_CFG's emb_dim=8, n_classes=0 (inherit) -> 8 type classes.
+    sec_type_idx = torch.zeros(8, K_MAX, dtype=torch.long)
+    sec_type_idx[:, :2] = 1
+    sec_mask = torch.ones(8, K_MAX, dtype=torch.bool)
+    return model, cond_cont, cond_cat, x1_s1, sec_type_idx, sec_mask, device
+
+
+def test_flow_ddpm_trainer_type_loss_none_leaves_weight_unset():
+    model, *_ = _onehot_flow_stage2_setup()
+    spec = StageSpec(
+        name="stage2",
+        is_stage2=True,
+        generator="flow",
+        particle_type=ParticleTypeConfig(target="onehot", class_weighting="none"),
+        particle_type_n_classes=8,
+        ema_decay=0.0,
+    )
+    trainer = FlowDDPMStageTrainer(spec, model, torch.device("cpu"))
+    assert trainer.type_class_weights is None
+
+
+def test_flow_ddpm_trainer_type_loss_matches_manual_weighted_cross_entropy():
+    model, cond_cont, cond_cat, x1_s1, sec_type_idx, sec_mask, device = _onehot_flow_stage2_setup()
+    class_counts = {0: 1000, 1: 10, 2: 1, 3: 1, 4: 1, 5: 1, 6: 1, 7: 1}
+    weights = _type_class_weight_vector(class_counts, n_classes=8, scheme="inverse_freq")
+    spec = StageSpec(
+        name="stage2",
+        is_stage2=True,
+        generator="flow",
+        particle_type=ParticleTypeConfig(target="onehot", class_weighting="inverse_freq"),
+        particle_type_n_classes=8,
+        type_class_weights=weights,
+        ema_decay=0.0,
+    )
+    trainer = FlowDDPMStageTrainer(spec, model, device)
+    assert trainer.type_class_weights is not None
+    stage1_ctx = trainer._stage1_context(x1_s1, cond_cont, cond_cat, epoch=None)
+
+    with torch.no_grad():
+        type_out = model.predict_type(cond_cont, cond_cat, stage1_ctx)
+        weight_t = torch.tensor(weights)
+        ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, weight=weight_t, reduction="none")
+        expected = (ce * sec_mask.float()).sum() / sec_mask.float().sum().clamp(min=1)
+
+        l_type, _ = trainer._type_loss(cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device)
+
+    assert torch.allclose(l_type, expected, atol=1e-6)
+
+    # Unweighted trainer, same model/batch — the two losses must differ
+    # (the batch mixes the common and rare classes, so weighting changes the
+    # per-slot contributions), confirming the weight is actually plumbed in.
+    spec_none = StageSpec(
+        name="stage2",
+        is_stage2=True,
+        generator="flow",
+        particle_type=ParticleTypeConfig(target="onehot", class_weighting="none"),
+        particle_type_n_classes=8,
+        ema_decay=0.0,
+    )
+    trainer_none = FlowDDPMStageTrainer(spec_none, model, device)
+    with torch.no_grad():
+        l_type_none, _ = trainer_none._type_loss(cond_cont, cond_cat, stage1_ctx, sec_type_idx, sec_mask, device)
+    assert not torch.allclose(l_type, l_type_none)
+
+
+def test_build_stage_trainers_threads_sec_type_class_counts_into_weights():
+    cfg = _base_cfg()
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"]["particle_type"] = {
+        "target": "onehot",
+        "lambda": 1.0,
+        "class_weighting": "inverse_freq",
+    }
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    class_counts = {i: 100 for i in range(8)}
+    class_counts[1] = 1  # one rare class
+    trainers = build_stage_trainers(
+        cfg, models, critics, torch.device("cpu"), total_train_batches=4, sec_type_class_counts=class_counts
+    )
+    stage2_trainer = trainers["stage2"]
+    assert isinstance(stage2_trainer, FlowDDPMStageTrainer)
+    weights = stage2_trainer.type_class_weights
+    assert weights is not None
+    assert weights[1] > weights[0]
+
+
+def test_build_stage_trainers_no_class_counts_with_none_weighting_is_fine():
+    """The overwhelmingly common case (class_weighting = 'none', the
+    default): build_stage_trainers must not require sec_type_class_counts at
+    all."""
+    cfg = _base_cfg()
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0})
+    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)
+    stage2_trainer = trainers["stage2"]
+    assert isinstance(stage2_trainer, FlowDDPMStageTrainer)
+    assert stage2_trainer.type_class_weights is None
+
+
+# --- gitea #42: freeze / init_from -------------------------------------------
+
+
+def _state_dict_clone(module):
+    return {k: v.clone() for k, v in module.state_dict().items()}
+
+
+def _assert_state_dicts_equal(before, after, label):
+    for key, value in before.items():
+        assert torch.equal(value, after[key]), f"{label}: {key} changed while frozen"
+
+
+def test_frozen_flow_stage_trainer_step_does_not_update_model_or_ema():
+    cfg = _base_cfg()
+    model_config = _model_config(cfg)
+    model = build_models(model_config)["stage1"]
+    assert model is not None
+    spec = StageSpec(name="stage1", is_stage2=False, generator="flow", freeze=True, ema_decay=0.999, steps_per_epoch=4)
+    trainer = FlowDDPMStageTrainer(spec, model, torch.device("cpu"))
+    assert trainer.ema_model is not None
+    model_before = _state_dict_clone(trainer.model)
+    ema_before = _state_dict_clone(trainer.ema_model)
+    for batch in _fake_batches(4, 8):
+        trainer.step(batch, torch.device("cpu"), global_step=1)
+    _assert_state_dicts_equal(model_before, trainer.model.state_dict(), "frozen flow model")
+    _assert_state_dicts_equal(ema_before, trainer.ema_model.state_dict(), "frozen flow ema")
+
+
+def test_frozen_wgan_stage_trainer_step_does_not_update_generator_or_critic():
+    cfg = _base_cfg()
+    cfg["stage1_model"]["generator"] = "wgan"
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    assert models["stage1"] is not None and critics["stage1"] is not None
+    spec = StageSpec(
+        name="stage1",
+        is_stage2=False,
+        generator="wgan",
+        freeze=True,
+        n_critic=1,  # a generator step every batch, so a bug would surface immediately
+        ema_decay=0.999,
+        steps_per_epoch=4,
+    )
+    trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu"))
+    assert trainer.ema_model is not None
+    model_before = _state_dict_clone(trainer.model)
+    critic_before = _state_dict_clone(trainer.critic)
+    ema_before = _state_dict_clone(trainer.ema_model)
+    for global_step, batch in enumerate(_fake_batches(4, 8)):
+        trainer.step(batch, torch.device("cpu"), global_step=global_step)
+    _assert_state_dicts_equal(model_before, trainer.model.state_dict(), "frozen wgan generator")
+    _assert_state_dicts_equal(critic_before, trainer.critic.state_dict(), "frozen wgan critic")
+    _assert_state_dicts_equal(ema_before, trainer.ema_model.state_dict(), "frozen wgan ema")
+
+
+@pytest.mark.parametrize("stage1_generator", ["flow", "wgan"])
+def test_train_end_to_end_frozen_stage1_unchanged_while_stage2_trains(stage1_generator):
+    cfg = _base_cfg()
+    cfg["stage1_model"]["generator"] = stage1_generator
+    cfg["stage1_model"]["freeze"] = True
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    assert models["stage1"] is not None and models["stage2"] is not None
+    stage1_before = _state_dict_clone(models["stage1"])
+    stage2_before = _state_dict_clone(models["stage2"])
+    with tempfile.TemporaryDirectory() as tmp:
+        train(
+            cfg=cfg,
+            models=models,
+            critics=critics,
+            train_loader=_fake_batches(4, cfg["train"]["batch_size"]),
+            val_loader=_fake_batches(2, cfg["train"]["batch_size"], seed=1),
+            device=torch.device("cpu"),
+            out_dir=Path(tmp) / "run",
+            normalizer_dict={"cond": {}, "target": {}, "sec_phys": {}},
+            pdg_map={"22": 0},
+            mat_map={"G4_AIR": 0},
+            proc_map=None,
+            model_config=model_config,
+            total_train_batches=4,
+        )
+    _assert_state_dicts_equal(stage1_before, models["stage1"].state_dict(), "frozen stage1")
+    stage2_after = models["stage2"].state_dict()
+    assert any(not torch.equal(v, stage2_after[k]) for k, v in stage2_before.items()), (
+        "unfrozen stage2 should have trained"
+    )
+
+
+def test_init_stages_from_checkpoints_loads_matching_stage_and_ema_weights(tmp_path):
+    cfg = _base_cfg()
+    model_config = _model_config(cfg)
+    source_models = build_models(model_config)
+    source_critics = build_critics(model_config)
+    source_trainers = build_stage_trainers(cfg, source_models, source_critics, torch.device("cpu"), 4)
+    source_stage1_ema = source_trainers["stage1"].ema_model
+    assert source_stage1_ema is not None
+    # Diverge the source's EMA from its raw weights so a same-vs-different
+    # check below actually distinguishes the two copy paths.
+    for p in source_stage1_ema.parameters():
+        p.data.add_(1.0)
+    ckpt_path = tmp_path / "source.pt"
+    ckpt = build_checkpoint(source_trainers, epoch=1, global_step=1, best_val_loss=0.0, extras={})
+    torch.save(ckpt, ckpt_path)
+
+    cfg2 = copy.deepcopy(cfg)
+    cfg2["stage1_model"]["init_from"] = str(ckpt_path)
+    dest_models = build_models(_model_config(cfg2))
+    dest_critics = build_critics(_model_config(cfg2))
+    dest_trainers = build_stage_trainers(cfg2, dest_models, dest_critics, torch.device("cpu"), 4)
+
+    loaded = init_stages_from_checkpoints(dest_trainers)
+    assert len(loaded) == 1 and "stage1" in loaded[0]
+    dest_stage1_ema = dest_trainers["stage1"].ema_model
+    assert dest_stage1_ema is not None
+
+    _assert_state_dicts_equal(
+        source_trainers["stage1"].model.state_dict(), dest_trainers["stage1"].model.state_dict(), "init_from raw"
+    )
+    _assert_state_dicts_equal(
+        source_stage1_ema.state_dict(),
+        dest_stage1_ema.state_dict(),
+        "init_from ema",
+    )
+    # stage2 has no init_from set -- untouched fresh init, not the source's.
+    stage2_matches_source = all(
+        torch.equal(v, dest_trainers["stage2"].model.state_dict()[k])
+        for k, v in source_trainers["stage2"].model.state_dict().items()
+    )
+    assert not stage2_matches_source
+
+
+def test_run_train_job_stage1_init_from_freeze_produces_rollout_capable_checkpoint(tmp_path):
+    """The exact scenario gitea #42 exists for: retrain stage 2 alone against
+    a fixed, known-good stage 1, and still get a checkpoint giant rollout can
+    load (checkpoint_io.load_for_inference with require_stage2=True)."""
+    normalizer_dict = {
+        "cond": Normalizer().fit(np.zeros((1, COND_DIM), dtype=np.float32)).to_dict(),
+        "target": Normalizer().fit(np.zeros((1, X_DIM), dtype=np.float32)).to_dict(),
+        "sec_phys": Normalizer().fit(np.zeros((1, 2), dtype=np.float32)).to_dict(),
+    }
+
+    cfg = _base_cfg()
+    source_out = tmp_path / "source"
+    _run_train(cfg, source_out, normalizer_dict=normalizer_dict)
+    source_ckpt = torch.load(source_out / "best.pt", weights_only=False)
+
+    cfg2 = copy.deepcopy(cfg)
+    cfg2["stage1_model"]["init_from"] = str(source_out / "best.pt")
+    cfg2["stage1_model"]["freeze"] = True
+    retrain_out = tmp_path / "retrain"
+    _run_train(cfg2, retrain_out, normalizer_dict=normalizer_dict)
+
+    ctx = load_for_inference(retrain_out / "best.pt", torch.device("cpu"), "rollout", require_stage2=True)
+    assert ctx.stage1 is not None and ctx.stage2 is not None
+
+    retrain_ckpt = torch.load(retrain_out / "best.pt", weights_only=False)
+    for key, value in source_ckpt["model"].items():
+        assert torch.equal(value, retrain_ckpt["model"][key]), f"frozen stage1 {key} drifted across the retrain"
+
+
 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
@@ -831,3 +1146,127 @@ def test_wgan_physical_omits_grad_norm_slice_columns():
         header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
         assert "stage2/train/grad_norm_type_slice" not in header
         assert "stage2/train/grad_norm_cont_slice" not in header
+
+
+# --- stage2_model.stage1_context = "sampled" (gitea #41) --------------------
+
+
+def _sampled_ctx_cfg(ema_decay=0.999):
+    cfg = _base_cfg()
+    cfg["stage1_model"]["generator"] = "flow"
+    cfg["stage2_model"]["generator"] = "flow"
+    cfg["stage2_model"]["stage1_context"] = "sampled"
+    cfg["stage2_model"]["ctx_p_start"] = 0.0
+    cfg["stage2_model"]["ctx_p_end"] = 0.0
+    cfg["train"]["ema_decay"] = ema_decay
+    return cfg
+
+
+def _build_sampled_trainers(cfg):
+    model_config = _model_config(cfg)
+    models = build_models(model_config)
+    critics = build_critics(model_config)
+    return build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
+
+
+def test_build_stage_trainers_attaches_stage1_only_under_sampled():
+    trainers = _build_sampled_trainers(_sampled_ctx_cfg())
+    assert trainers["stage2"].stage1_source is trainers["stage1"]
+    assert trainers["stage1"].stage1_source is None
+
+
+def test_build_stage_trainers_leaves_stage1_source_none_under_truth():
+    """Regression guard for the old silent no-op: 'truth' (the default) must
+    never attach a stage1_source, so _stage1_context short-circuits without
+    ever calling sample_stage1."""
+    cfg = _base_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    assert trainers["stage2"].stage1_source is None
+
+
+def test_stage1_context_sampled_calls_sample_stage1_and_differs_from_truth():
+    cfg = _sampled_ctx_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    stage1, stage2 = trainers["stage1"], trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy:
+        ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0)
+    assert spy.call_count == 1
+    assert spy.call_args.args[0] is stage1.sampling_model()
+    assert not torch.equal(ctx, x1_s1)
+
+
+def test_stage1_context_truth_default_never_calls_sample_stage1():
+    cfg = _base_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    stage2 = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy:
+        ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0)
+    assert spy.call_count == 0
+    assert torch.equal(ctx, x1_s1)
+
+
+def test_stage1_context_val_epoch_none_uses_ground_truth_even_under_sampled():
+    cfg = _sampled_ctx_cfg()
+    trainers = _build_sampled_trainers(cfg)
+    stage2 = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    with patch("giant.training.trainers.sample_stage1", wraps=trainers_sample_stage1) as spy:
+        ctx = stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=None)
+    assert spy.call_count == 0
+    assert torch.equal(ctx, x1_s1)
+
+
+def test_stage1_context_sampled_preserves_stage1_training_mode():
+    """Every sampler in giant/sample.py flips its model to .eval() as a side
+    effect with no restore of its own (see sample_flow). Sampling from the
+    RAW stage-1 model (ema_decay=0, so sampling_model() returns self.model,
+    the same weights the stage-1 trainer is actively training on) must not
+    silently leave it in eval mode for the rest of the epoch's stage-1
+    updates."""
+    cfg = _sampled_ctx_cfg(ema_decay=0.0)
+    trainers = _build_sampled_trainers(cfg)
+    stage1, stage2 = trainers["stage1"], trainers["stage2"]
+    stage1.train_mode()
+    assert stage1.model.training
+    batch = _fake_batches(1, 4)[0]
+    cond_cont, cond_cat, x1_s1 = batch.cond_cont, batch.cond_cat, batch.target_s1
+
+    stage2._stage1_context(x1_s1, cond_cont, cond_cat, epoch=0)
+    assert stage1.model.training
+
+
+@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
+def test_build_stage_trainers_sampled_step_runs(stage2_generator):
+    """Both trainer subclasses' call sites (FlowDDPMStageTrainer._compute,
+    WGANStageTrainer.step) must run end to end under 'sampled' and produce a
+    finite loss."""
+    cfg = _sampled_ctx_cfg()
+    cfg["stage2_model"]["generator"] = stage2_generator
+    trainers = _build_sampled_trainers(cfg)
+    trainer = trainers["stage2"]
+    batch = _fake_batches(1, 4)[0]
+    stats = trainer.step(batch, torch.device("cpu"), global_step=1)
+    loss_key = "g_loss" if stage2_generator == "wgan" else "loss"
+    assert math.isfinite(stats[loss_key])
+
+
+def test_train_end_to_end_stage1_context_sampled():
+    """Full train() run with stage1_context='sampled' must complete and
+    write a checkpoint + metrics.csv with finite losses throughout."""
+    cfg = _sampled_ctx_cfg()
+    with tempfile.TemporaryDirectory() as tmp:
+        out_dir = Path(tmp) / "run"
+        _run_train(cfg, out_dir)
+        assert (out_dir / "last.pt").exists()
+        with open(out_dir / "metrics.csv", newline="") as f:
+            rows = list(csv.DictReader(f))
+        assert len(rows) == cfg["train"]["epochs"]
+        assert all(math.isfinite(float(r["stage2/train/loss"])) for r in rows)

tests/test_training_plots.py

@@ -0,0 +1,333 @@
+"""Tests for giant.training.plots (gitea #75) — render smoke tests skipped
+where plotstyle/LaTeX is unavailable, plus pure-function column-classification
+coverage that needs neither."""
+
+from __future__ import annotations
+
+import csv
+from pathlib import Path
+
+import pytest
+
+pytest.importorskip("plotstyle")
+
+from giant.training import plots as plots_mod  # noqa: E402
+from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics  # noqa: E402
+
+# --- fixtures ----------------------------------------------------------
+
+
+_RICH_HEADER = [
+    "epoch",
+    "stage1/train/loss",
+    "stage1/train/loss_gen",
+    "stage1/train/nsec_acc",
+    "stage1/train/grad_norm",
+    "stage1/val/loss",
+    "stage1/val/loss_gen",
+    "stage1/val/nsec_acc",
+    "stage1/lr",
+    "stage1/router/entropy",
+    "stage1/router/util_min",
+    "stage1/router/util_max",
+    "stage1/router/util_std",
+    "stage2/train/d_loss",
+    "stage2/train/g_loss",
+    "stage2/train/wasserstein",
+    "stage2/train/gp_loss",
+    "stage2/train/loss_nsec",
+    "stage2/train/nsec_acc",
+    "stage2/train/grad_norm_d",
+    "stage2/train/grad_norm_g",
+    "stage2/lr",
+    "stage2/critic_lr",
+    "val/loss",
+    "val/marginal_kl",
+    "grad_norm",
+    "gpu_mem_mb",
+    "samples_per_sec",
+    "is_best",
+    "epoch_time_s",
+]
+
+_RICH_ROWS = [
+    [
+        1,
+        1.0,
+        0.8,
+        0.5,
+        1.2,
+        0.9,
+        0.7,
+        0.6,
+        3e-4,
+        1.5,
+        0.05,
+        0.3,
+        0.1,
+        -0.2,
+        0.3,
+        0.5,
+        0.1,
+        0.4,
+        0.4,
+        0.9,
+        1.1,
+        3e-4,
+        1e-4,
+        0.85,
+        0.4,
+        2.1,
+        512.0,
+        100.0,
+        1,
+        5.0,
+    ],
+    [
+        2,
+        0.8,
+        0.6,
+        0.6,
+        1.0,
+        0.7,
+        0.5,
+        0.7,
+        2e-4,
+        1.6,
+        0.06,
+        0.28,
+        0.09,
+        -0.1,
+        0.25,
+        0.4,
+        0.09,
+        0.3,
+        0.5,
+        0.8,
+        1.0,
+        2e-4,
+        8e-5,
+        0.7,
+        0.35,
+        1.9,
+        520.0,
+        105.0,
+        0,
+        5.1,
+    ],
+]
+
+_MINIMAL_HEADER = [
+    "epoch",
+    "stage1/train/loss",
+    "stage1/train/loss_gen",
+    "stage1/val/loss",
+    "stage1/val/loss_gen",
+    "stage1/lr",
+    "val/loss",
+    "grad_norm",
+    "gpu_mem_mb",
+    "samples_per_sec",
+    "is_best",
+    "epoch_time_s",
+]
+
+_MINIMAL_ROWS = [
+    [1, 1.0, 0.8, 0.9, 0.7, 3e-4, 0.85, 0.4, 0.0, 100.0, 0, 5.0],
+    [2, 0.8, 0.6, 0.7, 0.5, 2e-4, 0.7, 0.35, 0.0, 105.0, 1, 5.1],
+]
+
+
+def _write_csv(path: Path, header: list[str], rows: list[list]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with open(path, "w", newline="") as f:
+        writer = csv.writer(f)
+        writer.writerow(header)
+        writer.writerows(rows)
+
+
+# --- MetricsTable --------------------------------------------------------
+
+
+def test_metrics_table_load_round_trips(tmp_path: Path):
+    csv_path = tmp_path / "metrics.csv"
+    _write_csv(csv_path, _MINIMAL_HEADER, _MINIMAL_ROWS)
+    table = MetricsTable.load(csv_path)
+    assert table.epochs == [1, 2]
+    assert table.columns["stage1/train/loss"] == [1.0, 0.8]
+    assert "epoch" not in table.columns
+    assert table.best_epochs() == [2]
+
+
+def test_metrics_table_best_epochs_empty_without_is_best_column():
+    table = MetricsTable(epochs=[1, 2], columns={"stage1/train/loss": [1.0, 0.5]})
+    assert table.best_epochs() == []
+
+
+# --- column classification (pure functions, no matplotlib) --------------
+
+
+def _rich_columns() -> dict[str, list]:
+    return {name: [0.0] for name in _RICH_HEADER if name != "epoch"}
+
+
+def test_stages_detects_only_stages_present():
+    assert plots_mod._stages(_rich_columns()) == ["stage1", "stage2"]
+    assert plots_mod._stages({"stage2/train/loss": [0.0]}) == ["stage2"]
+    assert plots_mod._stages({"val/loss": [0.0]}) == []
+
+
+def test_split_matches_stage_and_split_prefix_only():
+    cols = _rich_columns()
+    train = plots_mod._split(cols, "stage1", "train")
+    assert train == {
+        "loss": "stage1/train/loss",
+        "loss_gen": "stage1/train/loss_gen",
+        "nsec_acc": "stage1/train/nsec_acc",
+        "grad_norm": "stage1/train/grad_norm",
+    }
+    assert plots_mod._split(cols, "stage2", "val") == {}
+
+
+def test_point_in_time_excludes_train_val_router():
+    cols = _rich_columns()
+    pit = plots_mod._point_in_time(cols, "stage1")
+    assert pit == {"lr": "stage1/lr"}
+    pit2 = plots_mod._point_in_time(cols, "stage2")
+    assert pit2 == {"lr": "stage2/lr", "critic_lr": "stage2/critic_lr"}
+
+
+def test_router_columns():
+    cols = _rich_columns()
+    assert plots_mod._router(cols, "stage1") == {
+        "entropy": "stage1/router/entropy",
+        "util_min": "stage1/router/util_min",
+        "util_max": "stage1/router/util_max",
+        "util_std": "stage1/router/util_std",
+    }
+    assert plots_mod._router(cols, "stage2") == {}
+
+
+def test_run_level_excludes_stage_prefixed_columns_including_val_loss_lookalike():
+    cols = _rich_columns()
+    run_level = plots_mod._run_level(cols)
+    assert set(run_level) == {
+        "val/loss",
+        "val/marginal_kl",
+        "grad_norm",
+        "gpu_mem_mb",
+        "samples_per_sec",
+        "is_best",
+        "epoch_time_s",
+    }
+    # stage-prefixed "val/loss" lookalike (stage1/val/loss) must not leak in
+    assert "stage1/val/loss" not in run_level
+
+
+def test_loss_keys_excludes_acc_and_wgan_and_grad_norm():
+    train = {"loss": "x", "loss_gen": "x", "nsec_acc": "x", "grad_norm": "x", "d_loss": "x"}
+    val = {"loss": "x", "loss_gen": "x"}
+    assert plots_mod._loss_keys(train, val) == ["loss", "loss_gen"]
+
+
+# --- derive_metrics_dir ---------------------------------------------------
+
+
+def test_derive_metrics_dir_explicit_out_dir_wins():
+    assert derive_metrics_dir("runs/my-run", out_dir="/somewhere") == Path("/somewhere")
+
+
+def test_derive_metrics_dir_default_base():
+    assert derive_metrics_dir("runs/my-run", default_base="/data/analysis_runs") == Path(
+        "/data/analysis_runs/metrics_my-run"
+    )
+
+
+def test_derive_metrics_dir_falls_back_to_cwd_analysis_runs(monkeypatch, tmp_path):
+    monkeypatch.chdir(tmp_path)
+    assert derive_metrics_dir("runs/my-run") == tmp_path / "analysis_runs" / "metrics_my-run"
+
+
+# --- render_metrics end to end -------------------------------------------
+
+
+def _try_render(run_dir: Path, out_dir: Path) -> list[Path]:
+    try:
+        return render_metrics(run_dir, out_dir)
+    except RuntimeError as e:  # LaTeX missing at render time
+        pytest.skip(f"LaTeX rendering unavailable: {e}")
+
+
+def test_render_metrics_rich_run_produces_expected_plots_outside_run_dir(tmp_path: Path):
+    run_dir = tmp_path / "run"
+    out_dir = tmp_path / "out"
+    _write_csv(run_dir / "metrics.csv", _RICH_HEADER, _RICH_ROWS)
+
+    paths = _try_render(run_dir, out_dir)
+
+    names = {p.stem for p in paths}
+    assert names == {
+        "overview",
+        "stage1_loss",
+        "stage2_loss",
+        "lr",
+        "stage1_accuracy",
+        "stage2_accuracy",
+        "grad_norm",
+        "stage1_router",
+        "stage2_wgan_balance",
+        "throughput",
+    }
+    assert all(p.exists() for p in paths)
+    assert all(p.is_relative_to(out_dir) for p in paths)
+    # nothing written into the training run directory itself
+    assert not any(run_dir.rglob("*.pdf"))
+
+
+def test_render_metrics_minimal_run_omits_router_wgan_accuracy(tmp_path: Path):
+    run_dir = tmp_path / "run"
+    out_dir = tmp_path / "out"
+    _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)
+
+    paths = _try_render(run_dir, out_dir)
+
+    names = {p.stem for p in paths}
+    assert names == {"overview", "stage1_loss", "lr", "grad_norm", "throughput"}
+    assert "stage1_accuracy" not in names
+    assert "stage1_router" not in names
+    assert "stage1_wgan_balance" not in names
+
+
+def test_render_metrics_default_out_dir_uses_analysis_runs_convention(tmp_path: Path):
+    run_dir = tmp_path / "runs" / "my-run"
+    _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)
+    default_base = tmp_path / "analysis_runs"
+
+    try:
+        paths = render_metrics(run_dir, default_base=default_base)
+    except RuntimeError as e:
+        pytest.skip(f"LaTeX rendering unavailable: {e}")
+
+    assert paths
+    assert all(p.is_relative_to(default_base / "metrics_my-run") for p in paths)
+
+
+# --- CLI -------------------------------------------------------------------
+
+
+def test_cli_analyze_metrics_smoke(tmp_path: Path):
+    from typer.testing import CliRunner
+
+    from giant.cli import app
+
+    run_dir = tmp_path / "run"
+    out_dir = tmp_path / "out"
+    _write_csv(run_dir / "metrics.csv", _MINIMAL_HEADER, _MINIMAL_ROWS)
+
+    runner = CliRunner()
+    result = runner.invoke(app, ["analyze", "metrics", str(run_dir), "--out", str(out_dir)])
+
+    if result.exit_code != 0 and "latex" in (str(result.output) + str(result.exception)).lower():
+        pytest.skip("LaTeX rendering unavailable")
+    assert result.exit_code == 0, result.output or result.exception
+    assert any(out_dir.glob("*.pdf"))

tests/test_type_embedding_distance.py

@@ -3,6 +3,9 @@
 
 from __future__ import annotations
 
+import polars as pl
+
+from giant.analysis.sources import RolloutSide
 from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
 
 
@@ -18,26 +21,47 @@ def _summary(n=100):
     }
 
 
+def _side(l1_dist: dict | None) -> RolloutSide:
+    empty = pl.LazyFrame()
+    return RolloutSide(all=empty, phys=empty, type_embedding_l1_dist=l1_dist)
+
+
 def test_none_is_unavailable():
-    r = compute_type_embedding_l1_distance(None)
+    r = compute_type_embedding_l1_distance({"rollout": _side(None)})
     assert r.kind == "unavailable"
     assert r.id == "type_embedding_l1_distance"
     assert r.payload["note"]
 
 
 def test_summary_produces_single_hist():
-    r = compute_type_embedding_l1_distance(_summary())
+    r = compute_type_embedding_l1_distance({"rollout": _side(_summary())})
     assert r.kind == "single_hist"
     assert r.id == "type_embedding_l1_distance"
     assert r.payload["edges"] == [0.0, 1.0, 2.0, 3.0]
-    assert r.payload["rollout"] == [30, 40, 30]
+    assert r.payload["series"]["rollout"] == [30, 40, 30]
     assert r.payload["log_x"] is True
     assert r.payload["log_y"] is True
     assert "n=100" in r.payload["note"]
 
 
 def test_single_hist_payload_shape_matches_render_contract():
-    """_render_single (giant.analysis.render) requires len(rollout) ==
+    """_render_single (giant.analysis.render) requires each series' length ==
     len(edges) - 1."""
-    r = compute_type_embedding_l1_distance(_summary())
-    assert len(r.payload["rollout"]) == len(r.payload["edges"]) - 1
+    r = compute_type_embedding_l1_distance({"rollout": _side(_summary())})
+    assert len(r.payload["series"]["rollout"]) == len(r.payload["edges"]) - 1
+
+
+def test_two_rollouts_both_populated():
+    r = compute_type_embedding_l1_distance({"flow": _side(_summary(50)), "wgan": _side(_summary(80))})
+    assert list(r.payload["series"]) == ["flow", "wgan"]
+    assert "n=50" in r.payload["note"] and "n=80" in r.payload["note"]
+
+
+def test_one_of_two_rollouts_populated_only_that_one_appears():
+    r = compute_type_embedding_l1_distance({"flow": _side(None), "wgan": _side(_summary())})
+    assert list(r.payload["series"]) == ["wgan"]
+
+
+def test_none_populated_across_rollouts_is_unavailable():
+    r = compute_type_embedding_l1_distance({"flow": _side(None), "wgan": _side(None)})
+    assert r.kind == "unavailable"

tests/test_wgan.py

@@ -2,7 +2,7 @@ import torch
 
 from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
-from giant.model.network import CriticModel, Stage1Model, Stage2OneShot
+from giant.model.network import CriticModel, LinearTrunk, Stage1Model, Stage2OneShot
 from giant.model.wgan import critic_loss, generator_loss, gradient_penalty
 from giant.sample import sample_secondaries_wgan, sample_wgan
 
@@ -115,6 +115,33 @@ def test_critic_output_shape():
     assert out.shape == (B,)
 
 
+def test_critic_model_honours_trunk_type_and_block_conditioning():
+    """gitea #57: CriticModel routes its body through build_trunk/build_block
+    like every generator stage model, instead of hand-rolling a plain
+    ResBlock stack."""
+    B = 8
+    critic = CriticModel(
+        pdg_vocab=3,
+        mat_vocab=2,
+        particle_cfg=PARTICLE_CFG,
+        material_cfg=MATERIAL_CFG,
+        in_dim=X_DIM,
+        hidden_dim=32,
+        n_res_blocks=2,
+        stage="stage1",
+        trunk_type="linear",
+        block_conditioning="adaln",
+    )
+    assert isinstance(critic.trunk, LinearTrunk)
+    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.backward()
+    for name, p in critic.named_parameters():
+        assert p.grad is not None, f"no grad for {name}"
+
+
 def test_sample_wgan_shape():
     B = 6
     model = _small_generator()

uv.lock

@@ -5,19 +5,19 @@ resolution-markers = [
     "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
     "python_full_version >= '3.14' and sys_platform == 'emscripten' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
-    "python_full_version < '3.14' and sys_platform == 'win32' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
-    "python_full_version < '3.14' and sys_platform == 'emscripten' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
-    "python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
     "python_full_version >= '3.14' and sys_platform == 'win32' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
     "python_full_version >= '3.14' and sys_platform == 'emscripten' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
     "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
+    "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
+    "python_full_version >= '3.14' and sys_platform == 'emscripten' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
+    "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
+    "python_full_version < '3.14' and sys_platform == 'win32' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
+    "python_full_version < '3.14' and sys_platform == 'emscripten' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
+    "python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda'",
     "python_full_version < '3.14' and sys_platform == 'win32' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
     "python_full_version < '3.14' and sys_platform == 'emscripten' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
     "(python_full_version == '3.13.*' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda') or (python_full_version < '3.14' and platform_machine != 'aarch64' and sys_platform == 'linux' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda') or (python_full_version < '3.14' and platform_python_implementation != 'CPython' and sys_platform == 'linux' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda') or (python_full_version < '3.14' and sys_platform != 'darwin' and sys_platform != 'emscripten' and sys_platform != 'linux' and sys_platform != 'win32' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda')",
     "(python_full_version < '3.13' and platform_machine == 'aarch64' and platform_python_implementation == 'CPython' and sys_platform == 'linux' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda') or (python_full_version < '3.14' and sys_platform == 'darwin' and extra == 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda')",
-    "python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
-    "python_full_version >= '3.14' and sys_platform == 'emscripten' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
-    "python_full_version >= '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
     "python_full_version < '3.14' and sys_platform == 'win32' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
     "python_full_version < '3.14' and sys_platform == 'emscripten' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
     "python_full_version < '3.14' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-5-giant-cpu' and extra != 'extra-5-giant-cuda'",
@@ -45,6 +45,19 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" },
 ]
 
+[[package]]
+name = "anyio"
+version = "4.14.2"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "idna", marker = "sys_platform != 'emscripten' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+    { name = "typing-extensions", marker = "(python_full_version < '3.13' and sys_platform != 'emscripten') or (python_full_version >= '3.13' and extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda') or (sys_platform == 'emscripten' and extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/61/cc/a381afa6efea9f496eff839d4a6a1aed3bfafc7b3ab4b0d1b243a12573dd/anyio-4.14.2.tar.gz", hash = "sha256:cfa139f3ed1a23ee8f88a145ddb5ac7605b8bbfd8592baacd7ce3d8bb4313c7f", size = 260176, upload-time = "2026-07-12T20:29:07.082Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/da/35/f2287558c17e29fafc8ef3daf819bb9834061cfa43bff8014f7df7f63bdc/anyio-4.14.2-py3-none-any.whl", hash = "sha256:9f505dda5ac9f0c8309b5e8bd445a8c2bf7246f3ce950121e45ea15bc41d1494", size = 125813, upload-time = "2026-07-12T20:29:05.763Z" },
+]
+
 [[package]]
 name = "appnope"
 version = "0.1.4"
@@ -134,6 +147,35 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/44/a1/70ebfffd6c6edc6034a547838ee46287c65ed89f710592ddc39c76b4a5a8/awkward_cpp-53-cp314-cp314t-win_arm64.whl", hash = "sha256:1be0c1d87d9f4fdf94b767a061df849f1bb21579d302b2996fb101527fc80a97", size = 551257, upload-time = "2026-06-08T12:31:56.319Z" },
 ]
 
+[[package]]
+name = "bracex"
+version = "3.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ac/01/5f394b8bcd6e5b92f73130990960423bbb19711f906bd9fe9ea5557c667c/bracex-3.0.1.tar.gz", hash = "sha256:4e38e32392e4a4780fe15d644bfc7c8514057cfc3861e060b11814ce829c25e4", size = 44019, upload-time = "2026-07-20T13:43:00.335Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b8/8f/6f7273a7adb8d73fc8d21ede4376a3e475e52f98435c6007f69100dec8ca/bracex-3.0.1-py3-none-any.whl", hash = "sha256:6523ad83aeb5098a4ee597cff0f964442ff74e460bd3fafaffab6a013ff2288c", size = 11940, upload-time = "2026-07-20T13:42:59.268Z" },
+]
+
+[[package]]
+name = "bump-my-version"
+version = "1.5.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "click" },
+    { name = "httpx2" },
+    { name = "pydantic" },
+    { name = "pydantic-settings" },
+    { name = "questionary" },
+    { name = "rich" },
+    { name = "rich-click" },
+    { name = "tomlkit" },
+    { name = "wcmatch" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/23/09/5b09ac74962eca809cbf7010a08ea6ad405852bdd53489209a9f473d775c/bump_my_version-1.5.1.tar.gz", hash = "sha256:5079e443ab8c9a9903f140b427ff9f6fe8dd54013a55a4cf48b89326f3a71c07", size = 1132060, upload-time = "2026-08-06T14:26:38.9Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b3/0b/5885530f79d4400368b9d4dcb9b39274c0d52e633f7871e7fc6feceea1e3/bump_my_version-1.5.1-py3-none-any.whl", hash = "sha256:df3e2989d0d7fe704718feb24a5880f089b6b6369e427a4445b89c3adebfcff1", size = 65090, upload-time = "2026-08-06T14:26:37.083Z" },
+]
+
 [[package]]
 name = "certifi"
 version = "2026.7.22"
@@ -633,7 +675,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.2"
+version = "0.3.9"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },
@@ -666,6 +708,8 @@ cuda = [
 ]
 dev = [
     { name = "awkward" },
+    { name = "bump-my-version" },
+    { name = "git-cliff" },
     { name = "ipykernel" },
     { name = "matplotlib" },
     { name = "plotstyle" },
@@ -688,7 +732,9 @@ wandb = [
 [package.metadata]
 requires-dist = [
     { name = "awkward", marker = "extra == 'convert'", specifier = ">=2.6,<3" },
+    { name = "bump-my-version", marker = "extra == 'dev'", specifier = ">=1.2,<2" },
     { name = "giant", extras = ["convert", "analysis", "geometry", "wandb"], marker = "extra == 'dev'" },
+    { name = "git-cliff", marker = "extra == 'dev'", specifier = ">=2,<3" },
     { name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
     { name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
     { name = "numpy", specifier = ">=1.26,<3" },
@@ -713,6 +759,35 @@ requires-dist = [
 ]
 provides-extras = ["cpu", "cuda", "dev", "geometry", "wandb", "convert", "analysis"]
 
+[[package]]
+name = "git-cliff"
+version = "2.13.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/62/57/b12494e2cbc3c9154c942e64659b5aec2b1ce9f12d07f6dc6167e2c63ae5/git_cliff-2.13.1.tar.gz", hash = "sha256:e949ea9c3951ba6037b99eec465162be2584f27f0836ace45f44d6f45650f8c6", size = 113119, upload-time = "2026-04-26T10:33:42.331Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8f/dd/24768c3c0030710d36706c17b997d06aee27cb76b27ab2abb058ae254175/git_cliff-2.13.1-py3-none-macosx_10_12_x86_64.whl", hash = "sha256:08a9cb0ec760e165210ed22fefa295b6549a3520b420db995ccbb3620cbb1fbe", size = 7260035, upload-time = "2026-04-26T10:33:16.269Z" },
+    { url = "https://files.pythonhosted.org/packages/22/df/842973ead79d27a58cd1eccd167191c0a71513e5c1e9dc30337dafdb7d36/git_cliff-2.13.1-py3-none-macosx_11_0_arm64.whl", hash = "sha256:e92cf470ecbe73f7d2963dfa80e8961a6b76888d0c19949b0f028a28a0a0470c", size = 6854384, upload-time = "2026-04-26T10:33:18.808Z" },
+    { url = "https://files.pythonhosted.org/packages/7c/4d/6d6efa7d61be8632563990ccd402e401695e4a85a0bb1002f28730d03268/git_cliff-2.13.1-py3-none-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:2e8d8e420adf6a36b97e0fbdbf2b07e47712199a9baf3675c9a759430243ea26", size = 7308164, upload-time = "2026-04-26T10:33:20.869Z" },
+    { url = "https://files.pythonhosted.org/packages/f0/4a/98b8d2f53a2d0b7d313e98ab363ad7e0d6a514e878a8d678f22402cb0ec7/git_cliff-2.13.1-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:1ab059d671565189faa4f3858b2fb42535aa39fe265ad95d84d5c116a2724fc7", size = 7687163, upload-time = "2026-04-26T10:33:23.096Z" },
+    { url = "https://files.pythonhosted.org/packages/c7/07/cdd149b3909644aa3f0be7960406d9bbb38598f8618e039ac48cbb43ded6/git_cliff-2.13.1-py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:a93db30da45967c42df607fbbc2092111fcd576b0ca9e2fbddd3f653d8c71be7", size = 7317670, upload-time = "2026-04-26T10:33:25.323Z" },
+    { url = "https://files.pythonhosted.org/packages/2e/11/6d377a7f3113f6e26d87a28a32013eb05bd62bce176d6f8ee808e4868c1f/git_cliff-2.13.1-py3-none-manylinux_2_5_i686.manylinux1_i686.whl", hash = "sha256:df9f5a2bd16e5225030c9c2362e6ac70b34a906c0fbb83f8cbf5ae46932ba0d2", size = 7502294, upload-time = "2026-04-26T10:33:27.315Z" },
+    { url = "https://files.pythonhosted.org/packages/80/6b/e1da9acf3aec99e6600be02b6ce0c9e8bd42d072e3120384514c905231e2/git_cliff-2.13.1-py3-none-manylinux_2_5_x86_64.manylinux1_x86_64.whl", hash = "sha256:c12276784d280aa6a7148d3e52ff139e891f4c97720cd9be581b19892ea39fe0", size = 7927258, upload-time = "2026-04-26T10:33:29.554Z" },
+    { url = "https://files.pythonhosted.org/packages/ed/ea/9f2188a5e474e5f02193d9c1cdf7028773d483a3f508a1df4f1d93c9cc80/git_cliff-2.13.1-py3-none-musllinux_1_2_aarch64.whl", hash = "sha256:17da93dc605cbc48c762770402067fa726437cedbd62514f9caef6e0ccb58a43", size = 7308153, upload-time = "2026-04-26T10:33:31.589Z" },
+    { url = "https://files.pythonhosted.org/packages/b3/70/5e2b2a0e42c07956f911e1eccd6dc6d79b96fc6c8a604a526c1ee0474c84/git_cliff-2.13.1-py3-none-musllinux_1_2_i686.whl", hash = "sha256:cd4a08cf3f638ec71d2ed451aa8673bef99e1107a36366153db97ca18d981655", size = 7502287, upload-time = "2026-04-26T10:33:33.894Z" },
+    { url = "https://files.pythonhosted.org/packages/05/50/cc4c1d3d360621c0235d66d2c74472d9993d8aadf23ffa311eaf29d7a3aa/git_cliff-2.13.1-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:c5bb87f6e1db18be09e50c9d59234dc949713e20c2700e94eca378d22ad79719", size = 7927253, upload-time = "2026-04-26T10:33:36.205Z" },
+    { url = "https://files.pythonhosted.org/packages/86/5d/717d30f37dad65a6cc5220b04a3ef1bc44c31fb56d0ce2895114e159df4f/git_cliff-2.13.1-py3-none-win32.whl", hash = "sha256:c8878972e0a6c26d9137fc406a611116239333578d95ac05064d2807920bd83c", size = 6718261, upload-time = "2026-04-26T10:33:38.1Z" },
+    { url = "https://files.pythonhosted.org/packages/99/b2/99fac50978b9a90bfec0f1b89354a667ec83f4990301f6c708abce05484e/git_cliff-2.13.1-py3-none-win_amd64.whl", hash = "sha256:856d831a0bede9c258229dbd4d4c2b1c0810d8fce3d3882729669e8dc09c72bf", size = 7714969, upload-time = "2026-04-26T10:33:40.163Z" },
+]
+
+[[package]]
+name = "h11"
+version = "0.16.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/01/ee/02a2c011bdab74c6fb3c75474d40b3052059d95df7e73351460c8588d963/h11-0.16.0.tar.gz", hash = "sha256:4e35b956cf45792e4caa5885e69fba00bdbc6ffafbfa020300e549b208ee5ff1", size = 101250, upload-time = "2025-04-24T03:35:25.427Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/04/4b/29cac41a4d98d144bf5f6d33995617b185d14b22401f75ca86f384e87ff1/h11-0.16.0-py3-none-any.whl", hash = "sha256:63cf8bbe7522de3bf65932fda1d9c2772064ffb3dae62d55932da54b31cb6c86", size = 37515, upload-time = "2025-04-24T03:35:24.344Z" },
+]
+
 [[package]]
 name = "hepunits"
 version = "2.4.6"
@@ -722,6 +797,45 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/85/10/7f9c58d1ec6a0b7f7783fe552f3593f39cda30c2e1d7a9d148ae711e748d/hepunits-2.4.6-py3-none-any.whl", hash = "sha256:089c52c3b84ef67a159b5e9ee9bdd50e1a442e3fd0c101303cc409c1e9011c4d", size = 17090, upload-time = "2026-06-16T09:23:35.35Z" },
 ]
 
+[[package]]
+name = "httpcore2"
+version = "2.11.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "h11", marker = "sys_platform != 'emscripten' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+    { name = "truststore", marker = "sys_platform != 'emscripten' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/30/7d/ee6787efd5fe675d7cfd5eb149e40ccb5bdfc7e7c9252edcf7825c38986f/httpcore2-2.11.0.tar.gz", hash = "sha256:82e6fc95d784e6ee22ebd4b2cb57df53a2efb13ad6a11260a236ecebbc5f50c7", size = 67532, upload-time = "2026-08-18T08:03:53.008Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f3/54/e84a5c82ac0959d5e55b3970d326fd95306446b2bf0702302888745f7e5c/httpcore2-2.11.0-py3-none-any.whl", hash = "sha256:c7c899fbc6b8abb6e747dda427aa6f52934c45e191eabb986965164ebb02a908", size = 83061, upload-time = "2026-08-18T08:03:50.894Z" },
+]
+
+[[package]]
+name = "httpx2"
+version = "2.11.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "anyio", marker = "sys_platform != 'emscripten' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+    { name = "httpcore2", marker = "sys_platform != 'emscripten' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+    { name = "httpx2-jsfetch", marker = "sys_platform == 'emscripten' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+    { name = "idna" },
+    { name = "truststore", marker = "sys_platform != 'emscripten' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+    { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d5/4d/b3fcae38f29bfb0f300517d085c488f41f65e5b0a73023976b2122f568cc/httpx2-2.11.0.tar.gz", hash = "sha256:ea01b2e8febfb026e2601814c77ecb1e64fff114a87bc789cb520e67f27e7809", size = 99617, upload-time = "2026-08-18T08:03:53.691Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/39/29/f60bcf54028601920c0ce3da81537bcd0af4ee19dd9672af85dc9dfd60e0/httpx2-2.11.0-py3-none-any.whl", hash = "sha256:c9790f62a327110f52a099f1e2030cbe32f78b28781ad68fe58bfd6f23e73ab0", size = 95043, upload-time = "2026-08-18T08:03:52.004Z" },
+]
+
+[[package]]
+name = "httpx2-jsfetch"
+version = "1.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/cd/c4/0e5636363151a2a1795e0a77617168b9ca438e1748ec05fc9b5687f93d64/httpx2_jsfetch-1.0.tar.gz", hash = "sha256:70a0e3eabfef7cce5ad9c629f7d01ca05e418f586646f4ddf14782e4c1454c60", size = 6872, upload-time = "2026-08-07T00:13:07.492Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9b/43/832f631d32e4f1211caa2ba368317739fe71f0b8530e4c9d15dc454bac2a/httpx2_jsfetch-1.0-py3-none-any.whl", hash = "sha256:cb916b707601e69a07721aabc8f3f6659be3a6893bc1ff5c6f9e02241df2da32", size = 6382, upload-time = "2026-08-07T00:13:06.567Z" },
+]
+
 [[package]]
 name = "idna"
 version = "3.18"
@@ -1780,6 +1894,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" },
 ]
 
+[[package]]
+name = "pydantic-settings"
+version = "2.15.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pydantic" },
+    { name = "python-dotenv" },
+    { name = "typing-inspection" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/68/ca/31c57507b13119d7d3cfa1576dad2911a4861e3be07b579395f4e9d393f9/pydantic_settings-2.15.0.tar.gz", hash = "sha256:694b793e84f766ba76a90ebdefc01d0a9a045dab0382bee70393da93712ad117", size = 261253, upload-time = "2026-08-07T09:24:57.419Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/30/a4/2bffa9f8e804325a09867f0e9d30795c80ea9f8d62560bd1b6ad6220eb2f/pydantic_settings-2.15.0-py3-none-any.whl", hash = "sha256:0ba092c291c94baceb5eff768aa0d56400a457585bc0175925a5a5510303da42", size = 69413, upload-time = "2026-08-07T09:24:55.839Z" },
+]
+
 [[package]]
 name = "pygments"
 version = "2.20.0"
@@ -1840,6 +1968,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" },
 ]
 
+[[package]]
+name = "python-dotenv"
+version = "1.2.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/6a/53/ed9d74092561d4b01a2ef1349d52cdbc135e526c245f366b089cfca6de49/python_dotenv-1.2.3.tar.gz", hash = "sha256:a20a594dabeaa385725aa239d5244871c143ecb356add8a20fcf23773a6c3a35", size = 58945, upload-time = "2026-08-16T16:54:54.067Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/0d/17/c5c6b53ddc18f297992099b3d9ec16c855c0ccc83263a21fe4d1c625ec6c/python_dotenv-1.2.3-py3-none-any.whl", hash = "sha256:904552145e8bfed22162c09dab1c2b9b54fefa7b23ba780f4f26ca0316b0f0d9", size = 22780, upload-time = "2026-08-16T16:54:52.473Z" },
+]
+
 [[package]]
 name = "pyyaml"
 version = "6.0.3"
@@ -1929,6 +2066,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/81/d6/4bfbb40c9a0b42fc53c7cf442f6385db70b40f74a783130c5d0a5aa62228/pyzmq-27.1.0-cp314-cp314t-win_arm64.whl", hash = "sha256:dc5dbf68a7857b59473f7df42650c621d7e8923fb03fa74a526890f4d33cc4d7", size = 575170, upload-time = "2025-09-08T23:09:01.418Z" },
 ]
 
+[[package]]
+name = "questionary"
+version = "2.1.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "prompt-toolkit" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f6/45/eafb0bba0f9988f6a2520f9ca2df2c82ddfa8d67c95d6625452e97b204a5/questionary-2.1.1.tar.gz", hash = "sha256:3d7e980292bb0107abaa79c68dd3eee3c561b83a0f89ae482860b181c8bd412d", size = 25845, upload-time = "2025-08-28T19:00:20.851Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/3c/26/1062c7ec1b053db9e499b4d2d5bc231743201b74051c973dadeac80a8f43/questionary-2.1.1-py3-none-any.whl", hash = "sha256:a51af13f345f1cdea62347589fbb6df3b290306ab8930713bfae4d475a7d4a59", size = 36753, upload-time = "2025-08-28T19:00:19.56Z" },
+]
+
 [[package]]
 name = "requests"
 version = "2.34.2"
@@ -1957,6 +2106,20 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" },
 ]
 
+[[package]]
+name = "rich-click"
+version = "1.9.8"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "click" },
+    { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-5-giant-cpu' and extra == 'extra-5-giant-cuda')" },
+    { name = "rich" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/f7/ea/21e4867ea0ef881ffd4c0550fc21a061435e50d6324bcd034396633cbc18/rich_click-1.9.8.tar.gz", hash = "sha256:4008f921da88b5d91646c134ec881c1500e5a6b3f093e90e8f29400e09608371", size = 75363, upload-time = "2026-05-28T19:54:59.144Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/6d/97/a87901aef6b7e7e4a34c6dd6cc17dca8594a592ef9d9dd765fca2b7facf7/rich_click-1.9.8-py3-none-any.whl", hash = "sha256:12873865396e6927835d4eabb1cc3996edcd65b7ac9b2391a29eca4f335a2f93", size = 72189, upload-time = "2026-05-28T19:54:57.867Z" },
+]
+
 [[package]]
 name = "ruff"
 version = "0.15.17"
@@ -2149,6 +2312,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" },
 ]
 
+[[package]]
+name = "tomlkit"
+version = "0.15.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/94/96/e07752635b98536177fa1f37671c8f3cdde2e724c6bcf6034b2cfb571565/tomlkit-0.15.1.tar.gz", hash = "sha256:e25bbf38843005246210a12982776f27f99cb9be67160e14434d0c0d21ee1e97", size = 180129, upload-time = "2026-07-17T01:48:04.562Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/13/bc/8c13eb66537dce1d2bd3a57132902f38d0e7f5bb46fa9f4daed9fe9d76ee/tomlkit-0.15.1-py3-none-any.whl", hash = "sha256:177a05aece5a8ca5266fd3c448abb47b8d352f09d477d3ca8332db4d89b24304", size = 49449, upload-time = "2026-07-17T01:48:05.728Z" },
+]
+
 [[package]]
 name = "torch"
 version = "2.3.1"
@@ -2270,6 +2442,15 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/96/8d/1080ee4c231f361b6ce4470d556c8c435b67c7e0753aaa641497ee92f88b/traitlets-5.15.1-py3-none-any.whl", hash = "sha256:770a53705f84b81ac107e83a1b3328ff2dae16094d8fc3cfc004e4b22dfd8e92", size = 85858, upload-time = "2026-06-03T12:26:04.395Z" },
 ]
 
+[[package]]
+name = "truststore"
+version = "0.10.4"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/53/a3/1585216310e344e8102c22482f6060c7a6ea0322b63e026372e6dcefcfd6/truststore-0.10.4.tar.gz", hash = "sha256:9d91bd436463ad5e4ee4aba766628dd6cd7010cf3e2461756b3303710eebc301", size = 26169, upload-time = "2025-08-12T18:49:02.73Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/19/97/56608b2249fe206a67cd573bc93cd9896e1efb9e98bce9c163bcdc704b88/truststore-0.10.4-py3-none-any.whl", hash = "sha256:adaeaecf1cbb5f4de3b1959b42d41f6fab57b2b1666adb59e89cb0b53361d981", size = 18660, upload-time = "2025-08-12T18:49:01.46Z" },
+]
+
 [[package]]
 name = "ty"
 version = "0.0.50"
@@ -2394,6 +2575,18 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/07/78/75b6827a6665337a715c5347c5edbd84eca660f7a0f48d8d6d24d1f66bee/wandb-0.28.1-py3-none-win_arm64.whl", hash = "sha256:4aa07f13dd3bcac2c0524c8d0f49f76e83ab5c1054fd09f3b1a436cfcde146a6", size = 22299006, upload-time = "2026-07-16T18:47:02.71Z" },
 ]
 
+[[package]]
+name = "wcmatch"
+version = "11.0.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "bracex" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/57/43/30e407989e313677dbb9d5f045f966549a7254834571e342eaa4b55cc67b/wcmatch-11.0.1.tar.gz", hash = "sha256:1ea2b4fa678b8ca268253798d5963935df39132d47c3e241c0a0732224005e7d", size = 144662, upload-time = "2026-08-14T15:20:40.477Z" }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ce/77/7a02b0f05b3ffcdbef9719ce3ee0b508d6a29b58e95299f1580055671db3/wcmatch-11.0.1-py3-none-any.whl", hash = "sha256:fd149ecddb9f0a88ea780017d6dde17c994e494e7f7303d4e3c9d6251f978f4b", size = 43449, upload-time = "2026-08-14T15:20:39.379Z" },
+]
+
 [[package]]
 name = "wcwidth"
 version = "0.8.1"