Add energy-conservation PoC ODE-step comparison scripts

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>

analysis/compare_ode_steps_energy_conservation.py

@@ -0,0 +1,147 @@
+"""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}")
+disp_ratio = lambda r: (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

@@ -0,0 +1,71 @@
+"""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

@@ -0,0 +1,101 @@
+"""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")