Merge branch 'analysis-streaming-rewrite' into 'master'

Rewrite analysis module as a lean, fully-streaming pipeline

See merge request lbogner/giant!4

CLAUDE.md

@@ -50,7 +50,7 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
 
 **Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
 
-**Validation** (`giant/validate.py`): step-level marginal comparisons. Shower-level (rollout) observables live in `giant/analysis.py` (`compute_rollout_observables` + `plot_rollout_*`), fed by `giant rollout` output.
+**Validation** (`giant/validate.py`): step-level marginal comparisons. `giant/analysis.py` is a fully-streaming (lazy polars) diagnostics module, sized for predict/rollout files larger than RAM, with no in-memory `SampleCollection` and no full-array materialization. It covers one-step-ahead `giant predict --coord local` output (`compute_event_observables_pl` + `plot_total_energy`/`plot_longitudinal_profile`/etc. for shower-level observables, plus the marginal/correlation/constraint tiers) and, via the `RolloutVsTruth` source type, a full autoregressive `giant rollout` shower compared against held-out truth data (`compute_rollout_vs_truth_observables_pl` for shower-level observables, reusing the same plot functions) — see the module docstring.
 
 **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.
 

analysis/compare_ode_steps_energy_conservation.py

@@ -1,167 +0,0 @@
-"""Compare energy-conservation PoC at 10 vs 20 ODE steps.
-
-Runs the same event-level energy-budget analysis as
-`analysis/export_energy_conservation_poc.py` on both predict outputs and prints a
-side-by-side table. Also regenerates the two incident-energy comparison plots for
-the 20-step run (prefix `giant-energy-conservation-poc-ode20-`).
-"""
-
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import polars as pl
-
-from giant.analysis import _edep_pl, _hist_edges
-
-FILES = {
-    "10-step (baseline)": "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet",
-    "20-step": "/home/lars/Programming/giant/0e236919-65ae-4b9b-9957-d31a8211aca4.parquet",
-}
-OUT = Path("/home/lars/knowledge-base/meta/attachments")
-PREFIX20 = "giant-energy-conservation-poc-ode20"
-
-
-def per_event(file: str) -> dict:
-    pe = (
-        pl.scan_parquet(file)
-        .group_by("event_id")
-        .agg(
-            pl.col("pre_E").max().alias("primary_E"),
-            _edep_pl("true").sum().alias("real_total_edep"),
-            _edep_pl("pred").sum().alias("gen_total_edep"),
-        )
-        .collect(engine="streaming")
-    )
-    primary_E = pe["primary_E"].to_numpy()
-    assert np.unique(primary_E).size == 1, "expected a single fixed incident energy"
-    E0 = float(primary_E[0])
-    real = pe["real_total_edep"].to_numpy()
-    gen = pe["gen_total_edep"].to_numpy()
-    n_steps = pl.scan_parquet(file).select(pl.len()).collect(engine="streaming").item()
-    return {
-        "E0": E0,
-        "n_events": pe.height,
-        "n_rows": n_steps,
-        "real": real,
-        "gen": gen,
-    }
-
-
-results = {name: per_event(f) for name, f in FILES.items()}
-
-
-def fmt_row(label, fn):
-    cells = "  ".join(f"{fn(r):>14}" for r in results.values())
-    print(f"{label:<32}{cells}")
-
-
-print("=" * 80)
-header = "  ".join(f"{name:>14}" for name in results)
-print(f"{'metric':<32}{header}")
-print("-" * 80)
-fmt_row("n_events", lambda r: r["n_events"])
-fmt_row("n_rows (steps)", lambda r: r["n_rows"])
-fmt_row("E0 [MeV]", lambda r: f"{r['E0']:.1f}")
-print("-- REAL --")
-fmt_row("real mean [MeV]", lambda r: f"{r['real'].mean():.3f}")
-fmt_row("real std [MeV]", lambda r: f"{r['real'].std():.3f}")
-fmt_row("real sigma/mu", lambda r: f"{r['real'].std() / r['real'].mean():.4f}")
-print("-- GENERATED --")
-fmt_row("gen mean [MeV]", lambda r: f"{r['gen'].mean():.3f}")
-fmt_row("gen std [MeV]", lambda r: f"{r['gen'].std():.3f}")
-fmt_row("gen sigma/mu", lambda r: f"{r['gen'].std() / r['gen'].mean():.4f}")
-fmt_row("gen max [MeV]", lambda r: f"{r['gen'].max():.3f}")
-fmt_row("gen mean/E0", lambda r: f"{r['gen'].mean() / r['E0']:.4f}")
-fmt_row("gen max/E0", lambda r: f"{r['gen'].max() / r['E0']:.4f}")
-fmt_row("gen p99/E0", lambda r: f"{np.quantile(r['gen'], 0.99) / r['E0']:.4f}")
-fmt_row("frac events gen>E0", lambda r: f"{np.mean(r['gen'] > r['E0']):.4f}")
-
-
-def disp_ratio(r):
-    return (r["gen"].std() / r["gen"].mean()) / (r["real"].std() / r["real"].mean())
-
-
-fmt_row("dispersion ratio gen/real", lambda r: f"{disp_ratio(r):.2f}x")
-print("=" * 80)
-
-# --- plots for the 20-step run (mirror the baseline export) ---
-r = results["20-step"]
-E0, real_tot, gen_tot = r["E0"], r["real"], r["gen"]
-
-fig, ax = plt.subplots(figsize=(6, 4))
-edges = _hist_edges(real_tot, gen_tot, bins=50).tolist()
-ax.hist(
-    real_tot,
-    bins=edges,
-    density=True,
-    histtype="step",
-    label=f"real (σ/μ={real_tot.std() / real_tot.mean():.3f})",
-)
-ax.hist(
-    gen_tot,
-    bins=edges,
-    density=True,
-    histtype="step",
-    label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, "
-    f"{np.mean(gen_tot > E0):.1%} > E0)",
-)
-ax.axvline(
-    E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV"
-)
-ax.set_yscale("log")
-ax.set_xlabel("total deposited energy per event [MeV]")
-ax.set_title("20 ODE steps")
-ax.legend(fontsize=8)
-fig.tight_layout()
-fig.savefig(
-    OUT / f"{PREFIX20}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight"
-)
-
-fig, ax = plt.subplots(figsize=(6, 4))
-ratio_real = real_tot / E0
-ratio_gen = gen_tot / E0
-edges = _hist_edges(ratio_real, ratio_gen, bins=60).tolist()
-ax.hist(ratio_real, bins=edges, density=True, histtype="step", label="real")
-ax.hist(ratio_gen, bins=edges, density=True, histtype="step", label="generated")
-ax.axvline(1.0, color="k", linestyle="--", linewidth=1, label="conservation limit (=1)")
-ax.set_yscale("log")
-ax.set_xlabel("total deposited energy / incident energy, per event")
-ax.set_title("20 ODE steps")
-ax.legend(fontsize=8)
-fig.tight_layout()
-fig.savefig(OUT / f"{PREFIX20}-event-energy-ratio.png", dpi=150, bbox_inches="tight")
-
-# --- overlay: generated total-edep, 10 vs 20 steps, against real+E0 ---
-fig, ax = plt.subplots(figsize=(6, 4))
-all_arrays = [results["10-step (baseline)"]["real"]] + [
-    r2["gen"] for r2 in results.values()
-]
-edges = _hist_edges(*all_arrays, bins=60).tolist()
-ax.hist(
-    results["20-step"]["real"],
-    bins=edges,
-    density=True,
-    histtype="step",
-    color="k",
-    label="real",
-)
-for name, r2 in results.items():
-    ax.hist(
-        r2["gen"],
-        bins=edges,
-        density=True,
-        histtype="step",
-        label=f"gen {name} ({np.mean(r2['gen'] > r2['E0']):.1%} > E0)",
-    )
-ax.axvline(E0, color="gray", linestyle="--", linewidth=1, label=f"E0={E0:.0f} MeV")
-ax.set_yscale("log")
-ax.set_xlabel("total deposited energy per event [MeV]")
-ax.set_title("Generated event energy: 10 vs 20 ODE steps")
-ax.legend(fontsize=8)
-fig.tight_layout()
-fig.savefig(
-    OUT / f"{PREFIX20}-compare-event-total-energy.png", dpi=150, bbox_inches="tight"
-)
-
-print("DONE")

