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

Reviewed-on: #88

.bumpversion.toml

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

CHANGELOG.md

@@ -1,5 +1,17 @@
 # Changelog
 
+## [0.3.12] - 2026-08-28
+
+### Added
+
+- Add sampled n_sec under n_sec.mode = 'head' [gitea #86](https://git.larsbogner.de/lars/giant/issues/86)
+
+## [0.3.11] - 2026-08-26
+
+### Changed
+
+- Feat(analysis): per-step secondary multiplicity plots
+
 ## [0.3.10] - 2026-08-26
 
 ### Changed

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"))
+    )

giant/config.py

@@ -424,11 +424,15 @@ class NSecConfig:
     # n_sec head was trained against Stage 1's own ConditionEncoder output and so has
     # to stay attached there, not just be labeled as such).
     owner: str = "stage2"
-    # mode="stop_token" only: how sample_secondaries_ar turns a slot's stop logit into a
-    # stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample":
-    # a Bernoulli draw at sigmoid(logit) (a real sample from the learned length
-    # distribution, at the cost of an extra RNG draw per slot).
-    stop_sampling: str = "greedy"
+    # How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an
+    # actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits
+    # (deterministic — the conditional mode, not a sample); "sample" is a categorical draw
+    # from softmax(logits) (a real sample from the learned count distribution). mode=
+    # "stop_token": "greedy" is sigmoid(stop_logit) >= 0.5 per slot (deterministic);
+    # "sample" is a Bernoulli draw at sigmoid(stop_logit) per slot. Renamed from
+    # "stop_sampling" (gitea #86), which is still accepted as a deprecated alias since it
+    # appears in existing checkpoints' model_config.
+    sampling: str = "greedy"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "NSecConfig":
@@ -437,7 +441,7 @@ class NSecConfig:
             mode=d.get("mode", "head"),
             lambda_weight=d.get("lambda", 0.1),
             owner=d.get("owner", "stage2"),
-            stop_sampling=d.get("stop_sampling", "greedy"),
+            sampling=d.get("sampling", d.get("stop_sampling", "greedy")),
         )
 
     def to_dict(self) -> dict:
@@ -445,7 +449,7 @@ class NSecConfig:
             "mode": self.mode,
             "lambda": self.lambda_weight,
             "owner": self.owner,
-            "stop_sampling": self.stop_sampling,
+            "sampling": self.sampling,
         }
 
 
@@ -1075,6 +1079,27 @@ def _set_path(d: dict, dotted: str, value) -> None:
     cur[parts[-1]] = value
 
 
+def _pop_path(d: dict, dotted: str) -> None:
+    """Remove a dotted path from a nested dict, if present. No-op if any
+    component along the path is missing."""
+    parts = dotted.split(".")
+    cur = d
+    for part in parts[:-1]:
+        if not isinstance(cur, dict) or part not in cur:
+            return
+        cur = cur[part]
+    if isinstance(cur, dict):
+        cur.pop(parts[-1], None)
+
+
+# Config keys renamed within v0.3 itself (not part of the v0.2->v0.3 migration
+# above) — normalized by migrate_config so a config.toml still using an older
+# v0.3 key name keeps passing validate_config_keys.
+_RENAMED_KEYS = {
+    "stage2_model.n_sec.stop_sampling": "stage2_model.n_sec.sampling",  # gitea #86
+}
+
+
 def _deep_merge(base: dict, override: dict) -> dict:
     """Recursively merge `override` onto a copy of `base`.
 
@@ -1271,11 +1296,21 @@ def migrate_config(cfg: dict) -> dict:
     (which additionally carries n_sec_head ownership and needs
     `network.build_models`'s cooperation) is a separate migration surface,
     deferred to the network.py refactor.
-    """
-    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
-        return copy.deepcopy(cfg)
 
