From 9fa6420183f648add8a1a91193fbb84f86e9490f Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Wed, 26 Aug 2026 14:13:18 +0200 Subject: [PATCH] feat(analysis): per-step secondary multiplicity plots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the event-level n_sec confusion matrix with two step-resolved secondary-multiplicity comparisons: - sec_count_per_step: overlay histogram of how many secondaries a single step emits, rollout series vs reference. - sec_count_per_step_by_species: heatmap of per-step multiplicity of one species (zero row included) against species, drawn as one panel per rollout plus a reference panel, raw counts on a log color scale. Both are backed by a new sources.secondaries_by_step view, which tags each secondary with its emitting step — (event_id, parent_id, birth position) on the rollout side, the row index on the reference side — so neither plot needs a join against the step frame. Steps that emitted nothing are recovered by subtraction from the chunk's step count, keeping both specs sum-mergeable across condor chunks. The rollout multiplicity is derived from the actual secondary birth rows rather than the n_sec_pred column, which records the predicted count before the per-event max-tracks cap. _render_heatmap gained reference-panel and log-color support; marginal_distance_summary sets neither key and is unchanged. Co-Authored-By: Claude Opus 5 --- CLAUDE.md | 2 +- giant/analysis/catalog.py | 249 +++++++++++++++++++++------------- giant/analysis/reduce.py | 17 --- giant/analysis/reduced.py | 2 +- giant/analysis/render.py | 14 +- giant/analysis/sources.py | 31 +++++ tests/test_analysis_reduce.py | 22 +-- tests/test_catalog.py | 64 ++++----- tests/test_render.py | 4 +- 9 files changed, 241 insertions(+), 164 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 190152c..156689f 100644 --- a/CLAUDE.md +++ b/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** (`/analysis_runs/analysis_/` 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/__.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 ` first joins every plot's chunk partials into `reduced/.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 ` is a separate, unrelated entry point: training-progress plots straight from a run's `metrics.csv`. diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index 18fbc1b..eaa81eb 100644 --- a/giant/analysis/catalog.py +++ b/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", diff --git a/giant/analysis/reduce.py b/giant/analysis/reduce.py index 8e1d48e..ce65c89 100644 --- a/giant/analysis/reduce.py +++ b/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 diff --git a/giant/analysis/reduced.py b/giant/analysis/reduced.py index bbf79ac..f697cbc 100644 --- a/giant/analysis/reduced.py +++ b/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) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index d4ae5c4..8f9b237 100644 --- a/giant/analysis/render.py +++ b/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") diff --git a/giant/analysis/sources.py b/giant/analysis/sources.py index 15b496a..42dccc1 100644 --- a/giant/analysis/sources.py +++ b/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")) + ) diff --git a/tests/test_analysis_reduce.py b/tests/test_analysis_reduce.py index 71cb4a8..0a456f0 100644 --- a/tests/test_analysis_reduce.py +++ b/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(): diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 377ac05..be50f27 100644 --- a/tests/test_catalog.py +++ b/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 diff --git a/tests/test_render.py b/tests/test_render.py index da580aa..c5f0fbf 100644 --- a/tests/test_render.py +++ b/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, }, ), ]