lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Rewrite analysis module as a lean, fully-streaming pipeline

Browse Source 
Replace giant/analysis.py's dual numpy-SampleCollection + polars paths with a
single polars-streaming implementation that produces the validation notebook's
plots directly from a `giant predict --coord local` parquet, sized for files
larger than RAM.

- Drop the numpy SampleCollection path (load_predicted_local, marginal_table,
  correlation_matrices, direction_alignment, constraint_report, plot_kl_bars)
  and the rollout observables; the 5 remaining plotters now take a parquet
  path / LazyFrame and stream internally.
- Rewrite compute_event_observables_pl to aggregate in parallel streaming
  polars (post-pos reconstruction as expressions) instead of a serial
  pyarrow-batch + numpy loop, fixing a pre-existing OOM (holistic median +
  323M-row join in the bin-edge sizing). Medians are approximated from a
  streaming log-bin histogram with within-bin interpolation.
- Keep every full-file scan narrow (few columns): on a file larger than RAM,
  peak mmap memory, not scan count, is the binding constraint. Marginals run
  one dim at a time (~15GB peak) rather than a combined all-dims pass (OOM).
- Update analysis/validation.ipynb to the path-based API; delete the
  analysis/export_*.py and compare_ode_steps_*.py one-off scripts.
- Rewrite tests/test_analysis.py around parquet fixtures with an inline numpy
  oracle; add correlation/streaming-plotter and approx-median coverage.