analysis/compare_ode_steps_kl.py

@@ -1,69 +0,0 @@
-"""Per-step marginal KL(real||gen) per target dim, 10 vs 20 ODE steps.
-
-Streaming histogram over shared bin edges (computed from the true distribution),
-so the two runs are directly comparable dim-by-dim.
-"""
-
-import numpy as np
-import polars as pl
-
-from giant.analysis import RAW_TARGET_NAMES, _kl_from_counts, _raw_dim_expr
-
-FILES = {
-    "10-step": "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet",
-    "20-step": "/home/lars/Programming/giant/0e236919-65ae-4b9b-9957-d31a8211aca4.parquet",
-}
-BINS = 100
-
-# Fixed shared edges from the true distribution (identical across both files), using
-# robust quantiles to avoid a few outliers dominating the range.
-base = FILES["10-step"]
-edges = {}
-for j, name in enumerate(RAW_TARGET_NAMES):
-    lo, hi = (
-        pl.scan_parquet(base)
-        .select(
-            _raw_dim_expr("true", j).quantile(0.001).alias("lo"),
-            _raw_dim_expr("true", j).quantile(0.999).alias("hi"),
-        )
-        .collect(engine="streaming")
-        .row(0)
-    )
-    if not (hi - lo > 1e-9):
-        lo, hi = lo - 0.5, hi + 0.5
-    edges[name] = np.linspace(lo, hi, BINS + 1)
-
-
-def counts(file, prefix, j, e):
-    vals = (
-        pl.scan_parquet(file)
-        .select(_raw_dim_expr(prefix, j).alias("v"))
-        .collect(engine="streaming")["v"]
-        .to_numpy()
-    )
-    c, _ = np.histogram(vals, bins=e)
-    return c
-
-
-kls = {name: {} for name in FILES}
-for j, name in enumerate(RAW_TARGET_NAMES):
-    e = edges[name]
-    real_c = counts(base, "true", j, e)  # identical true dist across files
-    for run, f in FILES.items():
-        gen_c = counts(f, "pred", j, e)
-        kls[run][name] = _kl_from_counts(real_c, gen_c)
-
-print(f"{'dim':<14}{'KL 10-step':>14}{'KL 20-step':>14}{'ratio 20/10':>14}")
-print("-" * 56)
-tot = {"10-step": 0.0, "20-step": 0.0}
-for name in RAW_TARGET_NAMES:
-    a, b = kls["10-step"][name], kls["20-step"][name]
-    tot["10-step"] += a
-    tot["20-step"] += b
-    print(f"{name:<14}{a:>14.5f}{b:>14.5f}{b / a if a else float('nan'):>14.2f}")
-print("-" * 56)
-print(
-    f"{'SUM':<14}{tot['10-step']:>14.5f}{tot['20-step']:>14.5f}"
-    f"{tot['20-step'] / tot['10-step']:>14.2f}"
-)
-print(f"{'MEAN':<14}{tot['10-step'] / 9:>14.5f}{tot['20-step'] / 9:>14.5f}")

