lars
1a22ae022c
Add analysis scripts for the 10-vs-20 flow-matching ODE-step ablation on the energy-conservation PoC predict outputs: - compare_ode_steps_energy_conservation.py: per-event energy-budget table + 20-step plots and the 10-vs-20 overlay. - compare_ode_steps_kl.py: per-step marginal KL(real||gen) per target dim over fixed shared bins, so the two runs are directly comparable dim-by-dim. Also commit export_energy_conservation_poc.py (the baseline event-level budget export) and repoint validation.ipynb at the PoC predict file at sample_frac=1.0. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
"""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")
|