giant/giant/analysis/catalog.py

"""The declarative plot catalog: one ``PlotSpec`` per figure.

Each spec knows its stable ``id`` (used for the reduced-data filename, the PDF
stem and the condor queue item), its gallery ``family`` (subdirectory), and a
``compute_partial(bundle) -> dict`` / ``finalize(parts, ctx) -> Reduced`` pair
that together run the streaming reduction. ``compute_partial`` runs once per
``(plot, chunk)`` condor job against a ``Bundle`` whose LazyFrames are
already filtered to that chunk (see ``Bundle.open``'s ``chunk`` argument); it
returns a small JSON-safe partial artifact — either a raw sum-mergeable count
dict (histograms/species sums against fixed edges) or a raw per-event/
per-secondary array to be concatenated (anything that derives its own edges or
a mean/std from the full dataset). ``finalize`` merges the per-chunk partials
(in chunk order) and does the actual histogramming/edge-selection/mean-std
collapse, once, over the merged data — for ``n_chunks=1`` this reproduces
exactly what a single unchunked pass would produce. Specs marked
``chunkable=False`` (the router ones) always run as a single chunk regardless
of the configured chunk count.

Every ``compute_partial`` here returns ``{"r": {rollout_name: <shape>}, "t":
<shape>}`` — one entry per rollout in ``Bundle.rollouts`` (insertion order,
which is the order rollouts were given on the CLI) plus the single reference.
``finalize`` merges each rollout's chunks independently and assembles a
``Reduced.payload`` keyed the same way: ``"series": {name: ...}`` for the
rollouts, ``"reference": ...`` as one distinguished entry (omitted on
rollout-only plots like ``leakage_fraction``). The heatmap-shaped specs
(``marginal_distance_summary``, ``sec_count_per_step_by_species``) and the
router diagnostics are inherently one-matrix/one-checkpoint per rollout, so their
``"series"`` entries are whole per-rollout artifacts (a matrix, a gating
dict) rather than a single number/array — ``render.py`` draws those as one
panel per rollout instead of one line/bar per rollout.

Rendering lives in ``render.py`` and dispatches on ``Reduced.kind`` — the
catalog itself never imports plotstyle, so ``compute-one`` jobs stay LaTeX-free.

The registry is built by expanding parametric families (marginals over
variable x grouping, secondaries, ...) into concrete specs.
"""

from __future__ import annotations

from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from typing import cast

import numpy as np
import polars as pl

from giant.analysis.context import Context
from giant.analysis.geant4_reference import GEANT4_REFERENCE, geant4_per_step_us
from giant.analysis.grouping import (
    energy_bin_labels,
    event_energy_bins,
    material_label,
    pdg_label,
)
from giant.analysis.prediction import (
    PAIRED_SCALARS,
    PAIRED_VARS,
    PredictionSide,
    PredictionSpec,
    dir_alignment_expr,
    open_prediction,
    paired_frame,
    paired_secondaries,
    paired_vars_for_coord,
)
from giant.analysis.reduce import (
    attach_entry_axis,
    binned_moments,
    depth_expr,
    entry_axis,
    event_scalars,
    hist1d,
    hist2d,
    leakage_fraction,
    profile_finalize,
    profile_partial,
    species_share,
    sum_merge,
    transverse_expr,
)
from giant.analysis.reduced import Reduced
from giant.analysis.router_gating import (
    compute_router_gating,
    compute_router_share_by_pdg,
    compute_router_share_by_process,
    compute_router_specialization,
)
from giant.analysis.sources import (
    RolloutSide,
    RolloutSpec,
    Side,
    open_side,
    physical_steps,
    secondaries,
    secondaries_by_step,
)
from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
from giant.analysis.variables import RANGED_VARS, cos_scatter_expr


@dataclass
class Bundle:
    """Everything a compute runs against — built once per ``compute-one`` job."""

    ctx: Context
    rollouts: dict[str, RolloutSide]  # name -> frames, insertion order = CLI order
    t_all: pl.LazyFrame  # reference, all rows
    t_phys: pl.LazyFrame  # reference, physical steps only
    predictions: dict[str, PredictionSide] = field(default_factory=dict)  # name -> paired frames, CLI order

    @classmethod
    def open(
        cls,
        rollouts: list[RolloutSpec],
        reference,
        ctx: Context,
        chunk: tuple[int, int] | None = None,
        predictions: list[PredictionSpec] | None = None,
    ) -> Bundle:
        """Open the reference + every rollout + every prediction, optionally
        restricted to one event-disjoint chunk.

        ``chunk = (chunk_index, n_chunks)`` filters every side to
        ``event_id % n_chunks == chunk_index`` *before* deriving the physical/
        secondary/paired views, so every downstream reduction (which is either
        row-local or a ``group_by("event_id")``) sees a self-contained,
        event-disjoint slice — no cross-chunk lookups are ever needed.
        """
        t_all = open_side(reference, Side.reference)
        pred = None
        if chunk is not None:
            idx, n = chunk
            pred = pl.col("event_id") % n == idx
            t_all = t_all.filter(pred)
        sides: dict[str, RolloutSide] = {}
        for rs in rollouts:
            r_all = open_side(rs.source, Side.rollout)
            if pred is not None:
                r_all = r_all.filter(pred)
            sides[rs.name] = RolloutSide(
                all=r_all,
                phys=physical_steps(r_all, Side.rollout),
                checkpoint=rs.checkpoint,
                type_embedding_l1_dist=rs.type_embedding_l1_dist,
                timing=rs.timing,
            )
        pred_sides: dict[str, PredictionSide] = {}
        for ps in predictions or []:
            opened = open_prediction(ps.source)
            lf = opened.lf if pred is None else opened.lf.filter(pred)
            pred_sides[ps.name] = PredictionSide(
                lf=lf,
                paired=paired_frame(lf, opened.coord, opened.has_truth),
                coord=opened.coord,
                has_truth=opened.has_truth,
                checkpoint=ps.checkpoint,
            )
        return cls(
            ctx=ctx,
            rollouts=sides,
            t_all=t_all,
            t_phys=physical_steps(t_all, Side.reference),
            predictions=pred_sides,
        )


@dataclass
class PlotSpec:
    id: str
    family: str
    compute_partial: Callable[[Bundle], dict]
    finalize: Callable[[list[dict], Context], Reduced]
    chunkable: bool = True


def _unchunkable(
    compute: Callable[[Bundle], Reduced],
) -> tuple[Callable[[Bundle], dict], Callable[[list[dict], Context], Reduced]]:
    """Wrap a whole-dataset ``compute(bundle) -> Reduced`` as a trivial
    ``(compute_partial, finalize)`` pair, for specs marked ``chunkable=False``
    (which always run as a single chunk, so ``parts`` is always one element).
    """

    def partial(b: Bundle) -> dict:
        return {"reduced": asdict(compute(b))}

    def finalize(parts: list[dict], ctx: Context) -> Reduced:
        return Reduced(**parts[0]["reduced"])

    return partial, finalize


# ---------------------------------------------------------------------------
# small numpy/hist helpers
# ---------------------------------------------------------------------------


def _per_rollout(b: Bundle, fn: Callable[[RolloutSide], object]) -> dict[str, object]:
    """``{name: fn(rollout_side)}`` over every rollout, preserving CLI order."""
    return {name: fn(rs) for name, rs in b.rollouts.items()}


def _counts(h: dict, key, nbins: int) -> list[int]:
    return h.get(key, np.zeros(nbins, dtype=np.int64)).astype(np.int64).tolist()


def _partial_hist(
    lf: pl.LazyFrame, value: pl.Expr, edges: np.ndarray, group: pl.Expr | None = None
) -> dict[str, list[int]]:
    """One chunk's raw ``hist1d`` result as a JSON-safe, sum-mergeable dict."""
    nb = len(edges) - 1
    h = hist1d(lf, value, edges, group=group)
    return {str(k): _counts(h, k, nb) for k in h}


def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]:
    """One group's merged counts (zero-filled if the group never appeared)."""
    return list(merged.get(str(key), [0] * nbins))


def _np_hist_shared_edges(arrays: list[np.ndarray], nbins: int) -> tuple[np.ndarray, list[np.ndarray]]:
    """Shared-edge histogram of several small per-event arrays (robust range).

    The edges are sized from the union of every array (reference + all
    rollouts), so every series in the resulting overlay is directly
    comparable on one axis.
    """
    non_empty = [a for a in arrays if len(a)]
    both = np.concatenate(non_empty) if non_empty else np.array([0.0, 1.0])
    lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
    if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
        lo, hi = lo - 0.5, hi + 0.5
    edges = np.linspace(lo, hi, nbins + 1)
    return edges, [np.histogram(a, edges)[0] for a in arrays]


def _ks_statistic(r_counts, t_counts) -> float:
    """KS statistic (max |CDF diff|) between two same-edge binned histograms.

    ``nan`` when neither side has any mass (nothing to compare); 1.0 (maximal
    mismatch) when exactly one side is entirely empty and the other isn't —
    correctly the worst score rather than an undefined one.
    """
    r_counts = np.asarray(r_counts, dtype=np.float64)
    t_counts = np.asarray(t_counts, dtype=np.float64)
    r_tot, t_tot = r_counts.sum(), t_counts.sum()
    if r_tot == 0 and t_tot == 0:
        return float("nan")
    if r_tot == 0 or t_tot == 0:
        return 1.0
    r_cdf = np.cumsum(r_counts) / r_tot
    t_cdf = np.cumsum(t_counts) / t_tot
    return float(np.max(np.abs(r_cdf - t_cdf)))


def _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 = {
    "step_length": "Step length",
    "edep": "Deposited energy per step",
    "delta_e": "Energy loss per step",
    "post_E": "Post-step energy",
    "cos_scatter": "Scattering cosine",
}


def _var(var: str):
    """(axis label, value expr) for a marginal variable name."""
    if var == "cos_scatter":
        return ("cos of scattering angle", cos_scatter_expr())
    label, expr = RANGED_VARS[var]
    return (label, expr)


def _marginal_edges(ctx: Context, var: str) -> np.ndarray:
    if var == "cos_scatter":
        return np.linspace(-1.0, 1.0, ctx.n_marginal_bins + 1)
    return ctx.marginal_edges(var)


# ---------------------------------------------------------------------------
# marginals: variable x {overall, energy, pdg, material}
# ---------------------------------------------------------------------------