analysis/export_energy_conservation_poc.py

@@ -1,105 +0,0 @@
-"""One-off export of the energy-budget-violation plot for the
-energy-conservation PoC checkpoint (ALR simplex output space, retrained on
-the patched-Geant4 regenerated dataset). Not part of the package; run
-manually.
-
-Unlike `plot_total_energy` (per-step conservation only, no reference to the
-fixed primary/incident energy), this adds an explicit comparison against
-`primary_E` (== max(pre_E) per event, since this PoC dataset uses a single
-fixed incident energy) to show event-level conservation violations that the
-per-step ALR simplex constraint does not prevent.
-"""
-
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import polars as pl
-
-from giant.analysis import _edep_pl, _hist_edges
-
-FILE = "/home/lars/Programming/giant/9879e806-5e88-4b06-b1fa-0e61de9cda6f.parquet"
-OUT = Path("/home/lars/knowledge-base/meta/attachments")
-PREFIX = "giant-energy-conservation-poc"
-
-print("=== per-event primary energy + total edep ===")
-per_event = (
-    pl.scan_parquet(FILE)
-    .group_by("event_id")
-    .agg(
-        pl.col("pre_E").max().alias("primary_E"),
-        _edep_pl("true").sum().alias("real_total_edep"),
-        _edep_pl("pred").sum().alias("gen_total_edep"),
-    )
-    .collect(engine="streaming")
-)
-print("n events:", per_event.height)
-
-primary_E = per_event["primary_E"].to_numpy()
-real_tot = per_event["real_total_edep"].to_numpy()
-gen_tot = per_event["gen_total_edep"].to_numpy()
-
-assert np.unique(primary_E).size == 1, "expected a single fixed incident energy"
-E0 = float(primary_E[0])
-print(f"fixed incident energy E0 = {E0} MeV")
-
-print()
-print(
-    f"real:  mean={real_tot.mean():.3f} std={real_tot.std():.3f} "
-    f"sigma/mu={real_tot.std() / real_tot.mean():.4f} max={real_tot.max():.3f} "
-    f"frac>E0={np.mean(real_tot > E0):.4f}"
-)
-print(
-    f"gen:   mean={gen_tot.mean():.3f} std={gen_tot.std():.3f} "
-    f"sigma/mu={gen_tot.std() / gen_tot.mean():.4f} max={gen_tot.max():.3f} "
-    f"frac>E0={np.mean(gen_tot > E0):.4f}"
-)
-print(
-    f"gen/E0 ratio: mean={np.mean(gen_tot / E0):.4f} max={np.max(gen_tot / E0):.4f} "
-    f"p99={np.quantile(gen_tot / E0, 0.99):.4f}"
-)
-
-print("=== plot: total edep per event, marked against incident energy ===")
-fig, ax = plt.subplots(figsize=(6, 4))
-edges = _hist_edges(real_tot, gen_tot, bins=50).tolist()
-ax.hist(
-    real_tot,
-    bins=edges,
-    density=True,
-    histtype="step",
-    label=f"real (σ/μ={real_tot.std() / real_tot.mean():.3f})",
-)
-ax.hist(
-    gen_tot,
-    bins=edges,
-    density=True,
-    histtype="step",
-    label=f"generated (σ/μ={gen_tot.std() / gen_tot.mean():.3f}, "
-    f"{np.mean(gen_tot > E0):.1%} > E0)",
-)
-ax.axvline(
-    E0, color="k", linestyle="--", linewidth=1, label=f"incident energy E0={E0:.0f} MeV"
-)
-ax.set_yscale("log")
-ax.set_xlabel("total deposited energy per event [MeV]")
-ax.legend(fontsize=8)
-fig.tight_layout()
-fig.savefig(
-    OUT / f"{PREFIX}-event-total-energy-vs-E0.png", dpi=150, bbox_inches="tight"
-)
-
-print("=== plot: total edep / incident energy ratio ===")
-fig, ax = plt.subplots(figsize=(6, 4))
-ratio_real = real_tot / E0
-ratio_gen = gen_tot / E0
-edges = _hist_edges(ratio_real, ratio_gen, bins=60).tolist()
-ax.hist(ratio_real, bins=edges, density=True, histtype="step", label="real")
-ax.hist(ratio_gen, bins=edges, density=True, histtype="step", label="generated")
-ax.axvline(1.0, color="k", linestyle="--", linewidth=1, label="conservation limit (=1)")
-ax.set_yscale("log")
-ax.set_xlabel("total deposited energy / incident energy, per event")
-ax.legend(fontsize=8)
-fig.tight_layout()
-fig.savefig(OUT / f"{PREFIX}-event-energy-ratio.png", dpi=150, bbox_inches="tight")
-
-print("DONE")

