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

Reviewed-on: #78

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.5"
+current_version = "0.3.8"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

CHANGELOG.md

@@ -1,5 +1,28 @@
 # Changelog
 
+## [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

giant/analysis/catalog.py

@@ -47,6 +47,7 @@ from giant.analysis.reduce import (
     leakage_fraction,
     profile_finalize,
     profile_partial,
+    sec_count_by_event,
     species_share,
     sum_merge,
     transverse_expr,
@@ -56,6 +57,7 @@ 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.type_embedding_distance import compute_type_embedding_l1_distance
@@ -176,6 +178,68 @@ def _np_hist_pair(r: np.ndarray, t: np.ndarray, nbins: int) -> tuple[np.ndarray,
     return edges, np.histogram(r, edges)[0], np.histogram(t, edges)[0]
 
 
+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) -> 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 = 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;
 # these read cleanly as a title without them).
 _TITLE_NAMES = {
@@ -299,6 +363,64 @@ def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis:
     )
 
 
+# ---------------------------------------------------------------------------
+# distance summary: a var x group-axis scorecard, 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]
+    matrix: list[list[float]] = []
+    for var in MARGINAL_VARS:
+        edges = _marginal_edges(ctx, var)
+        nb = len(edges) - 1
+        row: list[float] = []
+
+        r = sum_merge([p[var]["overall"]["r"] for p in parts])
+        t = sum_merge([p[var]["overall"]["t"] for p in parts])
+        row.append(_ks_statistic(_finalize_counts(r, 0, nb), _finalize_counts(t, 0, nb)))
+
+        for axis in GROUPING_AXES:
+            r = sum_merge([p[var][axis]["r"] for p in parts])
+            t = sum_merge([p[var][axis]["t"] for p in parts])
+            dists, weights = [], []
+            for k in _group_keys(ctx, axis):
+                rc, tc = _finalize_counts(r, k, nb), _finalize_counts(t, k, nb)
+                w = sum(rc) + sum(tc)
+                if w == 0:
+                    continue
+                dists.append(_ks_statistic(rc, tc))
+                weights.append(w)
+            row.append(float(np.average(dists, weights=weights)) if dists else float("nan"))
+        matrix.append(row)
+
+    return Reduced(
+        id="marginal_distance_summary",
+        family="quality",
+        kind="heatmap",
+        title="Marginal distance summary (KS statistic, rollout vs reference)",
+        xlabel="grouping axis",
+        payload={
+            "matrix": matrix,
+            "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
 # ---------------------------------------------------------------------------
@@ -438,6 +560,41 @@ def _profile_finalize(
     )
 
 
+# ---------------------------------------------------------------------------
+# 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
+    _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_full = np.concatenate([np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
+    t_full = np.concatenate([np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
+    r_depth = _containment_depths(r_full, edges, quantile)
+    t_depth = _containment_depths(t_full, edges, quantile)
+    hedges, rc, tc = _np_hist_pair(r_depth, t_depth, ctx.n_marginal_bins)
+    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(),
+            _ROLL: rc.astype(np.int64).tolist(),
+            _REF: tc.astype(np.int64).tolist(),
+            "log_y": False,
+        },
+    )
+
+
 # ---------------------------------------------------------------------------
 # species share + leakage
 # ---------------------------------------------------------------------------
@@ -627,6 +784,43 @@ def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
     )
 
 
+def _n_sec_confusion_partial(b: Bundle) -> dict:
+    r_sec, t_sec = _sec_frames(b)
+    r_ids, r_n = sec_count_by_event(b.r_phys, r_sec)
+    t_ids, t_n = sec_count_by_event(b.t_all, t_sec)
+    return {"r_ids": r_ids.tolist(), "r_n": r_n.tolist(), "t_ids": t_ids.tolist(), "t_n": t_n.tolist()}
+
+
+def _n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
+    r_ids = np.concatenate([np.asarray(p["r_ids"], dtype=np.int64) for p in parts])
+    r_n = np.concatenate([np.asarray(p["r_n"], dtype=np.int64) for p in parts])
+    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])
+    # 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.
+    r_map = dict(zip(r_ids.tolist(), r_n.tolist()))
+    t_map = dict(zip(t_ids.tolist(), t_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)
+    labels, mat = _integer_confusion(true_n, pred_n)
+    return Reduced(
+        id="n_sec_confusion",
+        family="secondaries",
+        kind="heatmap",
+        title="Predicted vs true secondary count per event",
+        xlabel="predicted secondaries (rollout)",
+        payload={
+            "matrix": mat.tolist(),
+            "row_labels": labels,
+            "col_labels": labels,
+            "ylabel": "true secondaries (reference)",
+            "cbar_label": "event count",
+            "vmin": 0.0,
+        },
+    )
+
+
 # ---------------------------------------------------------------------------
 # router diagnostics (not chunked — already bounded/subsampled)
 # ---------------------------------------------------------------------------
@@ -640,6 +834,9 @@ _router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
 _router_share_process_partial, _router_share_process_finalize = _unchunkable(
     lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
 )
+_router_specialization_partial, _router_specialization_finalize = _unchunkable(
+    lambda b: compute_router_specialization(b.checkpoint, b.r_phys, 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)
 )
@@ -676,6 +873,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 +951,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 +998,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 +1025,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/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

@@ -19,6 +19,10 @@ from pathlib import Path
 #   "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
+#   "router_specialization"  max gate weight vs energy, rollout + reference (one
+#                    scalar trend line summarizing "router_gating")
+#   "heatmap"        row x col matrix + colorbar (distance scorecard or a
+#                    predicted-vs-true confusion matrix)
 #   "unavailable"    plot not applicable to this run (e.g. non-MoE checkpoint)
 
 

giant/analysis/render.py

@@ -260,6 +260,47 @@ def _render_router_share(r: Reduced, params: dict):
     return fig
 
 
+def _render_router_specialization(r: Reduced, params: dict):
+    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
+    for key in ("reference", "rollout"):
+        side = r.payload.get(key)
+        if side and side["centers"]:
+            ax.plot(side["centers"], side["score"], label=_SERIES_LABELS[key], marker="o", markersize=3)
+    chance = r.payload.get("chance_level")
+    if chance is not None:
+        ax.axhline(chance, linestyle="--", color="gray", label="chance level (1/n_experts)")
+    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=f"{r.payload.get('router_type', '')} router")
+    return fig
+
+
+def _render_heatmap(r: Reduced, params: dict):
+    mat = np.asarray(r.payload["matrix"], dtype=float)
+    row_labels = r.payload["row_labels"]
+    col_labels = r.payload["col_labels"]
+    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
+    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)
+    ax.set_ylabel(r.payload.get("ylabel", ""))
+    fig.colorbar(im, ax=ax, label=r.payload.get("cbar_label", "value"))
+    return fig
+
+
 def _render_unavailable(r: Reduced, params: dict):
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
     ax.axis("off")
@@ -284,6 +325,8 @@ _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,
 }
 

