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

Reviewed-on: #79

.bumpversion.toml

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

CHANGELOG.md

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

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
-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

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.6"
+version = "0.3.7"
 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]]

uv.lock

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