analysis/export_event_observables.py

@@ -1,80 +0,0 @@
-"""One-off export of Tier 4 event-level/pdg-share plots for
-checkpoints/scan/h1024_n8_d0.1_lr0.0003/best.pt into the knowledge-base
-attachments folder. Not part of the package; run manually."""
-
-from pathlib import Path
-
-import giant.analysis as a
-
-FILE = "/home/lars/Programming/giant/pbwo4_10k_9_predicted_local.parquet"
-OUT = Path("/home/lars/knowledge-base/meta/attachments")
-PREFIX = "giant-h1024n8d0.1lr3e-4"
-
-print("=== computing event observables ===")
-obs = a.compute_event_observables_pl(FILE)
-table = obs.event_table
-print("n events:", table.height)
-
-print("=== total energy / total length ===")
-fig = a.plot_total_energy(obs)
-fig.savefig(OUT / f"{PREFIX}-event-total-energy.png", dpi=150, bbox_inches="tight")
-fig = a.plot_total_length(obs)
-fig.savefig(OUT / f"{PREFIX}-event-total-length.png", dpi=150, bbox_inches="tight")
-
-print("=== mean/median energy & length per step ===")
-fig = a.plot_mean_energy_per_step(obs)
-fig.savefig(
-    OUT / f"{PREFIX}-event-mean-energy-per-step.png", dpi=150, bbox_inches="tight"
-)
-fig = a.plot_mean_length_per_step(obs)
-fig.savefig(
-    OUT / f"{PREFIX}-event-mean-length-per-step.png", dpi=150, bbox_inches="tight"
-)
-
-print("=== longitudinal / transverse profiles ===")
-fig = a.plot_longitudinal_profile(obs)
-fig.savefig(
-    OUT / f"{PREFIX}-event-longitudinal-profile.png", dpi=150, bbox_inches="tight"
-)
-fig = a.plot_transverse_profile(obs)
-fig.savefig(
-    OUT / f"{PREFIX}-event-transverse-profile.png", dpi=150, bbox_inches="tight"
-)
-
-print("=== shower-max depth ===")
-fig = a.plot_shower_max_depth(obs)
-fig.savefig(OUT / f"{PREFIX}-event-shower-max-depth.png", dpi=150, bbox_inches="tight")
-
-print("=== pdg contribution shares ===")
-pdg_table = a.pdg_contribution_table_pl(FILE)
-fig = a.plot_pdg_energy_share(pdg_table)
-fig.savefig(OUT / f"{PREFIX}-pdg-energy-share.png", dpi=150, bbox_inches="tight")
-fig = a.plot_pdg_length_share(pdg_table)
-fig.savefig(OUT / f"{PREFIX}-pdg-length-share.png", dpi=150, bbox_inches="tight")
-
-print("=== summary stats ===")
-for label, real_col, gen_col in [
-    ("total_edep", "real_total_edep", "gen_total_edep"),
-    ("total_length", "real_total_length", "gen_total_length"),
-    ("mean_edep", "real_mean_edep", "gen_mean_edep"),
-    ("mean_length", "real_mean_length", "gen_mean_length"),
-    ("median_edep", "real_median_edep", "gen_median_edep"),
-    ("median_length", "real_median_length", "gen_median_length"),
-    ("centroid_depth", "real_centroid_depth", "gen_centroid_depth"),
-    ("transverse_rms", "real_transverse_rms", "gen_transverse_rms"),
-    ("max_depth", "real_max_depth", "gen_max_depth"),
-]:
-    real = table[real_col].to_numpy()
-    gen = table[gen_col].to_numpy()
-    print(
-        f"{label}: real mean={real.mean():.4g} std={real.std():.4g} sigma/mu={real.std() / real.mean():.4f} | "
-        f"gen mean={gen.mean():.4g} std={gen.std():.4g} sigma/mu={gen.std() / gen.mean():.4f} | "
-        f"mean_diff%={100 * (gen.mean() - real.mean()) / real.mean():.2f}"
-    )
-
-print("n_steps per event: mean", table["n_steps"].to_numpy().mean())
-
-print("=== pdg shares table ===")
-print(pdg_table.to_pandas().to_string())
-
-print("DONE")

analysis/export_photon_edep.py