Verified end-to-end on the 32GB predict file: full notebook completes at
~25GB peak (no OOM); event rollup runs at ~13 cores.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-10 14:47:46 +02:00
parent 740ebdf6b6
commit 55c676fb9b
13 changed files with 1298 additions and 2284 deletions
Download Patch File Download Diff File Expand all files Collapse all files
analysis/compare_ode_steps_energy_conservation.py
-167
View File
@@ -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
-69
View File
@@ -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
-105
View File
@@ -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
-80
View File
@@ -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
-76
View File
@@ -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
-72
View File
@@ -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
-63
View File
@@ -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
-145
View File
@@ -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
-48
View File
@@ -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
-80
View File
@@ -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")
analysis/validation.ipynb
+58 -41
View File
File diff suppressed because one or more lines are too long
giant/analysis.py
+913 -1095
View File
File diff suppressed because it is too large Load Diff
tests/test_analysis.py
+327 -243
View File
@@ -3,6 +3,7 @@ import matplotlib
matplotlib.use("Agg") # no display needed for plot smoke tests
import numpy as np
import pandas as pd
import polars as pl
import pyarrow as pa
import pyarrow.parquet as pq
@@ -10,20 +11,14 @@ import pytest
from giant.analysis import (
RAW_TARGET_NAMES,
SampleCollection,
compute_event_observables_pl,
constraint_report,
constraint_report_pl,
correlation_matrices,
direction_alignment,
load_predicted_local,
marginal_table,
correlation_matrices_pl,
marginal_table_pl,
pdg_contribution_table_pl,
plot_constraint_violations,
plot_correlation_matrices,
plot_direction_alignment,
plot_kl_bars,
plot_kl_bars_pl,
plot_longitudinal_profile,
plot_marginals,
@@ -48,146 +43,16 @@ from giant.data.transforms import (
reconstruct_post_pos,
)
# ---------------------------------------------------------------------------
# Fixtures: predict `--coord local` parquet writers
# ---------------------------------------------------------------------------
def _unit_vectors(rng, n):
v = rng.standard_normal((n, 3)).astype(np.float32)
return v / np.linalg.norm(v, axis=1, keepdims=True)
def _make_collection(n=200, seed=0, gen_offset=0.0) -> SampleCollection:
rng = np.random.default_rng(seed)
real = np.column_stack(
[
rng.uniform(0.1, 5.0, n), # step_length
rng.uniform(0.1, 5.0, n), # delta_e
rng.uniform(0.1, 5.0, n), # edep
_unit_vectors(rng, n), # post_dir
_unit_vectors(rng, n), # travel_dir
]
).astype(np.float32)
gen = real + gen_offset
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
cond_cont_raw = np.column_stack(
[
rng.standard_normal((n, 3)),
pre_E,
rng.standard_normal((n, 3)),
rng.integers(0, 5, n),
rng.integers(0, 3, n),
]
).astype(np.float32)
return SampleCollection(
cond_cont_raw=cond_cont_raw,
pdg=rng.choice([11, -11, 22], size=n),
material=rng.choice(["W", "Pb"], size=n),
real_raw=real,
gen_raw=gen,
)
def test_marginal_table_aggregate_has_all_dims():
table = marginal_table(_make_collection())
assert set(table["dim"]) == set(RAW_TARGET_NAMES)
assert (table["group"] == "all").all()
@pytest.mark.parametrize("group_by", ["pdg", "material", "energy"])
def test_marginal_table_grouped_covers_all_rows(group_by):
collection = _make_collection()
table = marginal_table(collection, group_by=group_by)
assert table["n"].groupby(table["group"]).first().sum() == len(collection.pdg)
def test_marginal_table_identical_distributions_have_zero_kl():
collection = _make_collection(gen_offset=0.0)
table = marginal_table(collection)
np.testing.assert_allclose(table["kl_real_gen"], 0.0, atol=1e-6)
def test_marginal_table_shifted_distribution_has_positive_kl():
collection = _make_collection(gen_offset=3.0)
table = marginal_table(collection)
assert (table["kl_real_gen"] > 0).all()
def test_correlation_matrices_are_symmetric_unit_diagonal():
real_corr, gen_corr = correlation_matrices(_make_collection())
for corr in (real_corr, gen_corr):
np.testing.assert_allclose(np.diag(corr), 1.0, atol=1e-5)
np.testing.assert_allclose(corr, corr.T, atol=1e-5)
def test_direction_alignment_real_data_is_unit_norm_dot_product():
real_cos, gen_cos = direction_alignment(_make_collection())
assert np.all(real_cos >= -1.0 - 1e-5) and np.all(real_cos <= 1.0 + 1e-5)
assert np.all(gen_cos >= -1.0 - 1e-5) and np.all(gen_cos <= 1.0 + 1e-5)
def test_constraint_report_clean_data_has_no_violations():
report = constraint_report(_make_collection(gen_offset=0.0))
assert (report["violation_rate"] == 0.0).all()
def test_constraint_report_flags_negative_log_dims_and_bad_norms():
collection = _make_collection(gen_offset=0.0)
collection.gen_raw[:, 0] = -1.0 # negative step_length
collection.gen_raw[:, 3:6] *= 2.0 # post_dir no longer unit norm
report = constraint_report(collection)
violations = dict(zip(report["check"], report["violation_rate"]))
assert violations["step_length >= 0"] == 1.0
assert violations["post_dir unit norm"] == 1.0
def test_plot_marginals_runs_without_error():
fig = plot_marginals(_make_collection())
assert fig is not None
def test_plot_marginals_grouped_runs_without_error():
fig = plot_marginals(_make_collection(), group_by="material")
assert fig is not None
def test_plot_correlation_matrices_runs_without_error():
fig = plot_correlation_matrices(_make_collection())
assert fig is not None
def test_plot_pairwise_runs_without_error():
fig = plot_pairwise(_make_collection())
assert fig is not None
def test_plot_direction_alignment_runs_without_error():
fig = plot_direction_alignment(_make_collection())
assert fig is not None
def test_plot_constraint_violations_runs_without_error():
fig = plot_constraint_violations(_make_collection())
assert fig is not None
def test_plot_kl_bars_runs_without_error():
fig = plot_kl_bars(_make_collection())
assert fig is not None
def test_plot_kl_bars_grouped_runs_without_error():
fig = plot_kl_bars(_make_collection(), group_by="material")
assert fig is not None
def test_plot_kl_bars_caps_groups_by_pdg():
collection = _make_collection(n=600)
collection.pdg = np.arange(600) % 8 # 8 distinct pdg values, > max_groups
fig = plot_kl_bars(collection, group_by="pdg", max_groups=3)
ax = fig.axes[0]
assert len({line.get_label() for line in ax.containers}) <= 3
def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
"""Mimic `giant predict --coord local`'s output schema for the loader tests.
@@ -231,81 +96,6 @@ def _write_predicted_local_parquet(path, n=50, metadata=None, rng=None):
return true_log_local, pred_log_local, pre_E
def test_load_predicted_local_round_trips_values(tmp_path):
path = tmp_path / "predicted_local.parquet"
true_log_local, pred_log_local, pre_E = _write_predicted_local_parquet(
path,
metadata={
PREDICT_COORD_METADATA_KEY: "local",
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
},
)
collection = load_predicted_local(path)
def expected_raw(log_local):
raw = log_local.copy()
raw[:, 0] = np.exp(log_local[:, 0]) - 1e-8
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(log_local[:, 1:3], pre_E)
raw[:, 1] = delta_e
raw[:, 2] = edep
return raw
np.testing.assert_allclose(
collection.real_raw, expected_raw(true_log_local), atol=1e-4
)
np.testing.assert_allclose(
collection.gen_raw, expected_raw(pred_log_local), atol=1e-4
)
def test_load_predicted_local_usable_by_downstream_plots(tmp_path):
path = tmp_path / "predicted_local.parquet"
_write_predicted_local_parquet(
path,
metadata={
PREDICT_COORD_METADATA_KEY: "local",
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
},
)
collection = load_predicted_local(path)
assert marginal_table(collection) is not None
assert plot_marginals(collection) is not None
def test_load_predicted_local_rejects_missing_metadata(tmp_path):
path = tmp_path / "no_metadata.parquet"
_write_predicted_local_parquet(path, metadata=None)
with pytest.raises(ValueError, match="no '.*' parquet metadata"):
load_predicted_local(path)
def test_load_predicted_local_rejects_global_coord(tmp_path):
path = tmp_path / "global.parquet"
_write_predicted_local_parquet(
path,
metadata={
PREDICT_COORD_METADATA_KEY: "global",
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
},
)
with pytest.raises(ValueError, match="coord=local"):
load_predicted_local(path)
def test_load_predicted_local_rejects_mismatched_schema_version(tmp_path):
path = tmp_path / "old_version.parquet"
_write_predicted_local_parquet(
path,
metadata={
PREDICT_COORD_METADATA_KEY: "local",
PREDICT_SCHEMA_VERSION_KEY: "999",
},
)
with pytest.raises(ValueError, match="schema version"):
load_predicted_local(path)
def _predicted_local_path(tmp_path, n=200, seed=0):
path = tmp_path / "predicted_local.parquet"
_write_predicted_local_parquet(
@@ -320,14 +110,133 @@ def _predicted_local_path(tmp_path, n=200, seed=0):
return path
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
path = _predicted_local_path(tmp_path)
collection = load_predicted_local(path)
# ---------------------------------------------------------------------------
# Inline numpy oracle: the ground truth the streaming functions are checked
# against. This is the raw-space decoding the old `SampleCollection` path did,
# recomputed directly from the small fixture rather than in the module.
# ---------------------------------------------------------------------------
expected = marginal_table(collection, group_by=group_by).sort_values(
["group", "dim"]
def _raw_from_parquet(path):
"""(real_raw, gen_raw, pdg, material, pre_E) in physical units, via numpy."""
df = pl.read_parquet(path)
pre_E = df["pre_E"].to_numpy().astype(np.float32)
def decode(prefix):
log_local = np.column_stack(
[df[f"{prefix}_{name}"].to_numpy() for name in LOCAL_TARGET_NAMES]
).astype(np.float32)
raw = log_local.copy()
raw[:, 0] = inv_log_transform(log_local[:, 0])
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(log_local[:, 1:3], pre_E)
raw[:, 1] = delta_e
raw[:, 2] = edep
return raw
return (
decode("true"),
decode("pred"),
df["pdg"].to_numpy(),
df["material"].to_numpy(),
pre_E,
)
def _histogram_kl(p_samples, q_samples, bins=50, eps=1e-8):
"""KL(P || Q) between two 1D samples via a shared histogram (numpy oracle)."""
lo = min(p_samples.min(), q_samples.min())
hi = max(p_samples.max(), q_samples.max())
if hi <= lo:
return 0.0
edges = np.linspace(lo, hi, bins + 1)
p_hist, _ = np.histogram(p_samples, bins=edges)
q_hist, _ = np.histogram(q_samples, bins=edges)
p = p_hist.astype(np.float64) + eps
q = q_hist.astype(np.float64) + eps
p /= p.sum()
q /= q.sum()
return float(np.sum(p * np.log(p / q)))
def _group_masks(pdg, material, pre_E, group_by, n_energy_bins=4):
"""Replicate the module's `_group` labelling so oracle labels line up."""
n = len(pdg)
if group_by is None:
return [("all", np.ones(n, dtype=bool))]
if group_by == "pdg":
return [(f"pdg={int(v)}", pdg == v) for v in np.unique(pdg)]
if group_by == "material":
return [(f"material={v}", material == v) for v in np.unique(material)]
if group_by == "energy":
edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
edges[-1] += 1e-6
bin_idx = np.digitize(pre_E, edges[1:-1])
return [
(f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", bin_idx == i)
for i in range(n_energy_bins)
]
raise ValueError(group_by)
def _numpy_marginal_table(real, gen, pdg, material, pre_E, group_by, bins=50):
rows = []
for label, mask in _group_masks(pdg, material, pre_E, group_by):
if mask.sum() < 2:
continue
r, g = real[mask], gen[mask]
for j, name in enumerate(RAW_TARGET_NAMES):
rows.append(
{
"group": label,
"dim": name,
"n": int(mask.sum()),
"real_mean": r[:, j].mean(),
"gen_mean": g[:, j].mean(),
"real_std": r[:, j].std(),
"gen_std": g[:, j].std(),
"kl_real_gen": _histogram_kl(r[:, j], g[:, j], bins=bins),
}
)
return pd.DataFrame(rows).sort_values(["group", "dim"]).reset_index(drop=True)
def _numpy_constraint_report(gen, norm_tol=0.05):
post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)
rows = [
{
"check": "post_dir unit norm",
"violation_rate": float(np.mean(np.abs(post_norm - 1) > norm_tol)),
"mean_abs_error": float(np.mean(np.abs(post_norm - 1))),
},
{
"check": "travel_dir unit norm",
"violation_rate": float(np.mean(np.abs(travel_norm - 1) > norm_tol)),
"mean_abs_error": float(np.mean(np.abs(travel_norm - 1))),
},
]
for j, name in enumerate(RAW_TARGET_NAMES[:3]):
rows.append(
{
"check": f"{name} >= 0",
"violation_rate": float(np.mean(gen[:, j] < 0)),
"mean_abs_error": float(np.mean(np.clip(-gen[:, j], 0, None))),
}
)
return pd.DataFrame(rows)
# ---------------------------------------------------------------------------
# Tier 1: stratified marginals
# ---------------------------------------------------------------------------
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_marginal_table_pl_matches_numpy_oracle(tmp_path, group_by):
path = _predicted_local_path(tmp_path)
real, gen, pdg, material, pre_E = _raw_from_parquet(path)
expected = _numpy_marginal_table(real, gen, pdg, material, pre_E, group_by)
actual = (
marginal_table_pl(path, group_by=group_by).sort(["group", "dim"]).to_pandas()
)
@@ -336,26 +245,29 @@ def test_marginal_table_pl_matches_numpy_version(tmp_path, group_by):
assert list(expected["n"]) == list(actual["n"])
for col in ["real_mean", "gen_mean", "real_std", "gen_std"]:
np.testing.assert_allclose(
expected[col].to_numpy(),
actual[col].to_numpy(),
atol=1e-4,
rtol=1e-4,
expected[col].to_numpy(), actual[col].to_numpy(), atol=1e-4, rtol=1e-4
)
# KL uses np.histogram (numpy path) vs polars Series.hist (lazy path); the two
# backends bin the boundary (min/max) sample differently, so allow a small
# absolute discrepancy rather than requiring bit-identical estimates.
# KL uses np.histogram (oracle) vs polars binning (lazy path); the two bin the
# boundary (min/max) sample differently, so allow a small absolute discrepancy
# rather than requiring bit-identical estimates.
np.testing.assert_allclose(
expected["kl_real_gen"].to_numpy(),
actual["kl_real_gen"].to_numpy(),
atol=2e-2,
expected["kl_real_gen"].to_numpy(), actual["kl_real_gen"].to_numpy(), atol=2e-2
)
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_plot_kl_bars_pl_runs_without_error(tmp_path, group_by):
def test_marginal_table_pl_aggregate_has_all_dims(tmp_path):
table = marginal_table_pl(_predicted_local_path(tmp_path))
assert set(table["dim"].to_list()) == set(RAW_TARGET_NAMES)
assert (table["group"] == "all").all()
def test_marginal_table_pl_accepts_lazyframe(tmp_path):
path = _predicted_local_path(tmp_path)
fig = plot_kl_bars_pl(path, group_by=group_by)
assert fig is not None
from_path = marginal_table_pl(path).sort(["group", "dim"])
from_lf = marginal_table_pl(pl.scan_parquet(path)).sort(["group", "dim"])
np.testing.assert_allclose(
from_lf["kl_real_gen"].to_numpy(), from_path["kl_real_gen"].to_numpy()
)
def test_marginal_table_pl_rejects_missing_metadata(tmp_path):
@@ -365,11 +277,105 @@ def test_marginal_table_pl_rejects_missing_metadata(tmp_path):
marginal_table_pl(path)
def test_constraint_report_pl_matches_numpy_version(tmp_path):
path = _predicted_local_path(tmp_path)
collection = load_predicted_local(path)
def test_marginal_table_pl_rejects_global_coord(tmp_path):
path = tmp_path / "global.parquet"
_write_predicted_local_parquet(
path,
metadata={
PREDICT_COORD_METADATA_KEY: "global",
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
},
)
with pytest.raises(ValueError, match="coord=local"):
marginal_table_pl(path)
expected = constraint_report(collection)
def test_marginal_table_pl_rejects_mismatched_schema_version(tmp_path):
path = tmp_path / "old_version.parquet"
_write_predicted_local_parquet(
path,
metadata={
PREDICT_COORD_METADATA_KEY: "local",
PREDICT_SCHEMA_VERSION_KEY: "999",
},
)
with pytest.raises(ValueError, match="schema version"):
marginal_table_pl(path)
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_plot_kl_bars_pl_runs_without_error(tmp_path, group_by):
fig = plot_kl_bars_pl(_predicted_local_path(tmp_path), group_by=group_by)
assert fig is not None
def test_plot_marginals_runs_without_error(tmp_path):
fig = plot_marginals(_predicted_local_path(tmp_path))
assert len(fig.axes) == len(RAW_TARGET_NAMES) # single "all" row
def test_plot_marginals_grouped_runs_without_error(tmp_path):
fig = plot_marginals(_predicted_local_path(tmp_path), group_by="material")
assert fig is not None
def test_plot_marginals_caps_groups(tmp_path):
path = _predicted_local_path(tmp_path, n=400)
df = pl.read_parquet(path).with_columns(
pl.Series("pdg", np.arange(400) % 8, dtype=pl.Int64)
)
fig = plot_marginals(df.lazy(), group_by="pdg", max_groups=3)
n_rows = len(fig.axes) // len(RAW_TARGET_NAMES)
assert n_rows <= 3
# ---------------------------------------------------------------------------
# Tier 2: joint structure
# ---------------------------------------------------------------------------
def test_correlation_matrices_pl_matches_numpy(tmp_path):
path = _predicted_local_path(tmp_path)
real, gen, *_ = _raw_from_parquet(path)
real_corr, gen_corr = correlation_matrices_pl(path)
np.testing.assert_allclose(
real_corr, np.corrcoef(real, rowvar=False), atol=1e-4, rtol=1e-4
)
np.testing.assert_allclose(
gen_corr, np.corrcoef(gen, rowvar=False), atol=1e-4, rtol=1e-4
)
for corr in (real_corr, gen_corr):
np.testing.assert_allclose(np.diag(corr), 1.0, atol=1e-6)
np.testing.assert_allclose(corr, corr.T, atol=1e-6)
def test_plot_correlation_matrices_runs_without_error(tmp_path):
fig = plot_correlation_matrices(_predicted_local_path(tmp_path))
assert len(fig.axes) >= 3 # real, generated, difference (+ colorbars)
def test_plot_pairwise_runs_without_error(tmp_path):
fig = plot_pairwise(_predicted_local_path(tmp_path), n_sample=100)
assert len(fig.axes) == 6 # 2 rows (real/gen) × 3 default pairs
def test_plot_direction_alignment_runs_without_error(tmp_path):
fig = plot_direction_alignment(_predicted_local_path(tmp_path))
assert fig is not None
# ---------------------------------------------------------------------------
# Tier 3: physical constraints
# ---------------------------------------------------------------------------
def test_constraint_report_pl_matches_numpy_oracle(tmp_path):
path = _predicted_local_path(tmp_path)
_real, gen, *_ = _raw_from_parquet(path)
expected = _numpy_constraint_report(gen)
actual = constraint_report_pl(path).to_pandas()
assert list(expected["check"]) == list(actual["check"])
@@ -385,6 +391,25 @@ def test_constraint_report_pl_matches_numpy_version(tmp_path):
)
def test_constraint_report_pl_flags_bad_direction_norms(tmp_path):
path = _predicted_local_path(tmp_path)
# Force the generated post_dir off the unit sphere for every row: the fixed
# (1, 1, 1) vector has norm √3 ≈ 1.73, well outside the tolerance.
df = pl.read_parquet(path).with_columns(
pl.lit(1.0).alias("pred_post_dx"),
pl.lit(1.0).alias("pred_post_dy"),
pl.lit(1.0).alias("pred_post_dz"),
)
report = constraint_report_pl(df.lazy()).to_pandas()
rate = dict(zip(report["check"], report["violation_rate"]))
assert rate["post_dir unit norm"] == 1.0
def test_plot_constraint_violations_runs_without_error(tmp_path):
fig = plot_constraint_violations(_predicted_local_path(tmp_path))
assert len(fig.axes) == 2 + 3 # 2 direction norms + 3 scalar dims
# ---------------------------------------------------------------------------
# Tier 4: event-level (shower) observables
# ---------------------------------------------------------------------------
@@ -394,8 +419,8 @@ def _make_event_level_arrays(rng):
"""3 events (3/2/4 steps), each with an unambiguous highest-pre_E row.
The forced max-pre_E rows (indices 1, 3, 7) fix a known shower axis/entry
point per event, so the expected event_table can be re-derived
independently in the test without depending on compute_event_observables_pl.
point per event, so the expected event_table can be re-derived independently
in the test without depending on compute_event_observables_pl.
"""
event_id = np.array([0, 0, 0, 1, 1, 2, 2, 2, 2], dtype=np.int64)
n = len(event_id)
@@ -471,7 +496,7 @@ def _expected_event_table(
entry_pos = pre_pos[entry_idx]
axis_dir = pre_dir[entry_idx]
def agg(log_local):
def agg(log_local, mask=mask, entry_pos=entry_pos, axis_dir=axis_dir):
step_length = inv_log_transform(log_local[mask, 0])
edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(
log_local[mask, 1:3], pre_E[mask]
@@ -584,6 +609,60 @@ def test_compute_event_observables_pl_accepts_lazyframe(tmp_path):
)
def test_compute_event_observables_pl_approx_median(tmp_path):
"""The streaming log-bin median approximates the true per-event median."""
rng = np.random.default_rng(11)
n = 4000 # one event, many steps → a well-defined median
true_log = rng.standard_normal((n, 9)).astype(np.float32)
true_log[:, 0] = log_transform(rng.uniform(0.1, 5.0, n).astype(np.float32))
for s in (3, 6):
v = rng.standard_normal((n, 3)).astype(np.float32)
true_log[:, s : s + 3] = v / np.linalg.norm(v, axis=1, keepdims=True)
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
pre = rng.standard_normal((n, 3)).astype(np.float32)
pdir = _unit_vectors(rng, n)
pre_E[0] = 1000.0 # entry step
path = tmp_path / "one_event.parquet"
table = pa.table(
{
"event_id": np.zeros(n, dtype=np.int64),
"pdg": rng.choice([11, 22], n),
"pre_x": pre[:, 0],
"pre_y": pre[:, 1],
"pre_z": pre[:, 2],
"pre_E": pre_E,
"pre_dx": pdir[:, 0],
"pre_dy": pdir[:, 1],
"pre_dz": pdir[:, 2],
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"n_sec": rng.integers(0, 3, n).astype(np.int32),
**{f"pred_{nm}": true_log[:, j] for j, nm in enumerate(LOCAL_TARGET_NAMES)},
**{f"true_{nm}": true_log[:, j] for j, nm in enumerate(LOCAL_TARGET_NAMES)},
}
).replace_schema_metadata(
{
PREDICT_COORD_METADATA_KEY: "local",
PREDICT_SCHEMA_VERSION_KEY: PREDICT_SCHEMA_VERSION,
}
)
pq.write_table(table, path)
true_edep = energy_simplex_decode(true_log[:, 1:3], pre_E)[0]
true_length = inv_log_transform(true_log[:, 0])
exact_edep_median = np.median(true_edep)
exact_length_median = np.median(true_length)
obs = compute_event_observables_pl(path)
approx_edep = obs.event_table["real_median_edep"][0]
approx_length = obs.event_table["real_median_length"][0]
# log-bin interpolation: expect within a few percent of the true median.
np.testing.assert_allclose(approx_edep, exact_edep_median, rtol=0.05)
np.testing.assert_allclose(approx_length, exact_length_median, rtol=0.05)
def test_event_level_plots_run_without_error(tmp_path):
path = tmp_path / "event_level.parquet"
_write_event_level_parquet(path)
@@ -596,6 +675,11 @@ def test_event_level_plots_run_without_error(tmp_path):
assert plot_shower_max_depth(obs) is not None
# ---------------------------------------------------------------------------
# Particle-species (pdg) contribution shares
# ---------------------------------------------------------------------------
def test_pdg_contribution_table_pl_matches_manual_sums(tmp_path):
path = tmp_path / "event_level.parquet"
_, _, _, pre_E, true_log_local, pred_log_local, pdg = _write_event_level_parquet(