+    Independently of the v0.2/v0.3 branch below, `_RENAMED_KEYS` normalizes
+    keys renamed within v0.3 itself (e.g. `stop_sampling` -> `sampling`,
+    gitea #86) so a config.toml written against an older v0.3 key name still
+    passes `validate_config_keys`.
+    """
     cfg = copy.deepcopy(cfg)
+    for old_path, new_path in _RENAMED_KEYS.items():
+        if _get_path(cfg, old_path) is not None and _get_path(cfg, new_path) is None:
+            _set_path(cfg, new_path, _get_path(cfg, old_path))
+        _pop_path(cfg, old_path)
+
+    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
+        return cfg
+
     old_train = cfg.pop("train", {})
     old_model = cfg.pop("model", {})
     old_router = dict(old_model.pop("router", {}))
@@ -1512,9 +1547,9 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
                 "conditioning to hang an EOS decision off"
             )
 
-    stop_sampling = _get_path(cfg, "stage2_model.n_sec.stop_sampling")
-    if stop_sampling not in ("greedy", "sample"):
-        raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
+    n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
+    if n_sec_sampling not in ("greedy", "sample"):
+        raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")
 
     precision = _get_path(cfg, "train.precision")
     if precision not in ("fp32", "bf16"):

giant/model/builders.py

@@ -138,7 +138,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
                 build_stop_head=stop_token,
-                stop_sampling=s2_spec.n_sec.stop_sampling,
+                n_sec_sampling=s2_spec.n_sec.sampling,
                 stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
             )
         else:
@@ -168,6 +168,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 cond_enc=shared_cond_enc,
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
+                n_sec_sampling=s2_spec.n_sec.sampling,
             )
 
     return result

giant/model/models.py

@@ -370,6 +370,7 @@ class Stage2OneShot(StageModel):
         cond_enc: ConditionEncoder | None = None,
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
+        n_sec_sampling: str = "greedy",
     ) -> None:
         super().__init__(
             pdg_vocab,
@@ -383,6 +384,7 @@ class Stage2OneShot(StageModel):
             particle_type_cfg=particle_type_cfg,
             cond_enc=cond_enc,
         )
+        self.n_sec_sampling = n_sec_sampling
         self._build_context_fusion(x_dim, context_dim, cond_out_dim)
         target = self.particle_type_cfg.target
         type_head_out_dim = None if target == "physical" else k_max * self.type_dim
@@ -498,7 +500,7 @@ class Stage2Autoregressive(StageModel):
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
         build_stop_head: bool = False,
-        stop_sampling: str = "greedy",
+        n_sec_sampling: str = "greedy",
         stop_head_cfg: dict | None = None,
     ) -> None:
         super().__init__(
@@ -514,7 +516,7 @@ class Stage2Autoregressive(StageModel):
             cond_enc=cond_enc,
         )
         self.history_kind = history
-        self.stop_sampling = stop_sampling
+        self.n_sec_sampling = n_sec_sampling
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
             nn.Linear(cond_out_dim + context_dim, cond_out_dim),

giant/model/summary.py

@@ -140,7 +140,7 @@ def _fingerprint(modules: dict[str, nn.Module]) -> list:
     exist, every parameter's/buffer's shape+dtype (never values — those are
     randomly initialized and irrelevant to *structure*), and every plain
     scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
-    .stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
+    .n_sec_sampling`, `EnergyRouter.temperature`) — this is what makes a
     non-parametric key's effect on construction observable."""
     sig = []
     for stage_name, module in modules.items():

giant/sample.py

@@ -261,7 +261,7 @@ def sample_secondaries_ar(
     slot's own stop logit (`predict_stop`, evaluated on the same prefix
     conditioning as the token itself — see `predict_type`'s docstring for
     why this needs no extra state) decides whether generation should have
-    already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
+    already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
     0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
     `n_sec_pred` is the first slot index where this fires; once every row in
     the batch has fired, the loop breaks before spending a model call on the
@@ -345,7 +345,7 @@ def sample_secondaries_ar(
                 slot_idx,
                 hist=hist,
             ).squeeze(1)
-            if sec_decoder.stop_sampling == "sample":
+            if sec_decoder.n_sec_sampling == "sample":
                 stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
             else:
                 stop_now = stop_logit >= 0.0
@@ -496,7 +496,12 @@ def resolve_n_sec(
 
     Raises if neither stage owns any n_sec mechanism at all — the only way
     that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
-    rollout-/predict-capable checkpoint."""
+    rollout-/predict-capable checkpoint.
+
+    `n_sec.mode = "head"` resolves the classifier logits per
+    `sec_decoder.n_sec_sampling`: "greedy" (default) takes the conditional
+    mode via argmax; "sample" draws a real sample from the learned count
+    distribution via `torch.multinomial` on the softmax — see gitea #86."""
     if n_sec_pred is not None:
         return n_sec_pred
     if getattr(sec_decoder, "stop_head", None) is not None:
@@ -508,4 +513,6 @@ def resolve_n_sec(
             "'truth' is standalone-evaluation-only"
         )
     logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
+    if sec_decoder.n_sec_sampling == "sample":
+        return torch.multinomial(logits.softmax(dim=-1), 1).squeeze(-1)
     return logits.argmax(dim=-1)

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.10"
+version = "0.3.12"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"

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_config.py

@@ -171,17 +171,40 @@ def test_n_sec_config_owner_defaults_to_stage2():
 def test_n_sec_config_owner_round_trips():
     n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
     assert n_sec.owner == "stage1"
-    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "sampling": "greedy"}
 
 
-def test_n_sec_config_stop_sampling_defaults_to_greedy():
-    assert gconfig.NSecConfig().stop_sampling == "greedy"
+def test_n_sec_config_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().sampling == "greedy"
 
 
-def test_n_sec_config_stop_sampling_round_trips():
+def test_n_sec_config_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "sampling": "sample"})
+    assert n_sec.sampling == "sample"
+    assert n_sec.to_dict()["sampling"] == "sample"
+
+
+def test_n_sec_config_stop_sampling_alias_still_honored():
+    """gitea #86: stop_sampling was renamed to sampling; old checkpoints'
+    model_config still carries the old key and must keep working."""
     n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