@@ -1,76 +0,0 @@
-"""Zoomed-in real-vs-generated edep histogram for photons only (pdg=22),
-to characterize the KL spike flagged by plot_kl_bars_pl(group_by='pdg')."""
-
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-import giant.analysis as a
-
-FILE = "/home/lars/Programming/giant/pbwo4_10k_9_predicted_local.parquet"
-OUT = Path("/home/lars/knowledge-base/meta/attachments")
-
-samples = a.load_predicted_local(FILE, sample_frac=0.15)
-mask = samples.pdg == 22
-edep_idx = a.RAW_TARGET_NAMES.index("edep")
-real = samples.real_raw[mask, edep_idx]
-gen = samples.gen_raw[mask, edep_idx]
-print("n photon rows:", mask.sum())
-print(
-    "real: mean",
-    real.mean(),
-    "std",
-    real.std(),
-    "max",
-    real.max(),
-    "frac==0",
-    (real == 0).mean(),
-)
-print(
-    "gen:  mean",
-    gen.mean(),
-    "std",
-    gen.std(),
-    "max",
-    gen.max(),
-    "frac==0",
-    (gen == 0).mean(),
-)
-for q in [0.5, 0.9, 0.99, 0.999]:
-    print(f"q={q}: real={np.quantile(real, q):.4f} gen={np.quantile(gen, q):.4f}")
-
-fig, axes = plt.subplots(1, 2, figsize=(10, 4))
-bins = np.linspace(0, np.quantile(real, 0.999), 80)
-axes[0].hist(real, bins=bins, alpha=0.6, label="real", density=True)
-axes[0].hist(gen, bins=bins, alpha=0.6, label="gen", density=True)
-axes[0].set_yscale("log")
-axes[0].set_xlabel("edep (photons, pdg=22)")
-axes[0].legend()
-
-axes[1].hist(
-    real,
-    bins=bins,
-    alpha=0.6,
-    label="real",
-    density=True,
-    cumulative=True,
-    histtype="step",
-)
-axes[1].hist(
-    gen,
-    bins=bins,
-    alpha=0.6,
-    label="gen",
-    density=True,
-    cumulative=True,
-    histtype="step",
-)
-axes[1].set_xlabel("edep (photons, pdg=22) - CDF")
-axes[1].legend()
-
-fig.tight_layout()
-fig.savefig(
-    OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-zoom.png", dpi=150, bbox_inches="tight"
-)
-print("saved photon-edep-zoom")

analysis/export_photon_edep_by_process.py

@@ -1,72 +0,0 @@
-"""Truth-only histogram of deposited photon (pdg=22) energy in eV, grouped by
-the Geant4 physics process that ended the step. Step-type histograms, one per
-process, in both linear and log energy space.
-
-Data is the raw miniCaloSim steps parquet (truth), loaded lazily with polars —
-only the photon rows and the (process, edep) columns are materialized.
-"""
-
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-import polars as pl
-
-FILE = "/home/lars/Programming/minicalo-data-exploration/pbwo4_10000events_hits.parquet"
-OUT = Path("/home/lars/knowledge-base/meta/attachments")
-
-MEV_TO_EV = 1e6
-
-# Lazy load: keep only photon steps and the two columns we need.
-df = (
-    pl.scan_parquet(FILE)
-    .filter(pl.col("pdg") == 22)
-    .select(
-        pl.col("process"),
-        (pl.col("edep") * MEV_TO_EV).alias("edep_ev"),
-    )
-    .collect()
-)
-print("n photon rows:", df.height)
-
-# Process order by abundance, so the legend is stable and the busiest on top.
-processes = (
-    df.group_by("process")
-    .len()
-    .sort("len", descending=True)
-    .get_column("process")
-    .to_list()
-)
-series = {
-    p: df.filter(pl.col("process") == p).get_column("edep_ev").to_numpy()
-    for p in processes
-}
-
-all_edep = df.get_column("edep_ev").to_numpy()
-lin_bins = np.linspace(0, np.quantile(all_edep, 0.999), 80)
-pos = all_edep[all_edep > 0]
-log_bins = np.geomspace(max(pos.min(), 1e-3), pos.max(), 80)
-
-
-def plot(bins, xscale: str, fname: str) -> None:
-    fig, ax = plt.subplots(figsize=(7, 4.5))
-    lo, hi = bins[0], bins[-1]
-    for p in processes:
-        v = series[p]
-        # Skip processes with nothing inside the bin range (e.g. all-zero edep
-        # processes on the log axis), which would add phantom legend entries.
-        if not np.any((v >= lo) & (v <= hi)):
-            continue
-        ax.hist(v, bins=bins, histtype="step", density=True, label=f"{p} (n={v.size})")
-    ax.set_xscale(xscale)
-    ax.set_yscale("log")
-    ax.set_xlabel("deposited energy [eV] (photons, pdg=22)")
-    ax.set_ylabel("density")
-    ax.legend(fontsize=8, title="process")
-    fig.tight_layout()
-    fig.savefig(OUT / fname, dpi=150, bbox_inches="tight")
-    print("saved", fname)
-
-
-plot(lin_bins, "linear", "giant-photon-edep-by-process-ev-lin.png")
-plot(log_bins, "log", "giant-photon-edep-by-process-ev-log.png")

analysis/export_photon_edep_ev.py