def _marginal_overall_partial(b: Bundle, var: str) -> dict:
    _, expr = _var(var)
    edges = _marginal_edges(b.ctx, var)
    return {
        "r": _per_rollout(b, lambda rs: _partial_hist(rs.phys, expr, edges)),
        "t": _partial_hist(b.t_phys, expr, edges),
    }


def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced:
    label, _ = _var(var)
    edges = _marginal_edges(ctx, var)
    nb = len(edges) - 1
    names = list(parts[0]["r"])
    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
    t = sum_merge([p["t"] for p in parts])
    return Reduced(
        id=f"marginal_{var}",
        family="marginals",
        kind="overlay_hist",
        title=_TITLE_NAMES[var],
        xlabel=label,
        payload={
            "edges": edges.tolist(),
            "series": series,
            "reference": _finalize_counts(t, 0, nb),
            "log_y": True,
        },
    )


def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
    ids, bins = event_energy_bins(lf, edges)
    return pl.col("event_id").replace_strict(ids, bins, default=-1, return_dtype=pl.Int64)


def _grouped_hist_dict(lf: pl.LazyFrame, expr: pl.Expr, edges: np.ndarray, axis: str, ctx: Context, nb: int) -> dict:
    if axis == "pdg":
        h = hist1d(lf, expr, edges, group=pl.col("pdg"))
    elif axis == "material":
        h = hist1d(lf, expr, edges, group=pl.col("material"))
    else:  # energy
        e_edges = np.asarray(ctx.energy_edges)
        h = hist1d(lf, expr, edges, group=_energy_group_expr(lf, e_edges))
    return {str(k): _counts(h, k, nb) for k in h}


def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
    _, expr = _var(var)
    edges = _marginal_edges(b.ctx, var)
    nb = len(edges) - 1
    return {
        "r": _per_rollout(b, lambda rs: _grouped_hist_dict(rs.phys, expr, edges, axis, b.ctx, nb)),
        "t": _grouped_hist_dict(b.t_phys, expr, edges, axis, b.ctx, nb),
    }


def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis: str) -> Reduced:
    label, _ = _var(var)
    edges = _marginal_edges(ctx, var)
    nb = len(edges) - 1
    names = list(parts[0]["r"])
    r_merged = {name: sum_merge([p["r"][name] for p in parts]) for name in names}
    t_merged = sum_merge([p["t"] for p in parts])

    if axis == "pdg":
        keys, labels = ctx.top_pdgs, [pdg_label(k) for k in ctx.top_pdgs]
    elif axis == "material":
        keys, labels = ctx.materials, [material_label(m) for m in ctx.materials]
    else:  # energy
        e_edges = np.asarray(ctx.energy_edges)
        keys, labels = list(range(len(e_edges) - 1)), energy_bin_labels(e_edges)

    groups: dict[str, dict] = {}
    for k, lbl in zip(keys, labels):
        groups[lbl] = {
            "series": {name: _finalize_counts(r_merged[name], k, nb) for name in names},
            "reference": _finalize_counts(t_merged, k, nb),
        }

    return Reduced(
        id=f"marginal_{var}_by_{axis}",
        family="marginals",
        kind="grouped_hist",
        title=f"{_TITLE_NAMES[var]} by {axis}",
        xlabel=label,
        payload={"edges": edges.tolist(), "groups": groups, "log_y": True},
    )


# ---------------------------------------------------------------------------
# distance summary: a var x group-axis scorecard per rollout, reusing the marginal hists
# ---------------------------------------------------------------------------


def _distance_summary_partial(b: Bundle) -> dict:
    out: dict[str, dict] = {}
    for var in MARGINAL_VARS:
        out[var] = {"overall": _marginal_overall_partial(b, var)}
        for axis in GROUPING_AXES:
            out[var][axis] = _marginal_grouped_partial(b, var, axis)
    return out


def _distance_summary_finalize(parts: list[dict], ctx: Context) -> Reduced:
    col_labels = ["overall", *GROUPING_AXES]
    names = list(parts[0][MARGINAL_VARS[0]]["overall"]["r"])
    matrices: dict[str, list[list[float]]] = {name: [] for name in names}

    for var in MARGINAL_VARS:
        edges = _marginal_edges(ctx, var)
        nb = len(edges) - 1

        t_overall = sum_merge([p[var]["overall"]["t"] for p in parts])
        r_overall = {name: sum_merge([p[var]["overall"]["r"][name] for p in parts]) for name in names}
        row: dict[str, list[float]] = {name: [] for name in names}
        for name in names:
            row[name].append(
                _ks_statistic(_finalize_counts(r_overall[name], 0, nb), _finalize_counts(t_overall, 0, nb))
            )

        for axis in GROUPING_AXES:
            t_grp = sum_merge([p[var][axis]["t"] for p in parts])
            r_grp = {name: sum_merge([p[var][axis]["r"][name] for p in parts]) for name in names}
            for name in names:
                dists, weights = [], []
                for k in _group_keys(ctx, axis):
                    rc, tc = _finalize_counts(r_grp[name], k, nb), _finalize_counts(t_grp, k, nb)
                    w = sum(rc) + sum(tc)
                    if w == 0:
                        continue
                    dists.append(_ks_statistic(rc, tc))
                    weights.append(w)
                row[name].append(float(np.average(dists, weights=weights)) if dists else float("nan"))

        for name in names:
            matrices[name].append(row[name])

    return Reduced(
        id="marginal_distance_summary",
        family="quality",
        kind="heatmap",
        title="Marginal distance summary (KS statistic, rollout vs reference)",
        xlabel="grouping axis",
        payload={
            "series": matrices,
            "row_labels": [_TITLE_NAMES[v] for v in MARGINAL_VARS],
            "col_labels": col_labels,
            "ylabel": "marginal variable",
            "cbar_label": "KS statistic (0 = identical, 1 = maximal mismatch)",
            "vmin": 0.0,
            "vmax": 1.0,
        },
    )


# ---------------------------------------------------------------------------
# per-event scalar observables
# ---------------------------------------------------------------------------


def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
    t_lf = b.t_all if use_all else b.t_phys

    def _vals(rs: RolloutSide) -> list[float]:
        lf = rs.all if use_all else rs.phys
        return event_scalars(lf)[col].to_numpy().tolist()

    return {
        "r": _per_rollout(b, _vals),
        "t": event_scalars(t_lf)[col].to_numpy().tolist(),
    }


def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str) -> Reduced:
    names = list(parts[0]["r"])
    r_arrays = {name: np.concatenate([np.asarray(p["r"][name], dtype=float) for p in parts]) for name in names}
    t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
    edges, counts = _np_hist_shared_edges([t, *(r_arrays[n] for n in names)], ctx.n_marginal_bins)
    t_counts, *r_counts = counts
    return Reduced(
        id=spec_id,
        family="event",
        kind="overlay_hist",
        title=title,
        xlabel=xlabel,
        payload={
            "edges": edges.tolist(),
            "series": {name: c.astype(np.int64).tolist() for name, c in zip(names, r_counts)},
            "reference": t_counts.astype(np.int64).tolist(),
            "log_y": False,
        },
    )


def _event_total_edep_by_energy_partial(b: Bundle) -> dict:
    t = event_scalars(b.t_all)

    def _vals(rs: RolloutSide) -> dict:
        r = event_scalars(rs.all)
        return {"incident": r["incident_E"].to_list(), "edep": r["total_edep"].to_list()}

    return {
        "r": _per_rollout(b, _vals),
        "t": {"incident": t["incident_E"].to_list(), "edep": t["total_edep"].to_list()},
    }


def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
    e_edges = np.asarray(ctx.energy_edges)
    names = list(parts[0]["r"])
    t_inc = np.concatenate([np.asarray(p["t"]["incident"], dtype=float) for p in parts])
    t_val = np.concatenate([np.asarray(p["t"]["edep"], dtype=float) for p in parts])
    r_inc = {n: np.concatenate([np.asarray(p["r"][n]["incident"], dtype=float) for p in parts]) for n in names}
    r_val = {n: np.concatenate([np.asarray(p["r"][n]["edep"], dtype=float) for p in parts]) for n in names}

    edges, _ = _np_hist_shared_edges([t_val, *(r_val[n] for n in names)], ctx.n_marginal_bins)
    t_bin = np.clip(np.digitize(t_inc, e_edges[1:-1]), 0, len(e_edges) - 2)
    r_bin = {n: np.clip(np.digitize(r_inc[n], e_edges[1:-1]), 0, len(e_edges) - 2) for n in names}

    groups: dict[str, dict] = {}
    for bi, lbl in enumerate(energy_bin_labels(e_edges)):
        tc = np.histogram(t_val[t_bin == bi], edges)[0]
        groups[lbl] = {
            "series": {n: np.histogram(r_val[n][r_bin[n] == bi], edges)[0].astype(np.int64).tolist() for n in names},
            "reference": tc.astype(np.int64).tolist(),
        }
    return Reduced(
        id="event_total_edep_by_energy",
        family="event",
        kind="grouped_hist",
        title="Total deposited energy per event by incident energy",
        xlabel="total deposited energy [MeV]",
        payload={"edges": edges.tolist(), "groups": groups, "log_y": False},
    )


# ---------------------------------------------------------------------------
# shower shape profiles
# ---------------------------------------------------------------------------


def _profile_partial(b: Bundle, coord_fn, edges_key: str) -> dict:
    edges = np.asarray(getattr(b.ctx, edges_key))

    def _mat(lf: pl.LazyFrame) -> dict:
        lf2 = attach_entry_axis(lf, entry_axis(lf))
        ids, mat = profile_partial(lf2, coord_fn(), edges, pl.col("edep"))
        return {"ids": ids.tolist(), "mat": mat.tolist()}

    return {
        "r": _per_rollout(b, lambda rs: _mat(rs.all)),
        "t": _mat(b.t_all),
    }


def _assert_event_disjoint(id_lists: list[list[int]], spec_id: str, side: str) -> None:
    """Guard the chunking invariant profiles depend on: no event in two chunks.

    A violation would silently double-count that event in the merged mean/RMS
    with no other symptom, so this is worth a loud failure rather than a
    quietly-wrong plot.
    """
    seen: set[int] = set()
    for ids in id_lists:
        overlap = seen & set(ids)
        if overlap:
            raise ValueError(
                f"{spec_id} ({side}): event_id(s) {sorted(overlap)[:5]} appear "
                "in more than one chunk — chunking must be event-disjoint"
            )
        seen.update(ids)


