From ffb7c0cc2af16062d183488c7d905324c9f33c7f Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 24 Aug 2026 11:12:16 +0200 Subject: [PATCH] Add rollout-quality distance, confusion, containment and router plots (gitea #76) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Picks 4 of the 7 catalog additions the issue proposed (the smaller-lift ones; 2D joint plots, PIT calibration, and the throughput/accuracy scatter are left for follow-up issues): - marginal_distance_summary: a var x grouping-axis KS-statistic heatmap, reusing the existing marginal hist1d compute and just adding a finalize — a single at-a-glance regression scorecard instead of N overlay plots. - n_sec_confusion: predicted (rollout) vs true (reference) secondary count per event, paired by event_id since a rollout is seeded from the same events as its reference file. Needed a new zero-filling primitive (reduce.sec_count_by_event) since a plain group_by over secondary rows silently drops zero-secondary events. - shower_containment_depth_{90,95}: per-event depth containing 90%/95% of deposited energy, derived from the same per-event depth-bin matrix the longitudinal profile already computes. - router_specialization: max gate weight vs energy per side, summarizing router_gating's full stacked area into the one trend line the roadmap's MoE writeup describes (the ~60-65% ceiling), to make a future lambda_balance>0 retrain's effect on specialization checkable at a glance. Both new heatmap-shaped plots (distance summary, confusion matrix) share one new "heatmap" Reduced kind/renderer rather than two near-identical ones. Co-Authored-By: Claude Sonnet 5 --- giant/analysis/catalog.py | 230 +++++++++++++++++++++++++++++ giant/analysis/reduce.py | 17 +++ giant/analysis/reduced.py | 4 + giant/analysis/render.py | 43 ++++++ giant/analysis/router_gating.py | 49 ++++++ giant/analysis/runtime_estimate.py | 6 +- tests/test_analysis_reduce.py | 14 ++ tests/test_catalog.py | 68 ++++++++- 8 files changed, 426 insertions(+), 5 deletions(-) diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index 458965b..e9b4157 100644 --- a/giant/analysis/catalog.py +++ b/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", diff --git a/giant/analysis/reduce.py b/giant/analysis/reduce.py index ce65c89..8e1d48e 100644 --- a/giant/analysis/reduce.py +++ b/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 diff --git a/giant/analysis/reduced.py b/giant/analysis/reduced.py index 2188800..d3b2c24 100644 --- a/giant/analysis/reduced.py +++ b/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) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index f5f57ea..41ad7dc 100644 --- a/giant/analysis/render.py +++ b/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, } diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py index 2a6a70f..8df6225 100644 --- a/giant/analysis/router_gating.py +++ b/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, diff --git a/giant/analysis/runtime_estimate.py b/giant/analysis/runtime_estimate.py index ae703be..5f639c6 100644 --- a/giant/analysis/runtime_estimate.py +++ b/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 diff --git a/tests/test_analysis_reduce.py b/tests/test_analysis_reduce.py index 1241716..71cb4a8 100644 --- a/tests/test_analysis_reduce.py +++ b/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 diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 92ddb24..ec5a134 100644 --- a/tests/test_catalog.py +++ b/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]]