@@ -1,63 +0,0 @@
-"""Recreate of the photon (pdg=22) edep histogram in transformed (local-frame)
-coordinates, with the deposited-energy axis in electron volts.
-
-Same data path as export_photon_edep.py, but edep is converted from the raw
-MeV units to eV (x1e6) so the small-deposition photon spike is readable.
-"""
-
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-import giant.analysis as a
-
-FILE = "/home/lars/Programming/giant/pbwo4_10k_9_predicted_local.parquet"
-OUT = Path("/home/lars/knowledge-base/meta/attachments")
-
-MEV_TO_EV = 1e6
-
-samples = a.load_predicted_local(FILE, sample_frac=0.15)
-mask = samples.pdg == 22
-edep_idx = a.RAW_TARGET_NAMES.index("edep")
-real = samples.real_raw[mask, edep_idx] * MEV_TO_EV  # MeV -> eV
-gen = samples.gen_raw[mask, edep_idx] * MEV_TO_EV
-print("n photon rows:", mask.sum())
-print("real [eV]: mean", real.mean(), "max", real.max(), "frac==0", (real == 0).mean())
-print("gen  [eV]: mean", gen.mean(), "max", gen.max(), "frac==0", (gen == 0).mean())
-
-fig, ax = plt.subplots(figsize=(6, 4))
-bins = np.linspace(0, np.quantile(real, 0.999), 80)
-ax.hist(real, bins=bins, alpha=0.6, label="real", density=True)
-ax.hist(gen, bins=bins, alpha=0.6, label="gen", density=True)
-ax.set_yscale("log")
-ax.set_xlabel("deposited energy [eV] (photons, pdg=22)")
-ax.set_ylabel("density")
-ax.legend()
-
-fig.tight_layout()
-fig.savefig(
-    OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-ev.png", dpi=150, bbox_inches="tight"
-)
-print("saved photon-edep-ev")
-
-# Second version: log-spaced energy axis to expose the low-deposition structure.
-fig, ax = plt.subplots(figsize=(6, 4))
-lo = max(min(real[real > 0].min(), gen[gen > 0].min()), 1.0)
-hi = max(real.max(), gen.max())
-log_bins = np.geomspace(lo, hi, 80)
-ax.hist(real, bins=log_bins, alpha=0.6, label="real", density=True)
-ax.hist(gen, bins=log_bins, alpha=0.6, label="gen", density=True)
-ax.set_xscale("log")
-ax.set_yscale("log")
-ax.set_xlabel("deposited energy [eV] (photons, pdg=22)")
-ax.set_ylabel("density")
-ax.legend()
-
-fig.tight_layout()
-fig.savefig(
-    OUT / "giant-h1024n8d0.1lr3e-4-photon-edep-ev-logx.png",
-    dpi=150,
-    bbox_inches="tight",
-)
-print("saved photon-edep-ev-logx")

analysis/export_presentation_plots.py