-    assert n_sec.stop_sampling == "sample"
-    assert n_sec.to_dict()["stop_sampling"] == "sample"
+    assert n_sec.sampling == "sample"
+    assert "stop_sampling" not in n_sec.to_dict()
+
+
+def test_n_sec_config_sampling_key_wins_over_stop_sampling_alias():
+    n_sec = gconfig.NSecConfig.from_dict({"sampling": "sample", "stop_sampling": "greedy"})
+    assert n_sec.sampling == "sample"
+
+
+def test_migrate_config_renames_stop_sampling_key():
+    cfg = {
+        "meta": {"config_version": gconfig.CONFIG_VERSION},
+        "stage2_model": {"n_sec": {"stop_sampling": "sample"}},
+    }
+    migrated = gconfig.migrate_config(cfg)
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.sampling") == "sample"
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.stop_sampling") is None
 
 
 # ---------------------------------------------------------------------------
@@ -851,13 +874,13 @@ def test_validate_config_stop_token_rejected_for_stage1_owner():
         assert "stop_token" in str(e) and "owner" in str(e)
 
 
-def test_validate_config_bad_stop_sampling_rejected():
-    cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
+def test_validate_config_bad_n_sec_sampling_rejected():
+    cfg = _cfg_with(**{"stage2_model.n_sec.sampling": "bogus"})
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "stop_sampling" in str(e)
+        assert "sampling" in str(e)
 
 
 def test_validate_config_default_precision_is_fp32():

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,
             },
         ),
     ]

tests/test_sample.py

@@ -14,6 +14,7 @@ from giant.model.network import (
     stage2_trunk_sec_dim,
 )
 from giant.sample import (
+    resolve_n_sec,
     sample_flow,
     sample_secondaries,
     sample_secondaries_ar,
@@ -72,6 +73,7 @@ def _stage2_ar(
     mat: int = 2,
     k_max: int = 5,
     history: str = "markov",
+    n_sec_sampling: str = "greedy",
 ) -> Stage2Autoregressive:
     particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     return Stage2Autoregressive(
@@ -89,6 +91,7 @@ def _stage2_ar(
         history=history,
         attn_n_heads=2,
         attn_n_layers=1,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -99,7 +102,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
 def _stage2_ar_stop_token(
     target: str,
     generator: str,
-    stop_sampling: str = "greedy",
+    n_sec_sampling: str = "greedy",
     emb_dim: int = 6,
     pdg: int = 3,
     mat: int = 2,
@@ -120,7 +123,7 @@ def _stage2_ar_stop_token(
         particle_type_cfg=ParticleTypeConfig(target=target),
         build_n_sec_head=False,
         build_stop_head=True,
-        stop_sampling=stop_sampling,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -267,14 +270,14 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
 # ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_sampling):
     """A stop_head pinned to a large positive logit fires at slot 0 for
     every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
     a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
     generating any token."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, 50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -283,13 +286,13 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(sto
     assert not sec_valid.any()
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_sampling):
     """A stop_head pinned to a large negative logit never fires under either
     policy, so every row is capped at k_max (the safety cap, not a modeling
     ceiling)."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, -50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -336,3 +339,60 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
     stage1_out = torch.randn(3, X_DIM)
     with pytest.raises(AssertionError):
         sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+
+
+# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
+
+
+def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
+    """Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
+    returns `bias` (broadcast over the batch) as logits regardless of
+    conditioning — mirrors `_force_stop_head_logit`."""
+    assert decoder.n_sec_head is not None
+    last_linear = decoder.n_sec_head[-1]
+    with torch.no_grad():
+        last_linear.weight.zero_()
+        last_linear.bias.copy_(bias)
+
+
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
+    """A logit vector overwhelmingly favoring one class gives the same
+    answer under both policies — greedy because it's the argmax, sample
+    because softmax puts ~all mass on it."""
+    B, k_max = 8, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
+    bias = torch.full((k_max + 1,), -50.0)
+    bias[2] = 50.0
+    _force_n_sec_head_bias(decoder, bias)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))
+
+
+def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
+    B, k_max = 32, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() == 1
+
+
+def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
+    """Under a flat logit vector, a categorical draw across a large batch
+    should hit more than one class — the whole point of gitea #86: greedy
+    always collapses to one, sample should not."""
+    torch.manual_seed(0)
+    B, k_max = 256, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() > 1

uv.lock

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