def _profile_finalize(
    parts: list[dict],
    ctx: Context,
    spec_id: str,
    title: str,
    xlabel: str,
    edges_key: str,
) -> Reduced:
    edges = np.asarray(getattr(ctx, edges_key))
    nb = len(edges) - 1
    names = list(parts[0]["r"])

    _assert_event_disjoint([p["t"]["ids"] for p in parts], spec_id, "reference")
    t_mats = [np.asarray(p["t"]["mat"], dtype=float).reshape(-1, nb) for p in parts]
    t_mean, t_std = profile_finalize(t_mats)

    series: dict[str, dict] = {}
    for name in names:
        _assert_event_disjoint([p["r"][name]["ids"] for p in parts], spec_id, name)
        mats = [np.asarray(p["r"][name]["mat"], dtype=float).reshape(-1, nb) for p in parts]
        mean, std = profile_finalize(mats)
        series[name] = {"mean": mean.tolist(), "std": std.tolist()}

    return Reduced(
        id=spec_id,
        family="shower",
        kind="profile",
        title=title,
        xlabel=xlabel,
        payload={
            "edges": edges.tolist(),
            "series": series,
            "reference": {"mean": t_mean.tolist(), "std": t_std.tolist()},
            "ylabel": "mean deposited energy per event [MeV]",
        },
    )


# ---------------------------------------------------------------------------
# shower containment depth (reuses the longitudinal profile's per-event matrix)
# ---------------------------------------------------------------------------

_CONTAINMENT_QUANTILES: list[tuple[float, str]] = [
    (0.90, "shower_containment_depth_90"),
    (0.95, "shower_containment_depth_95"),
]