@@ -1,145 +0,0 @@
-"""One-off export of presentation-specific plots for the 2026-06-24 ETP group
-update, written directly into the presentation's images/ folder. Not part of
-the package; run manually.
-
-All plots are saved as PDF (vector) except the pairwise scatter plot, which
-stays PNG/raster in the .tex (scatter plots with thousands of points blow up
-as vector files and gain nothing from being scalable).
-
-1. Aggregate marginals, split into a 3x3 grid (instead of plot_marginals'
-   single wide row of 9 columns) for a more manageable slide aspect ratio.
-2. post_dir/travel_dir unit-norm histograms only (subset of
-   plot_constraint_violations).
-3. KL bar plot stratified by pdg.
-4. Zoomed-in photon edep histogram + CDF.
-5. Event-level profiles/histograms/pdg-share plots.
-"""
-
-from pathlib import Path
-
-import matplotlib.pyplot as plt
-import numpy as np
-
-import giant.analysis as a
-from giant.analysis import RAW_TARGET_NAMES, _hist_edges
-
-FILE = "/home/lars/Programming/giant/pbwo4_10k_9_predicted_local.parquet"
-OUT = Path(
-    "/home/lars/Programming/thesis-presentations/presentations/2026-06-24-group-update/images"
-)
-PREFIX = "giant-h1024n8d0.1lr3e-4"
-
-
-def savefig(fig, name: str) -> None:
-    fig.savefig(OUT / f"{PREFIX}-{name}.pdf", bbox_inches="tight")
-    print("saved", name)
-
-
-print("=== loading sampled SampleCollection ===")
-samples = a.load_predicted_local(FILE, sample_frac=0.15)
-print("n rows:", len(samples.gen_raw))
-
-print("=== marginals, split 3x3 grid ===")
-n_rows, n_cols = 3, 3
-fig, axes = plt.subplots(n_rows, n_cols, figsize=(3.6 * n_cols, 2.8 * n_rows))
-for idx, name in enumerate(RAW_TARGET_NAMES):
-    row, col = divmod(idx, n_cols)
-    ax = axes[row][col]
-    real, gen = samples.real_raw[:, idx], samples.gen_raw[:, idx]
-    edges = _hist_edges(real, gen, bins=50)
-    ax.hist(real, bins=edges, density=True, histtype="step", label="real")
-    ax.hist(gen, bins=edges, density=True, histtype="step", label="generated")
-    ax.set_yscale("log")
-    ax.set_title(name, fontsize=9)
-    if idx == 0:
-        ax.legend(fontsize=7)
-fig.tight_layout()
-savefig(fig, "marginals-grid")
-
-print("=== post_dir / travel_dir unit-norm histograms ===")
-gen = samples.gen_raw
-post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
-travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)
-
-fig, axes = plt.subplots(1, 2, figsize=(8, 3.5))
-for ax, norm, title in [
-    (axes[0], post_norm, r"$\|\mathrm{post\_dir}\|$"),
-    (axes[1], travel_norm, r"$\|\mathrm{travel\_dir}\|$"),
-]:
-    ax.hist(norm, bins=_hist_edges(norm, bins=50), histtype="step")
-    ax.set_yscale("log")
-    ax.axvline(1.0, color="k", linestyle="--", linewidth=1)
-    ax.set_title(title)
-fig.tight_layout()
-savefig(fig, "direction-norms")
-
-print("=== KL bars by pdg ===")
-fig = a.plot_kl_bars_pl(FILE, group_by="pdg")
-savefig(fig, "kl-bars-pdg")
-
-print("=== photon edep zoom ===")
-mask = samples.pdg == 22
-edep_idx = a.RAW_TARGET_NAMES.index("edep")
-real = samples.real_raw[mask, edep_idx]
-gen = samples.gen_raw[mask, edep_idx]
-
-fig, axes = plt.subplots(1, 2, figsize=(10, 4))
-bins = np.linspace(0, np.quantile(real, 0.999), 80)
-axes[0].hist(real, bins=bins, alpha=0.6, label="real", density=True)
-axes[0].hist(gen, bins=bins, alpha=0.6, label="gen", density=True)
-axes[0].set_yscale("log")
-axes[0].set_xlabel("edep (photons, pdg=22)")
-axes[0].legend()
-axes[1].hist(
-    real,
-    bins=bins,
-    alpha=0.6,
-    label="real",
-    density=True,
-    cumulative=True,
-    histtype="step",
-)
-axes[1].hist(
-    gen,
-    bins=bins,
-    alpha=0.6,
-    label="gen",
-    density=True,
-    cumulative=True,
-    histtype="step",
-)
-axes[1].set_xlabel("edep (photons, pdg=22) - CDF")
-axes[1].legend()
-fig.tight_layout()
-savefig(fig, "photon-edep-zoom")
-
-print("=== event observables ===")
-obs = a.compute_event_observables_pl(FILE)
-
-fig = a.plot_total_energy(obs)
-savefig(fig, "event-total-energy")
-
-fig = a.plot_total_length(obs)
-savefig(fig, "event-total-length")
-
-fig = a.plot_mean_energy_per_step(obs)
-savefig(fig, "event-mean-energy-per-step")
-
-fig = a.plot_mean_length_per_step(obs)
-savefig(fig, "event-mean-length-per-step")
-
-fig = a.plot_longitudinal_profile(obs)
-savefig(fig, "event-longitudinal-profile")
-
-fig = a.plot_transverse_profile(obs)
-savefig(fig, "event-transverse-profile")
-
-fig = a.plot_shower_max_depth(obs)
-savefig(fig, "event-shower-max-depth")
-
-print("=== pdg length share ===")
-pdg_table = a.pdg_contribution_table_pl(FILE)
-fig = a.plot_pdg_length_share(pdg_table)
-savefig(fig, "pdg-length-share")
-
-print("DONE")

analysis/export_rollout_observables.py

@@ -1,48 +0,0 @@
-"""Export shower-level plots from a `giant rollout` steps parquet.
-
-Not part of the package; run manually. Optionally overlays the real showers
-seeded from the same events (a `giant predict --coord local` file) by passing a
-reference path. Usage::
-
-    python analysis/export_rollout_observables.py ROLLOUT.parquet [REFERENCE_local.parquet]
-"""
-
-import sys
-from pathlib import Path
-
-import giant.analysis as a
-
-rollout_file = sys.argv[1] if len(sys.argv) > 1 else "rollout.parquet"
-reference_file = (
-    sys.argv[2] if len(sys.argv) > 2 and sys.argv[2] not in ("", "-") else None
-)
-OUT = Path(sys.argv[3]) if len(sys.argv) > 3 else Path(".")
-
-print(f"=== computing rollout observables: {rollout_file} ===")
-obs = a.compute_rollout_observables(rollout_file)
-tbl = obs.event_table
-print(f"n events: {len(tbl)}")
-print(
-    f"total_edep/event: mean={tbl['total_edep'].mean():.4g} MeV  "
-    f"leaked_E/event: mean={tbl['leaked_E'].mean():.4g} MeV  "
-    f"n_tracks/event: mean={tbl['n_tracks'].mean():.1f}  "
-    f"n_steps/event: mean={tbl['n_steps'].mean():.1f}"
-)
-
-reference = None
-if reference_file is not None:
-    print(f"=== computing real reference: {reference_file} ===")
-    reference = a.compute_event_observables_pl(reference_file)
-
-print("=== plots ===")
-for name, fn in [
-    ("longitudinal", a.plot_rollout_longitudinal),
-    ("transverse", a.plot_rollout_transverse),
-    ("total-energy", a.plot_rollout_total_energy),
-]:
-    fig = fn(obs, reference=reference)
-    path = OUT / f"rollout-{name}.png"
-    fig.savefig(path, dpi=150, bbox_inches="tight")
-    print(f"wrote {path}")
-
-print("DONE")

