diff --git a/CLAUDE.md b/CLAUDE.md index 32a2617..bbeb296 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -72,7 +72,7 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep **Validation** (`giant/validate.py`): step-level marginal comparisons. -**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. 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 + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `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`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_/`) holding `shared.json`, `run_meta.json`, `reduced/`, `plots/`. **Compute/render split:** `giant analyze submit rollout.yaml` runs `prep` then submits one HTCondor job per plot (`compute-one --run-dir`, polars/numpy only — no LaTeX on workers), each writing a small `reduced/.json`; the local `giant analyze render ` turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`. +**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. 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 + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `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`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_/`) holding `shared.json`, `run_meta.json`, `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit rollout.yaml --chunks N` runs `prep` (recording the run's chunk count `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`) and writing a small `reduced_partial/__.json`; every `PlotSpec` (`catalog.py`) splits into a `compute_partial`/`finalize` pair so a plot's chunks can be summed/concatenated back together correctly (`chunkable=False` specs — the router diagnostics, already bounded/subsampled — always run as a single chunk regardless of `N`). The local `giant analyze render ` first joins every plot's chunk partials into `reduced/.json` (`merge_all`, a no-op join when `N=1`), then turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`. **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. diff --git a/giant/analysis/__init__.py b/giant/analysis/__init__.py index a646f80..8af48cb 100644 --- a/giant/analysis/__init__.py +++ b/giant/analysis/__init__.py @@ -2,7 +2,8 @@ Compares one autoregressive ``giant rollout`` against a held-out miniCaloSim reference file, producing publication-styled comparison plots generated in -parallel on HTCondor (one job per plot, compute/render split). +parallel on HTCondor (one job per plot x data chunk, compute/merge/render +split). Only ``render`` (and the ``render`` CLI path) imports plotstyle/LaTeX; everything re-exported here is plotstyle-free so it runs on a compute worker. Import @@ -17,11 +18,13 @@ from giant.analysis.condor import ( compute_reduced, derive_run_dir, load_rollout_yaml, + merge_all, + merge_one, prep, write_submit, ) from giant.analysis.context import Context, build_context -from giant.analysis.reduced import Reduced +from giant.analysis.reduced import Partial, Reduced from giant.analysis.sources import Side __all__ = [ @@ -34,10 +37,13 @@ __all__ = [ "compute_reduced", "derive_run_dir", "load_rollout_yaml", + "merge_all", + "merge_one", "prep", "write_submit", "Context", "build_context", + "Partial", "Reduced", "Side", ] diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index a0e5bc0..fc665e3 100644 --- a/giant/analysis/catalog.py +++ b/giant/analysis/catalog.py @@ -2,9 +2,22 @@ Each spec knows its stable ``id`` (used for the reduced-data filename, the PDF stem and the condor queue item), its gallery ``family`` (subdirectory), and a -``compute(bundle) -> Reduced`` that runs the streaming reduction. Rendering lives -in ``render.py`` and dispatches on ``Reduced.kind`` — the catalog itself never -imports plotstyle, so ``compute-one`` jobs stay LaTeX-free. +``compute_partial(bundle) -> dict`` / ``finalize(parts, ctx) -> Reduced`` pair +that together run the streaming reduction. ``compute_partial`` runs once per +``(plot, chunk)`` condor job against a ``Bundle`` whose four LazyFrames are +already filtered to that chunk (see ``Bundle.open``'s ``chunk`` argument); it +returns a small JSON-safe partial artifact — either a raw sum-mergeable count +dict (histograms/species sums against fixed edges) or a raw per-event/ +per-secondary array to be concatenated (anything that derives its own edges or +a mean/std from the full dataset). ``finalize`` merges the per-chunk partials +(in chunk order) and does the actual histogramming/edge-selection/mean-std +collapse, once, over the merged data — for ``n_chunks=1`` this reproduces +exactly what a single unchunked pass would produce. Specs marked +``chunkable=False`` (the router ones) always run as a single chunk regardless +of the configured chunk count. + +Rendering lives in ``render.py`` and dispatches on ``Reduced.kind`` — the +catalog itself never imports plotstyle, so ``compute-one`` jobs stay LaTeX-free. The registry is built by expanding parametric families (marginals over variable x grouping, secondaries, ...) into concrete specs. @@ -12,7 +25,7 @@ variable x grouping, secondaries, ...) into concrete specs. from __future__ import annotations -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import Callable import numpy as np @@ -32,9 +45,11 @@ from giant.analysis.reduce import ( event_scalars, hist1d, leakage_fraction, + profile_finalize, + profile_partial, species_share, + sum_merge, transverse_expr, - weighted_profile, ) from giant.analysis.reduced import Reduced from giant.analysis.router_gating import ( @@ -58,9 +73,29 @@ class Bundle: checkpoint: str | None = None # from the rollout YAML; router_gating only @classmethod - def open(cls, rollout, reference, ctx: Context, checkpoint=None) -> "Bundle": + def open( + cls, + rollout, + reference, + ctx: Context, + checkpoint=None, + chunk: tuple[int, int] | None = None, + ) -> "Bundle": + """Open both sides, optionally restricted to one event-disjoint chunk. + + ``chunk = (chunk_index, n_chunks)`` filters both sides to + ``event_id % n_chunks == chunk_index`` *before* deriving the physical/ + secondary views, so every downstream reduction (which is either + row-local or a ``group_by("event_id")``) sees a self-contained, + event-disjoint slice — no cross-chunk lookups are ever needed. + """ r_all = open_side(rollout, Side.rollout) t_all = open_side(reference, Side.reference) + if chunk is not None: + idx, n = chunk + pred = pl.col("event_id") % n == idx + r_all = r_all.filter(pred) + t_all = t_all.filter(pred) return cls( ctx=ctx, r_all=r_all, @@ -75,7 +110,26 @@ class Bundle: class PlotSpec: id: str family: str - compute: Callable[[Bundle], Reduced] + compute_partial: Callable[[Bundle], dict] + finalize: Callable[[list[dict], Context], Reduced] + chunkable: bool = True + + +def _unchunkable( + compute: Callable[[Bundle], Reduced], +) -> tuple[Callable[[Bundle], dict], Callable[[list[dict], Context], Reduced]]: + """Wrap a whole-dataset ``compute(bundle) -> Reduced`` as a trivial + ``(compute_partial, finalize)`` pair, for specs marked ``chunkable=False`` + (which always run as a single chunk, so ``parts`` is always one element). + """ + + def partial(b: Bundle) -> dict: + return {"reduced": asdict(compute(b))} + + def finalize(parts: list[dict], ctx: Context) -> Reduced: + return Reduced(**parts[0]["reduced"]) + + return partial, finalize # --------------------------------------------------------------------------- @@ -90,6 +144,20 @@ def _counts(h: dict, key, nbins: int) -> list[int]: return h.get(key, np.zeros(nbins, dtype=np.int64)).astype(np.int64).tolist() +def _partial_hist( + lf: pl.LazyFrame, value: pl.Expr, edges: np.ndarray, group: pl.Expr | None = None +) -> dict[str, list[int]]: + """One chunk's raw ``hist1d`` result as a JSON-safe, sum-mergeable dict.""" + nb = len(edges) - 1 + h = hist1d(lf, value, edges, group=group) + return {str(k): _counts(h, k, nb) for k in h} + + +def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]: + """One group's merged counts (zero-filled if the group never appeared).""" + return list(merged.get(str(key), [0] * nbins)) + + def _np_hist_pair( r: np.ndarray, t: np.ndarray, nbins: int ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: @@ -132,12 +200,21 @@ def _marginal_edges(ctx: Context, var: str) -> np.ndarray: # --------------------------------------------------------------------------- -def _marginal_overall(b: Bundle, var: str) -> Reduced: - label, expr = _var(var) +def _marginal_overall_partial(b: Bundle, var: str) -> dict: + _, expr = _var(var) edges = _marginal_edges(b.ctx, var) + return { + "r": _partial_hist(b.r_phys, expr, edges), + "t": _partial_hist(b.t_phys, expr, edges), + } + + +def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced: + label, _ = _var(var) + edges = _marginal_edges(ctx, var) nb = len(edges) - 1 - r = hist1d(b.r_phys, expr, edges) - t = hist1d(b.t_phys, expr, edges) + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) return Reduced( id=f"marginal_{var}", family="marginals", @@ -146,8 +223,8 @@ def _marginal_overall(b: Bundle, var: str) -> Reduced: xlabel=label, payload={ "edges": edges.tolist(), - _ROLL: _counts(r, 0, nb), - _REF: _counts(t, 0, nb), + _ROLL: _finalize_counts(r, 0, nb), + _REF: _finalize_counts(t, 0, nb), "log_y": True, }, ) @@ -160,31 +237,55 @@ def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr: ) -def _marginal_grouped(b: Bundle, var: str, axis: str) -> Reduced: - label, expr = _var(var) +def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict: + _, expr = _var(var) edges = _marginal_edges(b.ctx, var) - nb = len(edges) - 1 - groups: dict[str, dict] = {} - if axis == "pdg": r = hist1d(b.r_phys, expr, edges, group=pl.col("pdg")) t = hist1d(b.t_phys, expr, edges, group=pl.col("pdg")) - for k in b.ctx.top_pdgs: - groups[pdg_label(k)] = {_ROLL: _counts(r, k, nb), _REF: _counts(t, k, nb)} elif axis == "material": r = hist1d(b.r_phys, expr, edges, group=pl.col("material")) t = hist1d(b.t_phys, expr, edges, group=pl.col("material")) - for m in b.ctx.materials: - groups[material_label(m)] = { - _ROLL: _counts(r, m, nb), - _REF: _counts(t, m, nb), - } else: # energy e_edges = np.asarray(b.ctx.energy_edges) r = hist1d(b.r_phys, expr, edges, group=_energy_group_expr(b.r_phys, e_edges)) t = hist1d(b.t_phys, expr, edges, group=_energy_group_expr(b.t_phys, e_edges)) + nb = len(edges) - 1 + return { + "r": {str(k): _counts(r, k, nb) for k in r}, + "t": {str(k): _counts(t, k, nb) for k in t}, + } + + +def _marginal_grouped_finalize( + parts: list[dict], ctx: Context, var: str, axis: str +) -> Reduced: + label, _ = _var(var) + edges = _marginal_edges(ctx, var) + nb = len(edges) - 1 + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) + groups: dict[str, dict] = {} + + if axis == "pdg": + for k in ctx.top_pdgs: + groups[pdg_label(k)] = { + _ROLL: _finalize_counts(r, k, nb), + _REF: _finalize_counts(t, k, nb), + } + elif axis == "material": + for m in ctx.materials: + groups[material_label(m)] = { + _ROLL: _finalize_counts(r, m, nb), + _REF: _finalize_counts(t, m, nb), + } + else: # energy + e_edges = np.asarray(ctx.energy_edges) for bi, lbl in enumerate(energy_bin_labels(e_edges)): - groups[lbl] = {_ROLL: _counts(r, bi, nb), _REF: _counts(t, bi, nb)} + groups[lbl] = { + _ROLL: _finalize_counts(r, bi, nb), + _REF: _finalize_counts(t, bi, nb), + } return Reduced( id=f"marginal_{var}_by_{axis}", @@ -201,13 +302,19 @@ def _marginal_grouped(b: Bundle, var: str, axis: str) -> Reduced: # --------------------------------------------------------------------------- -def _event_scalar( - b: Bundle, spec_id: str, title: str, xlabel: str, col: str, use_all: bool -) -> Reduced: +def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict: r_lf, t_lf = (b.r_all, b.t_all) if use_all else (b.r_phys, b.t_phys) r = event_scalars(r_lf)[col].to_numpy() t = event_scalars(t_lf)[col].to_numpy() - edges, rc, tc = _np_hist_pair(r, t, b.ctx.n_marginal_bins) + return {"r": r.tolist(), "t": t.tolist()} + + +def _event_scalar_finalize( + parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str +) -> Reduced: + r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts]) + t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts]) + edges, rc, tc = _np_hist_pair(r, t, ctx.n_marginal_bins) return Reduced( id=spec_id, family="event", @@ -223,18 +330,26 @@ def _event_scalar( ) -def _event_total_edep_by_energy(b: Bundle) -> Reduced: - e_edges = np.asarray(b.ctx.energy_edges) +def _event_total_edep_by_energy_partial(b: Bundle) -> dict: r = event_scalars(b.r_all) t = event_scalars(b.t_all) - r_bin = np.clip( - np.digitize(r["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2 - ) - t_bin = np.clip( - np.digitize(t["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2 - ) - r_val, t_val = r["total_edep"].to_numpy(), t["total_edep"].to_numpy() - edges, _, _ = _np_hist_pair(r_val, t_val, b.ctx.n_marginal_bins) + return { + "r_incident": r["incident_E"].to_list(), + "r_edep": r["total_edep"].to_list(), + "t_incident": t["incident_E"].to_list(), + "t_edep": t["total_edep"].to_list(), + } + + +def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Reduced: + e_edges = np.asarray(ctx.energy_edges) + r_inc = np.concatenate([np.asarray(p["r_incident"], dtype=float) for p in parts]) + r_val = np.concatenate([np.asarray(p["r_edep"], dtype=float) for p in parts]) + t_inc = np.concatenate([np.asarray(p["t_incident"], dtype=float) for p in parts]) + t_val = np.concatenate([np.asarray(p["t_edep"], dtype=float) for p in parts]) + r_bin = np.clip(np.digitize(r_inc, e_edges[1:-1]), 0, len(e_edges) - 2) + t_bin = np.clip(np.digitize(t_inc, e_edges[1:-1]), 0, len(e_edges) - 2) + edges, _, _ = _np_hist_pair(r_val, t_val, ctx.n_marginal_bins) groups: dict[str, dict] = {} for bi, lbl in enumerate(energy_bin_labels(e_edges)): rc = np.histogram(r_val[r_bin == bi], edges)[0] @@ -258,15 +373,54 @@ def _event_total_edep_by_energy(b: Bundle) -> Reduced: # --------------------------------------------------------------------------- -def _profile( - b: Bundle, spec_id: str, title: str, xlabel: str, coord_fn, edges_key: str -) -> Reduced: +def _profile_partial(b: Bundle, coord_fn, edges_key: str) -> dict: edges = np.asarray(getattr(b.ctx, edges_key)) - r_ea, t_ea = entry_axis(b.r_all), entry_axis(b.t_all) - r_lf = attach_entry_axis(b.r_all, r_ea) - t_lf = attach_entry_axis(b.t_all, t_ea) - r_mean, r_std = weighted_profile(r_lf, coord_fn(), edges, pl.col("edep")) - t_mean, t_std = weighted_profile(t_lf, coord_fn(), edges, pl.col("edep")) + r_lf = attach_entry_axis(b.r_all, entry_axis(b.r_all)) + t_lf = attach_entry_axis(b.t_all, entry_axis(b.t_all)) + r_ids, r_mat = profile_partial(r_lf, coord_fn(), edges, pl.col("edep")) + t_ids, t_mat = profile_partial(t_lf, coord_fn(), edges, pl.col("edep")) + return { + "r_ids": r_ids.tolist(), + "r_mat": r_mat.tolist(), + "t_ids": t_ids.tolist(), + "t_mat": t_mat.tolist(), + } + + +def _assert_event_disjoint(id_lists: list[list[int]], spec_id: str, side: str) -> None: + """Guard the chunking invariant profiles depend on: no event in two chunks. + + A violation would silently double-count that event in the merged mean/RMS + with no other symptom, so this is worth a loud failure rather than a + quietly-wrong plot. + """ + seen: set[int] = set() + for ids in id_lists: + overlap = seen & set(ids) + if overlap: + raise ValueError( + f"{spec_id} ({side}): event_id(s) {sorted(overlap)[:5]} appear " + "in more than one chunk — chunking must be event-disjoint" + ) + seen.update(ids) + + +def _profile_finalize( + parts: list[dict], + ctx: Context, + spec_id: str, + title: str, + xlabel: str, + edges_key: str, +) -> Reduced: + edges = np.asarray(getattr(ctx, edges_key)) + nb = len(edges) - 1 + _assert_event_disjoint([p["r_ids"] for p in parts], spec_id, "rollout") + _assert_event_disjoint([p["t_ids"] for p in parts], spec_id, "reference") + r_mats = [np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts] + t_mats = [np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts] + r_mean, r_std = profile_finalize(r_mats) + t_mean, t_std = profile_finalize(t_mats) return Reduced( id=spec_id, family="shower", @@ -289,14 +443,21 @@ def _profile( # --------------------------------------------------------------------------- -def _species_share(b: Bundle) -> Reduced: +def _species_share_partial(b: Bundle) -> dict: r = species_share(b.r_all) t = species_share(b.t_all) - r_map = dict(zip(r["pdg"].to_list(), r["total_edep"].to_list())) - t_map = dict(zip(t["pdg"].to_list(), t["total_edep"].to_list())) + return { + "r": {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())}, + "t": {str(k): v for k, v in zip(t["pdg"].to_list(), t["total_edep"].to_list())}, + } + + +def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced: + r_map = sum_merge([p["r"] for p in parts]) + t_map = sum_merge([p["t"] for p in parts]) r_tot = sum(r_map.values()) or 1.0 t_tot = sum(t_map.values()) or 1.0 - labels = [pdg_label(k) for k in b.ctx.top_pdgs] + labels = [pdg_label(k) for k in ctx.top_pdgs] return Reduced( id="species_edep_share", family="species", @@ -305,19 +466,22 @@ def _species_share(b: Bundle) -> Reduced: xlabel="species", payload={ "labels": labels, - _ROLL: [r_map.get(k, 0.0) / r_tot for k in b.ctx.top_pdgs], - _REF: [t_map.get(k, 0.0) / t_tot for k in b.ctx.top_pdgs], + _ROLL: [r_map.get(str(k), 0.0) / r_tot for k in ctx.top_pdgs], + _REF: [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs], "ylabel": "fraction of total deposited energy", }, ) -def _leakage(b: Bundle) -> Reduced: +def _leakage_partial(b: Bundle) -> dict: frac = leakage_fraction(b.r_all) + return {"frac": frac.tolist()} + + +def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced: + frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts]) edges = np.linspace( - 0.0, - max(float(frac.max()) if len(frac) else 1.0, 1e-3), - b.ctx.n_marginal_bins + 1, + 0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1 ) counts = np.histogram(frac, edges)[0] return Reduced( @@ -347,7 +511,7 @@ def _sec_frames(b: Bundle): ) -def _sec_count_per_event(b: Bundle) -> Reduced: +def _sec_count_per_event_partial(b: Bundle) -> dict: r_sec, t_sec = _sec_frames(b) r = ( r_sec.group_by("event_id") @@ -361,9 +525,13 @@ def _sec_count_per_event(b: Bundle) -> Reduced: .collect(engine="streaming")["n"] .to_numpy() ) - edges, rc, tc = _np_hist_pair( - r.astype(float), t.astype(float), min(b.ctx.n_marginal_bins, 40) - ) + return {"r": r.tolist(), "t": t.tolist()} + + +def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced: + r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts]) + t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts]) + edges, rc, tc = _np_hist_pair(r, t, min(ctx.n_marginal_bins, 40)) return Reduced( id="sec_count_per_event", family="secondaries", @@ -379,32 +547,21 @@ def _sec_count_per_event(b: Bundle) -> Reduced: ) -def _sec_count_per_species(b: Bundle) -> Reduced: +def _counts_by_pdg(sec_lf: pl.LazyFrame) -> dict[str, int]: + df = sec_lf.group_by("pdg").agg(pl.len().alias("n")).collect(engine="streaming") + return {str(k): v for k, v in zip(df["pdg"].to_list(), df["n"].to_list())} + + +def _sec_count_per_species_partial(b: Bundle) -> dict: r_sec, t_sec = _sec_frames(b) - r = dict( - zip( - *[ - r_sec.group_by("pdg") - .agg(pl.len().alias("n")) - .collect(engine="streaming")[c] - .to_list() - for c in ("pdg", "n") - ] - ) - ) - t = dict( - zip( - *[ - t_sec.group_by("pdg") - .agg(pl.len().alias("n")) - .collect(engine="streaming")[c] - .to_list() - for c in ("pdg", "n") - ] - ) - ) + return {"r": _counts_by_pdg(r_sec), "t": _counts_by_pdg(t_sec)} + + +def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced: + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[ - : len(b.ctx.top_pdgs) + : len(ctx.top_pdgs) ] return Reduced( id="sec_count_per_species", @@ -413,7 +570,7 @@ def _sec_count_per_species(b: Bundle) -> Reduced: title="Secondary count by species", xlabel="species", payload={ - "labels": [pdg_label(k) for k in keys], + "labels": [pdg_label(int(k)) for k in keys], _ROLL: [float(r.get(k, 0)) for k in keys], _REF: [float(t.get(k, 0)) for k in keys], "ylabel": "secondary count", @@ -421,12 +578,20 @@ def _sec_count_per_species(b: Bundle) -> Reduced: ) -def _sec_energy(b: Bundle) -> Reduced: +def _sec_energy_partial(b: Bundle) -> dict: r_sec, t_sec = _sec_frames(b) edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1) - r = hist1d(r_sec, pl.col("energy"), edges) - t = hist1d(t_sec, pl.col("energy"), edges) + return { + "r": _partial_hist(r_sec, pl.col("energy"), edges), + "t": _partial_hist(t_sec, pl.col("energy"), edges), + } + + +def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced: + edges = np.linspace(*ctx.sec_energy_range, ctx.n_sec_bins + 1) nb = len(edges) - 1 + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) return Reduced( id="sec_energy", family="secondaries", @@ -435,27 +600,34 @@ def _sec_energy(b: Bundle) -> Reduced: xlabel="secondary energy [MeV]", payload={ "edges": edges.tolist(), - _ROLL: _counts(r, 0, nb), - _REF: _counts(t, 0, nb), + _ROLL: _finalize_counts(r, 0, nb), + _REF: _finalize_counts(t, 0, nb), "log_y": True, }, ) -def _sec_cos_angle(b: Bundle) -> Reduced: +def _sec_cos_angle_partial(b: Bundle) -> dict: edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1) - nb = len(edges) - 1 cos = ( pl.col("sdx") * pl.col("axis_x") + pl.col("sdy") * pl.col("axis_y") + pl.col("sdz") * pl.col("axis_z") ).clip(-1.0, 1.0) - def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> list[int]: + def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]: ea = entry_axis(steps_lf) - return _counts(hist1d(attach_entry_axis(sec_lf, ea), cos, edges), 0, nb) + return _partial_hist(attach_entry_axis(sec_lf, ea), cos, edges) r_sec, t_sec = _sec_frames(b) + return {"r": _side(r_sec, b.r_phys), "t": _side(t_sec, b.t_all)} + + +def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced: + edges = np.linspace(-1.0, 1.0, ctx.n_sec_bins + 1) + nb = len(edges) - 1 + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) return Reduced( id="sec_cos_angle", family="secondaries", @@ -464,13 +636,30 @@ def _sec_cos_angle(b: Bundle) -> Reduced: xlabel="cos of emission angle", payload={ "edges": edges.tolist(), - _ROLL: _side(r_sec, b.r_phys), - _REF: _side(t_sec, b.t_all), + _ROLL: _finalize_counts(r, 0, nb), + _REF: _finalize_counts(t, 0, nb), "log_y": False, }, ) +# --------------------------------------------------------------------------- +# router diagnostics (not chunked — already bounded/subsampled) +# --------------------------------------------------------------------------- + +_router_gating_partial, _router_gating_finalize = _unchunkable( + lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys) +) +_router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable( + lambda b: compute_router_share_by_pdg( + b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs + ) +) +_router_share_process_partial, _router_share_process_finalize = _unchunkable( + lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys) +) + + # --------------------------------------------------------------------------- # registry assembly # --------------------------------------------------------------------------- @@ -486,7 +675,12 @@ def build_catalog() -> list[PlotSpec]: for var in MARGINAL_VARS: specs.append( PlotSpec( - f"marginal_{var}", "marginals", lambda b, v=var: _marginal_overall(b, v) + f"marginal_{var}", + "marginals", + compute_partial=lambda b, v=var: _marginal_overall_partial(b, v), + finalize=lambda parts, ctx, v=var: _marginal_overall_finalize( + parts, ctx, v + ), ) ) for axis in GROUPING_AXES: @@ -494,7 +688,12 @@ def build_catalog() -> list[PlotSpec]: PlotSpec( f"marginal_{var}_by_{axis}", "marginals", - lambda b, v=var, a=axis: _marginal_grouped(b, v, a), + compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial( + b, v, a + ), + finalize=lambda parts, ctx, v=var, a=axis: ( + _marginal_grouped_finalize(parts, ctx, v, a) + ), ) ) @@ -502,86 +701,135 @@ def build_catalog() -> list[PlotSpec]: PlotSpec( "event_total_edep", "event", - lambda b: _event_scalar( - b, + compute_partial=lambda b: _event_scalar_partial( + b, "total_edep", use_all=True + ), + finalize=lambda parts, ctx: _event_scalar_finalize( + parts, + ctx, "event_total_edep", "Total deposited energy per event", "total deposited energy [MeV]", - "total_edep", - use_all=True, ), ), - PlotSpec("event_total_edep_by_energy", "event", _event_total_edep_by_energy), + PlotSpec( + "event_total_edep_by_energy", + "event", + compute_partial=_event_total_edep_by_energy_partial, + finalize=_event_total_edep_by_energy_finalize, + ), PlotSpec( "event_mean_length", "event", - lambda b: _event_scalar( - b, + compute_partial=lambda b: _event_scalar_partial( + b, "mean_length", use_all=False + ), + finalize=lambda parts, ctx: _event_scalar_finalize( + parts, + ctx, "event_mean_length", "Mean step length per event", "mean step length [mm]", - "mean_length", - use_all=False, ), ), PlotSpec( "event_n_steps", "event", - lambda b: _event_scalar( - b, + compute_partial=lambda b: _event_scalar_partial( + b, "n_steps", use_all=False + ), + finalize=lambda parts, ctx: _event_scalar_finalize( + parts, + ctx, "event_n_steps", "Number of steps per event", "steps per event", - "n_steps", - use_all=False, ), ), PlotSpec( "shower_longitudinal", "shower", - lambda b: _profile( - b, + compute_partial=lambda b: _profile_partial(b, depth_expr, "depth_edges"), + finalize=lambda parts, ctx: _profile_finalize( + parts, + ctx, "shower_longitudinal", "Longitudinal shower profile", "depth along shower axis [mm]", - depth_expr, "depth_edges", ), ), PlotSpec( "shower_transverse", "shower", - lambda b: _profile( - b, + compute_partial=lambda b: _profile_partial( + b, transverse_expr, "transverse_edges" + ), + finalize=lambda parts, ctx: _profile_finalize( + parts, + ctx, "shower_transverse", "Transverse shower profile", "radius from shower axis [mm]", - transverse_expr, "transverse_edges", ), ), - PlotSpec("species_edep_share", "species", _species_share), - PlotSpec("leakage_fraction", "species", _leakage), - PlotSpec("sec_count_per_event", "secondaries", _sec_count_per_event), - PlotSpec("sec_count_per_species", "secondaries", _sec_count_per_species), - PlotSpec("sec_energy", "secondaries", _sec_energy), - PlotSpec("sec_cos_angle", "secondaries", _sec_cos_angle), + PlotSpec( + "species_edep_share", + "species", + compute_partial=_species_share_partial, + finalize=_species_share_finalize, + ), + PlotSpec( + "leakage_fraction", + "species", + compute_partial=_leakage_partial, + finalize=_leakage_finalize, + ), + PlotSpec( + "sec_count_per_event", + "secondaries", + compute_partial=_sec_count_per_event_partial, + finalize=_sec_count_per_event_finalize, + ), + PlotSpec( + "sec_count_per_species", + "secondaries", + compute_partial=_sec_count_per_species_partial, + finalize=_sec_count_per_species_finalize, + ), + PlotSpec( + "sec_energy", + "secondaries", + compute_partial=_sec_energy_partial, + finalize=_sec_energy_finalize, + ), + PlotSpec( + "sec_cos_angle", + "secondaries", + compute_partial=_sec_cos_angle_partial, + finalize=_sec_cos_angle_finalize, + ), PlotSpec( "router_gating", "model", - lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys), + compute_partial=_router_gating_partial, + finalize=_router_gating_finalize, + chunkable=False, ), PlotSpec( "router_share_by_pdg", "model", - lambda b: compute_router_share_by_pdg( - b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs - ), + compute_partial=_router_share_pdg_partial, + finalize=_router_share_pdg_finalize, + chunkable=False, ), PlotSpec( "router_share_by_process", "model", - lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys), + compute_partial=_router_share_process_partial, + finalize=_router_share_process_finalize, + chunkable=False, ), ] return specs diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index b9ad621..606d93b 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -15,18 +15,25 @@ next to the rollout parquet, and lays everything out under it: /shared.json fixed bin edges / group sets (prep) /run_meta.json resolved rollout/reference paths + plot metadata - /reduced/.json one per compute job + /reduced_partial/__.json one per (plot, chunk) job + /reduced/.json merged, per plot /plots//.pdf rendered locally -Job model (one condor job per plot, compute/render split): +Job model (one condor job per (plot, chunk), compute/merge/render split): 1. ``prep`` runs once on the submit node — reads the YAML, resolves the shared - context from a subsample, writes ``shared.json`` + ``run_meta.json``. -2. one job per catalog id runs ``giant analyze compute-one --run-dir`` on a - worker — a single streaming pass writing ``reduced/.json`` (polars/numpy - only, no LaTeX). -3. a final *local* ``giant analyze render`` turns those into the styled PDF + - gallery tree (that step imports plotstyle/LaTeX). + context from a subsample, writes ``shared.json`` + ``run_meta.json`` + (including the run's configured ``n_chunks``). +2. one job per catalog id x chunk index runs ``giant analyze compute-one + --run-dir`` on a worker — a single streaming pass over that + ``event_id``-disjoint chunk, writing ``reduced_partial/__.json`` + (polars/numpy only, no LaTeX). Specs marked ``chunkable=False`` + (``PlotSpec``, ``catalog.py``) always run as a single chunk. +3. a *local* ``giant analyze render`` first merges every plot's chunk partials + (``merge_all`` — sums/concatenates them and re-derives any data-dependent + histogram edges or mean/std, per ``PlotSpec.finalize``) into + ``reduced/.json``, then renders those into the styled PDF + gallery tree + (that step imports plotstyle/LaTeX). Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no HTCondor file transfer of the multi-GB inputs. @@ -42,6 +49,7 @@ import yaml from giant.analysis.catalog import Bundle, catalog_ids, get_spec from giant.analysis.context import Context, build_context +from giant.analysis.reduced import Partial # Keys copied verbatim from a rollout YAML into each plot's gallery metadata. _PLOT_META_KEYS = ( @@ -104,6 +112,7 @@ class RunMeta: run_dir: str title: str plot_meta: dict + n_chunks: int = 1 def save(self, path: str | Path) -> None: Path(path).write_text(json.dumps(self.__dict__, indent=2)) @@ -116,11 +125,16 @@ class RunMeta: def prep( rollout_yaml: str | Path, run_dir: str | Path | None = None, + n_chunks: int = 1, **ctx_kwargs, ) -> Path: """Read the rollout YAML, build the shared context, and lay out the run dir. Writes ``shared.json`` + ``run_meta.json`` and returns the run directory. + ``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one`` + job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is + resolved once here rather than re-passed (and risking disagreement) at every + later step. """ y = load_rollout_yaml(rollout_yaml) run_path = derive_run_dir(y, run_dir) @@ -137,12 +151,13 @@ def prep( run_dir=str(run_path), title=f"GIANT rollout analysis — {ckpt}", plot_meta=_plot_meta(y), + n_chunks=n_chunks, ).save(run_path / "run_meta.json") return run_path # --------------------------------------------------------------------------- -# per-plot compute (what each condor job runs) +# per-(plot, chunk) compute (what each condor job runs) # --------------------------------------------------------------------------- @@ -153,18 +168,41 @@ def compute_reduced( shared: str | Path, out: str | Path, checkpoint: str | None = None, + chunk_index: int = 0, + n_chunks: int = 1, ) -> Path: - """Core: run one plot's reduction against explicit paths → ``Reduced`` JSON.""" + """Core: run one (plot, chunk)'s partial reduction against explicit paths. + + Writes a ``Partial`` JSON — the raw, not-yet-merged output of + ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one`` + is what combines every chunk's ``Partial`` for a plot into the final + ``Reduced``. Specs with ``chunkable=False`` always run as a single chunk + regardless of ``n_chunks``. + """ ctx = Context.load(shared) - bundle = Bundle.open(rollout, reference, ctx, checkpoint=checkpoint) - reduced = get_spec(spec_id).compute(bundle) + spec = get_spec(spec_id) + effective_n = n_chunks if spec.chunkable else 1 + if not (0 <= chunk_index < effective_n): + raise ValueError( + f"{spec_id}: chunk_index={chunk_index} out of range for " + f"n_chunks={effective_n} (chunkable={spec.chunkable})" + ) + bundle = Bundle.open( + rollout, reference, ctx, checkpoint=checkpoint, chunk=(chunk_index, effective_n) + ) + partial = Partial( + id=spec_id, + family=spec.family, + chunk=chunk_index, + data=spec.compute_partial(bundle), + ) out = Path(out) - reduced.save(out) + partial.save(out) return out -def compute_one(spec_id: str, run_dir: str | Path) -> Path: - """Run one plot's reduction from a prepped run directory.""" +def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path: + """Run one (plot, chunk)'s partial reduction from a prepped run directory.""" run_path = Path(run_dir) meta = RunMeta.load(run_path / "run_meta.json") return compute_reduced( @@ -172,11 +210,56 @@ def compute_one(spec_id: str, run_dir: str | Path) -> Path: meta.rollout, meta.reference, run_path / "shared.json", - run_path / "reduced" / f"{spec_id}.json", + run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json", checkpoint=meta.plot_meta.get("checkpoint"), + chunk_index=chunk_index, + n_chunks=meta.n_chunks, ) +# --------------------------------------------------------------------------- +# per-plot merge (the join step ``render`` runs before rendering) +# --------------------------------------------------------------------------- + + +def merge_one(spec_id: str, run_dir: str | Path) -> Path: + """Merge every chunk's partial for one plot into the final ``Reduced`` JSON. + + Fails loudly if fewer partials exist than the run's configured chunk count + for this plot — that is what catches an incomplete/failed condor job + instead of silently rendering a plot from partial data. Idempotent: safe + to call again (e.g. from ``render_run``) once all chunks are in. + """ + run_path = Path(run_dir) + meta = RunMeta.load(run_path / "run_meta.json") + ctx = Context.load(run_path / "shared.json") + spec = get_spec(spec_id) + effective_n = meta.n_chunks if spec.chunkable else 1 + + partial_dir = run_path / "reduced_partial" + found = { + p.chunk: p + for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json")) + } + missing = sorted(set(range(effective_n)) - set(found)) + if missing: + raise FileNotFoundError( + f"{spec_id}: missing chunk partial(s) {missing} of {effective_n} " + f"under {partial_dir} — did every compute-one job finish?" + ) + + parts = [found[k].data for k in range(effective_n)] + reduced = spec.finalize(parts, ctx) + out = run_path / "reduced" / f"{spec_id}.json" + reduced.save(out) + return out + + +def merge_all(run_dir: str | Path) -> list[Path]: + """Merge every catalog plot's chunk partials into ``reduced/.json``.""" + return [merge_one(spec_id, run_dir) for spec_id in catalog_ids()] + + # --------------------------------------------------------------------------- # submit description # --------------------------------------------------------------------------- @@ -192,16 +275,17 @@ class SubmitConfig: request_cpus: int = 1 request_walltime_s: int = 3600 remote: bool = False # +RemoteJob (grid I/O) vs ProvidesETPResources (local files) + n_chunks: int = 1 # per-plot data chunks; ignored for chunkable=False specs _WRAPPER = """#!/bin/bash set -euo pipefail cd {repo_dir} -exec uv run giant analyze compute-one --id "$1" --run-dir {run_dir} +exec uv run giant analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir} """ -def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str: +def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str: reqs_attrs = ( "+RemoteJob = True\n" if cfg.remote @@ -211,7 +295,7 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str "universe = docker\n" f"docker_image = {cfg.docker_image}\n" f"executable = {wrapper}\n" - "arguments = $(plotid)\n" + "arguments = $(plotid) $(chunk)\n" "should_transfer_files = YES\n" "when_to_transfer_output = ON_EXIT\n" f"request_memory = {cfg.request_memory_mb}\n" @@ -219,31 +303,41 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str f"+RequestWalltime = {cfg.request_walltime_s}\n" f"accounting_group = {cfg.accounting_group}\n" f"{reqs_attrs}" - f"output = {cfg.run_dir}/logs/$(plotid).out\n" - f"error = {cfg.run_dir}/logs/$(plotid).err\n" + f"output = {cfg.run_dir}/logs/$(plotid)__$(chunk).out\n" + f"error = {cfg.run_dir}/logs/$(plotid)__$(chunk).err\n" f"log = {cfg.run_dir}/logs/condor.log\n" - f"queue plotid from {ids_file}\n" + f"queue plotid,chunk from {jobs_file}\n" ) def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path: - """Write the wrapper script, plot-id list, and HTCondor submit description. + """Write the wrapper script, (plot, chunk) job list, and HTCondor submit + description. - Returns the submit description path (``/analyze.sub``). Does not - submit — call ``condor_submit`` on the returned file. + Each catalog id gets ``cfg.n_chunks`` jobs, except ``chunkable=False`` + specs (the router diagnostics), which always get exactly one regardless of + ``cfg.n_chunks``. Returns the submit description path + (``/analyze.sub``). Does not submit — call ``condor_submit`` on + the returned file. """ ids = ids or catalog_ids() run_dir = cfg.run_dir (run_dir / "logs").mkdir(parents=True, exist_ok=True) (run_dir / "reduced").mkdir(parents=True, exist_ok=True) + (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True) wrapper = run_dir / "run_compute.sh" wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, run_dir=run_dir)) wrapper.chmod(0o755) - ids_file = run_dir / "plotids.txt" - ids_file.write_text("\n".join(ids) + "\n") + jobs = [ + (spec_id, chunk) + for spec_id in ids + for chunk in range(cfg.n_chunks if get_spec(spec_id).chunkable else 1) + ] + jobs_file = run_dir / "jobs.txt" + jobs_file.write_text("\n".join(f"{i},{k}" for i, k in jobs) + "\n") sub = run_dir / "analyze.sub" - sub.write_text(_submit_description(cfg, wrapper, ids_file)) + sub.write_text(_submit_description(cfg, wrapper, jobs_file)) return sub diff --git a/giant/analysis/reduce.py b/giant/analysis/reduce.py index 51bf035..15089f6 100644 --- a/giant/analysis/reduce.py +++ b/giant/analysis/reduce.py @@ -16,6 +16,8 @@ efficiency" section): from __future__ import annotations +from typing import Any + import numpy as np import polars as pl @@ -58,6 +60,24 @@ def hist1d( return out +def sum_merge(dicts: list[dict[str, Any]]) -> dict[str, Any]: + """Elementwise-sum a list of sum-mergeable count/total dicts (JSON-safe keys). + + Used to merge chunked ``hist1d``/``species_share``-style partials, whose + values bin/group against edges or keys fixed by ``Context`` — a chunk's raw + count dict is exactly a partial sum, so merging is a plain elementwise sum + over the union of keys (a key absent from some chunk is all-zero there). + Values may be per-bin count lists or plain scalar totals; both round-trip + through ``np.asarray``/``.tolist()`` unchanged in shape. + """ + out: dict[str, np.ndarray] = {} + for d in dicts: + for k, v in d.items(): + arr = np.asarray(v) + out[k] = arr.copy() if k not in out else out[k] + arr + return {k: v.tolist() for k, v in out.items()} + + # --------------------------------------------------------------------------- # Per-event scalar observables (one bounded group_by pass) # --------------------------------------------------------------------------- @@ -149,18 +169,19 @@ def transverse_expr() -> pl.Expr: return (tx**2 + ty**2 + tz**2).sqrt() -def weighted_profile( +def profile_partial( lf: pl.LazyFrame, coord: pl.Expr, edges: np.ndarray, weight: pl.Expr, ) -> tuple[np.ndarray, np.ndarray]: - """Event-averaged, ``weight``-summed profile of ``coord``, with an event-RMS band. + """One chunk's per-event x bin ``weight``-sum matrix: ``(event_ids, matrix)``. - One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin); - collapsed in numpy to the per-bin mean over events and its event-to-event std - (the band). ``coord``/``weight`` require the entry/axis columns attached. - Returns ``(mean, std)``, each length ``len(edges)-1``. + One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin). + A chunk's matrix rows are only the events present in that chunk, so chunks' + matrices stack cleanly with no cross-chunk lookup — this requires chunking + to be event-disjoint (every row of an event lands in one chunk). + ``coord``/``weight`` require the entry/axis columns attached. """ lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1 grid = ( @@ -177,7 +198,35 @@ def weighted_profile( uniq, inv = np.unique(ev, return_inverse=True) mat = np.zeros((len(uniq), nbins), dtype=np.float64) np.add.at(mat, (inv, grid["_b"].to_numpy()), grid["_ws"].to_numpy()) - return mat.mean(axis=0), mat.std(axis=0) + return uniq, mat + + +def profile_finalize(mats: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Collapse per-chunk per-event x bin matrices into the final mean/std profile. + + Chunks are event-disjoint, so row-wise concatenation of their matrices + reconstructs the full per-event matrix; the mean/event-RMS collapse must + happen once over that full matrix — an average of per-chunk means/stds + would be wrong (chunks generally hold different numbers of events). + Returns ``(mean, std)``, each length ``nbins``. + """ + full = np.concatenate(mats, axis=0) + return full.mean(axis=0), full.std(axis=0) + + +def weighted_profile( + lf: pl.LazyFrame, + coord: pl.Expr, + edges: np.ndarray, + weight: pl.Expr, +) -> tuple[np.ndarray, np.ndarray]: + """Single-pass profile of ``coord`` (mean +/- event-RMS band over events). + + Convenience wrapper for the unchunked (whole-dataset) case; ``mean_std + + profile_partial`` is what a chunked compute/finalize split uses instead. + """ + _, mat = profile_partial(lf, coord, edges, weight) + return profile_finalize([mat]) # --------------------------------------------------------------------------- diff --git a/giant/analysis/reduced.py b/giant/analysis/reduced.py index 836c923..2188800 100644 --- a/giant/analysis/reduced.py +++ b/giant/analysis/reduced.py @@ -39,3 +39,27 @@ class Reduced: @classmethod def load(cls, path: str | Path) -> "Reduced": return cls(**json.loads(Path(path).read_text())) + + +@dataclass +class Partial: + """The raw, not-yet-finalized output of one ``(plot, chunk)`` compute job. + + ``data`` holds whatever shape that plot's ``PlotSpec.compute_partial`` + returns — a raw sum-mergeable count dict, or a raw per-event/per-secondary + array to be concatenated across chunks — never a finished histogram/profile. + ``PlotSpec.finalize`` is the only thing that knows how to interpret it. + """ + + id: str + family: str + chunk: int + data: dict + + def save(self, path: str | Path) -> None: + Path(path).parent.mkdir(parents=True, exist_ok=True) + Path(path).write_text(json.dumps(asdict(self))) + + @classmethod + def load(cls, path: str | Path) -> "Partial": + return cls(**json.loads(Path(path).read_text())) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index 35cd3d6..76d2cd4 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -303,12 +303,16 @@ def render_all( def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]: """Render a prepped run directory: ``/reduced`` → ``/plots``. - Pulls the rollout provenance (checkpoint, paths, cutoffs) from - ``run_meta.json`` into every plot's gallery metadata. + First joins every plot's chunk partials (``reduced_partial/__*.json``) + into ``reduced/.json`` via ``merge_all`` — a no-op merge when the run + wasn't chunked (``n_chunks=1``) — then pulls the rollout provenance + (checkpoint, paths, cutoffs) from ``run_meta.json`` into every plot's + gallery metadata and renders. """ - from giant.analysis.condor import RunMeta + from giant.analysis.condor import RunMeta, merge_all run_dir = Path(run_dir) + merge_all(run_dir) meta = RunMeta.load(run_dir / "run_meta.json") run_meta = { "title": meta.title, diff --git a/giant/cli.py b/giant/cli.py index e371120..1f9769f 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -1150,6 +1150,12 @@ def analyze_prep( n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4, n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50, top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6, + chunks: Annotated[ + int, + typer.Option( + "--chunks", help="Split each plot's data into this many event_id chunks" + ), + ] = 1, ) -> None: """Read the rollout YAML → shared.json + run_meta.json in the run directory.""" from giant.analysis import prep @@ -1157,6 +1163,7 @@ def analyze_prep( path = prep( rollout_yaml, run_dir, + n_chunks=chunks, n_energy_bins=n_energy_bins, n_marginal_bins=n_marginal_bins, top_k_pdg=top_k_pdg, @@ -1172,11 +1179,34 @@ def analyze_compute_one( run_dir: Annotated[ Path, typer.Option("--run-dir", help="Run directory from `analyze prep`") ], + chunk: Annotated[ + int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)") + ] = 0, ) -> None: - """Run one plot's streaming reduction (this is what each condor job runs).""" + """Run one (plot, chunk)'s streaming reduction (this is what each condor job runs).""" from giant.analysis import compute_one - path = compute_one(id, run_dir) + path = compute_one(id, run_dir, chunk_index=chunk) + typer.echo(f"wrote {path}") + + +@analyze_app.command("merge-one") +def analyze_merge_one( + id: Annotated[ + str, typer.Option("--id", help="Catalog plot id (see `analyze list`)") + ], + run_dir: Annotated[ + Path, typer.Option("--run-dir", help="Run directory from `analyze prep`") + ], +) -> None: + """Merge one plot's chunk partials into its final reduced JSON. + + Runs automatically as part of `analyze render`; useful standalone to + debug a specific plot without re-rendering everything. + """ + from giant.analysis import merge_one + + path = merge_one(id, run_dir) typer.echo(f"wrote {path}") @@ -1224,16 +1254,23 @@ def analyze_submit( bool, typer.Option("--remote/--local", help="+RemoteJob vs ProvidesETPResources"), ] = False, + chunks: Annotated[ + int, + typer.Option( + "--chunks", + help="Split each plot's data into this many event_id chunks/jobs", + ), + ] = 1, dry_run: Annotated[ bool, typer.Option("--dry-run", help="Write files but don't condor_submit") ] = False, ) -> None: - """prep + write the HTCondor submit description (one job per plot), then submit.""" + """prep + write the HTCondor submit description (one job per plot x chunk), then submit.""" import subprocess from giant.analysis import SubmitConfig, prep, write_submit - path = prep(rollout_yaml, run_dir) + path = prep(rollout_yaml, run_dir, n_chunks=chunks) cfg = SubmitConfig( run_dir=path, accounting_group=accounting_group, @@ -1241,6 +1278,7 @@ def analyze_submit( docker_image=docker_image, request_memory_mb=request_memory, remote=remote, + n_chunks=chunks, ) sub = write_submit(cfg) typer.echo(f"run directory: {path}") diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 90447d3..c8a47f0 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -2,21 +2,30 @@ from __future__ import annotations +import numpy as np import pytest from giant.analysis import build_catalog, catalog_ids, get_spec -from giant.analysis.catalog import Bundle -from giant.analysis.context import build_context +from giant.analysis.catalog import Bundle, PlotSpec +from giant.analysis.context import Context, build_context from tests.test_analysis_reduce import _reference_frame, _rollout_frame -@pytest.fixture(scope="module") -def bundle() -> Bundle: +def _build_ctx() -> Context: r, t = _rollout_frame(), _reference_frame() - ctx = build_context( + return build_context( r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000 ) - return Bundle.open(r, t, ctx) + + +@pytest.fixture(scope="module") +def ctx() -> Context: + return _build_ctx() + + +@pytest.fixture(scope="module") +def bundle(ctx: Context) -> Bundle: + return Bundle.open(_rollout_frame(), _reference_frame(), ctx) def test_catalog_ids_unique_and_nonempty(): @@ -36,7 +45,7 @@ def test_get_spec_roundtrip_and_unknown(): def test_every_spec_computes_valid_reduced(bundle: Bundle): for spec in build_catalog(): - r = spec.compute(bundle) + r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx) assert r.id == spec.id assert r.kind in { "overlay_hist", @@ -81,3 +90,64 @@ def _validate_payload(r) -> None: for side in ("rollout", "reference"): if side in p: assert cat in p[side] + + +# --------------------------------------------------------------------------- +# chunked (compute_partial x N -> finalize) must match the unchunked (N=1) result +# --------------------------------------------------------------------------- + +# One representative id per merge shape: sum-mergeable (marginal_edep, +# sec_count_per_species via pdg-keyed sums), concat-then-finalize with +# data-dependent edges (event_total_edep), concat-then-mean/std (shower_ +# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a +# ratio (species_edep_share), and a chunkable=False passthrough (router_gating). +_CHUNK_EQUIVALENCE_IDS = [ + "marginal_edep", + "species_edep_share", + "event_total_edep", + "shower_longitudinal", + "leakage_fraction", + "sec_count_per_species", + "router_gating", +] + + +def _assert_payload_close(a, b, path: str = "payload") -> None: + """Recursively compare two JSON-shaped payloads (float-tolerant).""" + assert type(a) is type(b), f"{path}: {type(a)} != {type(b)}" + if isinstance(a, dict): + assert set(a) == set(b), f"{path}: key mismatch {set(a)} != {set(b)}" + for k in a: + _assert_payload_close(a[k], b[k], f"{path}.{k}") + elif isinstance(a, list): + assert len(a) == len(b), f"{path}: length mismatch" + for i, (x, y) in enumerate(zip(a, b)): + _assert_payload_close(x, y, f"{path}[{i}]") + elif isinstance(a, float): + assert np.isclose(a, b, atol=1e-9), f"{path}: {a} != {b}" + else: + assert a == b, f"{path}: {a} != {b}" + + +@pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS) +def test_chunked_matches_unchunked(ctx: Context, spec_id: str): + """A plot computed over N event-disjoint chunks then merged must equal the + same plot computed in one unchunked pass — the core chunking correctness + guarantee (see the analysis-rollout-plots chunking plan).""" + spec: PlotSpec = get_spec(spec_id) + r, t = _rollout_frame(), _reference_frame() + + unchunked_bundle = Bundle.open(r, t, ctx) + unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx) + + # 4 chunks over only 2 distinct event_ids also exercises empty chunks. + n_chunks = 4 if spec.chunkable else 1 + parts = [ + spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks))) + for k in range(n_chunks) + ] + chunked = spec.finalize(parts, ctx) + + assert chunked.id == unchunked.id + assert chunked.kind == unchunked.kind + _assert_payload_close(unchunked.payload, chunked.payload) diff --git a/tests/test_condor.py b/tests/test_condor.py index 1d6ff3e..6634bd8 100644 --- a/tests/test_condor.py +++ b/tests/test_condor.py @@ -16,11 +16,13 @@ from giant.analysis import ( compute_reduced, derive_run_dir, load_rollout_yaml, + merge_one, prep, write_submit, ) +from giant.analysis.catalog import get_spec from giant.analysis.condor import Context -from giant.analysis.reduced import Reduced +from giant.analysis.reduced import Partial, Reduced from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE from tests.test_analysis_reduce import _reference_frame, _rollout_frame @@ -51,11 +53,14 @@ def _write_inputs(tmp_path: Path) -> Path: return yaml_path -def _prep(rollout_yaml: Path, run_dir: str | Path | None = None) -> Path: +def _prep( + rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1 +) -> Path: """``prep`` with small test-sized context bins/sampling.""" return prep( rollout_yaml, run_dir, + n_chunks=chunks, n_energy_bins=2, n_marginal_bins=8, top_k_pdg=3, @@ -87,15 +92,16 @@ def test_prep_lays_out_run_dir(tmp_path: Path): assert meta.reference.endswith("reference.parquet") assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt" assert "best.pt" in meta.title + assert meta.n_chunks == 1 def test_compute_one_from_run_dir(tmp_path: Path): run_dir = _prep(_write_inputs(tmp_path)) out = compute_one("marginal_edep", run_dir) - assert out == run_dir / "reduced" / "marginal_edep.json" - reduced = Reduced.load(out) - assert reduced.id == "marginal_edep" - assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1 + assert out == run_dir / "reduced_partial" / "marginal_edep__0.json" + partial = Partial.load(out) + assert partial.id == "marginal_edep" and partial.chunk == 0 + assert "r" in partial.data and "t" in partial.data def test_compute_reduced_explicit_paths(tmp_path: Path): @@ -108,7 +114,45 @@ def test_compute_reduced_explicit_paths(tmp_path: Path): run_dir / "shared.json", tmp_path / "r.json", ) - assert Reduced.load(out).id == "marginal_step_length" + assert Partial.load(out).id == "marginal_step_length" + + +def test_merge_one_produces_reduced(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path)) + compute_one("marginal_edep", run_dir) + out = merge_one("marginal_edep", run_dir) + assert out == run_dir / "reduced" / "marginal_edep.json" + reduced = Reduced.load(out) + assert reduced.id == "marginal_edep" + assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1 + + +def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path), chunks=2) + compute_one("marginal_edep", run_dir, chunk_index=0) # chunk 1 never computed + with pytest.raises(FileNotFoundError, match="missing chunk"): + merge_one("marginal_edep", run_dir) + + +def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path): + (tmp_path / "a").mkdir() + (tmp_path / "b").mkdir() + unchunked_dir = _prep(_write_inputs(tmp_path / "a")) + compute_one("marginal_step_length", unchunked_dir) + unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir)) + + chunked_dir = _prep(_write_inputs(tmp_path / "b"), chunks=2) + for k in range(2): + compute_one("marginal_step_length", chunked_dir, chunk_index=k) + chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir)) + + assert chunked.payload == unchunked.payload + + +def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path)) # n_chunks=1 (default) + with pytest.raises(ValueError, match="out of range"): + compute_one("marginal_edep", run_dir, chunk_index=1) def test_write_submit_description(tmp_path: Path): @@ -119,12 +163,15 @@ def test_write_submit_description(tmp_path: Path): assert "docker_image = mschnepf/slc7-condocker" in txt assert "requirements = TARGET.ProvidesETPResources" in txt assert "accounting_group = cms" in txt - assert "queue plotid from" in txt - assert (run_dir / "plotids.txt").read_text().split() == catalog_ids() + assert "queue plotid,chunk from" in txt + jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()] + assert [i for i, _ in jobs] == catalog_ids() + assert all(k == "0" for _, k in jobs) # n_chunks=1 default wrapper = run_dir / "run_compute.sh" assert wrapper.exists() and (wrapper.stat().st_mode & 0o111) body = wrapper.read_text() - assert "giant analyze compute-one --id" in body and "--run-dir" in body + assert "giant analyze compute-one --id" in body + assert "--chunk" in body and "--run-dir" in body def test_write_submit_remote_flag(tmp_path: Path): @@ -135,3 +182,18 @@ def test_write_submit_remote_flag(tmp_path: Path): txt = write_submit(cfg).read_text() assert "+RemoteJob = True" in txt assert "ProvidesETPResources" not in txt + + +def test_write_submit_chunks_respect_chunkable(tmp_path: Path): + assert get_spec("router_gating").chunkable is False + run_dir = _prep(_write_inputs(tmp_path)) + cfg = SubmitConfig( + run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4 + ) + write_submit(cfg) + jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()] + counts: dict[str, int] = {} + for spec_id, _ in jobs: + counts[spec_id] = counts.get(spec_id, 0) + 1 + assert counts["marginal_edep"] == 4 + assert counts["router_gating"] == 1 # chunkable=False, ignores n_chunks