lars
55c676fb9b
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>
1720 lines
66 KiB
Python
1720 lines
66 KiB
Python
"""Notebook diagnostics for a trained model's sample quality, fully streaming.
|
||
|
||
Every check here reads directly from a `giant predict --coord local` parquet
|
||
file (`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled)
|
||
and is computed with **lazy, streaming polars** so peak memory stays bounded no
|
||
matter how large the file is (production predict output runs to tens of GB).
|
||
There is no in-memory `SampleCollection`, no full-array materialization, and no
|
||
on-the-fly (checkpoint + live sampler) path: generate predictions once via the
|
||
CLI, then run every diagnostic below against that file.
|
||
|
||
Each public function takes `source: str | Path | pl.LazyFrame` — a path to the
|
||
predict parquet, or a pre-built `LazyFrame` with the same columns (for tests) —
|
||
and returns either a matplotlib figure or a small reduced table/dataclass.
|
||
|
||
Four tiers of checks::
|
||
|
||
from giant.analysis import plot_kl_bars_pl, plot_marginals
|
||
from giant.analysis import plot_correlation_matrices, plot_pairwise
|
||
from giant.analysis import plot_direction_alignment, plot_constraint_violations
|
||
|
||
FILE = "path/to/steps_predicted_local.parquet"
|
||
|
||
# Tier 1 — stratified marginals (per-dimension real-vs-generated, sliced by
|
||
# pdg / material / energy so failures hidden by the aggregate show)
|
||
for grouping in (None, "energy", "pdg", "material"):
|
||
plot_kl_bars_pl(FILE, group_by=grouping)
|
||
plot_marginals(FILE)
|
||
plot_marginals(FILE, group_by="energy")
|
||
|
||
# Tier 2 — joint structure (correlations + physically-coupled pairs + the
|
||
# post/travel direction alignment marginals can't see)
|
||
plot_correlation_matrices(FILE)
|
||
plot_pairwise(FILE)
|
||
plot_direction_alignment(FILE)
|
||
|
||
# Tier 3 — physical constraints (unit-norm directions, non-negative scalars;
|
||
# any violation is a pure generation artifact of the unconstrained MLP)
|
||
plot_constraint_violations(FILE)
|
||
|
||
For event-level (shower) observables — total deposited energy, total length,
|
||
longitudinal/transverse profiles, shower-max depth — aggregated per `event_id`
|
||
in the world frame with physical units (mm, MeV)::
|
||
|
||
from giant.analysis import compute_event_observables_pl
|
||
from giant.analysis import plot_total_energy, plot_total_length
|
||
from giant.analysis import plot_mean_energy_per_step, plot_mean_length_per_step
|
||
from giant.analysis import plot_longitudinal_profile, plot_transverse_profile
|
||
from giant.analysis import plot_shower_max_depth
|
||
|
||
obs = compute_event_observables_pl(FILE)
|
||
plot_total_energy(obs)
|
||
plot_total_length(obs)
|
||
plot_longitudinal_profile(obs)
|
||
plot_transverse_profile(obs)
|
||
plot_shower_max_depth(obs)
|
||
|
||
This re-aggregates one-step-ahead generations (each row generated conditioned on
|
||
the *real* preceding state) grouped by event — not a full autoregressive shower
|
||
rollout — so it won't surface covariate-shift failures that only appear under
|
||
true rollout, only how well one-step generation reconstructs aggregate shower
|
||
structure when fed real conditioning throughout.
|
||
|
||
For the dataset-wide breakdown of which particle species (pdg) contributed how
|
||
much of the total energy/length::
|
||
|
||
from giant.analysis import pdg_contribution_table_pl
|
||
from giant.analysis import plot_pdg_energy_share, plot_pdg_length_share
|
||
|
||
table = pdg_contribution_table_pl(FILE)
|
||
plot_pdg_energy_share(table)
|
||
plot_pdg_length_share(table)
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import dataclass
|
||
from pathlib import Path
|
||
|
||
import matplotlib.pyplot as plt
|
||
import numpy as np
|
||
import pandas as pd
|
||
import polars as pl
|
||
|
||
import pyarrow.parquet as pq
|
||
|
||
from giant.constants import (
|
||
LOCAL_TARGET_NAMES,
|
||
PREDICT_COORD_METADATA_KEY,
|
||
PREDICT_SCHEMA_VERSION,
|
||
PREDICT_SCHEMA_VERSION_KEY,
|
||
)
|
||
|
||
# The first 3 target dims are the scalar (non-direction) outputs. In raw/physical
|
||
# space they are step_length, delta_e, edep; in the model's native target space the
|
||
# first is log_step_length and the next two are the deposit/secondary ALR energy
|
||
# logits (see giant.constants.LOCAL_TARGET_NAMES and energy_simplex_encode).
|
||
_N_SCALAR_DIMS = 3
|
||
RAW_TARGET_NAMES = [
|
||
"step_length",
|
||
"delta_e",
|
||
"edep",
|
||
"post_dx",
|
||
"post_dy",
|
||
"post_dz",
|
||
"travel_dx",
|
||
"travel_dy",
|
||
"travel_dz",
|
||
]
|
||
_LOG_EPS = (
|
||
1e-8 # mirrors giant.data.transforms._EPS, duplicated for use in polars exprs
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Predict-parquet scanning, metadata verification, and physical-value exprs
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _check_predict_metadata(path: Path) -> None:
|
||
"""Verify a parquet file's giant-predict tag before trusting its column layout.
|
||
|
||
Raises rather than warns: a wrong or missing tag means the column-layout
|
||
assumptions below don't hold, so silently proceeding could mix up which
|
||
columns are predictions vs. ground truth.
|
||
"""
|
||
metadata = pq.read_schema(path).metadata or {}
|
||
coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
|
||
if coord is None:
|
||
raise ValueError(
|
||
f"{path} has no '{PREDICT_COORD_METADATA_KEY}' parquet metadata — it wasn't "
|
||
"written by `giant predict` (or predates schema tagging), so its column "
|
||
"layout can't be verified"
|
||
)
|
||
if coord.decode() != "local":
|
||
raise ValueError(
|
||
f"{path} was written with --coord {coord.decode()!r}, not 'local' — "
|
||
"giant.analysis only supports coord=local predict output, which is the "
|
||
"only mode that also writes ground-truth columns"
|
||
)
|
||
version = metadata.get(PREDICT_SCHEMA_VERSION_KEY.encode())
|
||
if version is not None and version.decode() != PREDICT_SCHEMA_VERSION:
|
||
raise ValueError(
|
||
f"{path} has predict schema version {version.decode()!r}, but "
|
||
f"giant.analysis expects {PREDICT_SCHEMA_VERSION!r} — column layout may "
|
||
"have changed; update giant.analysis to match"
|
||
)
|
||
|
||
|
||
def _scan_predicted_local(source: str | Path | pl.LazyFrame) -> pl.LazyFrame:
|
||
if isinstance(source, pl.LazyFrame):
|
||
return source
|
||
path = Path(source)
|
||
_check_predict_metadata(path)
|
||
return pl.scan_parquet(path)
|
||
|
||
|
||
def _edep_pl(prefix: str) -> pl.Expr:
|
||
"""Physical edep from a `{prefix}_edep_logit`/`{prefix}_sec_logit` pair + `pre_E`.
|
||
|
||
Polars equivalent of `energy_simplex_decode(...)[0]` (the deposit component):
|
||
a softmax over `[z_edep, z_sec, 0]` times pre_E.
|
||
"""
|
||
z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
|
||
m = pl.max_horizontal(z1, z2, pl.lit(0.0))
|
||
e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
|
||
return (e1 / (e1 + e2 + e3)) * pl.col("pre_E")
|
||
|
||
|
||
def _delta_e_pl(prefix: str) -> pl.Expr:
|
||
"""Physical delta_e (= edep + e_sec = pre_E - post_E) from the ALR logits + pre_E."""
|
||
z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
|
||
m = pl.max_horizontal(z1, z2, pl.lit(0.0))
|
||
e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
|
||
return ((e1 + e2) / (e1 + e2 + e3)) * pl.col("pre_E")
|
||
|
||
|
||
def _raw_dim_expr(prefix: str, j: int) -> pl.Expr:
|
||
"""Physical raw value of target dim j (`RAW_TARGET_NAMES[j]`) from a predict parquet.
|
||
|
||
Dim 0 is de-logged step_length; dims 1–2 are the physical `delta_e`/`edep`
|
||
decoded from the deposit/secondary ALR logits against `pre_E`; dims ≥3 are
|
||
direction components, used as-is. `prefix` is "true" or "pred".
|
||
"""
|
||
if j == 0:
|
||
return pl.col(f"{prefix}_log_step_length").exp() - _LOG_EPS
|
||
if j == 1:
|
||
return _delta_e_pl(prefix)
|
||
if j == 2:
|
||
return _edep_pl(prefix)
|
||
return pl.col(f"{prefix}_{LOCAL_TARGET_NAMES[j]}")
|
||
|
||
|
||
def _hist_edges(*arrays: np.ndarray, bins: int) -> np.ndarray:
|
||
"""Bin edges that don't blow up on near-constant data (e.g. a tight unit-norm cluster).
|
||
|
||
Plain `np.linspace(lo, hi, bins+1)` raises when `hi - lo` is too small relative
|
||
to float precision to support `bins` distinct edges, which is a real failure
|
||
mode here (a well-trained model can push direction norms to within float32
|
||
epsilon of 1.0), not just a test artifact. Used only on small per-event
|
||
reduced arrays, never on raw file rows.
|
||
"""
|
||
lo = min(a.min() for a in arrays)
|
||
hi = max(a.max() for a in arrays)
|
||
if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
|
||
lo, hi = lo - 0.5, hi + 0.5
|
||
return np.linspace(lo, hi, bins + 1)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Streaming histogram / KL primitives
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _kl_from_counts(
|
||
real_counts: np.ndarray, gen_counts: np.ndarray, eps: float = 1e-8
|
||
) -> float:
|
||
"""KL(real || gen) from two aligned histogram bin-count arrays.
|
||
|
||
Add-eps smoothing then normalize both to distributions, as is standard for a
|
||
histogram-estimated KL; the counts here come from a polars `group_by`
|
||
aggregation rather than `np.histogram`.
|
||
"""
|
||
p = real_counts.astype(np.float64) + eps
|
||
q = gen_counts.astype(np.float64) + eps
|
||
p /= p.sum()
|
||
q /= q.sum()
|
||
return float(np.sum(p * np.log(p / q)))
|
||
|
||
|
||
def _resolve_hist_range(
|
||
lf: pl.LazyFrame, expr: pl.Expr, lo: float | None, hi: float | None
|
||
) -> tuple[float, float]:
|
||
"""Fill in any missing (lo, hi) from a streaming min/max pass over `expr`.
|
||
|
||
Mirrors `_hist_edges`'s degenerate-range widening (±0.5 when the span is too
|
||
tight to support distinct edges) so fixed-edge and data-ranged histograms
|
||
agree on edge cases. Only runs the extra streaming pass when a bound is
|
||
actually unknown.
|
||
"""
|
||
if lo is not None and hi is not None:
|
||
return lo, hi
|
||
row = (
|
||
lf.select(expr.min().alias("lo"), expr.max().alias("hi"))
|
||
.collect(engine="streaming")
|
||
.row(0, named=True)
|
||
)
|
||
data_lo, data_hi = row["lo"], row["hi"]
|
||
if not (data_hi - data_lo > 1e-6 * max(abs(data_hi), 1.0)):
|
||
data_lo, data_hi = data_lo - 0.5, data_hi + 0.5
|
||
return (data_lo if lo is None else lo, data_hi if hi is None else hi)
|
||
|
||
|
||
def _streaming_hist1d(
|
||
lf: pl.LazyFrame,
|
||
expr: pl.Expr,
|
||
bins: int,
|
||
lo: float | None = None,
|
||
hi: float | None = None,
|
||
) -> tuple[np.ndarray, np.ndarray]:
|
||
"""Bounded-memory 1D histogram of `expr` over `lf`, returning (counts, edges).
|
||
|
||
A single streaming `group_by(bin).len()` pass — a proper single hash-pass
|
||
histogram — after (optionally) a streaming min/max pass to fix the range
|
||
when `lo`/`hi` aren't supplied. Never materializes the underlying column.
|
||
"""
|
||
lo, hi = _resolve_hist_range(lf, expr, lo, hi)
|
||
width = hi - lo
|
||
bin_expr = ((expr - lo) / width * bins).floor().cast(pl.Int64).clip(0, bins - 1)
|
||
hist = (
|
||
lf.select(bin_expr.alias("bin"))
|
||
.group_by("bin")
|
||
.agg(pl.len().alias("count"))
|
||
.collect(engine="streaming")
|
||
)
|
||
counts = np.zeros(bins, dtype=np.int64)
|
||
for bin_idx, count in hist.iter_rows():
|
||
counts[bin_idx] = count
|
||
return counts, np.linspace(lo, hi, bins + 1)
|
||
|
||
|
||
def _step_density(counts: np.ndarray, edges: np.ndarray) -> np.ndarray:
|
||
"""Density-normalized bin heights (integrate to 1), matching `hist(density=True)`."""
|
||
total = counts.sum()
|
||
if total == 0:
|
||
return counts.astype(np.float64)
|
||
width = edges[1] - edges[0]
|
||
return counts.astype(np.float64) / (total * width)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 1: stratified marginals
|
||
# ---------------------------------------------------------------------------
|
||
#
|
||
# One native `group_by` per dim (covering every group's mean/std/n/histogram
|
||
# range at once) plus one pass for per-bin counts — runtime independent of group
|
||
# cardinality. Filtering and re-collecting once per (group, dim) pair instead
|
||
# would make runtime scale with the number of groups: on a 114M-row file with
|
||
# 138 distinct pdg codes that extrapolated to ~an hour vs ~1 minute here.
|
||
|
||
|
||
def _add_group_label(
|
||
lf: pl.LazyFrame, group_by: str | None, n_energy_bins: int
|
||
) -> pl.LazyFrame:
|
||
"""Add a `_group` string column identifying each row's stratum.
|
||
|
||
For "pdg"/"material" the label is a direct string expr over the existing
|
||
column — `group_by("_group")` downstream discovers the distinct values
|
||
itself. "energy" needs one streaming pass over `pre_E` to fix quantile bin
|
||
edges before a label can be assigned per row.
|
||
"""
|
||
if group_by is None:
|
||
return lf.with_columns(pl.lit("all").alias("_group"))
|
||
if group_by == "pdg":
|
||
return lf.with_columns(
|
||
(pl.lit("pdg=") + pl.col("pdg").cast(pl.Int64).cast(pl.Utf8)).alias(
|
||
"_group"
|
||
)
|
||
)
|
||
if group_by == "material":
|
||
return lf.with_columns(
|
||
(pl.lit("material=") + pl.col("material")).alias("_group")
|
||
)
|
||
if group_by == "energy":
|
||
pre_E = lf.select("pre_E").collect(engine="streaming").to_series().to_numpy()
|
||
edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
|
||
edges[-1] += 1e-6
|
||
labels = [
|
||
f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})" for i in range(n_energy_bins)
|
||
]
|
||
expr = pl.when(pl.col("pre_E") < edges[1]).then(pl.lit(labels[0]))
|
||
for i in range(1, n_energy_bins - 1):
|
||
expr = expr.when(pl.col("pre_E") < edges[i + 1]).then(pl.lit(labels[i]))
|
||
expr = expr.otherwise(pl.lit(labels[-1]))
|
||
return lf.with_columns(expr.alias("_group"))
|
||
raise ValueError(f"unknown group_by={group_by!r}")
|
||
|
||
|
||
# On a file larger than RAM (never page-cached) the binding constraint is *peak
|
||
# memory*, not scan count: each streaming scan mmaps the columns it reads, so a
|
||
# scan that touches only ~3 columns keeps resident memory low, while one that
|
||
# reads all 9 dims' ~19 source columns at once mmaps almost the whole 32GB file
|
||
# and OOMs. So every scan here is kept *narrow* — stats and histogram are
|
||
# computed one dim at a time over just that dim's real/gen columns (2 scans per
|
||
# dim, real+gen histograms merged into one). More scans than a combined pass, but
|
||
# each is cheap in memory, which is what actually matters at this file size.
|
||
|
||
|
||
def _pad_range(lo: float, hi: float) -> tuple[float, float]:
|
||
"""Widen (lo, hi) by ±0.5 when too tight to support `bins` distinct edges.
|
||
|
||
Scalar duplicate of `_hist_edges`'s degenerate-range handling, applied per
|
||
(group, dim).
|
||
"""
|
||
if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
|
||
return lo - 0.5, hi + 0.5
|
||
return lo, hi
|
||
|
||
|
||
def _dim_stats(lf: pl.LazyFrame, j: int) -> pl.DataFrame:
|
||
"""One narrow streaming pass: per-group n + mean/std/range for dim `j`.
|
||
|
||
Reads only this dim's real/gen columns (see the note above on why scans stay
|
||
narrow). `n` (= `pl.len()`) is the group size; ddof=0 matches numpy's
|
||
population std.
|
||
"""
|
||
r, g = _raw_dim_expr("true", j), _raw_dim_expr("pred", j)
|
||
return (
|
||
lf.select("_group", r.alias("real"), g.alias("gen"))
|
||
.group_by("_group")
|
||
.agg(
|
||
pl.len().alias("n"),
|
||
pl.col("real").mean().alias("real_mean"),
|
||
pl.col("gen").mean().alias("gen_mean"),
|
||
pl.col("real").std(ddof=0).alias("real_std"),
|
||
pl.col("gen").std(ddof=0).alias("gen_std"),
|
||
pl.min_horizontal(pl.col("real").min(), pl.col("gen").min()).alias("lo"),
|
||
pl.max_horizontal(pl.col("real").max(), pl.col("gen").max()).alias("hi"),
|
||
)
|
||
.collect(engine="streaming")
|
||
)
|
||
|
||
|
||
def _dim_hist_counts(
|
||
lf: pl.LazyFrame, j: int, lo_hi: pl.DataFrame, bins: int
|
||
) -> tuple[dict[str, np.ndarray], dict[str, np.ndarray]]:
|
||
"""One streaming pass: real & gen histogram bin counts for dim `j`, by group.
|
||
|
||
Projects to just `(_group, real, gen)` for this dim (an explicit up-front
|
||
`.select` — projection pushdown doesn't reliably prune across the join in
|
||
this polars version), joins each row to its group's padded `(lo, hi)` range,
|
||
bins with the same floor/clip formula the original per-dim path used (so
|
||
counts are bit-identical), and counts real and gen together by unpivoting only
|
||
the two bin-index columns — a bounded 2× explosion, unlike unpivoting all 18
|
||
raw values at once (which OOMs on a >RAM file). The inner join also drops rows
|
||
whose group had n < 2.
|
||
"""
|
||
narrow = lf.select(
|
||
"_group",
|
||
_raw_dim_expr("true", j).alias("real"),
|
||
_raw_dim_expr("pred", j).alias("gen"),
|
||
)
|
||
width = pl.col("hi") - pl.col("lo")
|
||
|
||
def _bin(col: str) -> pl.Expr:
|
||
return (
|
||
((pl.col(col) - pl.col("lo")) / width * bins)
|
||
.floor()
|
||
.cast(pl.Int64)
|
||
.clip(0, bins - 1)
|
||
)
|
||
|
||
hist = (
|
||
narrow.join(lo_hi.lazy(), on="_group")
|
||
.select("_group", _bin("real").alias("real"), _bin("gen").alias("gen"))
|
||
.unpivot(index="_group", variable_name="kind", value_name="bin")
|
||
.group_by(["_group", "kind", "bin"])
|
||
.agg(pl.len().alias("count"))
|
||
.collect(engine="streaming")
|
||
)
|
||
|
||
real_by: dict[str, np.ndarray] = {}
|
||
gen_by: dict[str, np.ndarray] = {}
|
||
for group, kind, bin_idx, count in hist.iter_rows():
|
||
target = real_by if kind == "real" else gen_by
|
||
target.setdefault(group, np.zeros(bins, dtype=np.int64))[bin_idx] = count
|
||
return real_by, gen_by
|
||
|
||
|
||
# Per (group, dim) histogram payload: (real_counts, gen_counts, lo, hi).
|
||
_DimHists = dict[str, dict[str, tuple[np.ndarray, np.ndarray, float, float]]]
|
||
|
||
|
||
def _marginal_histograms(
|
||
source: str | Path | pl.LazyFrame,
|
||
group_by: str | None,
|
||
n_energy_bins: int,
|
||
bins: int,
|
||
) -> tuple[pl.DataFrame, _DimHists]:
|
||
"""Stratified per-dim marginal stats table + the per-(group, dim) histograms.
|
||
|
||
The single source of truth for both `marginal_table_pl` (which wants only the
|
||
table) and `plot_marginals` (which also wants the bin counts). Two narrow
|
||
scans per dim (stats, then merged real+gen histogram) — see the note above
|
||
these helpers on why scans stay narrow rather than combined.
|
||
"""
|
||
lf = _add_group_label(_scan_predicted_local(source), group_by, n_energy_bins)
|
||
|
||
empty = np.zeros(bins, dtype=np.int64)
|
||
rows = []
|
||
hists: _DimHists = {name: {} for name in RAW_TARGET_NAMES}
|
||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||
stats = _dim_stats(lf, j).filter(pl.col("n") >= 2)
|
||
range_by = {
|
||
row["_group"]: _pad_range(row["lo"], row["hi"])
|
||
for row in stats.iter_rows(named=True)
|
||
}
|
||
lo_hi = pl.DataFrame(
|
||
[{"_group": g, "lo": lo, "hi": hi} for g, (lo, hi) in range_by.items()],
|
||
schema={"_group": pl.Utf8, "lo": pl.Float64, "hi": pl.Float64},
|
||
)
|
||
real_by, gen_by = _dim_hist_counts(lf, j, lo_hi, bins)
|
||
for row in stats.iter_rows(named=True):
|
||
group = row["_group"]
|
||
real_counts = real_by.get(group, empty)
|
||
gen_counts = gen_by.get(group, empty)
|
||
lo, hi = range_by[group]
|
||
hists[name][group] = (real_counts, gen_counts, lo, hi)
|
||
rows.append(
|
||
{
|
||
"group": group,
|
||
"dim": name,
|
||
"n": row["n"],
|
||
"real_mean": row["real_mean"],
|
||
"gen_mean": row["gen_mean"],
|
||
"real_std": row["real_std"],
|
||
"gen_std": row["gen_std"],
|
||
"kl_real_gen": _kl_from_counts(real_counts, gen_counts),
|
||
}
|
||
)
|
||
|
||
table = pl.DataFrame(rows).sort("kl_real_gen", descending=True)
|
||
return table, hists
|
||
|
||
|
||
def marginal_table_pl(
|
||
source: str | Path | pl.LazyFrame,
|
||
group_by: str | None = None,
|
||
n_energy_bins: int = 4,
|
||
bins: int = 50,
|
||
) -> pl.DataFrame:
|
||
"""Per-dimension real-vs-generated summary stats + KL(real||gen), in raw units.
|
||
|
||
`source` is a path to a `giant predict --coord local` parquet file, or an
|
||
already-built LazyFrame with the same `pred_*`/`true_*`/`pdg`/`material`/
|
||
`pre_E` columns (e.g. for testing). `group_by`: None for an aggregate table,
|
||
or one of "pdg", "material", "energy" to stratify. Sorted worst-KL first, so
|
||
failure modes hidden by the aggregate surface at the top.
|
||
"""
|
||
table, _ = _marginal_histograms(source, group_by, n_energy_bins, bins)
|
||
return table
|
||
|
||
|
||
def _limit_groups_by_kl_n(table: pd.DataFrame, max_groups: int) -> pd.DataFrame:
|
||
"""Keep only the `max_groups` groups with the largest max(kl) * n.
|
||
|
||
Ranks by how badly wrong *and* how common a stratum is, rather than by KL
|
||
alone, so a rare pdg/material with a noisy, high-variance KL estimate from a
|
||
handful of samples doesn't crowd out groups that actually matter.
|
||
"""
|
||
by_group = table.groupby("group").agg(
|
||
kl_max=("kl_real_gen", "max"), n=("n", "first")
|
||
)
|
||
worst_first = (by_group["kl_max"] * by_group["n"]).sort_values(ascending=False)
|
||
keep = set(worst_first.index[:max_groups])
|
||
return table[table["group"].isin(keep)]
|
||
|
||
|
||
def _worst_first_groups(table: pd.DataFrame, max_groups: int) -> list[str]:
|
||
"""Group labels ranked by max(kl) * n (worst first), capped at `max_groups`."""
|
||
by_group = table.groupby("group").agg(
|
||
kl_max=("kl_real_gen", "max"), n=("n", "first")
|
||
)
|
||
worst_first = (by_group["kl_max"] * by_group["n"]).sort_values(ascending=False)
|
||
return list(worst_first.index[:max_groups])
|
||
|
||
|
||
def plot_marginals(
|
||
source: str | Path | pl.LazyFrame,
|
||
dims: list[str] | None = None,
|
||
group_by: str | None = None,
|
||
n_energy_bins: int = 4,
|
||
bins: int = 50,
|
||
max_groups: int = 6,
|
||
figsize_per_axis: tuple[float, float] = (3.5, 2.8),
|
||
):
|
||
"""Overlaid real-vs-generated histograms: one row per group, one column per dim.
|
||
|
||
Without `group_by`, a single row over the whole file. With "pdg", "material",
|
||
or "energy", one row per stratum, ranked by max(kl) * n (capped at
|
||
`max_groups`) so groups that are both badly wrong and common surface first,
|
||
rather than rare groups with a noisy, high-variance KL estimate. Every
|
||
histogram is computed with streaming bin counts (`_marginal_histograms`) — no
|
||
row is ever materialized.
|
||
"""
|
||
dims = dims or RAW_TARGET_NAMES
|
||
table, hists = _marginal_histograms(source, group_by, n_energy_bins, bins)
|
||
|
||
if group_by is None:
|
||
groups = ["all"]
|
||
else:
|
||
groups = _worst_first_groups(table.to_pandas(), max_groups)
|
||
|
||
n_rows, n_cols = len(groups), len(dims)
|
||
fig, axes = plt.subplots(
|
||
n_rows,
|
||
n_cols,
|
||
squeeze=False,
|
||
figsize=(figsize_per_axis[0] * n_cols, figsize_per_axis[1] * n_rows),
|
||
)
|
||
for row, label in enumerate(groups):
|
||
for col, name in enumerate(dims):
|
||
ax = axes[row][col]
|
||
entry = hists[name].get(label)
|
||
if entry is not None:
|
||
real_counts, gen_counts, lo, hi = entry
|
||
edges = np.linspace(lo, hi, bins + 1)
|
||
ax.stairs(_step_density(real_counts, edges), edges, label="real")
|
||
ax.stairs(_step_density(gen_counts, edges), edges, label="generated")
|
||
ax.set_yscale("log")
|
||
if row == 0:
|
||
ax.set_title(name, fontsize=9)
|
||
if col == 0:
|
||
ax.set_ylabel(label, fontsize=8)
|
||
if row == 0 and col == n_cols - 1:
|
||
ax.legend(fontsize=7)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def _plot_kl_bars(table: pd.DataFrame, figsize: tuple[float, float]):
|
||
"""Shared bar-plot body for `plot_kl_bars_pl`.
|
||
|
||
`table` is a `marginal_table_pl` result converted to pandas with
|
||
`group`/`dim`/`kl_real_gen` columns.
|
||
"""
|
||
pivot = table.pivot(index="dim", columns="group", values="kl_real_gen")
|
||
pivot = pivot.reindex(RAW_TARGET_NAMES)
|
||
groups = sorted(pivot.columns)
|
||
pivot = pivot[groups]
|
||
|
||
n_dims, n_groups = len(RAW_TARGET_NAMES), len(groups)
|
||
x = np.arange(n_dims)
|
||
width = 0.8 / n_groups
|
||
|
||
fig, ax = plt.subplots(figsize=figsize)
|
||
for i, group in enumerate(groups):
|
||
offset = (i - (n_groups - 1) / 2) * width
|
||
ax.bar(x + offset, pivot[group].to_numpy(), width=width, label=group)
|
||
ax.set_xticks(x)
|
||
ax.set_xticklabels(RAW_TARGET_NAMES, rotation=45, ha="right")
|
||
ax.set_ylabel("KL(real || gen)")
|
||
if n_groups > 1:
|
||
ax.legend(fontsize=7)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def plot_kl_bars_pl(
|
||
source: str | Path | pl.LazyFrame,
|
||
group_by: str | None = None,
|
||
n_energy_bins: int = 4,
|
||
bins: int = 50,
|
||
max_groups: int = 6,
|
||
figsize: tuple[float, float] = (8, 4),
|
||
):
|
||
"""Bar plot of KL(real||gen) per target dimension, optionally stratified.
|
||
|
||
One bar cluster per dimension; with `group_by` set, one bar per stratum
|
||
within each cluster, so which dimension/stratum combination drives a KL
|
||
regression is visible at a glance. With "pdg" or "material" — open-ended
|
||
vocabularies — groups are capped at `max_groups`, ranked by max(kl) * n;
|
||
"energy" is already bounded by `n_energy_bins` and isn't capped.
|
||
"""
|
||
table = marginal_table_pl(
|
||
source, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
|
||
).to_pandas()
|
||
if group_by in ("pdg", "material"):
|
||
table = _limit_groups_by_kl_n(table, max_groups)
|
||
return _plot_kl_bars(table, figsize=figsize)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 2: joint structure
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def correlation_matrices_pl(
|
||
source: str | Path | pl.LazyFrame,
|
||
) -> tuple[np.ndarray, np.ndarray]:
|
||
"""9×9 Pearson correlation matrices of the raw targets, (real, generated).
|
||
|
||
Computed from streaming sufficient statistics — one pass accumulating
|
||
per-dim sums and all i≤j cross-products (in float64) — rather than loading
|
||
the target arrays and calling `np.corrcoef`. `cov_ij = E[x_i x_j] -
|
||
E[x_i]E[x_j]` (population/N normalization, matching `np.corrcoef`), then
|
||
`corr = cov / sqrt(cov_ii cov_jj)`.
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
n_dims = len(RAW_TARGET_NAMES)
|
||
narrow = lf.select(
|
||
*[
|
||
_raw_dim_expr("true", j).cast(pl.Float64).alias(f"r{j}")
|
||
for j in range(n_dims)
|
||
],
|
||
*[
|
||
_raw_dim_expr("pred", j).cast(pl.Float64).alias(f"g{j}")
|
||
for j in range(n_dims)
|
||
],
|
||
)
|
||
|
||
aggs: list[pl.Expr] = [pl.len().alias("n")]
|
||
for j in range(n_dims):
|
||
aggs.append(pl.col(f"r{j}").sum().alias(f"rs{j}"))
|
||
aggs.append(pl.col(f"g{j}").sum().alias(f"gs{j}"))
|
||
for i in range(n_dims):
|
||
for j in range(i, n_dims):
|
||
aggs.append((pl.col(f"r{i}") * pl.col(f"r{j}")).sum().alias(f"rp{i}_{j}"))
|
||
aggs.append((pl.col(f"g{i}") * pl.col(f"g{j}")).sum().alias(f"gp{i}_{j}"))
|
||
row = narrow.select(aggs).collect(engine="streaming").row(0, named=True)
|
||
n = row["n"]
|
||
|
||
def _build(sum_prefix: str, prod_prefix: str) -> np.ndarray:
|
||
mean = np.array([row[f"{sum_prefix}{j}"] for j in range(n_dims)]) / n
|
||
cov = np.empty((n_dims, n_dims), dtype=np.float64)
|
||
for i in range(n_dims):
|
||
for j in range(i, n_dims):
|
||
c = row[f"{prod_prefix}{i}_{j}"] / n - mean[i] * mean[j]
|
||
cov[i, j] = cov[j, i] = c
|
||
std = np.sqrt(np.diag(cov))
|
||
return cov / np.outer(std, std)
|
||
|
||
return _build("rs", "rp"), _build("gs", "gp")
|
||
|
||
|
||
def plot_correlation_matrices(source: str | Path | pl.LazyFrame):
|
||
"""Side-by-side real/generated correlation heatmaps, plus their difference."""
|
||
real_corr, gen_corr = correlation_matrices_pl(source)
|
||
diff = gen_corr - real_corr
|
||
|
||
fig, axes = plt.subplots(1, 3, figsize=(13, 4))
|
||
for ax, mat, title, cmap, vlim in [
|
||
(axes[0], real_corr, "real", "coolwarm", (-1, 1)),
|
||
(axes[1], gen_corr, "generated", "coolwarm", (-1, 1)),
|
||
(axes[2], diff, "generated − real", "PuOr", (-0.5, 0.5)),
|
||
]:
|
||
im = ax.imshow(mat, vmin=vlim[0], vmax=vlim[1], cmap=cmap)
|
||
ax.set_xticks(range(len(RAW_TARGET_NAMES)))
|
||
ax.set_xticklabels(RAW_TARGET_NAMES, rotation=90, fontsize=7)
|
||
ax.set_yticks(range(len(RAW_TARGET_NAMES)))
|
||
ax.set_yticklabels(RAW_TARGET_NAMES, fontsize=7)
|
||
ax.set_title(title)
|
||
fig.colorbar(im, ax=ax, fraction=0.046)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
_DEFAULT_PAIRS = [
|
||
("step_length", "delta_e"),
|
||
("delta_e", "edep"),
|
||
("step_length", "edep"),
|
||
]
|
||
|
||
|
||
def plot_pairwise(
|
||
source: str | Path | pl.LazyFrame,
|
||
pairs: list[tuple[str, str]] | None = None,
|
||
n_sample: int = 3000,
|
||
seed: int = 0,
|
||
):
|
||
"""Real-vs-generated scatter for physically coupled target pairs.
|
||
|
||
Marginals matching doesn't imply the joint does — these pairs are coupled by
|
||
the underlying physics (energy loss tracks distance, edep is part of
|
||
delta_e), so a model that decorrelates them shows up here even with clean
|
||
per-dimension marginals. Only ~`n_sample` rows are ever materialized: a cheap
|
||
streaming count fixes a hash-subsample fraction, and only the dims referenced
|
||
by `pairs` are read for the kept rows.
|
||
"""
|
||
pairs = pairs or _DEFAULT_PAIRS
|
||
names = sorted({name for pair in pairs for name in pair})
|
||
lf = _scan_predicted_local(source)
|
||
|
||
n_total = lf.select(pl.len()).collect(engine="streaming").item()
|
||
if n_total == 0:
|
||
raise ValueError("no rows in predict source")
|
||
frac = min(1.0, n_sample / n_total)
|
||
threshold = int(frac * 2**32)
|
||
|
||
select_exprs = []
|
||
for name in names:
|
||
j = RAW_TARGET_NAMES.index(name)
|
||
select_exprs.append(_raw_dim_expr("true", j).alias(f"real_{name}"))
|
||
select_exprs.append(_raw_dim_expr("pred", j).alias(f"gen_{name}"))
|
||
|
||
sampled = (
|
||
lf.select(*select_exprs)
|
||
.with_row_index("_ri")
|
||
.filter((pl.col("_ri").hash(seed=seed) % 2**32) < threshold)
|
||
.drop("_ri")
|
||
.collect(engine="streaming")
|
||
)
|
||
if sampled.height > n_sample:
|
||
idx = np.random.default_rng(seed).choice(
|
||
sampled.height, size=n_sample, replace=False
|
||
)
|
||
sampled = sampled[idx]
|
||
|
||
fig, axes = plt.subplots(2, len(pairs), squeeze=False, figsize=(4 * len(pairs), 7))
|
||
for col, (a, b) in enumerate(pairs):
|
||
for row, (prefix, title) in enumerate([("real", "real"), ("gen", "generated")]):
|
||
ax = axes[row][col]
|
||
ax.scatter(
|
||
sampled[f"{prefix}_{a}"].to_numpy(),
|
||
sampled[f"{prefix}_{b}"].to_numpy(),
|
||
s=3,
|
||
alpha=0.3,
|
||
)
|
||
ax.set_xlabel(a)
|
||
ax.set_ylabel(b)
|
||
if col == 0:
|
||
ax.set_title(title, loc="left", fontsize=9)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def _cos_alignment_expr(prefix: str) -> pl.Expr:
|
||
"""cos(angle) between `{prefix}` post_dir and travel_dir, per row."""
|
||
post = [pl.col(f"{prefix}_{n}") for n in ("post_dx", "post_dy", "post_dz")]
|
||
travel = [pl.col(f"{prefix}_{n}") for n in ("travel_dx", "travel_dy", "travel_dz")]
|
||
dot = sum(p * t for p, t in zip(post, travel))
|
||
post_norm = pl.sum_horizontal([p**2 for p in post]).sqrt()
|
||
travel_norm = pl.sum_horizontal([t**2 for t in travel]).sqrt()
|
||
return dot / (post_norm * travel_norm + 1e-8)
|
||
|
||
|
||
def plot_direction_alignment(source: str | Path | pl.LazyFrame, bins: int = 50):
|
||
"""cos(angle) between post_dir and travel_dir, real vs generated.
|
||
|
||
These two unit vectors are coupled through the scattering physics, so their
|
||
joint alignment is a check the per-dimension marginals can't see. Fixed
|
||
`[-1, 1]` edges, one streaming bin-count pass per side — no materialization.
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
real_counts, edges = _streaming_hist1d(
|
||
lf, _cos_alignment_expr("true"), bins, lo=-1.0, hi=1.0
|
||
)
|
||
gen_counts, _ = _streaming_hist1d(
|
||
lf, _cos_alignment_expr("pred"), bins, lo=-1.0, hi=1.0
|
||
)
|
||
|
||
fig, ax = plt.subplots(figsize=(5, 4))
|
||
ax.stairs(_step_density(real_counts, edges), edges, label="real")
|
||
ax.stairs(_step_density(gen_counts, edges), edges, label="generated")
|
||
ax.set_yscale("log")
|
||
ax.set_xlabel("cos(angle) between post_dir and travel_dir")
|
||
ax.legend()
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 3: physical constraints
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _pred_norm_expr(start: int) -> pl.Expr:
|
||
"""||pred direction|| for the 3 components starting at target dim `start`."""
|
||
cols = [pl.col(f"pred_{LOCAL_TARGET_NAMES[k]}") for k in range(start, start + 3)]
|
||
return pl.sum_horizontal([c**2 for c in cols]).sqrt()
|
||
|
||
|
||
def constraint_report_pl(
|
||
source: str | Path | pl.LazyFrame, norm_tol: float = 0.05
|
||
) -> pl.DataFrame:
|
||
"""Rate of physical-constraint violations in the generated raw-space samples.
|
||
|
||
The model is an unconstrained MLP, so nothing forces post_dir / travel_dir to
|
||
stay unit-norm or step_length/delta_e/edep to stay non-negative — all hold by
|
||
construction in the real data, so any violation rate here is purely a
|
||
generation artifact. A single streaming aggregation over the `pred_*` columns;
|
||
never materializes the target arrays.
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
post_norm = _pred_norm_expr(3)
|
||
travel_norm = _pred_norm_expr(6)
|
||
raw_scalar_dims = [_raw_dim_expr("pred", j) for j in range(_N_SCALAR_DIMS)]
|
||
|
||
agg = (
|
||
lf.select(
|
||
[
|
||
((post_norm - 1).abs() > norm_tol)
|
||
.mean()
|
||
.alias("post_dir_violation_rate"),
|
||
(post_norm - 1).abs().mean().alias("post_dir_mean_abs_error"),
|
||
((travel_norm - 1).abs() > norm_tol)
|
||
.mean()
|
||
.alias("travel_dir_violation_rate"),
|
||
(travel_norm - 1).abs().mean().alias("travel_dir_mean_abs_error"),
|
||
*[
|
||
(raw < 0).mean().alias(f"{name}_violation_rate")
|
||
for raw, name in zip(
|
||
raw_scalar_dims, RAW_TARGET_NAMES[:_N_SCALAR_DIMS]
|
||
)
|
||
],
|
||
*[
|
||
raw.clip(upper_bound=0).abs().mean().alias(f"{name}_mean_abs_error")
|
||
for raw, name in zip(
|
||
raw_scalar_dims, RAW_TARGET_NAMES[:_N_SCALAR_DIMS]
|
||
)
|
||
],
|
||
]
|
||
)
|
||
.collect(engine="streaming")
|
||
.row(0, named=True)
|
||
)
|
||
|
||
rows = [
|
||
{
|
||
"check": "post_dir unit norm",
|
||
"violation_rate": agg["post_dir_violation_rate"],
|
||
"mean_abs_error": agg["post_dir_mean_abs_error"],
|
||
},
|
||
{
|
||
"check": "travel_dir unit norm",
|
||
"violation_rate": agg["travel_dir_violation_rate"],
|
||
"mean_abs_error": agg["travel_dir_mean_abs_error"],
|
||
},
|
||
]
|
||
for name in RAW_TARGET_NAMES[:_N_SCALAR_DIMS]:
|
||
rows.append(
|
||
{
|
||
"check": f"{name} >= 0",
|
||
"violation_rate": agg[f"{name}_violation_rate"],
|
||
"mean_abs_error": agg[f"{name}_mean_abs_error"],
|
||
}
|
||
)
|
||
return pl.DataFrame(rows)
|
||
|
||
|
||
def plot_constraint_violations(source: str | Path | pl.LazyFrame, bins: int = 50):
|
||
"""Histograms backing `constraint_report_pl`: direction norms and sign of the scalars.
|
||
|
||
Each panel is a streaming histogram (`_streaming_hist1d`) over the generated
|
||
columns — the direction norms and the three physical scalar dims — so nothing
|
||
is materialized. The dashed line marks the constraint boundary (1.0 for norms,
|
||
0.0 for the non-negative scalars).
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
|
||
panels: list[tuple[pl.Expr, str, float]] = [
|
||
(_pred_norm_expr(3), "||post_dir||", 1.0),
|
||
(_pred_norm_expr(6), "||travel_dir||", 1.0),
|
||
]
|
||
for k, name in enumerate(RAW_TARGET_NAMES[:_N_SCALAR_DIMS]):
|
||
panels.append((_raw_dim_expr("pred", k), f"generated {name}", 0.0))
|
||
|
||
fig, axes = plt.subplots(1, len(panels), figsize=(4 * len(panels), 3.5))
|
||
for ax, (expr, title, boundary) in zip(axes, panels):
|
||
counts, edges = _streaming_hist1d(lf, expr, bins)
|
||
ax.stairs(counts, edges)
|
||
ax.set_yscale("log")
|
||
ax.axvline(boundary, color="k", linestyle="--", linewidth=1)
|
||
ax.set_title(title)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 4: event-level (shower) observables
|
||
#
|
||
# Built directly on `giant predict --coord local` parquet output, streamed in
|
||
# two passes rather than materialized: per-event sums (total deposited energy,
|
||
# etc.) run into the tens of millions of rows.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_PRE_COLS = ["pre_x", "pre_y", "pre_z", "pre_dx", "pre_dy", "pre_dz", "pre_E"]
|
||
|
||
|
||
@dataclass
|
||
class EventObservables:
|
||
event_table: pl.DataFrame # one row per event_id; real_*/gen_* columns, mm/MeV
|
||
depth_edges: np.ndarray # (depth_bins+1,) mm, along shower axis
|
||
transverse_edges: np.ndarray # (transverse_bins+1,) mm, perpendicular to axis
|
||
real_depth_profile: np.ndarray # (depth_bins,) mean edep/event/bin, MeV
|
||
gen_depth_profile: np.ndarray
|
||
real_depth_profile_std: np.ndarray # event-to-event RMS per bin
|
||
gen_depth_profile_std: np.ndarray
|
||
real_transverse_profile: np.ndarray
|
||
gen_transverse_profile: np.ndarray
|
||
real_transverse_profile_std: np.ndarray
|
||
gen_transverse_profile_std: np.ndarray
|
||
|
||
|
||
# Every per-row use of the per-event entry point / shower axis attaches those
|
||
# six values with `replace_strict` (a 10⁴-entry hash map applied as an
|
||
# expression) rather than a `.join`: polars 1.4's streaming join buffers the
|
||
# whole 323M-row left side and OOMs on a file larger than RAM, whereas
|
||
# `replace_strict` streams in bounded memory.
|
||
|
||
|
||
def _entry_axis_exprs(entry_df: pl.DataFrame) -> list[pl.Expr]:
|
||
"""`replace_strict` expressions mapping event_id → each entry_*/axis_* value."""
|
||
event_ids = entry_df["event_id"].to_numpy()
|
||
return [
|
||
pl.col("event_id")
|
||
.replace_strict(event_ids, entry_df[col].to_numpy(), return_dtype=pl.Float64)
|
||
.alias(col)
|
||
for col in ("entry_x", "entry_y", "entry_z", "axis_x", "axis_y", "axis_z")
|
||
]
|
||
|
||
|
||
def _entry_axis_and_bin_edges(
|
||
lf: pl.LazyFrame,
|
||
depth_bins: int,
|
||
transverse_bins: int,
|
||
sample_rows: int = 1_000_000,
|
||
seed: int = 0,
|
||
) -> tuple[pl.DataFrame, np.ndarray, np.ndarray]:
|
||
"""Per-event shower axis (highest-pre_E row) plus depth/transverse bin edges.
|
||
|
||
The per-event entry point / axis is a `group_by(event_id)` (a bounded
|
||
streaming aggregation, no join). Bin edges are sized from `pre_pos` alone (no
|
||
post_pos reconstruction needed) via robust quantiles of the depth/transverse
|
||
proxies — estimated on a hash-subsample of ~`sample_rows` rows rather than the
|
||
whole file, because an exact `quantile` is holistic (materializes every row
|
||
and OOMs on a >RAM file) and outlier-clipping bin bounds don't need more than
|
||
a sample.
|
||
"""
|
||
entry_df = (
|
||
lf.select(["event_id", *_PRE_COLS])
|
||
.group_by("event_id")
|
||
.agg(
|
||
pl.col("pre_x").get(pl.col("pre_E").arg_max()).alias("entry_x"),
|
||
pl.col("pre_y").get(pl.col("pre_E").arg_max()).alias("entry_y"),
|
||
pl.col("pre_z").get(pl.col("pre_E").arg_max()).alias("entry_z"),
|
||
pl.col("pre_dx").get(pl.col("pre_E").arg_max()).alias("axis_x"),
|
||
pl.col("pre_dy").get(pl.col("pre_E").arg_max()).alias("axis_y"),
|
||
pl.col("pre_dz").get(pl.col("pre_E").arg_max()).alias("axis_z"),
|
||
)
|
||
.collect(engine="streaming")
|
||
.sort("event_id")
|
||
)
|
||
|
||
n_total = lf.select(pl.len()).collect(engine="streaming").item()
|
||
frac = min(1.0, sample_rows / max(n_total, 1))
|
||
threshold = int(frac * 2**32)
|
||
|
||
dx = pl.col("pre_x") - pl.col("entry_x")
|
||
dy = pl.col("pre_y") - pl.col("entry_y")
|
||
dz = pl.col("pre_z") - pl.col("entry_z")
|
||
depth = dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
|
||
tx = dx - depth * pl.col("axis_x")
|
||
ty = dy - depth * pl.col("axis_y")
|
||
tz = dz - depth * pl.col("axis_z")
|
||
transverse = (tx**2 + ty**2 + tz**2).sqrt()
|
||
|
||
sample = (
|
||
lf.select(["event_id", "pre_x", "pre_y", "pre_z"])
|
||
.filter((pl.col("pre_x").hash(seed=seed) % 2**32) < threshold)
|
||
.with_columns(*_entry_axis_exprs(entry_df))
|
||
.select(depth.alias("depth_proxy"), transverse.alias("transverse_proxy"))
|
||
.collect(engine="streaming")
|
||
)
|
||
depth_proxy = sample["depth_proxy"].to_numpy()
|
||
transverse_proxy = sample["transverse_proxy"].to_numpy()
|
||
|
||
depth_lo = float(np.quantile(depth_proxy, 0.001))
|
||
depth_hi = float(np.quantile(depth_proxy, 0.999))
|
||
if not (depth_hi - depth_lo > 1e-6 * max(abs(depth_hi), 1.0)):
|
||
depth_lo, depth_hi = depth_lo - 0.5, depth_hi + 0.5
|
||
depth_edges = np.linspace(depth_lo, depth_hi, depth_bins + 1)
|
||
transverse_hi = max(float(np.quantile(transverse_proxy, 0.999)), 1e-6)
|
||
transverse_edges = np.linspace(0.0, transverse_hi, transverse_bins + 1)
|
||
return entry_df, depth_edges, transverse_edges
|
||
|
||
|
||
# Approximate per-event medians are estimated from fixed log10-spaced value bins
|
||
# (median = the bin the running count crosses half at). Exact medians need a
|
||
# holistic per-group aggregation that can't stream and materializes every row
|
||
# (~18 GB / OOM on a >RAM file), whereas this histogram is a bounded streaming
|
||
# group_by. The range is generous enough to cover physical edep/step values.
|
||
_MED_LOG10_LO = -8.0
|
||
_MED_LOG10_HI = 6.0
|
||
_MED_BINS = 280
|
||
_MED_QUANTITIES = (
|
||
"real_median_edep",
|
||
"gen_median_edep",
|
||
"real_median_length",
|
||
"gen_median_length",
|
||
)
|
||
|
||
|
||
def _uniform_bin(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
|
||
"""Bin `value` into `[0, nbins)` over uniform edges [lo, hi].
|
||
|
||
Matches `np.digitize(value, linspace(lo, hi, nbins+1)[1:-1])` for uniform
|
||
edges — floor of the scaled offset, clipped to the valid range.
|
||
"""
|
||
return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int64).clip(0, nbins - 1)
|
||
|
||
|
||
def _geometry_exprs(src: str, out: str) -> list[pl.Expr]:
|
||
"""Per-row edep / step_length / depth / transverse for one side (`true`/`pred`).
|
||
|
||
Reconstructs world-frame post_pos as polars expressions — the inverse
|
||
local-frame (Rodrigues) rotation of `giant.data.transforms`, closed-form for
|
||
the constant ẑ axis — then projects `post_pos - entry` onto the per-event
|
||
shower axis (`depth`) and its perpendicular (`transverse`). Requires the
|
||
per-event `entry_*`/`axis_*` columns to be joined on already. Doing this in
|
||
expressions (rather than a numpy loop over pyarrow batches) lets the whole
|
||
per-row pass parallelize across the streaming engine's cores.
|
||
"""
|
||
norm = (
|
||
pl.col("pre_dx") ** 2 + pl.col("pre_dy") ** 2 + pl.col("pre_dz") ** 2
|
||
).sqrt()
|
||
ux, uy, uz = (
|
||
pl.col("pre_dx") / norm,
|
||
pl.col("pre_dy") / norm,
|
||
pl.col("pre_dz") / norm,
|
||
)
|
||
cos_t = uz.clip(-1.0, 1.0)
|
||
sin_t = (1.0 - cos_t**2).clip(lower_bound=0.0).sqrt()
|
||
|
||
# Rodrigues axis pre_dir × ẑ = [uy, -ux, 0], normalized; x̂ when pre_dir ∥ ẑ.
|
||
axis_norm = (uy**2 + ux**2).sqrt()
|
||
degenerate = axis_norm < 1e-7
|
||
ax = pl.when(degenerate).then(pl.lit(1.0)).otherwise(uy / axis_norm)
|
||
ay = pl.when(degenerate).then(pl.lit(0.0)).otherwise(-ux / axis_norm)
|
||
|
||
vx, vy, vz = (
|
||
pl.col(f"{src}_travel_dx"),
|
||
pl.col(f"{src}_travel_dy"),
|
||
pl.col(f"{src}_travel_dz"),
|
||
)
|
||
# axis × v (axis_z = 0); axis · v; inverse rotation applies R^T (−angle).
|
||
kxv_x, kxv_y, kxv_z = ay * vz, -(ax * vz), ax * vy - ay * vx
|
||
kdv = ax * vx + ay * vy
|
||
omc = 1.0 - cos_t
|
||
wx = vx * cos_t - kxv_x * sin_t + ax * kdv * omc
|
||
wy = vy * cos_t - kxv_y * sin_t + ay * kdv * omc
|
||
wz = vz * cos_t - kxv_z * sin_t # axis_z * kdv * omc == 0
|
||
|
||
step = pl.col(f"{src}_log_step_length").exp() - _LOG_EPS
|
||
post_x = pl.col("pre_x") + step * wx
|
||
post_y = pl.col("pre_y") + step * wy
|
||
post_z = pl.col("pre_z") + step * wz
|
||
|
||
dx = post_x - pl.col("entry_x")
|
||
dy = post_y - pl.col("entry_y")
|
||
dz = post_z - pl.col("entry_z")
|
||
depth = dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
|
||
tx = dx - depth * pl.col("axis_x")
|
||
ty = dy - depth * pl.col("axis_y")
|
||
tz = dz - depth * pl.col("axis_z")
|
||
transverse = (tx**2 + ty**2 + tz**2).sqrt()
|
||
|
||
return [
|
||
_edep_pl(src).alias(f"{out}_edep"),
|
||
step.alias(f"{out}_step"),
|
||
depth.alias(f"{out}_depth"),
|
||
transverse.alias(f"{out}_transverse"),
|
||
]
|
||
|
||
|
||
def _spatial_grid(
|
||
geo: pl.LazyFrame,
|
||
out: str,
|
||
depth_edges: np.ndarray,
|
||
transverse_edges: np.ndarray,
|
||
depth_bins: int,
|
||
transverse_bins: int,
|
||
) -> pl.DataFrame:
|
||
"""One streaming pass: per-(event, depth-bin, transverse-bin) edep/length sums.
|
||
|
||
The reduced grid (at most n_events × depth_bins × transverse_bins rows) carries
|
||
everything the event-level observables need: summing its cells recovers exact
|
||
per-event totals and the edep-weighted centroid/RMS numerators (since a sum of
|
||
per-cell sums is the full per-event sum), and marginalizing one axis gives each
|
||
longitudinal/transverse profile — so the whole 323M-row reduction happens in a
|
||
parallel `group_by`, not a serial python loop.
|
||
"""
|
||
depth_bin = _uniform_bin(
|
||
pl.col(f"{out}_depth"), depth_edges[0], depth_edges[-1], depth_bins
|
||
)
|
||
transverse_bin = _uniform_bin(
|
||
pl.col(f"{out}_transverse"),
|
||
transverse_edges[0],
|
||
transverse_edges[-1],
|
||
transverse_bins,
|
||
)
|
||
edep = pl.col(f"{out}_edep")
|
||
return (
|
||
geo.with_columns(depth_bin.alias("db"), transverse_bin.alias("tb"))
|
||
.group_by(["event_id", "db", "tb"])
|
||
.agg(
|
||
pl.len().alias("cnt"),
|
||
edep.sum().alias("s_edep"),
|
||
(edep * pl.col(f"{out}_depth")).sum().alias("s_edep_depth"),
|
||
(edep * pl.col(f"{out}_transverse") ** 2).sum().alias("s_edep_t2"),
|
||
pl.col(f"{out}_step").sum().alias("s_step"),
|
||
)
|
||
.collect(engine="streaming")
|
||
)
|
||
|
||
|
||
def _reduce_spatial_grid(
|
||
grid: pl.DataFrame,
|
||
event_ids: np.ndarray,
|
||
depth_bins: int,
|
||
transverse_bins: int,
|
||
depth_edges: np.ndarray,
|
||
) -> dict[str, np.ndarray]:
|
||
"""Collapse a `_spatial_grid` table into per-event arrays (aligned to event_ids).
|
||
|
||
Pure numpy over the small reduced grid: `np.bincount` scatter-adds recover the
|
||
per-event totals, centroid/RMS, dense (event, bin) profile matrices, and the
|
||
shower-max depth — identical to the previous full-file numpy accumulation, but
|
||
fed pre-summed grid cells instead of every row.
|
||
"""
|
||
n_events = len(event_ids)
|
||
idx = np.searchsorted(event_ids, grid["event_id"].to_numpy())
|
||
db = grid["db"].to_numpy()
|
||
tb = grid["tb"].to_numpy()
|
||
cnt = grid["cnt"].to_numpy().astype(np.float64)
|
||
s_edep = grid["s_edep"].to_numpy().astype(np.float64)
|
||
s_edep_depth = grid["s_edep_depth"].to_numpy().astype(np.float64)
|
||
s_edep_t2 = grid["s_edep_t2"].to_numpy().astype(np.float64)
|
||
s_step = grid["s_step"].to_numpy().astype(np.float64)
|
||
|
||
total_edep = np.bincount(idx, weights=s_edep, minlength=n_events)
|
||
total_length = np.bincount(idx, weights=s_step, minlength=n_events)
|
||
n_steps = np.bincount(idx, weights=cnt, minlength=n_events)
|
||
sum_edep_depth = np.bincount(idx, weights=s_edep_depth, minlength=n_events)
|
||
sum_edep_t2 = np.bincount(idx, weights=s_edep_t2, minlength=n_events)
|
||
depth_mat = np.bincount(
|
||
idx * depth_bins + db, weights=s_edep, minlength=n_events * depth_bins
|
||
).reshape(n_events, depth_bins)
|
||
transverse_mat = np.bincount(
|
||
idx * transverse_bins + tb, weights=s_edep, minlength=n_events * transverse_bins
|
||
).reshape(n_events, transverse_bins)
|
||
|
||
safe_total = np.where(total_edep > 0, total_edep, 1.0)
|
||
safe_n = np.where(n_steps > 0, n_steps, 1)
|
||
depth_centers = 0.5 * (depth_edges[:-1] + depth_edges[1:])
|
||
return {
|
||
"total_edep": total_edep,
|
||
"total_length": total_length,
|
||
"n_steps": n_steps.astype(np.int64),
|
||
"mean_edep": total_edep / safe_n,
|
||
"mean_length": total_length / safe_n,
|
||
"centroid_depth": sum_edep_depth / safe_total,
|
||
"transverse_rms": np.sqrt(sum_edep_t2 / safe_total),
|
||
"max_depth": depth_centers[np.argmax(depth_mat, axis=1)],
|
||
"depth_mat": depth_mat,
|
||
"transverse_mat": transverse_mat,
|
||
}
|
||
|
||
|
||
def _approx_medians(geo: pl.LazyFrame, event_ids: np.ndarray) -> dict[str, np.ndarray]:
|
||
"""Per-event approximate median edep/step for real & gen, via a streaming pass.
|
||
|
||
Histograms each quantity into fixed log10 bins per event (one streaming
|
||
`group_by` after unpivoting only the four bin-index columns — a bounded 4×
|
||
expansion), then reads off the bin its running count crosses half at. See the
|
||
`_MED_*` note for why exact medians are avoided.
|
||
"""
|
||
|
||
def _log_bin(col: str) -> pl.Expr:
|
||
log10 = pl.col(col).clip(lower_bound=1e-30).log10()
|
||
return (
|
||
((log10 - _MED_LOG10_LO) / (_MED_LOG10_HI - _MED_LOG10_LO) * _MED_BINS)
|
||
.floor()
|
||
.cast(pl.Int64)
|
||
.clip(0, _MED_BINS - 1)
|
||
)
|
||
|
||
hist = (
|
||
geo.select(
|
||
"event_id",
|
||
_log_bin("real_edep").alias("real_median_edep"),
|
||
_log_bin("gen_edep").alias("gen_median_edep"),
|
||
_log_bin("real_step").alias("real_median_length"),
|
||
_log_bin("gen_step").alias("gen_median_length"),
|
||
)
|
||
.unpivot(index="event_id", variable_name="q", value_name="bin")
|
||
.group_by(["q", "event_id", "bin"])
|
||
.agg(pl.len().alias("c"))
|
||
.collect(engine="streaming")
|
||
)
|
||
|
||
n_events = len(event_ids)
|
||
log_width = (_MED_LOG10_HI - _MED_LOG10_LO) / _MED_BINS
|
||
rows = np.arange(n_events)
|
||
out: dict[str, np.ndarray] = {}
|
||
for name in _MED_QUANTITIES:
|
||
sub = hist.filter(pl.col("q") == name)
|
||
idx = np.searchsorted(event_ids, sub["event_id"].to_numpy())
|
||
counts = np.bincount(
|
||
idx * _MED_BINS + sub["bin"].to_numpy(),
|
||
weights=sub["c"].to_numpy().astype(np.float64),
|
||
minlength=n_events * _MED_BINS,
|
||
).reshape(n_events, _MED_BINS)
|
||
cum = np.cumsum(counts, axis=1)
|
||
total = cum[:, -1]
|
||
half = total / 2.0
|
||
median_bin = (cum >= half[:, None]).argmax(axis=1)
|
||
# Linear interpolation of the log-CDF within the median bin, for sub-bin
|
||
# resolution (a bare bin center is only ~12% granular).
|
||
cum_before = np.where(median_bin > 0, cum[rows, median_bin - 1], 0.0)
|
||
count_in = counts[rows, median_bin]
|
||
frac = np.where(count_in > 0, (half - cum_before) / count_in, 0.5)
|
||
log_val = _MED_LOG10_LO + (median_bin + frac) * log_width
|
||
median = 10.0**log_val
|
||
median[total == 0] = 0.0
|
||
out[name] = median
|
||
return out
|
||
|
||
|
||
def compute_event_observables_pl(
|
||
source: str | Path | pl.LazyFrame,
|
||
depth_bins: int = 20,
|
||
transverse_bins: int = 20,
|
||
) -> EventObservables:
|
||
"""Stream a `giant predict --coord local` parquet file into event-level observables.
|
||
|
||
For each event, the highest-`pre_E` row is the primary's entry step
|
||
(secondaries always carry less energy than their parent), fixing a shower
|
||
axis/entry point shared by real and generated rows. Every row's `post_pos`/
|
||
`edep` is reconstructed into the world frame in physical units (mm, MeV) and
|
||
projected onto depth-along-axis / transverse-distance-from-axis.
|
||
|
||
Everything is computed with parallel streaming polars aggregations — no serial
|
||
python-over-pyarrow-batch loop, and no holistic per-event median (medians are
|
||
approximated from a streaming log-bin histogram) — so the 323M-row reduction
|
||
both uses all cores and stays within a bounded memory budget on a file larger
|
||
than RAM. `event_table` carries `real_total_length`/`gen_total_length`
|
||
(`sum(step_length)` per event) alongside the deposited-energy observables.
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
entry_df, depth_edges, transverse_edges = _entry_axis_and_bin_edges(
|
||
lf, depth_bins, transverse_bins
|
||
)
|
||
event_ids = entry_df["event_id"].to_numpy()
|
||
|
||
# Project to just the geometry inputs *before* the join, and to just the
|
||
# geometry outputs right after: projection pushdown doesn't reliably prune
|
||
# across a join in this polars version, so without the explicit narrowing the
|
||
# join buffers all ~30 columns × 323M rows and OOMs (the same trap the
|
||
# marginal path documents). Narrowed, only ~9 columns cross the join.
|
||
geo_inputs = [
|
||
"event_id",
|
||
"pre_x",
|
||
"pre_y",
|
||
"pre_z",
|
||
"pre_dx",
|
||
"pre_dy",
|
||
"pre_dz",
|
||
"pre_E",
|
||
*[
|
||
f"{prefix}_{name}"
|
||
for prefix in ("true", "pred")
|
||
for name in (
|
||
"log_step_length",
|
||
"edep_logit",
|
||
"sec_logit",
|
||
"travel_dx",
|
||
"travel_dy",
|
||
"travel_dz",
|
||
)
|
||
],
|
||
]
|
||
geo = (
|
||
lf.select(geo_inputs)
|
||
.join(entry_df.lazy(), on="event_id")
|
||
.select(
|
||
"event_id",
|
||
*_geometry_exprs("true", "real"),
|
||
*_geometry_exprs("pred", "gen"),
|
||
)
|
||
)
|
||
|
||
real = _reduce_spatial_grid(
|
||
_spatial_grid(
|
||
geo, "real", depth_edges, transverse_edges, depth_bins, transverse_bins
|
||
),
|
||
event_ids,
|
||
depth_bins,
|
||
transverse_bins,
|
||
depth_edges,
|
||
)
|
||
gen = _reduce_spatial_grid(
|
||
_spatial_grid(
|
||
geo, "gen", depth_edges, transverse_edges, depth_bins, transverse_bins
|
||
),
|
||
event_ids,
|
||
depth_bins,
|
||
transverse_bins,
|
||
depth_edges,
|
||
)
|
||
medians = _approx_medians(geo, event_ids)
|
||
|
||
event_table = pl.DataFrame(
|
||
{
|
||
"event_id": event_ids,
|
||
"n_steps": real["n_steps"],
|
||
"real_total_edep": real["total_edep"],
|
||
"gen_total_edep": gen["total_edep"],
|
||
"real_total_length": real["total_length"],
|
||
"gen_total_length": gen["total_length"],
|
||
"real_mean_edep": real["mean_edep"],
|
||
"gen_mean_edep": gen["mean_edep"],
|
||
"real_mean_length": real["mean_length"],
|
||
"gen_mean_length": gen["mean_length"],
|
||
"real_median_edep": medians["real_median_edep"],
|
||
"gen_median_edep": medians["gen_median_edep"],
|
||
"real_median_length": medians["real_median_length"],
|
||
"gen_median_length": medians["gen_median_length"],
|
||
"real_centroid_depth": real["centroid_depth"],
|
||
"gen_centroid_depth": gen["centroid_depth"],
|
||
"real_transverse_rms": real["transverse_rms"],
|
||
"gen_transverse_rms": gen["transverse_rms"],
|
||
"real_max_depth": real["max_depth"],
|
||
"gen_max_depth": gen["max_depth"],
|
||
}
|
||
)
|
||
|
||
return EventObservables(
|
||
event_table=event_table,
|
||
depth_edges=depth_edges,
|
||
transverse_edges=transverse_edges,
|
||
real_depth_profile=real["depth_mat"].mean(axis=0),
|
||
gen_depth_profile=gen["depth_mat"].mean(axis=0),
|
||
real_depth_profile_std=real["depth_mat"].std(axis=0),
|
||
gen_depth_profile_std=gen["depth_mat"].std(axis=0),
|
||
real_transverse_profile=real["transverse_mat"].mean(axis=0),
|
||
gen_transverse_profile=gen["transverse_mat"].mean(axis=0),
|
||
real_transverse_profile_std=real["transverse_mat"].std(axis=0),
|
||
gen_transverse_profile_std=gen["transverse_mat"].std(axis=0),
|
||
)
|
||
|
||
|
||
def plot_total_energy(observables: EventObservables, bins: int = 50):
|
||
"""Real-vs-generated histogram of total deposited energy per event, with resolution."""
|
||
table = observables.event_table
|
||
real = table["real_total_edep"].to_numpy()
|
||
gen = table["gen_total_edep"].to_numpy()
|
||
|
||
fig, ax = plt.subplots(figsize=(6, 4))
|
||
edges = _hist_edges(real, gen, bins=bins).tolist()
|
||
ax.hist(
|
||
real,
|
||
bins=edges,
|
||
density=True,
|
||
histtype="step",
|
||
label=f"real (σ/μ={real.std() / real.mean():.3f})",
|
||
)
|
||
ax.hist(
|
||
gen,
|
||
bins=edges,
|
||
density=True,
|
||
histtype="step",
|
||
label=f"generated (σ/μ={gen.std() / gen.mean():.3f})",
|
||
)
|
||
ax.set_yscale("log")
|
||
ax.set_xlabel("total deposited energy per event [MeV]")
|
||
ax.legend(fontsize=8)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def plot_total_length(observables: EventObservables, bins: int = 50):
|
||
"""Real-vs-generated histogram of total length traveled per event (sum of step_length)."""
|
||
table = observables.event_table
|
||
real = table["real_total_length"].to_numpy()
|
||
gen = table["gen_total_length"].to_numpy()
|
||
|
||
fig, ax = plt.subplots(figsize=(6, 4))
|
||
edges = _hist_edges(real, gen, bins=bins).tolist()
|
||
ax.hist(
|
||
real,
|
||
bins=edges,
|
||
density=True,
|
||
histtype="step",
|
||
label=f"real (σ/μ={real.std() / real.mean():.3f})",
|
||
)
|
||
ax.hist(
|
||
gen,
|
||
bins=edges,
|
||
density=True,
|
||
histtype="step",
|
||
label=f"generated (σ/μ={gen.std() / gen.mean():.3f})",
|
||
)
|
||
ax.set_yscale("log")
|
||
ax.set_xlabel("total length traveled per event [mm]")
|
||
ax.legend(fontsize=8)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def _plot_mean_median_per_step(
|
||
real_mean: np.ndarray,
|
||
gen_mean: np.ndarray,
|
||
real_median: np.ndarray,
|
||
gen_median: np.ndarray,
|
||
mean_xlabel: str,
|
||
median_xlabel: str,
|
||
bins: int,
|
||
median_bins: int,
|
||
):
|
||
"""Side-by-side real-vs-generated histograms: per-event mean (left), median (right).
|
||
|
||
The mean is pulled down by a compressed/under-sampled right tail (rare large
|
||
values), while the median is robust to that tail. Splitting them into separate
|
||
panels keeps each comparison legible; if a tail-compression bias explains a
|
||
low generated mean, the median panel should overlap far more closely. Each
|
||
panel gets its own bin edges (`median_bins` larger, since the median is far
|
||
less spread than the mean).
|
||
"""
|
||
mean_edges = _hist_edges(real_mean, gen_mean, bins=bins)
|
||
median_edges = _hist_edges(real_median, gen_median, bins=median_bins)
|
||
prop_colors = plt.rcParams["axes.prop_cycle"].by_key()["color"]
|
||
real_color, gen_color = prop_colors[0], prop_colors[1]
|
||
|
||
fig, (ax_mean, ax_median) = plt.subplots(1, 2, figsize=(12, 4))
|
||
for ax, real, gen, edges, title, xlabel in [
|
||
(ax_mean, real_mean, gen_mean, mean_edges, "mean", mean_xlabel),
|
||
(ax_median, real_median, gen_median, median_edges, "median", median_xlabel),
|
||
]:
|
||
ax.hist(
|
||
real,
|
||
bins=edges,
|
||
density=True,
|
||
histtype="step",
|
||
color=real_color,
|
||
label=f"real (σ/μ={real.std() / real.mean():.3f})",
|
||
)
|
||
ax.hist(
|
||
gen,
|
||
bins=edges,
|
||
density=True,
|
||
histtype="step",
|
||
color=gen_color,
|
||
label=f"generated (σ/μ={gen.std() / gen.mean():.3f})",
|
||
)
|
||
ax.set_yscale("log")
|
||
ax.set_title(title)
|
||
ax.set_xlabel(xlabel)
|
||
ax.legend(fontsize=8)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def plot_mean_energy_per_step(
|
||
observables: EventObservables, bins: int = 50, median_bins: int = 150
|
||
):
|
||
"""Real-vs-generated histograms of mean/median deposited energy per step, per event.
|
||
|
||
Per event: `total_edep / n_steps` (mean) and the per-step median edep —
|
||
distinct from `plot_total_energy`, which histograms the per-event *total*;
|
||
these instead ask whether the typical step's energy deposit is right,
|
||
independent of how many steps the event happened to have.
|
||
"""
|
||
table = observables.event_table
|
||
return _plot_mean_median_per_step(
|
||
table["real_mean_edep"].to_numpy(),
|
||
table["gen_mean_edep"].to_numpy(),
|
||
table["real_median_edep"].to_numpy(),
|
||
table["gen_median_edep"].to_numpy(),
|
||
mean_xlabel="mean deposited energy per step, per event [MeV]",
|
||
median_xlabel="median deposited energy per step, per event [MeV]",
|
||
bins=bins,
|
||
median_bins=median_bins,
|
||
)
|
||
|
||
|
||
def plot_mean_length_per_step(
|
||
observables: EventObservables, bins: int = 50, median_bins: int = 150
|
||
):
|
||
"""Real-vs-generated histograms of mean/median step length per step, per event.
|
||
|
||
Per event: `total_length / n_steps` (mean) and the per-step median step length
|
||
— distinct from `plot_total_length`, which histograms the per-event *total*;
|
||
these instead ask whether the typical step length is right, independent of how
|
||
many steps the event happened to have.
|
||
"""
|
||
table = observables.event_table
|
||
return _plot_mean_median_per_step(
|
||
table["real_mean_length"].to_numpy(),
|
||
table["gen_mean_length"].to_numpy(),
|
||
table["real_median_length"].to_numpy(),
|
||
table["gen_median_length"].to_numpy(),
|
||
mean_xlabel="mean step length per step, per event [mm]",
|
||
median_xlabel="median step length per step, per event [mm]",
|
||
bins=bins,
|
||
median_bins=median_bins,
|
||
)
|
||
|
||
|
||
def _plot_profile(
|
||
centers: np.ndarray,
|
||
real_mean: np.ndarray,
|
||
gen_mean: np.ndarray,
|
||
real_std: np.ndarray,
|
||
gen_std: np.ndarray,
|
||
xlabel: str,
|
||
):
|
||
fig, ax = plt.subplots(figsize=(6, 4))
|
||
ax.errorbar(centers, real_mean, yerr=real_std, fmt="o-", label="real", capsize=2)
|
||
ax.errorbar(centers, gen_mean, yerr=gen_std, fmt="s-", label="generated", capsize=2)
|
||
ax.set_xlabel(xlabel)
|
||
ax.set_ylabel("mean edep per event per bin [MeV]")
|
||
ax.legend(fontsize=8)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def plot_longitudinal_profile(observables: EventObservables):
|
||
"""E_dep(depth) mean ± event-to-event RMS, real vs generated."""
|
||
centers = 0.5 * (observables.depth_edges[:-1] + observables.depth_edges[1:])
|
||
return _plot_profile(
|
||
centers,
|
||
observables.real_depth_profile,
|
||
observables.gen_depth_profile,
|
||
observables.real_depth_profile_std,
|
||
observables.gen_depth_profile_std,
|
||
"depth along shower axis [mm]",
|
||
)
|
||
|
||
|
||
def plot_transverse_profile(observables: EventObservables):
|
||
"""E_dep(transverse distance) mean ± event-to-event RMS, real vs generated (Molière-style)."""
|
||
centers = 0.5 * (
|
||
observables.transverse_edges[:-1] + observables.transverse_edges[1:]
|
||
)
|
||
return _plot_profile(
|
||
centers,
|
||
observables.real_transverse_profile,
|
||
observables.gen_transverse_profile,
|
||
observables.real_transverse_profile_std,
|
||
observables.gen_transverse_profile_std,
|
||
"transverse distance from shower axis [mm]",
|
||
)
|
||
|
||
|
||
def plot_shower_max_depth(observables: EventObservables, bins: int = 30):
|
||
"""Real-vs-generated histogram of per-event shower-maximum depth."""
|
||
table = observables.event_table
|
||
real = table["real_max_depth"].to_numpy()
|
||
gen = table["gen_max_depth"].to_numpy()
|
||
|
||
fig, ax = plt.subplots(figsize=(6, 4))
|
||
edges = _hist_edges(real, gen, bins=bins).tolist()
|
||
ax.hist(real, bins=edges, density=True, histtype="step", label="real")
|
||
ax.hist(gen, bins=edges, density=True, histtype="step", label="generated")
|
||
ax.set_xlabel("depth of shower maximum [mm]")
|
||
ax.legend(fontsize=8)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Particle-species (pdg) contribution shares
|
||
#
|
||
# Dataset-wide (not per-event) breakdown of which pdg species contributed how
|
||
# much of the total deposited energy / total length traveled. Only needs scalar
|
||
# sums — no post_pos reconstruction, no shower axis — so it's a single lazy
|
||
# polars group_by.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
_PDG_NAMES = {11: "e-", -11: "e+", 22: "gamma", 2112: "n", 2212: "p"}
|
||
|
||
|
||
def _pdg_label(pdg: int) -> str:
|
||
if pdg in _PDG_NAMES:
|
||
return _PDG_NAMES[pdg]
|
||
if abs(pdg) > 1_000_000_000:
|
||
return f"ion{pdg}"
|
||
return str(pdg)
|
||
|
||
|
||
def pdg_contribution_table_pl(source: str | Path | pl.LazyFrame) -> pl.DataFrame:
|
||
"""Total edep / step_length contributed by each pdg species, real vs generated.
|
||
|
||
One row per pdg code, sorted by pdg. Pure lazy polars `group_by` over the
|
||
whole file — `edep`/`step_length` are scalars unaffected by the local-frame
|
||
rotation. `step_length` is a simple `exp(...) - eps` de-log; `edep` is decoded
|
||
from the deposit/secondary energy logits against pre_E (`_edep_pl`).
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
|
||
def _delog(col: str) -> pl.Expr:
|
||
return pl.col(col).exp() - _LOG_EPS
|
||
|
||
return (
|
||
lf.group_by("pdg")
|
||
.agg(
|
||
_edep_pl("true").sum().alias("real_total_edep"),
|
||
_edep_pl("pred").sum().alias("gen_total_edep"),
|
||
_delog("true_log_step_length").sum().alias("real_total_length"),
|
||
_delog("pred_log_step_length").sum().alias("gen_total_length"),
|
||
)
|
||
.collect(engine="streaming")
|
||
.sort("pdg")
|
||
)
|
||
|
||
|
||
def _pdg_pie_shares(
|
||
table: pl.DataFrame, real_col: str, gen_col: str, max_slices: int
|
||
) -> tuple[list[str], np.ndarray, np.ndarray]:
|
||
"""Pie-ready (labels, real_values, gen_values), lumping small contributors into 'other'.
|
||
|
||
Ranked by combined real+gen contribution so the same species end up in the
|
||
same slice position in both pies, making them easier to compare.
|
||
"""
|
||
pdg = table["pdg"].to_numpy()
|
||
real = table[real_col].to_numpy()
|
||
gen = table[gen_col].to_numpy()
|
||
|
||
order = np.argsort(-(real + gen))
|
||
pdg, real, gen = pdg[order], real[order], gen[order]
|
||
|
||
if len(pdg) > max_slices:
|
||
keep = max_slices - 1
|
||
labels = [_pdg_label(int(p)) for p in pdg[:keep]] + ["other"]
|
||
real = np.append(real[:keep], real[keep:].sum())
|
||
gen = np.append(gen[:keep], gen[keep:].sum())
|
||
else:
|
||
labels = [_pdg_label(int(p)) for p in pdg]
|
||
|
||
return labels, real, gen
|
||
|
||
|
||
def _plot_pdg_pie(
|
||
table: pl.DataFrame, real_col: str, gen_col: str, suptitle: str, max_slices: int
|
||
):
|
||
labels, real_vals, gen_vals = _pdg_pie_shares(table, real_col, gen_col, max_slices)
|
||
|
||
fig, axes = plt.subplots(1, 2, figsize=(9, 4.5))
|
||
for ax, vals, title in [
|
||
(axes[0], real_vals, "real"),
|
||
(axes[1], gen_vals, "generated"),
|
||
]:
|
||
ax.pie(vals, labels=labels, autopct="%1.1f%%", startangle=90)
|
||
ax.set_title(title)
|
||
fig.suptitle(suptitle)
|
||
fig.tight_layout()
|
||
return fig
|
||
|
||
|
||
def plot_pdg_energy_share(table: pl.DataFrame, max_slices: int = 6):
|
||
"""Real-vs-generated pies of total deposited energy share by pdg species."""
|
||
return _plot_pdg_pie(
|
||
table,
|
||
"real_total_edep",
|
||
"gen_total_edep",
|
||
"deposited energy share by particle type",
|
||
max_slices,
|
||
)
|
||
|
||
|
||
def plot_pdg_length_share(table: pl.DataFrame, max_slices: int = 6):
|
||
"""Real-vs-generated pies of total length-traveled share by pdg species."""
|
||
return _plot_pdg_pie(
|
||
table,
|
||
"real_total_length",
|
||
"gen_total_length",
|
||
"length traveled share by particle type",
|
||
max_slices,
|
||
)
|