giant/analysis/router_gating.py

@@ -203,6 +203,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)",
 }
 
 
@@ -250,6 +251,54 @@ def compute_router_gating(
     )
 
 
+def compute_router_specialization(
+    checkpoint: str | Path | None,
+    r_phys: pl.LazyFrame,
+    t_phys: pl.LazyFrame,
+    seed: int = 0,
+) -> Reduced:
+    """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_specialization")
+
+    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 Reduced(
+        id="router_specialization",
+        family="model",
+        kind="router_specialization",
+        title=_TITLES["router_specialization"],
+        xlabel="pre-step energy [MeV]",
+        payload={
+            "router_type": handle.router_type,
+            "n_experts": handle.router.n_experts,
+            "log_x": True,
+            "chance_level": 1.0 / handle.router.n_experts,
+            **sides,
+        },
+    )
+
+
 def compute_router_share_by_pdg(
     checkpoint: str | Path | None,
     r_phys: pl.LazyFrame,

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
+`chunkable=False` router specs in `_ROUTER_IDS` (they always run as a single
-regardless of chunk count).
+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/cli.py

@@ -1644,6 +1644,25 @@ 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")],

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/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._build_context_fusion(x_dim, context_dim, cond_out_dim)
-        self.fuse = nn.Sequential(
-            nn.Linear(cond_out_dim + context_dim, cond_out_dim),
-            nn.SiLU(),
-        )
         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._build_context_fusion(context_in_dim, context_dim, cond_out_dim)
-            self.fuse = nn.Sequential(
+        self.trunk = build_trunk(
-                nn.Linear(cond_out_dim + context_dim, cond_out_dim),
+            None, trunk_type, in_dim, 1, hidden_dim, n_res_blocks, cond_out_dim, dropout, block_conditioning
-                nn.SiLU(),
+        )
-            )
-        self.input_proj = nn.Linear(in_dim, hidden_dim)
-        self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)])
-        self.out_norm = nn.LayerNorm(hidden_dim)
-        self.out_proj = nn.Linear(hidden_dim, 1)
 
     def forward(
         self,
@@ -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)
+            assert stage1_out is not None, "stage='stage2' CriticModel requires stage1_out"
-            cond = self.fuse(torch.cat([base, ctx], dim=-1))
+            cond = self._cond_embed(cond_cont, cond_cat, stage1_out)
         else:
-            cond = base
+            cond = self.cond_enc(cond_cont, cond_cat)
-        h = self.input_proj(x)
+        return self.trunk(x, cond, cond_cont, cond_cat).squeeze(-1)
-        for block in self.blocks:
-            h = block(h, cond)
-        return self.out_proj(self.out_norm(h)).squeeze(-1)

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

pyproject.toml

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

tests/test_analysis_reduce.py

@@ -159,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

tests/test_catalog.py

@@ -6,7 +6,13 @@ 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 tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
@@ -53,6 +59,8 @@ 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
@@ -88,6 +96,14 @@ def _validate_payload(r) -> None:
             for side in ("rollout", "reference"):
                 if side in p:
                     assert cat in p[side]
+    elif r.kind == "router_specialization":
+        for side in ("rollout", "reference"):
+            if side in p:
+                assert len(p[side]["centers"]) == len(p[side]["score"])
+    elif r.kind == "heatmap":
+        assert len(p["matrix"]) == len(p["row_labels"])
+        for row in p["matrix"]:
+            assert len(row) == len(p["col_labels"])
 
 
 # ---------------------------------------------------------------------------
@@ -98,7 +114,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 +126,9 @@ _CHUNK_EQUIVALENCE_IDS = [
     "leakage_fraction",
     "sec_count_per_species",
     "router_gating",
+    "marginal_distance_summary",
+    "n_sec_confusion",
+    "shower_containment_depth_90",
 ]
 
 
@@ -146,3 +168,45 @@ def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
     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_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["matrix"] == [[0, 0], [1, 1]]

tests/test_network.py

@@ -8,7 +8,10 @@ from giant.model.network import (
     HISTORY_REGISTRY,
     AttentionHistory,
     ConditionEncoder,
+    CriticModel,
+    FilmResBlock,
     HistoryEncoder,
+    LinearTrunk,
     MarkovHistory,
     NoHistory,
     SinusoidalEmbedding,
@@ -944,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) ─
@@ -1021,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
@@ -1046,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 inherited.trunk.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    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 overridden.trunk.input_proj.out_features == 16
-    assert len(overridden.blocks) == 3
+    assert len(overridden.trunk.blocks) == 3
 
 
 def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
@@ -1065,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 inherited.trunk.input_proj.out_features == 8
-    assert len(inherited.blocks) == 1
+    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 overridden.trunk.input_proj.out_features == 16
-    assert len(overridden.blocks) == 3
+    assert len(overridden.trunk.blocks) == 3
 
 
 # ── StageModel base (gitea #39): Stage1Model/Stage2OneShot/Stage2Autoregressive
@@ -1234,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_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_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

@@ -675,7 +675,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.5"
+version = "0.3.8"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },