Merge pull request 'feat(analysis): per-step secondary multiplicity plots' (#85) from analysis/per-step-secondary-multiplicity into master

Reviewed-on: #85

CLAUDE.md

@@ -97,7 +97,7 @@ Secondary energies are a **stick-breaking partition of the `e_sec` budget** from
 
 **Validation** (`giant/validate.py`): step-level marginal + KL-divergence comparisons during training (`--validate-every`).
 
-**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive `giant rollout` runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + `RolloutSpec`/`Side` — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `variables.py` (the per-step value expressions shared by range sizing and the plot registry), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json` over the union of the reference and every rollout, so every compute job is one pass with no range scan), `reduced.py` (`Partial`/`Reduced` — the compact self-describing JSON a compute job emits), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles/containment, species/leakage, secondaries, distance/confusion summaries, router and type-embedding diagnostics; `giant analyze list` prints every id), `runtime_estimate.py` (per-(plot, chunk) walltime estimates for the submit description), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`; each rollout gets a stable `ps.get_color(i)` slot by its position in `series`, the reference always draws in one fixed dashed-ink style). `Bundle.rollouts` is a name-keyed dict of `Side`, and every `compute_partial`/`finalize` builds a `Reduced.payload["series"]` dict keyed the same way, with `payload["reference"]` as the one distinguished non-rollout entry. The heatmap-shaped specs (`marginal_distance_summary`, `n_sec_confusion`) and the checkpoint-bound diagnostics (`router_gating.py`, `type_embedding_distance.py`) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout.
+**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive `giant rollout` runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + `RolloutSpec`/`Side` — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `variables.py` (the per-step value expressions shared by range sizing and the plot registry), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json` over the union of the reference and every rollout, so every compute job is one pass with no range scan), `reduced.py` (`Partial`/`Reduced` — the compact self-describing JSON a compute job emits), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles/containment, species/leakage, secondaries, distance summaries, router and type-embedding diagnostics; `giant analyze list` prints every id), `runtime_estimate.py` (per-(plot, chunk) walltime estimates for the submit description), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`; each rollout gets a stable `ps.get_color(i)` slot by its position in `series`, the reference always draws in one fixed dashed-ink style). `Bundle.rollouts` is a name-keyed dict of `Side`, and every `compute_partial`/`finalize` builds a `Reduced.payload["series"]` dict keyed the same way, with `payload["reference"]` as the one distinguished non-rollout entry. The heatmap-shaped specs (`marginal_distance_summary`, `sec_count_per_step_by_species` — the latter also drawing the reference as its own panel) and the checkpoint-bound diagnostics (`router_gating.py`, `type_embedding_distance.py`) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout.
 
 **Input is one or more `giant rollout` YAML sidecars** (`condor.py:load_rollout_yamls`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML). `prep` creates a **run directory** (`<cwd>/analysis_runs/analysis_<id>/` by default, `--run-dir` to override) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit a.yaml [b.yaml ...] --chunks N` runs `prep` (recording `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` splits into a `compute_partial`/`finalize` pair so chunks can be summed/concatenated back per rollout (`chunkable=False` specs — the checkpoint-bound diagnostics, already bounded/subsampled — always run as a single chunk). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`; `merge-one` does a single plot for debugging), then turns those into the styled PDF/gallery tree. `giant analyze metrics <train_run_dir>` is a separate, unrelated entry point: training-progress plots straight from a run's `metrics.csv`.
 

giant/analysis/catalog.py

@@ -22,9 +22,9 @@ which is the order rollouts were given on the CLI) plus the single reference.
 ``finalize`` merges each rollout's chunks independently and assembles a
 ``Reduced.payload`` keyed the same way: ``"series": {name: ...}`` for the
 rollouts, ``"reference": ...`` as one distinguished entry (omitted on
-rollout-only plots like ``leakage_fraction``). The two heatmap-shaped specs
-(``marginal_distance_summary``, ``n_sec_confusion``) and the router
-diagnostics are inherently one-matrix/one-checkpoint per rollout, so their
+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.
@@ -60,7 +60,6 @@ from giant.analysis.reduce import (
     leakage_fraction,
     profile_finalize,
     profile_partial,
-    sec_count_by_event,
     species_share,
     sum_merge,
     transverse_expr,
@@ -72,7 +71,15 @@ from giant.analysis.router_gating import (
     compute_router_share_by_process,
     compute_router_specialization,
 )
-from giant.analysis.sources import RolloutSide, RolloutSpec, Side, open_side, physical_steps, secondaries
+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
 
@@ -211,31 +218,6 @@ def _ks_statistic(r_counts, t_counts) -> float:
     return float(np.max(np.abs(r_cdf - t_cdf)))
 
 
-def _integer_confusion(
-    t: np.ndarray, r: np.ndarray, max_bins: int = 21, cap: int | None = None
-) -> tuple[list[str], np.ndarray]:
-    """Confusion matrix of two paired small-integer arrays (e.g. secondary counts).
-
-    Bins are consecutive integers ``0..cap``, with the last bin an overflow
-    ``"cap+"`` bucket, so an occasional pathological count doesn't blow up the
-    heatmap. Returns ``(labels, matrix)`` with ``matrix[i, j]`` counting pairs
-    with ``t == i`` and ``r == j`` (both clipped into ``[0, cap]``).
-
-    ``cap``, if given, is used as-is instead of being derived from ``t``/``r``
-    — lets a multi-rollout caller fix one shared cap (and so one shared label
-    set) across every rollout's matrix rather than each panel picking its own.
-    """
-    if cap is None:
-        cap = min(max(int(t.max()) if len(t) else 0, int(r.max()) if len(r) else 0, 1), max_bins - 1)
-    t_c = np.clip(t.astype(np.int64), 0, cap)
-    r_c = np.clip(r.astype(np.int64), 0, cap)
-    n = cap + 1
-    mat = np.zeros((n, n), dtype=np.int64)
-    np.add.at(mat, (t_c, r_c), 1)
-    labels = [str(i) for i in range(cap)] + [f"{cap}+"]
-    return labels, mat
-
-
 def _containment_depths(mat: np.ndarray, edges: np.ndarray, quantile: float) -> np.ndarray:
     """Per-event depth containing ``quantile`` of that event's deposited energy.
 
@@ -802,6 +784,139 @@ def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
     )
 
 
+# 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 {
@@ -858,62 +973,6 @@ def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
     )
 
 
-def _n_sec_confusion_partial(b: Bundle) -> dict:
-    t_ids, t_n = sec_count_by_event(b.t_all, _t_sec(b))
-
-    def _r(rs: RolloutSide) -> dict:
-        ids, n = sec_count_by_event(rs.phys, _r_sec(rs))
-        return {"ids": ids.tolist(), "n": n.tolist()}
-
-    return {"r": _per_rollout(b, _r), "t": {"ids": t_ids.tolist(), "n": t_n.tolist()}}
-
-
-def _n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    names = list(parts[0]["r"])
-    # event-disjoint chunking (see Bundle.open) means each event_id appears in
-    # exactly one part on each side, so a plain dict build is a safe merge.
-    t_ids = np.concatenate([np.asarray(p["t"]["ids"], dtype=np.int64) for p in parts])
-    t_n = np.concatenate([np.asarray(p["t"]["n"], dtype=np.int64) for p in parts])
-    t_map = dict(zip(t_ids.tolist(), t_n.tolist()))
-
-    pairs: dict[str, tuple[np.ndarray, np.ndarray]] = {}
-    max_val = 0
-    for name in names:
-        r_ids = np.concatenate([np.asarray(p["r"][name]["ids"], dtype=np.int64) for p in parts])
-        r_n = np.concatenate([np.asarray(p["r"][name]["n"], dtype=np.int64) for p in parts])
-        r_map = dict(zip(r_ids.tolist(), r_n.tolist()))
-        common = sorted(set(r_map) & set(t_map))
-        true_n = np.array([t_map[e] for e in common], dtype=np.int64)
-        pred_n = np.array([r_map[e] for e in common], dtype=np.int64)
-        pairs[name] = (true_n, pred_n)
-        if len(true_n):
-            max_val = max(max_val, int(true_n.max()), int(pred_n.max()))
-
-    cap = min(max(max_val, 1), 20)
-    matrices: dict[str, list[list[int]]] = {}
-    labels: list[str] = []
-    for name in names:
-        true_n, pred_n = pairs[name]
-        labels, mat = _integer_confusion(true_n, pred_n, cap=cap)
-        matrices[name] = mat.tolist()
-
-    return Reduced(
-        id="n_sec_confusion",
-        family="secondaries",
-        kind="heatmap",
-        title="Predicted vs true secondary count per event",
-        xlabel="predicted secondaries (rollout)",
-        payload={
-            "series": matrices,
-            "row_labels": labels,
-            "col_labels": labels,
-            "ylabel": "true secondaries (reference)",
-            "cbar_label": "event count",
-            "vmin": 0.0,
-        },
-    )
-
-
 # ---------------------------------------------------------------------------
 # router diagnostics (not chunked — already bounded/subsampled)
 # ---------------------------------------------------------------------------
@@ -1077,6 +1136,18 @@ def build_catalog() -> list[PlotSpec]:
             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",
@@ -1089,12 +1160,6 @@ 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",

giant/analysis/reduce.py

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

giant/analysis/reduced.py

@@ -27,7 +27,7 @@ from pathlib import Path
 #   "router_specialization"  max gate weight vs energy (one scalar trend line
 #                    summarizing "router_gating"), per rollout with an enabled router
 #   "heatmap"        row x col matrix + colorbar, one panel per rollout (a
-#                    distance scorecard or a predicted-vs-true confusion matrix)
+#                    distance scorecard)
 #   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint)
 
 

giant/analysis/render.py

@@ -27,6 +27,7 @@ from pathlib import Path
 
 import numpy as np
 import plotstyle as ps
+from matplotlib.colors import LogNorm
 import yaml
 
 from giant.analysis.reduced import Reduced
@@ -376,10 +377,16 @@ def _render_router_specialization(r: Reduced, params: dict):
 
 
 def _render_heatmap(r: Reduced, params: dict):
-    series = r.payload["series"]
+    series = dict(r.payload["series"])
     row_labels = r.payload["row_labels"]
     col_labels = r.payload["col_labels"]
+    # A heatmap-shaped plot is one matrix per rollout, so the reference (when the
+    # comparison has one — the distance scorecard doesn't) becomes one more panel
+    # rather than another line.
+    if r.payload.get("reference") is not None:
+        series["reference"] = r.payload["reference"]
     names = list(series)
+    norm = LogNorm(vmin=1) if r.payload.get("log_color") else None
     fig, axes = ps.new_figure(
         "slide-16x9" if len(names) > 1 else "thesis-single",
         title=r.title,
@@ -397,8 +404,9 @@ def _render_heatmap(r: Reduced, params: dict):
             origin="upper",
             aspect="auto",
             cmap=r.payload.get("cmap", "viridis"),
-            vmin=r.payload.get("vmin"),
-            vmax=r.payload.get("vmax"),
+            norm=norm,
+            vmin=None if norm else r.payload.get("vmin"),
+            vmax=None if norm else r.payload.get("vmax"),
         )
         ax.set_xticks(range(len(col_labels)))
         ax.set_xticklabels(col_labels, rotation=45, ha="right")

giant/analysis/sources.py

@@ -236,3 +236,34 @@ def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
             pl.col("sec_dz_list").alias("sdz"),
         )
     )
+
+
+def secondaries_by_step(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
+    """One row per produced secondary, tagged with the step that produced it.
+
+    Canonical columns: ``step_key`` (an opaque struct identifying the emitting
+    step) and ``pdg``. ``secondaries`` deliberately drops that link; the
+    per-step multiplicity plots need it, so this is a separate view rather than
+    extra columns every other consumer would pay for.
+
+    - rollout: a secondary's birth row carries ``parent_id`` and a birth
+      position copied verbatim from the parent step's ``post_pos``, so
+      ``(event_id, parent_id, pre_pos)`` identifies the emitting step exactly —
+      no join against the (large) step frame is needed.
+    - reference: secondaries already live on their parent step's row, so the
+      row index *is* the step key. It is only ever used as a group key inside
+      one chunk's own aggregation, so indices repeating across chunks is
+      harmless.
+    """
+    if side is Side.rollout:
+        return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
+            pl.struct("event_id", "parent_id", "pre_x", "pre_y", "pre_z").alias("step_key"),
+            "pdg",
+        )
+    return (
+        lf.select("sec_pdg_list")
+        .with_row_index("_row")
+        .explode("sec_pdg_list")
+        .drop_nulls("sec_pdg_list")
+        .select(pl.struct("_row").alias("step_key"), pl.col("sec_pdg_list").cast(pl.Int64).alias("pdg"))
+    )

tests/test_analysis_reduce.py

@@ -13,6 +13,7 @@ from giant.analysis.sources import (
     open_side,
     physical_steps,
     secondaries,
+    secondaries_by_step,
 )
 from giant.data.loader import EVENT_ID_FILE_STRIDE
 
@@ -159,18 +160,17 @@ 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}
+def test_secondaries_by_step_keys_each_secondary_to_its_emitting_step():
+    r = secondaries_by_step(_rollout_frame(), Side.rollout).collect()
+    assert r["pdg"].to_list() == [22]
+    # the rollout key is (event_id, parent_id, birth position) — the parent
+    # step's post_pos, copied verbatim onto the child's birth row.
+    assert r["step_key"][0] == {"event_id": 1, "parent_id": 0, "pre_x": 0.0, "pre_y": 0.0, "pre_z": 1.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}
+    t = secondaries_by_step(_reference_frame(), Side.reference).collect()
+    assert t["pdg"].to_list() == [22, 22]
+    # one row per emitting step; the empty-list step drops out entirely
+    assert [k["_row"] for k in t["step_key"]] == [0, 2]
 
 
 def test_leakage_fraction():

tests/test_catalog.py

@@ -10,10 +10,10 @@ from giant.analysis.catalog import (
     Bundle,
     PlotSpec,
     _containment_depths,
-    _integer_confusion,
     _ks_statistic,
 )
 from giant.analysis.context import Context, build_context
+from giant.analysis.grouping import pdg_label
 from giant.analysis.sources import RolloutSpec
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
@@ -160,8 +160,9 @@ def _validate_payload(r, names: list[str]) -> None:
 # 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), 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
+# sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
+# nested sum-merge into a scorecard (marginal_distance_summary), and
+# concat-then-per-event-derived-quantity
 # (shower_containment_depth_90, reusing the profile matrix's own merge shape).
 _CHUNK_EQUIVALENCE_IDS = [
     "marginal_edep",
@@ -170,9 +171,10 @@ _CHUNK_EQUIVALENCE_IDS = [
     "shower_longitudinal",
     "leakage_fraction",
     "sec_count_per_species",
+    "sec_count_per_step",
+    "sec_count_per_step_by_species",
     "router_gating",
     "marginal_distance_summary",
-    "n_sec_confusion",
     "shower_containment_depth_90",
 ]
 
@@ -217,7 +219,7 @@ def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
 
 
 # ---------------------------------------------------------------------------
-# new (gitea #76) reductions: KS distance, confusion matrix, containment depth
+# new (gitea #76) reductions: KS distance and containment depth
 # ---------------------------------------------------------------------------
 
 
@@ -228,28 +230,6 @@ def test_ks_statistic():
     assert _ks_statistic([10, 0], [0, 0]) == 1.0  # one side empty, other isn't -> maximal mismatch
 
 
-def test_integer_confusion_matches_event_pairing():
-    # true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
-    labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
-    assert labels == ["0", "1+"]
-    assert mat.tolist() == [[0, 0], [1, 1]]  # row=true, col=pred
-
-
-def test_integer_confusion_caps_pathological_outliers():
-    labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
-    assert labels[-1] == "4+"
-    assert mat.shape == (5, 5)
-    assert mat.sum() == 2
-
-
-def test_integer_confusion_explicit_cap_overrides_local_range():
-    # Even though this pair's own max is 1, an explicit shared cap forces a
-    # wider (and so cross-rollout-consistent) label set.
-    labels, mat = _integer_confusion(np.array([1, 1]), np.array([0, 1]), cap=3)
-    assert labels == ["0", "1", "2", "3+"]
-    assert mat.shape == (4, 4)
-
-
 def test_containment_depths_simple_ramp():
     # one event, edep concentrated in the first bin -> 90%/95% containment
     # depth is the first bin's right edge; a zero-energy event is dropped.
@@ -259,17 +239,25 @@ def test_containment_depths_simple_ramp():
     assert depths.tolist() == [1.0]
 
 
-def test_n_sec_confusion_spec(bundle):
-    spec = get_spec("n_sec_confusion")
+def test_sec_count_per_step_counts_empty_steps(bundle):
+    spec = get_spec("sec_count_per_step")
     r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
-    assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
-    assert r.payload["series"]["rollout"] == [[0, 0], [1, 1]]
+    # reference: 3 steps, two of which emit exactly one secondary
+    assert r.payload["reference"][:2] == [1, 2]
+    # rollout: 4 physical steps, one of which emits a single secondary
+    assert r.payload["series"]["rollout"][:2] == [3, 1]
+    assert sum(r.payload["reference"]) == 3
 
 
-def test_n_sec_confusion_shares_one_cap_across_rollouts(two_bundle):
-    spec = get_spec("n_sec_confusion")
-    r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
-    assert list(r.payload["series"]) == ["flow", "wgan"]
-    # both rollouts share the same fixture data here, so their matrices (and
-    # the shared label set) must be identical.
-    assert r.payload["series"]["flow"] == r.payload["series"]["wgan"]
+def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
+    spec = get_spec("sec_count_per_step_by_species")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    cols = r.payload["col_labels"]
+    ref = r.payload["reference"]
+    g = cols.index(pdg_label(22))
+    # two reference steps emit one photon each; the third emits none
+    assert [row[g] for row in ref][:2] == [1, 2]
+    # every other species column is "no such secondary" on all 3 steps
+    for j, _ in enumerate(cols):
+        if j != g:
+            assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0

tests/test_render.py

@@ -304,13 +304,15 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "hm1",
             "secondaries",
             "heatmap",
-            "Confusion (single rollout)",
+            "Heatmap (single rollout)",
             "predicted",
             {
                 "series": {"flow": [[1, 0], [0, 1]]},
+                "reference": [[2, 0], [0, 1]],
                 "row_labels": ["0", "1+"],
                 "col_labels": ["0", "1+"],
                 "cbar_label": "count",
+                "log_color": True,
             },
         ),
     ]