def _containment_finalize(parts: list[dict], ctx: Context, spec_id: str, quantile: float) -> Reduced:
    edges = np.asarray(ctx.depth_edges)
    nb = len(edges) - 1
    names = list(parts[0]["r"])

    _assert_event_disjoint([p["t"]["ids"] for p in parts], spec_id, "reference")
    t_full = np.concatenate([np.asarray(p["t"]["mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
    t_depth = _containment_depths(t_full, edges, quantile)

    r_depths: dict[str, np.ndarray] = {}
    for name in names:
        _assert_event_disjoint([p["r"][name]["ids"] for p in parts], spec_id, name)
        full = np.concatenate([np.asarray(p["r"][name]["mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
        r_depths[name] = _containment_depths(full, edges, quantile)

    hedges, counts = _np_hist_shared_edges([t_depth, *(r_depths[n] for n in names)], ctx.n_marginal_bins)
    t_counts, *r_counts = counts
    return Reduced(
        id=spec_id,
        family="shower",
        kind="overlay_hist",
        title=f"Shower containment depth ({quantile:.0%} of deposited energy)",
        xlabel=f"depth containing {quantile:.0%} of deposited energy [mm]",
        payload={
            "edges": hedges.tolist(),
            "series": {name: c.astype(np.int64).tolist() for name, c in zip(names, r_counts)},
            "reference": t_counts.astype(np.int64).tolist(),
            "log_y": False,
        },
    )


# ---------------------------------------------------------------------------
# species share + leakage
# ---------------------------------------------------------------------------


def _species_share_partial(b: Bundle) -> dict:
    t = species_share(b.t_all)

    def _map(rs: RolloutSide) -> dict[str, float]:
        r = species_share(rs.all)
        return {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())}

    return {
        "r": _per_rollout(b, _map),
        "t": {str(k): v for k, v in zip(t["pdg"].to_list(), t["total_edep"].to_list())},
    }


def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced:
    names = list(parts[0]["r"])
    r_maps = {n: sum_merge([p["r"][n] for p in parts]) for n in names}
    t_map = sum_merge([p["t"] for p in parts])
    t_tot = sum(t_map.values()) or 1.0
    labels = [pdg_label(k) for k in ctx.top_pdgs]

    series: dict[str, list[float]] = {}
    for n in names:
        r_tot = sum(r_maps[n].values()) or 1.0
        series[n] = [r_maps[n].get(str(k), 0.0) / r_tot for k in ctx.top_pdgs]

    return Reduced(
        id="species_edep_share",
        family="species",
        kind="bar",
        title="Deposited-energy share by species",
        xlabel="species",
        payload={
            "labels": labels,
            "series": series,
            "reference": [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs],
            "ylabel": "fraction of total deposited energy",
        },
    )


def _leakage_partial(b: Bundle) -> dict:
    return {"r": _per_rollout(b, lambda rs: leakage_fraction(rs.all).tolist())}


def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
    names = list(parts[0]["r"])
    arrays = {n: np.concatenate([np.asarray(p["r"][n], dtype=float) for p in parts]) for n in names}
    max_val = max((float(a.max()) for a in arrays.values() if len(a)), default=1e-3)
    edges = np.linspace(0.0, max(max_val, 1e-3), ctx.n_marginal_bins + 1)
    series = {n: np.histogram(arrays[n], edges)[0].astype(np.int64).tolist() for n in names}
    return Reduced(
        id="leakage_fraction",
        family="species",
        kind="single_hist",
        title="Escaped (leakage) energy fraction per shower",
        xlabel="escaped energy fraction",
        payload={
            "edges": edges.tolist(),
            "series": series,
            "log_y": True,
            "note": "rollout only; the reference has no detector-escape concept",
        },
    )


# ---------------------------------------------------------------------------
# secondaries
# ---------------------------------------------------------------------------


def _t_sec(b: Bundle) -> pl.LazyFrame:
    return secondaries(b.t_all, Side.reference)


def _r_sec(rs: RolloutSide) -> pl.LazyFrame:
    return secondaries(rs.phys, Side.rollout)


def _sec_count_per_event_partial(b: Bundle) -> dict:
    t = _t_sec(b).group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()

    def _r(rs: RolloutSide) -> list[float]:
        return (
            _r_sec(rs)
            .group_by("event_id")
            .agg(pl.len().alias("n"))
            .collect(engine="streaming")["n"]
            .to_numpy()
            .tolist()
        )

    return {"r": _per_rollout(b, _r), "t": t.tolist()}


def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced:
    names = list(parts[0]["r"])
    t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
    r = {n: np.concatenate([np.asarray(p["r"][n], dtype=float) for p in parts]) for n in names}
    edges, counts = _np_hist_shared_edges([t, *(r[n] for n in names)], min(ctx.n_marginal_bins, 40))
    t_c, *r_cs = counts
    return Reduced(
        id="sec_count_per_event",
        family="secondaries",
        kind="overlay_hist",
        title="Number of secondaries per event",
        xlabel="secondaries per event",
        payload={
            "edges": edges.tolist(),
            "series": {n: c.astype(np.int64).tolist() for n, c in zip(names, r_cs)},
            "reference": t_c.astype(np.int64).tolist(),
            "log_y": False,
        },
    )


def _counts_by_pdg(sec_lf: pl.LazyFrame) -> dict[str, int]:
    df = sec_lf.group_by("pdg").agg(pl.len().alias("n")).collect(engine="streaming")
    return {str(k): v for k, v in zip(df["pdg"].to_list(), df["n"].to_list())}


def _sec_count_per_species_partial(b: Bundle) -> dict:
    return {"r": _per_rollout(b, lambda rs: _counts_by_pdg(_r_sec(rs))), "t": _counts_by_pdg(_t_sec(b))}


def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
    names = list(parts[0]["r"])
    r_maps = {n: sum_merge([p["r"][n] for p in parts]) for n in names}
    t = sum_merge([p["t"] for p in parts])

    all_keys = set(t)
    for m in r_maps.values():
        all_keys |= set(m)

    def _total(k: str) -> float:
        return t.get(k, 0) + sum(m.get(k, 0) for m in r_maps.values())

    keys = sorted(all_keys, key=lambda k: -_total(k))[: len(ctx.top_pdgs)]
    return Reduced(
        id="sec_count_per_species",
        family="secondaries",
        kind="bar",
        title="Secondary count by species",
        xlabel="species",
        payload={
            "labels": [pdg_label(int(k)) for k in keys],
            "series": {n: [float(r_maps[n].get(k, 0)) for k in keys] for n in names},
            "reference": [float(t.get(k, 0)) for k in keys],
            "ylabel": "secondary count",
        },
    )


# Per-step secondary multiplicity. Fixed integer edges (bin i == exactly i
# secondaries, the top bin an overflow bucket) keep both plots sum-mergeable
# across chunks — no shared-range pass needed. The species heatmap gets a
# shorter row axis because a single step rarely emits many of *one* species.
_N_SEC_STEP_CAP = 20
_N_SEC_SPECIES_CAP = 10
_OTHER_KEY = "other"


def _n_sec_edges(cap: int) -> np.ndarray:
    return np.arange(-0.5, cap + 1.5)


def _sec_step_key_lf(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
    """Secondaries with their emitting-step key.

    The rollout side reads *all* rows, not just physical ones: a secondary
    whose very first row is a synthetic termination row (born, then immediately
    escaped or cut) was still produced by its parent step, and dropping it would
    undercount that step's multiplicity.
    """
    return secondaries_by_step(lf, side)


def _n_steps(lf: pl.LazyFrame) -> int:
    """Number of (physical) step rows — the denominator the zero rows come from."""
    return int(lf.select(pl.len()).collect(engine="streaming").item())


def _sec_count_per_step_partial(b: Bundle) -> dict:
    edges = _n_sec_edges(_N_SEC_STEP_CAP)

    def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict:
        per_step = sec_lf.group_by("step_key").agg(pl.len().alias("n"))
        return {
            "h": _partial_hist(per_step, pl.col("n").clip(0, _N_SEC_STEP_CAP), edges),
            "n_steps": _n_steps(steps_lf),
        }

    return {
        "r": _per_rollout(b, lambda rs: _side(_sec_step_key_lf(rs.all, Side.rollout), rs.phys)),
        "t": _side(_sec_step_key_lf(b.t_all, Side.reference), b.t_phys),
    }


def _zero_filled(part_hists: list[dict], n_steps: int, key, nbins: int) -> list[int]:
    """Merged counts for one series, with bin 0 (= steps that emitted none) filled in.

    The reduction only ever sees steps that produced at least one secondary, so
    the empty ones are recovered by subtraction from the total step count.
    """
    counts = _finalize_counts(sum_merge(part_hists), key, nbins)
    counts[0] = max(n_steps - int(sum(counts)), 0)
    return [int(c) for c in counts]


def _sec_count_per_step_finalize(parts: list[dict], ctx: Context) -> Reduced:
    edges = _n_sec_edges(_N_SEC_STEP_CAP)
    nb = len(edges) - 1
    names = list(parts[0]["r"])
    series = {
        name: _zero_filled([p["r"][name]["h"] for p in parts], sum(p["r"][name]["n_steps"] for p in parts), 0, nb)
        for name in names
    }
    return Reduced(
        id="sec_count_per_step",
        family="secondaries",
        kind="overlay_hist",
        title="Number of secondaries per step",
        xlabel="secondaries per step",
        payload={
            "edges": edges.tolist(),
            "series": series,
            "reference": _zero_filled([p["t"]["h"] for p in parts], sum(p["t"]["n_steps"] for p in parts), 0, nb),
            "log_y": True,
        },
    )


def _species_key_expr(top_pdgs: list[int]) -> pl.Expr:
    """``pdg`` bucketed into the shared top-K columns plus one ``other`` bin."""
    return pl.when(pl.col("pdg").is_in(list(top_pdgs))).then(pl.col("pdg").cast(pl.Utf8)).otherwise(pl.lit(_OTHER_KEY))


def _sec_count_per_step_by_species_partial(b: Bundle) -> dict:
    edges = _n_sec_edges(_N_SEC_SPECIES_CAP)
    group = _species_key_expr(b.ctx.top_pdgs)

    def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict:
        per_step_species = sec_lf.group_by("step_key", "pdg").agg(pl.len().alias("n"))
        return {
            "h": _partial_hist(per_step_species, pl.col("n").clip(0, _N_SEC_SPECIES_CAP), edges, group=group),
            "n_steps": _n_steps(steps_lf),
        }

    return {
        "r": _per_rollout(b, lambda rs: _side(_sec_step_key_lf(rs.all, Side.rollout), rs.phys)),
        "t": _side(_sec_step_key_lf(b.t_all, Side.reference), b.t_phys),
    }


def _sec_count_per_step_by_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
    edges = _n_sec_edges(_N_SEC_SPECIES_CAP)
    nb = len(edges) - 1
    names = list(parts[0]["r"])
    keys = [str(p) for p in ctx.top_pdgs] + [_OTHER_KEY]

    def _matrix(hists: list[dict], n_steps: int) -> list[list[int]]:
        # columns = species, rows = multiplicity; every species gets its own
        # zero row (steps that produced none of *that* species).
        cols = [_zero_filled(hists, n_steps, k, nb) for k in keys]
        return [[cols[j][i] for j in range(len(keys))] for i in range(nb)]

    return Reduced(
        id="sec_count_per_step_by_species",
        family="secondaries",
        kind="heatmap",
        title="Per-step secondary multiplicity by species",
        xlabel="species",
        payload={
            "series": {
                n: _matrix([p["r"][n]["h"] for p in parts], sum(p["r"][n]["n_steps"] for p in parts)) for n in names
            },
            "reference": _matrix([p["t"]["h"] for p in parts], sum(p["t"]["n_steps"] for p in parts)),
            "row_labels": [str(i) for i in range(_N_SEC_SPECIES_CAP)] + [f"{_N_SEC_SPECIES_CAP}+"],
            "col_labels": [pdg_label(k) for k in ctx.top_pdgs] + [_OTHER_KEY],
            "ylabel": "secondaries of this species per step",
            "cbar_label": "step count",
            "log_color": True,
        },
    )


def _sec_energy_partial(b: Bundle) -> dict:
    edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1)
    return {
        "r": _per_rollout(b, lambda rs: _partial_hist(_r_sec(rs), pl.col("energy"), edges)),
        "t": _partial_hist(_t_sec(b), pl.col("energy"), edges),
    }


def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
    edges = np.linspace(*ctx.sec_energy_range, ctx.n_sec_bins + 1)
    nb = len(edges) - 1
    names = list(parts[0]["r"])
    t = sum_merge([p["t"] for p in parts])
    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
    return Reduced(
        id="sec_energy",
        family="secondaries",
        kind="overlay_hist",
        title="Secondary birth energy",
        xlabel="secondary energy [MeV]",
        payload={"edges": edges.tolist(), "series": series, "reference": _finalize_counts(t, 0, nb), "log_y": True},
    )


def _sec_cos_angle_partial(b: Bundle) -> dict:
    edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1)
    cos = (pl.col("sdx") * pl.col("axis_x") + pl.col("sdy") * pl.col("axis_y") + pl.col("sdz") * pl.col("axis_z")).clip(
        -1.0, 1.0
    )

    def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]:
        ea = entry_axis(steps_lf)
        return _partial_hist(attach_entry_axis(sec_lf, ea), cos, edges)

    return {
        "r": _per_rollout(b, lambda rs: _side(_r_sec(rs), rs.phys)),
        "t": _side(_t_sec(b), b.t_all),
    }


def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
    edges = np.linspace(-1.0, 1.0, ctx.n_sec_bins + 1)
    nb = len(edges) - 1
    names = list(parts[0]["r"])
    t = sum_merge([p["t"] for p in parts])
    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
    return Reduced(
        id="sec_cos_angle",
        family="secondaries",
        kind="overlay_hist",
        title="Secondary emission angle relative to the shower axis",
        xlabel="cos of emission angle",
        payload={"edges": edges.tolist(), "series": series, "reference": _finalize_counts(t, 0, nb), "log_y": False},
    )


# ---------------------------------------------------------------------------
# eval cost (not chunked — metadata-only, no row scan)
# ---------------------------------------------------------------------------

_EVAL_COST_LABELS = ["sampling / simulation", "parquet write / convert", "total"]
_EVAL_COST_NOTE = (
    "no rollout in this run carries a `timing` block — re-run `giant rollout` "
    "(timing instrumentation added after this checkpoint's rollout run) to "
    "populate this plot"
)


def _eval_cost_per_step(b: Bundle) -> Reduced:
    """Per-rollout µs/physical-step vs the measured Geant4 reference.

    ``timing`` (``giant.cli``'s ``rollout`` command) is metadata carried on
    the rollout YAML, not derived from the row data, so this needs no chunked
    scan — same shape as the router diagnostics above.
    """
    series: dict[str, list[float]] = {}
    speedup: dict[str, float] = {}
    for name, rs in b.rollouts.items():
        t = rs.timing
        if not t or t.get("us_per_step") is None:
            continue
        sample_us = t["us_per_step"]
        write_us = t.get("write_us_per_step") or 0.0
        series[name] = [sample_us, write_us, sample_us + write_us]

    if not series:
        return Reduced(
            id="eval_cost_per_step",
            family="cost",
            kind="unavailable",
            title="Eval cost per step: surrogate vs Geant4",
            xlabel="n/a",
            payload={"note": _EVAL_COST_NOTE},
        )

    g4 = geant4_per_step_us()
    reference = [g4["sim_us_per_step"], g4["convert_us_per_step"], g4["total_us_per_step"]]
    for name, vals in series.items():
        speedup[name] = reference[-1] / vals[-1] if vals[-1] else float("inf")

    return Reduced(
        id="eval_cost_per_step",
        family="cost",
        kind="bar",
        title="Eval cost per step: surrogate vs Geant4",
        xlabel="phase",
        payload={
            "labels": _EVAL_COST_LABELS,
            "series": series,
            "reference": reference,
            "ylabel": "µs per physical step",
            "log_y": True,
        },
        meta={
            "speedup_vs_geant4_total": speedup,
            "geant4_provenance": GEANT4_REFERENCE["provenance"],
            "caveat": (
                "The Geant4 reference is measured single-threaded on one CPU core "
                "(see giant.analysis.geant4_reference); a rollout's timing is "
                "whatever device it actually ran on (see each series' device in "
                "run_meta.json's plot_meta). This is a deployment-speedup ratio, "
                "not a same-hardware or per-FLOP comparison."
            ),
        },
    )


_eval_cost_per_step_partial, _eval_cost_per_step_finalize = _unchunkable(_eval_cost_per_step)


# ---------------------------------------------------------------------------
# router diagnostics (not chunked — already bounded/subsampled)
# ---------------------------------------------------------------------------

_router_gating_partial, _router_gating_finalize = _unchunkable(lambda b: compute_router_gating(b.rollouts, b.t_phys))
_router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
    lambda b: compute_router_share_by_pdg(b.rollouts, b.t_phys, b.ctx.top_pdgs)
)
_router_share_process_partial, _router_share_process_finalize = _unchunkable(
    lambda b: compute_router_share_by_process(b.rollouts, b.t_phys)
)
_router_specialization_partial, _router_specialization_finalize = _unchunkable(
    lambda b: compute_router_specialization(b.rollouts, b.t_phys)
)
_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = _unchunkable(
    lambda b: compute_type_embedding_l1_distance(b.rollouts)
)


# ---------------------------------------------------------------------------
# registry assembly
# ---------------------------------------------------------------------------

MARGINAL_VARS = ["step_length", "edep", "delta_e", "post_E", "cos_scatter"]
GROUPING_AXES = ["energy", "pdg", "material"]


# ---------------------------------------------------------------------------
# giant predict: paired truth/pred family ("prediction")
#
# Unlike every spec above (rollout series vs one shared reference — an
# *unpaired* distribution comparison), a prediction has a truth row for every
# output row. These specs compare true_<var> against pred_<var> from
# `giant.analysis.prediction.paired_frame` directly — no reference series,
# and (unlike a rollout) a prediction whose --coord is "local" never has
# secondaries or a material/pdg breakdown. `b.predictions` is empty on a run
# with no --prediction given, in which case every spec here degrades to
# kind="unavailable" rather than raising.
# ---------------------------------------------------------------------------

_PRED_UNAVAILABLE_NOTE = "no --prediction given to `analyze prep`/`submit`"

_PRED_TITLE_NAMES = {
    "step_length": "Step length",
    "edep": "Deposited energy per step",
    "delta_e": "Energy loss per step",
    "post_E": "Post-step energy",
    "cos_scatter": "Scattering cosine (pre_dir . post_dir)",
    "cos_travel": "Travel-direction cosine (pre_dir . (post_pos - pre_pos))",
}

_PRED_VAR_LABELS = {
    "step_length": "step length [mm]",
    "edep": "deposited energy [MeV]",
    "delta_e": "energy loss [MeV]",
    "post_E": "post-step energy [MeV]",
    "cos_scatter": "cos(scattering angle)",
    "cos_travel": "cos(travel-direction angle)",
}


def _per_prediction(b: Bundle, fn: Callable[[PredictionSide], object]) -> dict[str, object]:
    """``{name: fn(prediction_side)}`` over every prediction, preserving CLI order."""
    return {name: fn(ps) for name, ps in b.predictions.items()}


def _pred_unavailable(
    spec_id: str, family: str, title: str, xlabel: str, note: str = _PRED_UNAVAILABLE_NOTE
) -> Reduced:
    return Reduced(id=spec_id, family=family, kind="unavailable", title=title, xlabel=xlabel, payload={"note": note})


# ---- marginals + 2D truth-vs-pred -----------------------------------------


def _pred_marginal_partial(b: Bundle, var: str) -> dict:
    if not b.predictions or var not in b.ctx.pred_var_ranges:
        return {"available": False}
    edges = b.ctx.pred_marginal_edges(var)
    nb = len(edges) - 1

    def _one(ps: PredictionSide) -> dict | None:
        if var not in paired_vars_for_coord(ps.coord):
            return None
        entry = {"pred": _counts(hist1d(ps.paired, pl.col(f"pred_{var}"), edges), 0, nb)}
        if ps.has_truth:
            entry["true"] = _counts(hist1d(ps.paired, pl.col(f"true_{var}"), edges), 0, nb)
        return entry

    return {"available": True, "p": _per_prediction(b, _one)}


def _pred_marginal_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced:
    if not parts[0]["available"]:
        return _pred_unavailable(f"pred_marginal_{var}", "prediction", _PRED_TITLE_NAMES[var], _PRED_VAR_LABELS[var])
    edges = ctx.pred_marginal_edges(var)
    names = list(parts[0]["p"])
    series: dict[str, dict] = {}
    for name in names:
        entries = [p["p"][name] for p in parts]
        if entries[0] is None:
            continue
        merged = sum_merge(entries)
        series[name] = {k: [int(x) for x in v] for k, v in merged.items()}
    return Reduced(
        id=f"pred_marginal_{var}",
        family="prediction",
        kind="paired_hist",
        title=_PRED_TITLE_NAMES[var],
        xlabel=_PRED_VAR_LABELS[var],
        payload={"edges": edges.tolist(), "series": series, "log_y": var in PAIRED_SCALARS},
    )


def _pred_scatter_partial(b: Bundle, var: str) -> dict:
    if not b.predictions or var not in b.ctx.pred_var_ranges:
        return {"available": False}
    lo, hi = b.ctx.pred_var_ranges[var]
    edges = np.linspace(lo, hi, b.ctx.n_marginal_bins + 1)

    def _one(ps: PredictionSide) -> list[list[int]] | None:
        if var not in paired_vars_for_coord(ps.coord) or not ps.has_truth:
            return None
        return hist2d(ps.paired, pl.col(f"true_{var}"), pl.col(f"pred_{var}"), edges, edges).tolist()

    return {"available": True, "edges": edges.tolist(), "p": _per_prediction(b, _one)}


def _pred_scatter_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced:
    if not parts[0]["available"]:
        return _pred_unavailable(
            f"pred_scatter_{var}", "prediction", f"{_PRED_TITLE_NAMES[var]}: truth vs prediction", _PRED_VAR_LABELS[var]
        )
    edges = np.asarray(parts[0]["edges"])
    names = list(parts[0]["p"])
    series: dict[str, list] = {}
    for name in names:
        mats = [p["p"][name] for p in parts]
        if mats[0] is None:
            continue
        series[name] = np.sum([np.asarray(m, dtype=np.int64) for m in mats], axis=0).tolist()
    return Reduced(
        id=f"pred_scatter_{var}",
        family="prediction",
        kind="heatmap2d",
        title=f"{_PRED_TITLE_NAMES[var]}: truth vs prediction",
        xlabel=f"true {_PRED_VAR_LABELS[var]}",
        payload={
            "x_edges": edges.tolist(),
            "y_edges": edges.tolist(),
            "series": series,
            "ylabel": f"predicted {_PRED_VAR_LABELS[var]}",
            "cbar_label": "step count",
            "log_color": True,
            "diagonal": True,
        },
    )


# ---- residuals + calibration -----------------------------------------------


def _pred_residual_partial(b: Bundle, var: str) -> dict:
    if not b.predictions or var not in b.ctx.pred_residual_ranges:
        return {"available": False}
    edges = b.ctx.pred_residual_edges(var)
    nb = len(edges) - 1

    def _one(ps: PredictionSide) -> list[int] | None:
        if var not in paired_vars_for_coord(ps.coord) or not ps.has_truth:
            return None
        resid = pl.col(f"pred_{var}") - pl.col(f"true_{var}")
        return _counts(hist1d(ps.paired, resid, edges), 0, nb)

    return {"available": True, "p": _per_prediction(b, _one)}


def _pred_residual_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced:
    title, xlabel = f"{_PRED_TITLE_NAMES[var]} residual", f"pred - true {_PRED_VAR_LABELS[var]}"
    if not parts[0]["available"]:
        return _pred_unavailable(f"pred_residual_{var}", "prediction", title, xlabel)
    edges = ctx.pred_residual_edges(var)
    names = list(parts[0]["p"])
    series: dict[str, list[int]] = {}
    for name in names:
        entries = [p["p"][name] for p in parts]
        if entries[0] is None:
            continue
        series[name] = [int(x) for x in sum_merge([{"c": e} for e in entries])["c"]]
    return Reduced(
        id=f"pred_residual_{var}",
        family="prediction",
        kind="single_hist",
        title=title,
        xlabel=xlabel,
        payload={"edges": edges.tolist(), "series": series, "log_y": True},
    )


_REL_RESIDUAL_LO, _REL_RESIDUAL_HI = -2.0, 2.0


def _pred_relative_residual_partial(b: Bundle, var: str) -> dict:
    if not b.predictions:
        return {"available": False}
    edges = np.linspace(_REL_RESIDUAL_LO, _REL_RESIDUAL_HI, b.ctx.n_marginal_bins + 1)
    nb = len(edges) - 1

    def _one(ps: PredictionSide) -> list[int] | None:
        if var not in paired_vars_for_coord(ps.coord) or not ps.has_truth:
            return None
        rel = (pl.col(f"pred_{var}") - pl.col(f"true_{var}")) / pl.col(f"true_{var}")
        return _counts(hist1d(ps.paired, rel, edges), 0, nb)

    return {"available": True, "edges": edges.tolist(), "p": _per_prediction(b, _one)}


def _pred_relative_residual_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced:
    title, xlabel = f"{_PRED_TITLE_NAMES[var]} relative residual", "(pred - true) / true"
    if not parts[0]["available"]:
        return _pred_unavailable(f"pred_relative_residual_{var}", "prediction", title, xlabel)
    edges = np.asarray(parts[0]["edges"])
    names = list(parts[0]["p"])
    series: dict[str, list[int]] = {}
    for name in names:
        entries = [p["p"][name] for p in parts]
        if entries[0] is None:
            continue
        series[name] = [int(x) for x in sum_merge([{"c": e} for e in entries])["c"]]
    return Reduced(
        id=f"pred_relative_residual_{var}",
        family="prediction",
        kind="single_hist",
        title=title,
        xlabel=xlabel,
        payload={"edges": edges.tolist(), "series": series, "log_y": True},
    )


def _pred_residual_profile_partial(b: Bundle, var: str) -> dict:
    if not b.predictions or var not in b.ctx.pred_var_ranges:
        return {"available": False}
    edges = b.ctx.pred_marginal_edges(var)  # bin by truth value

    def _one(ps: PredictionSide) -> dict | None:
        if var not in paired_vars_for_coord(ps.coord) or not ps.has_truth:
            return None
        resid = pl.col(f"pred_{var}") - pl.col(f"true_{var}")
        return binned_moments(ps.paired, pl.col(f"true_{var}"), resid, edges)

    return {"available": True, "p": _per_prediction(b, _one)}


def _moments_to_mean_std(merged: dict[str, list]) -> tuple[np.ndarray, np.ndarray]:
    n = np.asarray(merged["n"], dtype=np.float64)
    s = np.asarray(merged["sum"], dtype=np.float64)
    ss = np.asarray(merged["sumsq"], dtype=np.float64)
    with np.errstate(invalid="ignore", divide="ignore"):
        mean = np.where(n > 0, s / n, 0.0)
        var = np.where(n > 0, ss / n - mean**2, 0.0)
    return mean, np.sqrt(np.clip(var, 0.0, None))


def _pred_residual_profile_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced:
    title, xlabel = f"{_PRED_TITLE_NAMES[var]} residual vs truth", f"true {_PRED_VAR_LABELS[var]}"
    if not parts[0]["available"]:
        return _pred_unavailable(f"pred_residual_profile_{var}", "prediction", title, xlabel)
    edges = ctx.pred_marginal_edges(var)
    names = list(parts[0]["p"])
    series: dict[str, dict] = {}
    for name in names:
        entries = [p["p"][name] for p in parts]
        if entries[0] is None:
            continue
        mean, std = _moments_to_mean_std(sum_merge(entries))
        series[name] = {"mean": mean.tolist(), "std": std.tolist()}
    return Reduced(
        id=f"pred_residual_profile_{var}",
        family="prediction",
        kind="profile",
        title=title,
        xlabel=xlabel,
        payload={"edges": edges.tolist(), "series": series},
    )


_PRED_GROUP_AXES = ("overall", "energy", "pdg", "material")


def _pred_group_expr(lf: pl.LazyFrame, axis: str, ctx: Context) -> pl.Expr:
    if axis == "pdg":
        return pl.col("pdg")
    if axis == "material":
        return pl.col("material")
    if axis == "energy":
        return _energy_group_expr(lf, np.asarray(ctx.energy_edges))
    return pl.lit(0, dtype=pl.Int64)


def _pred_group_keys(ctx: Context, axis: str) -> list:
    if axis == "pdg":
        return list(ctx.top_pdgs)
    if axis == "material":
        return list(ctx.materials)
    if axis == "energy":
        return list(range(len(ctx.energy_edges) - 1))
    return [0]


def _grouped_moments(lf: pl.LazyFrame, group: pl.Expr, value: pl.Expr) -> dict[str, dict[str, float]]:
    """``{str(group_key): {"n", "sum", "sumsq"}}`` — one streaming pass, sum-mergeable."""
    res = (
        lf.select(group.alias("_g"), value.alias("_v"))
        .drop_nulls(["_g", "_v"])
        .group_by("_g")
        .agg(pl.len().alias("_n"), pl.col("_v").sum().alias("_s"), (pl.col("_v") ** 2).sum().alias("_ss"))
        .collect(engine="streaming")
    )
    return {str(g): {"n": float(n), "sum": float(s), "sumsq": float(ss)} for g, n, s, ss in res.iter_rows()}


def _pred_scorecard_partial(b: Bundle) -> dict:
    """Per-(var, prediction, axis) group-keyed truth/pred histograms + residual
    moments — the shared input to the KS/bias/RMSE scorecards below (three
    separate specs, each doing its own finalize math over this one compute)."""
    if not b.predictions:
        return {"available": False}
    out: dict[str, dict] = {}
    for var in PAIRED_SCALARS:
        if var not in b.ctx.pred_var_ranges:
            continue
        edges = b.ctx.pred_marginal_edges(var)

        def _one(ps: PredictionSide, var=var, edges=edges) -> dict | None:
            if var not in paired_vars_for_coord(ps.coord) or not ps.has_truth:
                return None
            resid = pl.col(f"pred_{var}") - pl.col(f"true_{var}")
            per_axis = {}
            for axis in _PRED_GROUP_AXES:
                group = _pred_group_expr(ps.paired, axis, b.ctx)
                per_axis[axis] = {
                    "true_hist": {
                        str(k): v.tolist() for k, v in hist1d(ps.paired, pl.col(f"true_{var}"), edges, group).items()
                    },
                    "pred_hist": {
                        str(k): v.tolist() for k, v in hist1d(ps.paired, pl.col(f"pred_{var}"), edges, group).items()
                    },
                    "moments": _grouped_moments(ps.paired, group, resid),
                    "true_moments": _grouped_moments(ps.paired, group, pl.col(f"true_{var}")),
                }
            return per_axis

        out[var] = _per_prediction(b, _one)
    return {"available": True, "vars": out}


def _pred_scorecard_matrix(
    parts: list[dict], ctx: Context, cell: Callable[[dict, dict, str, str], float]
) -> dict[str, list[list[float]]]:
    """Shared finalize skeleton for the three scorecards: rows = ``PAIRED_SCALARS``,
    cols = ``_PRED_GROUP_AXES``, one matrix per prediction. ``cell(true_h_or_moments,
    pred_h_or_moments, axis, group_key)`` computes one entry from that axis'
    merged group-keyed dict pair (weighted-averaged over the axis' groups)."""
    names = list(parts[0]["vars"][next(iter(parts[0]["vars"]))])
    matrices: dict[str, list[list[float]]] = {name: [] for name in names}
    for var in PAIRED_SCALARS:
        var_parts = [p["vars"].get(var) for p in parts]
        for name in names:
            row: list[float] = []
            for axis in _PRED_GROUP_AXES:
                raw_entries = [vp[name][axis] if vp is not None and vp[name] is not None else None for vp in var_parts]
                if raw_entries[0] is None:
                    row.append(float("nan"))
                    continue
                # non-None for every chunk: (var, name, axis) availability is
                # a static fact of the prediction's coord, not chunk-dependent.
                entries = cast("list[dict]", raw_entries)
                merged_true_hist = sum_merge([{k: v for k, v in e["true_hist"].items()} for e in entries])
                merged_pred_hist = sum_merge([{k: v for k, v in e["pred_hist"].items()} for e in entries])
                merged_moments: dict[str, dict[str, float]] = {}
                for e in entries:
                    for k, m in e["moments"].items():
                        acc = merged_moments.setdefault(k, {"n": 0.0, "sum": 0.0, "sumsq": 0.0})
                        for f in ("n", "sum", "sumsq"):
                            acc[f] += m[f]
                merged_true_moments: dict[str, dict[str, float]] = {}
                for e in entries:
                    for k, m in e["true_moments"].items():
                        acc = merged_true_moments.setdefault(k, {"n": 0.0, "sum": 0.0, "sumsq": 0.0})
                        for f in ("n", "sum", "sumsq"):
                            acc[f] += m[f]
                dists, weights = [], []
                for k in _pred_group_keys(ctx, axis):
                    key = str(k)
                    if key not in merged_moments:
                        continue
                    val = cell(
                        {"hist": merged_true_hist.get(key), "moments": merged_true_moments.get(key)},
                        {"hist": merged_pred_hist.get(key), "moments": merged_moments.get(key)},
                        axis,
                        key,
                    )
                    w = merged_moments[key]["n"]
                    if w <= 0 or not np.isfinite(val):
                        continue
                    dists.append(val)
                    weights.append(w)
                row.append(float(np.average(dists, weights=weights)) if dists else float("nan"))
            matrices[name].append(row)
    return matrices


def _pred_ks_finalize(parts: list[dict], ctx: Context) -> Reduced:
    if not parts[0]["available"] or not parts[0]["vars"]:
        return _pred_unavailable(
            "pred_ks_summary", "prediction", "Truth/pred distance summary (KS statistic)", "grouping axis"
        )

    def _cell(true_side: dict, pred_side: dict, axis: str, key: str) -> float:
        if pred_side["hist"] is None or true_side["hist"] is None:
            return float("nan")
        return _ks_statistic(pred_side["hist"], true_side["hist"])

    matrices = _pred_scorecard_matrix(parts, ctx, _cell)
    return Reduced(
        id="pred_ks_summary",
        family="prediction",
        kind="heatmap",
        title="Truth/pred distance summary (KS statistic)",
        xlabel="grouping axis",
        payload={
            "series": matrices,
            "row_labels": [_PRED_TITLE_NAMES[v] for v in PAIRED_SCALARS],
            "col_labels": list(_PRED_GROUP_AXES),
            "ylabel": "variable",
            "cbar_label": "KS statistic (0 = identical, 1 = maximal mismatch)",
            "vmin": 0.0,
            "vmax": 1.0,
        },
    )


def _pred_bias_finalize(parts: list[dict], ctx: Context) -> Reduced:
    if not parts[0]["available"] or not parts[0]["vars"]:
        return _pred_unavailable("pred_bias_summary", "prediction", "Relative bias summary", "grouping axis")

    def _cell(true_side: dict, pred_side: dict, axis: str, key: str) -> float:
        n, s = pred_side["moments"]["n"], pred_side["moments"]["sum"]
        if n <= 0:
            return float("nan")
        bias = s / n
        tm = true_side["moments"]
        denom = abs(tm["sum"] / tm["n"]) if tm and tm["n"] > 0 else 0.0
        return bias / denom if denom > 1e-12 else float("nan")

    matrices = _pred_scorecard_matrix(parts, ctx, _cell)
    return Reduced(
        id="pred_bias_summary",
        family="prediction",
        kind="heatmap",
        title="Relative bias summary (mean(pred - true) / mean|true|)",
        xlabel="grouping axis",
        payload={
            "series": matrices,
            "row_labels": [_PRED_TITLE_NAMES[v] for v in PAIRED_SCALARS],
            "col_labels": list(_PRED_GROUP_AXES),
            "ylabel": "variable",
            "cbar_label": "relative bias",
            "cmap": "RdBu_r",
            "vmin": -0.5,
            "vmax": 0.5,
        },
    )


def _pred_rmse_finalize(parts: list[dict], ctx: Context) -> Reduced:
    if not parts[0]["available"] or not parts[0]["vars"]:
        return _pred_unavailable("pred_rmse_summary", "prediction", "Relative RMSE summary", "grouping axis")

    def _cell(true_side: dict, pred_side: dict, axis: str, key: str) -> float:
        n, s, ss = pred_side["moments"]["n"], pred_side["moments"]["sum"], pred_side["moments"]["sumsq"]
        if n <= 0:
            return float("nan")
        rmse = float(np.sqrt(max(ss / n, 0.0)))
        tm = true_side["moments"]
        denom = abs(tm["sum"] / tm["n"]) if tm and tm["n"] > 0 else 0.0
        del s
        return rmse / denom if denom > 1e-12 else float("nan")

    matrices = _pred_scorecard_matrix(parts, ctx, _cell)
    return Reduced(
        id="pred_rmse_summary",
        family="prediction",
        kind="heatmap",
        title="Relative RMSE summary",
        xlabel="grouping axis",
        payload={
            "series": matrices,
            "row_labels": [_PRED_TITLE_NAMES[v] for v in PAIRED_SCALARS],
            "col_labels": list(_PRED_GROUP_AXES),
            "ylabel": "variable",
            "cbar_label": "relative RMSE",
            "vmin": 0.0,
        },
    )


# ---- confusion matrices -----------------------------------------------


def _n_sec_confusion_edges(cap: int) -> np.ndarray:
    return np.arange(-0.5, cap + 1.5)


def _pred_n_sec_confusion_partial(b: Bundle) -> dict:
    if not b.predictions:
        return {"available": False}
    cap = b.ctx.pred_n_sec_cap
    edges = _n_sec_confusion_edges(cap)

    def _one(ps: PredictionSide) -> list[list[int]] | None:
        if ps.coord != "global":
            return None
        mat = hist2d(ps.paired, pl.col("n_sec").clip(0, cap), pl.col("n_sec_pred").clip(0, cap), edges, edges)
        return mat.tolist()

    return {"available": True, "p": _per_prediction(b, _one)}


def _row_normalize(mat: np.ndarray) -> np.ndarray:
    row_sums = mat.sum(axis=1, keepdims=True)
    return np.divide(mat, row_sums, out=np.zeros_like(mat), where=row_sums > 0)


def _pred_n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
    title, xlabel = "Secondary-count confusion matrix (row-normalised)", "predicted n_sec"
    if not parts[0]["available"]:
        return _pred_unavailable("pred_n_sec_confusion", "prediction", title, xlabel)
    cap = ctx.pred_n_sec_cap
    labels = [str(i) for i in range(cap)] + [f"{cap}+"]
    names = list(parts[0]["p"])
    series: dict[str, list] = {}
    for name in names:
        mats = [p["p"][name] for p in parts]
        if mats[0] is None:
            continue
        mat = np.sum([np.asarray(m, dtype=np.float64) for m in mats], axis=0)
        series[name] = _row_normalize(mat).tolist()
    return Reduced(
        id="pred_n_sec_confusion",
        family="prediction",
        kind="heatmap",
        title=title,
        xlabel=xlabel,
        payload={
            "series": series,
            "row_labels": labels,
            "col_labels": labels,
            "ylabel": "true n_sec",
            "cbar_label": "fraction of true-count rows",
            "vmin": 0.0,
            "vmax": 1.0,
        },
    )


_OTHER_SEC_KEY = "other"


def _sec_species_key_expr(col: str, top_pdgs: list[int]) -> pl.Expr:
    return pl.when(pl.col(col).is_in(list(top_pdgs))).then(pl.col(col).cast(pl.Utf8)).otherwise(pl.lit(_OTHER_SEC_KEY))


def _pred_sec_species_confusion_partial(b: Bundle) -> dict:
    if not b.predictions or not b.ctx.pred_top_sec_pdgs:
        return {"available": False}
    top = b.ctx.pred_top_sec_pdgs

    def _one(ps: PredictionSide) -> dict | None:
        if ps.coord != "global" or not ps.has_truth:
            return None
        pf = paired_secondaries(ps.lf)
        counts = (
            pf.select(
                _sec_species_key_expr("true_pdg", top).alias("_t"),
                _sec_species_key_expr("pred_pdg", top).alias("_p"),
            )
            .group_by("_t", "_p")
            .agg(pl.len().alias("_n"))
            .collect(engine="streaming")
        )
        out: dict[str, dict[str, int]] = {}
        for t, p, n in counts.iter_rows():
            out.setdefault(t, {})[p] = out.get(t, {}).get(p, 0) + n
        return out

    return {"available": True, "p": _per_prediction(b, _one)}


def _pred_sec_species_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
    title, xlabel = "Secondary-species confusion matrix (row-normalised)", "predicted species"
    if not parts[0]["available"]:
        return _pred_unavailable("pred_sec_species_confusion", "prediction", title, xlabel)
    keys = [str(k) for k in ctx.pred_top_sec_pdgs] + [_OTHER_SEC_KEY]
    labels = [pdg_label(k) for k in ctx.pred_top_sec_pdgs] + [_OTHER_SEC_KEY]
    names = list(parts[0]["p"])
    series: dict[str, list] = {}
    for name in names:
        entries = [p["p"][name] for p in parts]
        if entries[0] is None:
            continue
        merged: dict[str, dict[str, int]] = {}
        for e in entries:
            for t, row in e.items():
                acc = merged.setdefault(t, {})
                for p_key, n in row.items():
                    acc[p_key] = acc.get(p_key, 0) + n
        mat = []
        for t in keys:
            row_counts = merged.get(t, {})
            row_vals = [row_counts.get(p_key, 0) for p_key in keys]
            total = sum(row_vals)
            mat.append([v / total if total > 0 else 0.0 for v in row_vals])
        series[name] = mat
    return Reduced(
        id="pred_sec_species_confusion",
        family="prediction",
        kind="heatmap",
        title=title,
        xlabel=xlabel,
        payload={
            "series": series,
            "row_labels": labels,
            "col_labels": labels,
            "ylabel": "true species",
            "cbar_label": "fraction of true-species rows",
            "vmin": 0.0,
            "vmax": 1.0,
        },
    )


# ---- physics consistency -----------------------------------------------


def _pred_dir_alignment_partial(b: Bundle, kind: str) -> dict:
    if not b.predictions:
        return {"available": False}
    edges = np.linspace(-1.0, 1.0, b.ctx.n_marginal_bins + 1)
    nb = len(edges) - 1

    def _one(ps: PredictionSide) -> list[int] | None:
        if not ps.has_truth:
            return None
        return _counts(hist1d(ps.lf, dir_alignment_expr(ps.coord, kind), edges), 0, nb)

    return {"available": True, "edges": edges.tolist(), "p": _per_prediction(b, _one)}


def _pred_dir_alignment_finalize(parts: list[dict], ctx: Context, kind: str, title: str) -> Reduced:
    xlabel = f"cos(angle) between true and predicted {kind}_dir"
    spec_id = f"pred_dir_alignment_{kind}"
    if not parts[0]["available"]:
        return _pred_unavailable(spec_id, "prediction", title, xlabel)
    edges = np.asarray(parts[0]["edges"])
    names = list(parts[0]["p"])
    series: dict[str, list[int]] = {}
    for name in names:
        entries = [p["p"][name] for p in parts]
        if entries[0] is None:
            continue
        series[name] = [int(x) for x in sum_merge([{"c": e} for e in entries])["c"]]
    return Reduced(
        id=spec_id,
        family="prediction",
        kind="single_hist",
        title=title,
        xlabel=xlabel,
        payload={"edges": edges.tolist(), "series": series, "log_y": True},
    )


def _rate_partial(lf: pl.LazyFrame, cond: pl.Expr) -> tuple[int, int]:
    res = lf.select(cond.cast(pl.Int64).sum().alias("v"), pl.len().alias("n")).collect(engine="streaming")
    return int(res["v"][0]), int(res["n"][0])


_CONSTRAINT_LABELS = ["post_dir_norm", "travel_dir_norm", "step_length_neg", "edep_neg", "delta_e_neg"]


def _pred_constraint_partial(b: Bundle) -> dict:
    """Rate of physical-constraint violations in the *predicted* values.

    Direction unit-norm is only a meaningful check for `--coord local`
    (`giant predict --coord global` already renormalises before writing, so
    it is compliant by construction there — reported as 0/1 rather than
    skipped, since "no violations" is still the correct answer).
    """
    if not b.predictions:
        return {"available": False}

    def _one(ps: PredictionSide) -> dict:
        out: dict[str, tuple[int, int]] = {}
        if ps.coord == "local":
            for label, cols in (
                ("post_dir_norm", ["pred_post_dx", "pred_post_dy", "pred_post_dz"]),
                ("travel_dir_norm", ["pred_travel_dx", "pred_travel_dy", "pred_travel_dz"]),
            ):
                norm = pl.sum_horizontal([pl.col(c) ** 2 for c in cols]).sqrt()
                out[label] = _rate_partial(ps.lf, (norm - 1).abs() > 0.05)
        else:
            out["post_dir_norm"] = (0, 1)
            out["travel_dir_norm"] = (0, 1)
        for label, col in (
            ("step_length_neg", "pred_step_length"),
            ("edep_neg", "pred_edep"),
            ("delta_e_neg", "pred_delta_e"),
        ):
            out[label] = _rate_partial(ps.paired, pl.col(col) < 0)
        return {k: list(v) for k, v in out.items()}

    return {"available": True, "p": _per_prediction(b, _one)}


def _pred_constraint_finalize(parts: list[dict], ctx: Context) -> Reduced:
    title, xlabel = "Physical-constraint violation rate", "check"
    if not parts[0]["available"]:
        return _pred_unavailable("pred_constraint_violations", "prediction", title, xlabel)
    names = list(parts[0]["p"])
    series: dict[str, list[float]] = {}
    for name in names:
        entries = [p["p"][name] for p in parts]
        rates = []
        for check in _CONSTRAINT_LABELS:
            v = sum(e[check][0] for e in entries)
            n = sum(e[check][1] for e in entries)
            rates.append(v / n if n > 0 else 0.0)
        series[name] = rates
    return Reduced(
        id="pred_constraint_violations",
        family="prediction",
        kind="bar",
        title=title,
        xlabel=xlabel,
        payload={"labels": _CONSTRAINT_LABELS, "series": series, "ylabel": "violation rate"},
    )


_CORR_PAIRS = [(a, b) for i, a in enumerate(PAIRED_SCALARS) for b in PAIRED_SCALARS[i + 1 :]]


def _pred_correlation_partial(b: Bundle) -> dict:
    if not b.predictions:
        return {"available": False}

    def _one(ps: PredictionSide) -> dict | None:
        if not ps.has_truth:
            return None
        vs = [v for v in PAIRED_SCALARS if v in paired_vars_for_coord(ps.coord)]
        pairs = [(a, c) for a, c in _CORR_PAIRS if a in vs and c in vs]
        agg = [pl.len().alias("_n")]
        for prefix in ("true", "pred"):
            for v in vs:
                agg.append(pl.col(f"{prefix}_{v}").sum().alias(f"s_{prefix}_{v}"))
                agg.append((pl.col(f"{prefix}_{v}") ** 2).sum().alias(f"ss_{prefix}_{v}"))
            for a, c in pairs:
                agg.append((pl.col(f"{prefix}_{a}") * pl.col(f"{prefix}_{c}")).sum().alias(f"sxy_{prefix}_{a}_{c}"))
        res = ps.paired.select(agg).collect(engine="streaming")
        return {"vars": vs, "pairs": pairs, "row": {k: float(res[k][0]) for k in res.columns}}

    return {"available": True, "p": _per_prediction(b, _one)}


def _corr_matrix(vs: list[str], pairs: list[tuple[str, str]], row: dict[str, float], prefix: str) -> np.ndarray:
    n = row["_n"]
    mat = np.eye(len(vs))
    idx = {v: i for i, v in enumerate(vs)}
    for a, c in pairs:
        sa, sc = row[f"s_{prefix}_{a}"], row[f"s_{prefix}_{c}"]
        ssa, ssc = row[f"ss_{prefix}_{a}"], row[f"ss_{prefix}_{c}"]
        sxy = row[f"sxy_{prefix}_{a}_{c}"]
        cov = sxy / n - (sa / n) * (sc / n)
        var_a = ssa / n - (sa / n) ** 2
        var_c = ssc / n - (sc / n) ** 2
        denom = np.sqrt(max(var_a, 0.0) * max(var_c, 0.0))
        corr = cov / denom if denom > 1e-12 else float("nan")
        mat[idx[a], idx[c]] = mat[idx[c], idx[a]] = corr
    return mat


def _pred_correlation_finalize(parts: list[dict], ctx: Context) -> Reduced:
    title, xlabel = "Correlation delta (corr(pred) - corr(true))", "variable"
    if not parts[0]["available"]:
        return _pred_unavailable("pred_correlation_delta", "prediction", title, xlabel)
    names = list(parts[0]["p"])
    series: dict[str, list] = {}
    labels: list[str] = []
    for name in names:
        entries = [p["p"][name] for p in parts]
        if entries[0] is None:
            continue
        vs, pairs = entries[0]["vars"], [tuple(p) for p in entries[0]["pairs"]]
        labels = [_PRED_TITLE_NAMES[v] for v in vs]  # same for every prediction: all share one coord (see condor.py)
        merged_row: dict[str, float] = {}
        for e in entries:
            for k, v in e["row"].items():
                merged_row[k] = merged_row.get(k, 0.0) + v
        delta = _corr_matrix(vs, pairs, merged_row, "pred") - _corr_matrix(vs, pairs, merged_row, "true")
        series[name] = delta.tolist()
    return Reduced(
        id="pred_correlation_delta",
        family="prediction",
        kind="heatmap",
        title=title,
        xlabel=xlabel,
        payload={
            "series": series,
            "row_labels": labels,
            "col_labels": labels,
            "ylabel": "variable",
            "cbar_label": "corr(pred) - corr(true)",
            "cmap": "RdBu_r",
            "vmin": -1.0,
            "vmax": 1.0,
        },
    )


def build_catalog() -> list[PlotSpec]:
    """All concrete plot specs, each with a unique id."""
    specs: list[PlotSpec] = []

    for var in MARGINAL_VARS:
        specs.append(
            PlotSpec(
                f"marginal_{var}",
                "marginals",
                compute_partial=lambda b, v=var: _marginal_overall_partial(b, v),
                finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(parts, ctx, v),
            )
        )
        for axis in GROUPING_AXES:
            specs.append(
                PlotSpec(
                    f"marginal_{var}_by_{axis}",
                    "marginals",
                    compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(b, v, a),
                    finalize=lambda parts, ctx, v=var, a=axis: _marginal_grouped_finalize(parts, ctx, v, a),
                )
            )

    specs.append(
        PlotSpec(
            "marginal_distance_summary",
            "quality",
            compute_partial=_distance_summary_partial,
            finalize=_distance_summary_finalize,
        )
    )

    specs += [
        PlotSpec(
            "event_total_edep",
            "event",
            compute_partial=lambda b: _event_scalar_partial(b, "total_edep", use_all=True),
            finalize=lambda parts, ctx: _event_scalar_finalize(
                parts,
                ctx,
                "event_total_edep",
                "Total deposited energy per event",
                "total deposited energy [MeV]",
            ),
        ),
        PlotSpec(
            "event_total_edep_by_energy",
            "event",
            compute_partial=_event_total_edep_by_energy_partial,
            finalize=_event_total_edep_by_energy_finalize,
        ),
        PlotSpec(
            "event_mean_length",
            "event",
            compute_partial=lambda b: _event_scalar_partial(b, "mean_length", use_all=False),
            finalize=lambda parts, ctx: _event_scalar_finalize(
                parts,
                ctx,
                "event_mean_length",
                "Mean step length per event",
                "mean step length [mm]",
            ),
        ),
        PlotSpec(
            "event_n_steps",
            "event",
            compute_partial=lambda b: _event_scalar_partial(b, "n_steps", use_all=False),
            finalize=lambda parts, ctx: _event_scalar_finalize(
                parts,
                ctx,
                "event_n_steps",
                "Number of steps per event",
                "steps per event",
            ),
        ),
        PlotSpec(
            "shower_longitudinal",
            "shower",
            compute_partial=lambda b: _profile_partial(b, depth_expr, "depth_edges"),
            finalize=lambda parts, ctx: _profile_finalize(
                parts,
                ctx,
                "shower_longitudinal",
                "Longitudinal shower profile",
                "depth along shower axis [mm]",
                "depth_edges",
            ),
        ),
        PlotSpec(
            "shower_transverse",
            "shower",
            compute_partial=lambda b: _profile_partial(b, transverse_expr, "transverse_edges"),
            finalize=lambda parts, ctx: _profile_finalize(
                parts,
                ctx,
                "shower_transverse",
                "Transverse shower profile",
                "radius from shower axis [mm]",
                "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",
            compute_partial=_species_share_partial,
            finalize=_species_share_finalize,
        ),
        PlotSpec(
            "leakage_fraction",
            "species",
            compute_partial=_leakage_partial,
            finalize=_leakage_finalize,
        ),
        PlotSpec(
            "sec_count_per_event",
            "secondaries",
            compute_partial=_sec_count_per_event_partial,
            finalize=_sec_count_per_event_finalize,
        ),
        PlotSpec(
            "sec_count_per_species",
            "secondaries",
            compute_partial=_sec_count_per_species_partial,
            finalize=_sec_count_per_species_finalize,
        ),
        PlotSpec(
            "sec_count_per_step",
            "secondaries",
            compute_partial=_sec_count_per_step_partial,
            finalize=_sec_count_per_step_finalize,
        ),
        PlotSpec(
            "sec_count_per_step_by_species",
            "secondaries",
            compute_partial=_sec_count_per_step_by_species_partial,
            finalize=_sec_count_per_step_by_species_finalize,
        ),
        PlotSpec(
            "sec_energy",
            "secondaries",
            compute_partial=_sec_energy_partial,
            finalize=_sec_energy_finalize,
        ),
        PlotSpec(
            "sec_cos_angle",
            "secondaries",
            compute_partial=_sec_cos_angle_partial,
            finalize=_sec_cos_angle_finalize,
        ),
        PlotSpec(
            "eval_cost_per_step",
            "cost",
            compute_partial=_eval_cost_per_step_partial,
            finalize=_eval_cost_per_step_finalize,
            chunkable=False,
        ),
        PlotSpec(
            "router_gating",
            "model",
            compute_partial=_router_gating_partial,
            finalize=_router_gating_finalize,
            chunkable=False,
        ),
        PlotSpec(
            "router_share_by_pdg",
            "model",
            compute_partial=_router_share_pdg_partial,
            finalize=_router_share_pdg_finalize,
            chunkable=False,
        ),
        PlotSpec(
            "router_share_by_process",
            "model",
            compute_partial=_router_share_process_partial,
            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",
            compute_partial=_type_embedding_l1_distance_partial,
            finalize=_type_embedding_l1_distance_finalize,
            chunkable=False,
        ),
    ]

    for var in PAIRED_VARS:
        specs.append(
            PlotSpec(
                f"pred_marginal_{var}",
                "prediction",
                compute_partial=lambda b, v=var: _pred_marginal_partial(b, v),
                finalize=lambda parts, ctx, v=var: _pred_marginal_finalize(parts, ctx, v),
            )
        )
        specs.append(
            PlotSpec(
                f"pred_scatter_{var}",
                "prediction",
                compute_partial=lambda b, v=var: _pred_scatter_partial(b, v),
                finalize=lambda parts, ctx, v=var: _pred_scatter_finalize(parts, ctx, v),
            )
        )
        specs.append(
            PlotSpec(
                f"pred_residual_{var}",
                "prediction",
                compute_partial=lambda b, v=var: _pred_residual_partial(b, v),
                finalize=lambda parts, ctx, v=var: _pred_residual_finalize(parts, ctx, v),
            )
        )
    for var in PAIRED_SCALARS:
        specs.append(
            PlotSpec(
                f"pred_relative_residual_{var}",
                "prediction",
                compute_partial=lambda b, v=var: _pred_relative_residual_partial(b, v),
                finalize=lambda parts, ctx, v=var: _pred_relative_residual_finalize(parts, ctx, v),
            )
        )
        specs.append(
            PlotSpec(
                f"pred_residual_profile_{var}",
                "prediction",
                compute_partial=lambda b, v=var: _pred_residual_profile_partial(b, v),
                finalize=lambda parts, ctx, v=var: _pred_residual_profile_finalize(parts, ctx, v),
            )
        )

    specs += [
        PlotSpec(
            "pred_ks_summary",
            "prediction",
            compute_partial=_pred_scorecard_partial,
            finalize=_pred_ks_finalize,
        ),
        PlotSpec(
            "pred_bias_summary",
            "prediction",
            compute_partial=_pred_scorecard_partial,
            finalize=_pred_bias_finalize,
        ),
        PlotSpec(
            "pred_rmse_summary",
            "prediction",
            compute_partial=_pred_scorecard_partial,
            finalize=_pred_rmse_finalize,
        ),
        PlotSpec(
            "pred_n_sec_confusion",
            "prediction",
            compute_partial=_pred_n_sec_confusion_partial,
            finalize=_pred_n_sec_confusion_finalize,
        ),
        PlotSpec(
            "pred_sec_species_confusion",
            "prediction",
            compute_partial=_pred_sec_species_confusion_partial,
            finalize=_pred_sec_species_confusion_finalize,
        ),
        PlotSpec(
            "pred_dir_alignment_post",
            "prediction",
            compute_partial=lambda b: _pred_dir_alignment_partial(b, "post"),
            finalize=lambda parts, ctx: _pred_dir_alignment_finalize(
                parts, ctx, "post", "Post-direction alignment (true vs predicted)"
            ),
        ),
        PlotSpec(
            "pred_dir_alignment_travel",
            "prediction",
            compute_partial=lambda b: _pred_dir_alignment_partial(b, "travel"),
            finalize=lambda parts, ctx: _pred_dir_alignment_finalize(
                parts, ctx, "travel", "Travel-direction alignment (true vs predicted)"
            ),
        ),
        PlotSpec(
            "pred_constraint_violations",
            "prediction",
            compute_partial=_pred_constraint_partial,
            finalize=_pred_constraint_finalize,
        ),
        PlotSpec(
            "pred_correlation_delta",
            "prediction",
            compute_partial=_pred_correlation_partial,
            finalize=_pred_correlation_finalize,
        ),
    ]
    return specs


def catalog_ids() -> list[str]:
    return [s.id for s in build_catalog()]


def get_spec(spec_id: str) -> PlotSpec:
    for s in build_catalog():
        if s.id == spec_id:
            return s
    raise KeyError(f"unknown plot id: {spec_id!r}")