analysis/export_validation_plots.py

@@ -1,80 +0,0 @@
-"""One-off export of validation plots for checkpoints/scan/h1024_n8_d0.1_lr0.0003/best.pt
-into the knowledge-base attachments folder. Not part of the package; run manually."""
-
-from pathlib import Path
-
-import giant.analysis as a
-
-FILE = "/home/lars/Programming/giant/pbwo4_10k_9_predicted_local.parquet"
-OUT = Path("/home/lars/knowledge-base/meta/attachments")
-PREFIX = "giant-h1024n8d0.1lr3e-4"
-
-PDG_NAMES = {
-    11: "e-",
-    -11: "e+",
-    22: "gamma",
-    2112: "n",
-    2212: "p",
-}
-
-
-def pdg_label(code: int) -> str:
-    if code in PDG_NAMES:
-        return PDG_NAMES[code]
-    if code > 1000000000:
-        return f"ion{code}"
-    return str(code)
-
-
-print("=== KL bar plots (lazy, full dataset) ===")
-for grouping in [None, "energy", "pdg", "material"]:
-    fig = a.plot_kl_bars_pl(FILE, group_by=grouping)
-    name = f"{PREFIX}-kl-bars-{grouping or 'all'}.png"
-    fig.savefig(OUT / name, dpi=150, bbox_inches="tight")
-    print("saved", name)
-
-print("=== loading sampled SampleCollection ===")
-samples = a.load_predicted_local(FILE, sample_frac=0.15)
-print("n rows:", len(samples.gen_raw))
-
-print("=== marginals ===")
-fig = a.plot_marginals(samples)
-fig.savefig(OUT / f"{PREFIX}-marginals-all.png", dpi=150, bbox_inches="tight")
-
-fig = a.plot_marginals(samples, group_by="energy")
-fig.savefig(OUT / f"{PREFIX}-marginals-energy.png", dpi=150, bbox_inches="tight")
-
-fig = a.plot_marginals(samples, group_by="pdg")
-fig.savefig(OUT / f"{PREFIX}-marginals-pdg.png", dpi=150, bbox_inches="tight")
-print("saved marginals")
-
-print("=== correlation matrices ===")
-fig = a.plot_correlation_matrices(samples)
-fig.savefig(OUT / f"{PREFIX}-correlation.png", dpi=150, bbox_inches="tight")
-
-print("=== pairwise ===")
-fig = a.plot_pairwise(samples, n_sample=5000)
-fig.savefig(OUT / f"{PREFIX}-pairwise.png", dpi=150, bbox_inches="tight")
-
-print("=== direction alignment ===")
-fig = a.plot_direction_alignment(samples)
-fig.savefig(OUT / f"{PREFIX}-direction-alignment.png", dpi=150, bbox_inches="tight")
-
-print("=== constraint violations ===")
-fig = a.plot_constraint_violations(samples)
-fig.savefig(OUT / f"{PREFIX}-constraint-violations.png", dpi=150, bbox_inches="tight")
-
-print("=== marginal_table aggregate ===")
-agg = a.marginal_table(samples)
-print(agg.to_string())
-
-print("=== marginal_table by pdg (top rows incl. photon) ===")
-by_pdg = a.marginal_table(samples, group_by="pdg")
-by_pdg["particle"] = by_pdg["group"].astype(str)
-print(by_pdg.to_string())
-
-print("=== photon-only rows ===")
-photon_rows = by_pdg[by_pdg["group"].astype(str) == "pdg=22"]
-print(photon_rows.to_string())
-
-print("DONE")

giant/rollout.py

@@ -17,7 +17,7 @@ treated as detector leakage and not deposited.
 from __future__ import annotations
 
 from collections import Counter
-from typing import Callable
+from typing import Callable, TypedDict
 
 import numpy as np
 import torch
@@ -127,6 +127,13 @@ _RECORD_DTYPES: dict[str, type] = {
 }
 
 
+class RolloutSummary(TypedDict):
+    """`rollout()`'s return shape when streaming to `on_chunk` instead of materializing rows."""
+
+    n_rows: int
+    termination_reason_counts: dict[str, int]
+
+
 class _Recorder:
     """Accumulates per-step rows into column lists, materialised at the end —
     or, when `sink` is given, streams each non-empty chunk to it immediately
@@ -276,7 +283,7 @@ def rollout(
     max_tracks_per_event: int | None = None,
     escape_threshold: float | None = None,
     on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
-) -> dict[str, np.ndarray]:
+) -> dict[str, np.ndarray] | RolloutSummary:
     """Run showers to completion.
 
     By default, returns a step-record dict (see _RECORD_KEYS) with the whole