lars
25718f175e
Removes unused imports and an ambiguous variable name, narrows Optional types before use so ty's flow analysis is satisfied, swaps sum() over polars expressions for pl.sum_horizontal to avoid the Literal[0] fallback type, and converts numpy bin edges to plain lists before passing to matplotlib's hist (whose stub only accepts Sequence[float]). Also applies ruff format across the repo, which had drifted out of sync with the formatter. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
1696 lines
67 KiB
Python
1696 lines
67 KiB
Python
"""Notebook-friendly diagnostics for a trained model's sample quality.
|
||
|
||
All checks here read from `giant predict --coord local` parquet output
|
||
(`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled) —
|
||
there is no on-the-fly (checkpoint + live sampler) path; generate predictions
|
||
once via the CLI, then run every diagnostic below against that file::
|
||
|
||
from giant.analysis import load_predicted_local
|
||
from giant.analysis import plot_marginals, plot_kl_bars, plot_correlation_matrices
|
||
from giant.analysis import plot_pairwise, plot_direction_alignment
|
||
from giant.analysis import plot_constraint_violations
|
||
|
||
samples = load_predicted_local("path/to/steps_predicted_local.parquet")
|
||
|
||
plot_marginals(samples, group_by="energy")
|
||
plot_kl_bars(samples, group_by="energy")
|
||
plot_correlation_matrices(samples)
|
||
plot_pairwise(samples)
|
||
plot_direction_alignment(samples)
|
||
plot_constraint_violations(samples)
|
||
|
||
(`--coord global` output isn't supported here — it has no ground-truth columns
|
||
to compare against.)
|
||
|
||
For event-level (shower) observables on the same `--coord local` predict
|
||
output, see `compute_event_observables_pl` — it streams the file directly
|
||
(no row subsampling, no `SampleCollection`), since per-event sums would be
|
||
corrupted by partial events::
|
||
|
||
from giant.analysis import compute_event_observables_pl
|
||
from giant.analysis import plot_total_energy, plot_total_length
|
||
from giant.analysis import plot_longitudinal_profile, plot_transverse_profile
|
||
from giant.analysis import plot_shower_max_depth
|
||
|
||
obs = compute_event_observables_pl("path/to/steps_predicted_local.parquet")
|
||
plot_total_energy(obs)
|
||
plot_total_length(obs)
|
||
plot_mean_energy_per_step(obs)
|
||
plot_mean_length_per_step(obs)
|
||
plot_longitudinal_profile(obs)
|
||
plot_transverse_profile(obs)
|
||
plot_shower_max_depth(obs)
|
||
|
||
For the dataset-wide breakdown of which particle species (pdg) contributed
|
||
how much of the total energy/length, see `pdg_contribution_table_pl`::
|
||
|
||
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("path/to/steps_predicted_local.parquet")
|
||
plot_pdg_energy_share(table)
|
||
plot_pdg_length_share(table)
|
||
|
||
Four tiers of checks, building on the aggregate marginal/KL check in
|
||
`giant.validate.validate_marginals`:
|
||
|
||
1. stratified marginals — per-dimension real-vs-generated comparison, sliced by
|
||
pdg / material / energy so failures hidden by the aggregate don't go unnoticed.
|
||
2. joint structure — correlation matrices, physically-coupled pairwise
|
||
plots, and post/travel direction alignment, since marginals can match while
|
||
the model decorrelates targets that are coupled by the underlying physics.
|
||
3. physical constraints — unit-norm direction vectors and non-negative raw
|
||
step_length/delta_e/edep, checked in denormalized physical units; nothing in
|
||
the unconstrained MLP output enforces these, so violations are a pure
|
||
generation artifact.
|
||
4. event-level observables — total deposited energy, total length traveled,
|
||
longitudinal/transverse shower profiles, and shower-max depth, aggregated
|
||
per `event_id` in the world frame with physical units (mm, MeV). 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.
|
||
"""
|
||
|
||
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,
|
||
)
|
||
from giant.data.transforms import (
|
||
energy_simplex_decode,
|
||
inv_log_transform,
|
||
reconstruct_post_pos,
|
||
)
|
||
from giant.validate import _histogram_kl
|
||
|
||
# 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
|
||
)
|
||
|
||
|
||
def _decode_raw_targets(target_local: np.ndarray, pre_E: np.ndarray) -> np.ndarray:
|
||
"""Map a model-native target array (N, 9) to physical raw units.
|
||
|
||
Column 0 (log_step_length) is de-logged; columns 1–2 (the deposit/secondary
|
||
ALR energy logits) are decoded against `pre_E` into physical `delta_e` and
|
||
`edep` (energy_simplex_decode); the six direction components pass through
|
||
unchanged. Output columns therefore line up with `RAW_TARGET_NAMES`.
|
||
"""
|
||
raw = target_local.astype(np.float32).copy()
|
||
raw[:, 0] = inv_log_transform(target_local[:, 0])
|
||
edep, _e_sec, _post_E, delta_e = energy_simplex_decode(target_local[:, 1:3], pre_E)
|
||
raw[:, 1] = delta_e
|
||
raw[:, 2] = edep
|
||
return raw
|
||
|
||
|
||
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 _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.
|
||
"""
|
||
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)
|
||
|
||
|
||
@dataclass
|
||
class SampleCollection:
|
||
cond_cont_raw: np.ndarray # (N, 9) denormalized conditioning (pre_E delogged)
|
||
pdg: np.ndarray # (N,) raw PDG codes
|
||
material: np.ndarray # (N,) raw material names
|
||
real_raw: np.ndarray # (N, 9) denormalized + delogged real targets
|
||
gen_raw: np.ndarray # (N, 9) denormalized + delogged generated targets
|
||
|
||
|
||
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' — "
|
||
"load_predicted_local 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 load_predicted_local to match"
|
||
)
|
||
|
||
|
||
_COND_CONT_COLS = [
|
||
"pre_x",
|
||
"pre_y",
|
||
"pre_z",
|
||
"pre_E",
|
||
"pre_dx",
|
||
"pre_dy",
|
||
"pre_dz",
|
||
"layer_id",
|
||
"n_sec",
|
||
]
|
||
|
||
|
||
def load_predicted_local(
|
||
path: str | Path,
|
||
sample_frac: float = 1.0,
|
||
seed: int = 0,
|
||
batch_size: int = 1_000_000,
|
||
) -> SampleCollection:
|
||
"""Build a SampleCollection from a `giant predict --coord local` parquet file.
|
||
|
||
Reads the `pred_*`/`true_*` columns directly — no checkpoint or model needed,
|
||
since the predict CLI already denormalized them into this log-scaled,
|
||
local-frame space. Requires the file to carry the `giant predict` metadata
|
||
tag (see `_check_predict_metadata`); raises if it's missing or from
|
||
--coord global, rather than guessing from column names.
|
||
|
||
Reads via `_iter_predicted_local_batches` (pyarrow's row-batch reader)
|
||
with the needed columns selected, so at most `batch_size` rows are ever
|
||
materialized at once — a plain lazy-polars `.collect()` with the sampling
|
||
filter applied afterward looks lazy but doesn't push the row reduction
|
||
into the scan (see its `.explain()`), so it still peaks at the full
|
||
file's memory footprint even for a small `sample_frac`; streaming keeps
|
||
peak memory to one batch regardless of file size or `sample_frac`.
|
||
|
||
`sample_frac` (0 < sample_frac <= 1) randomly keeps only that fraction of
|
||
rows after the column projection — useful for files too large to
|
||
comfortably hold as numpy arrays in `real_raw`/`gen_raw`. The kept/dropped
|
||
decision is a hash of each row's position in the file (`seed`-dependent),
|
||
computed against a running offset across batches so it's equivalent to
|
||
hashing a single row index over the whole file rather than restarting at
|
||
each batch boundary.
|
||
"""
|
||
if not (0 < sample_frac <= 1):
|
||
raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
|
||
_check_predict_metadata(Path(path))
|
||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||
true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
|
||
columns = pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"]
|
||
|
||
threshold = int(sample_frac * 2**32) if sample_frac < 1.0 else None
|
||
frames = []
|
||
offset = 0
|
||
for batch_df in _iter_predicted_local_batches(path, columns, batch_size):
|
||
n = batch_df.height
|
||
if threshold is not None:
|
||
row_idx = pl.arange(offset, offset + n, eager=True).cast(pl.UInt32)
|
||
batch_df = batch_df.filter((row_idx.hash(seed=seed) % 2**32) < threshold)
|
||
offset += n
|
||
frames.append(batch_df)
|
||
df = pl.concat(frames) if len(frames) != 1 else frames[0]
|
||
|
||
gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
|
||
real_log_local = df.select(true_cols).to_numpy().astype(np.float32)
|
||
|
||
cond_cont_raw = df.select(_COND_CONT_COLS).to_numpy().astype(np.float32)
|
||
pre_E = cond_cont_raw[:, 3] # pre_E is already physical in predict output
|
||
|
||
return SampleCollection(
|
||
cond_cont_raw=cond_cont_raw,
|
||
pdg=df["pdg"].to_numpy(),
|
||
material=df["material"].to_numpy(),
|
||
real_raw=_decode_raw_targets(real_log_local, pre_E),
|
||
gen_raw=_decode_raw_targets(gen_log_local, pre_E),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Tier 1: stratified marginals
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _group_labels(
|
||
collection: SampleCollection,
|
||
group_by: str | None,
|
||
n_energy_bins: int,
|
||
) -> list[tuple[str, np.ndarray]]:
|
||
n = len(collection.pdg)
|
||
if group_by is None:
|
||
return [("all", np.ones(n, dtype=bool))]
|
||
if group_by == "pdg":
|
||
return [(f"pdg={v}", collection.pdg == v) for v in np.unique(collection.pdg)]
|
||
if group_by == "material":
|
||
return [
|
||
(f"material={v}", collection.material == v)
|
||
for v in np.unique(collection.material)
|
||
]
|
||
if group_by == "energy":
|
||
pre_E = collection.cond_cont_raw[:, 3]
|
||
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(f"unknown group_by={group_by!r}")
|
||
|
||
|
||
def marginal_table(
|
||
collection: SampleCollection,
|
||
group_by: str | None = None,
|
||
n_energy_bins: int = 4,
|
||
bins: int = 50,
|
||
) -> pd.DataFrame:
|
||
"""Per-dimension real-vs-generated summary stats + KL(real||gen), in raw units.
|
||
|
||
`group_by`: None for an aggregate table, or one of "pdg", "material", "energy"
|
||
to stratify rows by that conditioning variable. Sorted worst-KL first, so
|
||
failure modes hidden by the aggregate surface at the top.
|
||
"""
|
||
rows = []
|
||
for label, mask in _group_labels(collection, group_by, n_energy_bins):
|
||
if mask.sum() < 2:
|
||
continue
|
||
real, gen = collection.real_raw[mask], collection.gen_raw[mask]
|
||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||
rows.append(
|
||
{
|
||
"group": label,
|
||
"dim": name,
|
||
"n": int(mask.sum()),
|
||
"real_mean": real[:, j].mean(),
|
||
"gen_mean": gen[:, j].mean(),
|
||
"real_std": real[:, j].std(),
|
||
"gen_std": gen[:, j].std(),
|
||
"kl_real_gen": _histogram_kl(real[:, j], gen[:, j], bins=bins),
|
||
}
|
||
)
|
||
return (
|
||
pd.DataFrame(rows)
|
||
.sort_values("kl_real_gen", ascending=False)
|
||
.reset_index(drop=True)
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Polars-native duplicates
|
||
#
|
||
# `marginal_table`/`constraint_report` above require a `SampleCollection`
|
||
# with the full real/gen arrays already materialized in numpy. The functions
|
||
# below instead take a parquet path (or LazyFrame) directly and stay lazy,
|
||
# reading the file a small constant number of times — one native `group_by`
|
||
# (covering every group and every dim's mean/std/n/histogram-range at once)
|
||
# plus one more pass per dim for histogram bin counts — rather than filtering
|
||
# and re-collecting once per (group, dim) pair. The latter (the previous
|
||
# implementation) meant runtime scaled with the number of *groups*: on a
|
||
# 114M-row file with 138 distinct pdg codes, `group_by="pdg"` extrapolated to
|
||
# roughly an hour, against ~1 minute for `group_by=None`. This version's
|
||
# runtime is independent of group cardinality.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
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.
|
||
|
||
Same smoothing/normalization as `giant.validate._histogram_kl`, just
|
||
taking counts directly instead of raw samples (the counts here come from
|
||
a polars `group_by` aggregation, not `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 _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 _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 _add_group_label(
|
||
lf: pl.LazyFrame, group_by: str | None, n_energy_bins: int
|
||
) -> pl.LazyFrame:
|
||
"""Add a `_group` string column matching `_group_labels`'s label format.
|
||
|
||
For "pdg"/"material" the label is a direct string expr over the existing
|
||
column — no upfront pass to enumerate distinct values needed, since
|
||
`group_by("_group")` downstream discovers them itself. "energy" still
|
||
needs one 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}")
|
||
|
||
|
||
def _dim_narrow_lf(lf: pl.LazyFrame, j: int) -> pl.LazyFrame:
|
||
"""Project down to just `(_group, real, gen)` for one dim, before any join/group_by.
|
||
|
||
Polars' projection pushdown doesn't reliably prune columns across a
|
||
`.join()` in this version — without this explicit `.select()` up front,
|
||
`_dim_hist_counts`'s join ends up materializing every column of the file
|
||
(~30 columns × 114M rows ≈ 13GB) instead of just the ~2-3 this dim needs,
|
||
even though nothing downstream references the others. Selecting first
|
||
guarantees the narrow projection regardless of what the optimizer would
|
||
otherwise infer.
|
||
"""
|
||
return lf.select(
|
||
"_group",
|
||
_raw_dim_expr("true", j).alias("real"),
|
||
_raw_dim_expr("pred", j).alias("gen"),
|
||
)
|
||
|
||
|
||
def _dim_stats(narrow: pl.LazyFrame) -> pl.DataFrame:
|
||
"""Per-group mean/std/n plus histogram-range (min/max of real+gen) for one dim.
|
||
|
||
A single `group_by("_group")` pass over `narrow` (see `_dim_narrow_lf`),
|
||
mirroring the proven-cheap shape of `pdg_contribution_table_pl`, rather
|
||
than folding all 9 dims' source columns into one query — the latter reads
|
||
~19 columns and builds every dim's intermediate (softmax, log, etc.)
|
||
arrays for the whole file at once, which is memory-heavy enough to OOM on
|
||
a 114M-row file even though it's still just one pass.
|
||
"""
|
||
return (
|
||
narrow.group_by("_group")
|
||
.agg(
|
||
pl.len().alias("n"),
|
||
pl.col("real").mean().alias("real_mean"),
|
||
pl.col("gen").mean().alias("gen_mean"),
|
||
# ddof=0 to match numpy's (population-std) default used by marginal_table
|
||
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 _pad_degenerate_range(lo_hi: pl.DataFrame) -> pl.DataFrame:
|
||
"""Widen a (lo, hi) pair by ±0.5 when too tight to support `bins` distinct edges.
|
||
|
||
Polars duplicate of `_hist_edges`'s degenerate-range handling, applied
|
||
per group row instead of per call.
|
||
"""
|
||
ok = (pl.col("hi") - pl.col("lo")) > 1e-6 * pl.max_horizontal(
|
||
pl.col("hi").abs(), pl.lit(1.0)
|
||
)
|
||
return lo_hi.with_columns(
|
||
pl.when(ok).then(pl.col("lo")).otherwise(pl.col("lo") - 0.5).alias("lo"),
|
||
pl.when(ok).then(pl.col("hi")).otherwise(pl.col("hi") + 0.5).alias("hi"),
|
||
)
|
||
|
||
|
||
def _dim_hist_counts(
|
||
narrow: pl.LazyFrame, lo_hi: pl.DataFrame, bins: int
|
||
) -> tuple[pl.DataFrame, pl.DataFrame]:
|
||
"""Per-group real/gen histogram bin counts for one dim, as two long tables.
|
||
|
||
Joins each row (from `narrow`, see `_dim_narrow_lf`) to its group's
|
||
(lo, hi) range (from `lo_hi`, already collected and tiny — one row per
|
||
group), bins real/gen into `[0, bins)`, and counts via a two-key
|
||
`group_by(["_group", "bin"])` — a proper single hash-pass histogram.
|
||
(An earlier version aggregated with one `(bin == k).sum()` expression per
|
||
bin, i.e. `bins` separate boolean-compare-and-reduce passes over every
|
||
row; that's O(N × bins) work — 50 bins meant ~50x more comparisons than
|
||
necessary — and was the dominant cost, not the join.)
|
||
"""
|
||
width = pl.col("hi") - pl.col("lo")
|
||
real_bin = (
|
||
((pl.col("real") - pl.col("lo")) / width * bins)
|
||
.floor()
|
||
.cast(pl.Int64)
|
||
.clip(0, bins - 1)
|
||
)
|
||
gen_bin = (
|
||
((pl.col("gen") - pl.col("lo")) / width * bins)
|
||
.floor()
|
||
.cast(pl.Int64)
|
||
.clip(0, bins - 1)
|
||
)
|
||
|
||
joined = narrow.join(lo_hi.lazy(), on="_group")
|
||
real_hist = (
|
||
joined.select("_group", real_bin.alias("bin"))
|
||
.group_by(["_group", "bin"])
|
||
.agg(pl.len().alias("count"))
|
||
.collect(engine="streaming")
|
||
)
|
||
gen_hist = (
|
||
joined.select("_group", gen_bin.alias("bin"))
|
||
.group_by(["_group", "bin"])
|
||
.agg(pl.len().alias("count"))
|
||
.collect(engine="streaming")
|
||
)
|
||
return real_hist, gen_hist
|
||
|
||
|
||
def _hist_counts_by_group(hist: pl.DataFrame, bins: int) -> dict[str, np.ndarray]:
|
||
"""Long `(_group, bin, count)` table -> `{group: dense (bins,) count array}`."""
|
||
out: dict[str, np.ndarray] = {}
|
||
for group, bin_idx, count in hist.iter_rows():
|
||
out.setdefault(group, np.zeros(bins, dtype=np.int64))[bin_idx] = count
|
||
return out
|
||
|
||
|
||
def marginal_table_pl(
|
||
source: str | Path | pl.LazyFrame,
|
||
group_by: str | None = None,
|
||
n_energy_bins: int = 4,
|
||
bins: int = 50,
|
||
) -> pl.DataFrame:
|
||
"""Polars duplicate of `marginal_table`, reading straight from a predict parquet.
|
||
|
||
`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). Never builds a `SampleCollection` —
|
||
see the module-level note above on why this stays lazy, and on why this
|
||
reads the file three passes per dim (stats, real histogram, gen
|
||
histogram) rather than once per (group, dim) pair.
|
||
"""
|
||
lf = _add_group_label(_scan_predicted_local(source), group_by, n_energy_bins)
|
||
|
||
rows = []
|
||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||
narrow = _dim_narrow_lf(lf, j)
|
||
stats = _dim_stats(narrow)
|
||
stats = stats.filter(pl.col("n") >= 2)
|
||
lo_hi = _pad_degenerate_range(stats.select("_group", "lo", "hi"))
|
||
real_hist, gen_hist = _dim_hist_counts(narrow, lo_hi, bins)
|
||
real_by_group = _hist_counts_by_group(real_hist, bins)
|
||
gen_by_group = _hist_counts_by_group(gen_hist, bins)
|
||
for row in stats.iter_rows(named=True):
|
||
group = row["_group"]
|
||
real_counts = real_by_group.get(group, np.zeros(bins, dtype=np.int64))
|
||
gen_counts = gen_by_group.get(group, np.zeros(bins, dtype=np.int64))
|
||
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),
|
||
}
|
||
)
|
||
return pl.DataFrame(rows).sort("kl_real_gen", descending=True)
|
||
|
||
|
||
def plot_marginals(
|
||
collection: SampleCollection,
|
||
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 val set. 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
|
||
in the dataset surface first, rather than rare groups with a noisy,
|
||
high-variance KL estimate from just one or two samples.
|
||
"""
|
||
dims = dims or RAW_TARGET_NAMES
|
||
dim_idx = [RAW_TARGET_NAMES.index(d) for d in dims]
|
||
groups = _group_labels(collection, group_by, n_energy_bins)
|
||
|
||
if group_by is not None:
|
||
table = marginal_table(
|
||
collection, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
|
||
)
|
||
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)
|
||
order = {label: rank for rank, label in enumerate(worst_first.index)}
|
||
groups = [g for g in groups if g[0] in order]
|
||
groups = sorted(groups, key=lambda g: order[g[0]])[: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, mask) in enumerate(groups):
|
||
real, gen = collection.real_raw[mask], collection.gen_raw[mask]
|
||
for col, j in enumerate(dim_idx):
|
||
ax = axes[row][col]
|
||
edges = _hist_edges(real[:, j], gen[:, j], bins=bins)
|
||
ax.hist(real[:, j], bins=edges, density=True, histtype="step", label="real")
|
||
ax.hist(
|
||
gen[:, j], bins=edges, density=True, histtype="step", label="generated"
|
||
)
|
||
ax.set_yscale("log")
|
||
if row == 0:
|
||
ax.set_title(dims[col], 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 _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, like `plot_marginals`.
|
||
|
||
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 _plot_kl_bars(table: pd.DataFrame, figsize: tuple[float, float]):
|
||
"""Shared bar-plot body for `plot_kl_bars`/`plot_kl_bars_pl`.
|
||
|
||
`table` is a `marginal_table`/`marginal_table_pl` result (already converted
|
||
to pandas in the polars case) 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(
|
||
collection: SampleCollection,
|
||
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 rather than buried in `marginal_table`'s
|
||
sorted rows. With "pdg" or "material" — open-ended vocabularies that can
|
||
run to many distinct values — groups are capped at `max_groups`, ranked by
|
||
max(kl) * n as in `plot_marginals`; "energy" is already bounded by
|
||
`n_energy_bins` and isn't capped. Built on top of `marginal_table`; see
|
||
`plot_kl_bars_pl` for the polars-native, parquet-direct equivalent.
|
||
"""
|
||
table = marginal_table(
|
||
collection, group_by=group_by, n_energy_bins=n_energy_bins, bins=bins
|
||
)
|
||
if group_by in ("pdg", "material"):
|
||
table = _limit_groups_by_kl_n(table, max_groups)
|
||
return _plot_kl_bars(table, figsize=figsize)
|
||
|
||
|
||
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),
|
||
):
|
||
"""Polars duplicate of `plot_kl_bars`, reading straight from a predict parquet.
|
||
|
||
Built on top of `marginal_table_pl`; see that function for the `source`
|
||
argument and why it stays lazy until the final per-(group, dim) collect.
|
||
See `plot_kl_bars` for the `max_groups` capping behavior.
|
||
"""
|
||
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(collection: SampleCollection) -> tuple[np.ndarray, np.ndarray]:
|
||
"""Pearson correlation matrices of the raw targets, real vs generated."""
|
||
return (
|
||
np.corrcoef(collection.real_raw, rowvar=False),
|
||
np.corrcoef(collection.gen_raw, rowvar=False),
|
||
)
|
||
|
||
|
||
def plot_correlation_matrices(collection: SampleCollection):
|
||
"""Side-by-side real/generated correlation heatmaps, plus their difference."""
|
||
real_corr, gen_corr = correlation_matrices(collection)
|
||
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(
|
||
collection: SampleCollection,
|
||
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.
|
||
"""
|
||
pairs = pairs or _DEFAULT_PAIRS
|
||
rng = np.random.default_rng(seed)
|
||
n = len(collection.pdg)
|
||
idx = rng.choice(n, size=min(n_sample, n), replace=False)
|
||
|
||
fig, axes = plt.subplots(2, len(pairs), squeeze=False, figsize=(4 * len(pairs), 7))
|
||
for col, (a, b) in enumerate(pairs):
|
||
ia, ib = RAW_TARGET_NAMES.index(a), RAW_TARGET_NAMES.index(b)
|
||
for row, (data, title) in enumerate(
|
||
[(collection.real_raw, "real"), (collection.gen_raw, "generated")]
|
||
):
|
||
ax = axes[row][col]
|
||
ax.scatter(data[idx, ia], data[idx, ib], 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 direction_alignment(collection: SampleCollection) -> tuple[np.ndarray, np.ndarray]:
|
||
"""cos(angle) between post_dir_local and travel_dir_local, 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.
|
||
"""
|
||
|
||
def cos_angle(raw: np.ndarray) -> np.ndarray:
|
||
post, travel = raw[:, 3:6], raw[:, 6:9]
|
||
return np.sum(post * travel, axis=1) / (
|
||
np.linalg.norm(post, axis=1) * np.linalg.norm(travel, axis=1) + 1e-8
|
||
)
|
||
|
||
return cos_angle(collection.real_raw), cos_angle(collection.gen_raw)
|
||
|
||
|
||
def plot_direction_alignment(collection: SampleCollection, bins: int = 50):
|
||
real_cos, gen_cos = direction_alignment(collection)
|
||
fig, ax = plt.subplots(figsize=(5, 4))
|
||
edges = np.linspace(-1, 1, bins + 1).tolist()
|
||
ax.hist(real_cos, bins=edges, density=True, histtype="step", label="real")
|
||
ax.hist(gen_cos, bins=edges, density=True, histtype="step", 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 constraint_report(
|
||
collection: SampleCollection, norm_tol: float = 0.05
|
||
) -> pd.DataFrame:
|
||
"""Rate of physical-constraint violations in the generated raw-space samples.
|
||
|
||
The model is an unconstrained MLP, so nothing forces post_dir_local /
|
||
travel_dir_local to stay unit-norm or step_length/delta_e/edep to stay
|
||
non-negative — both hold by construction in the real data, so any
|
||
violation rate here is purely a generation artifact.
|
||
"""
|
||
gen = collection.gen_raw
|
||
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[:_N_SCALAR_DIMS]):
|
||
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)
|
||
|
||
|
||
def constraint_report_pl(
|
||
source: str | Path | pl.LazyFrame, norm_tol: float = 0.05
|
||
) -> pl.DataFrame:
|
||
"""Polars duplicate of `constraint_report`, reading straight from a predict parquet.
|
||
|
||
Computed as a single set of lazy aggregation expressions over the
|
||
`pred_*` columns — never materializes `gen_raw`. See `marginal_table_pl`
|
||
for the `source` argument and why this stays lazy.
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||
|
||
post_norm = pl.sum_horizontal(
|
||
[pl.col(pred_cols[k]) ** 2 for k in range(3, 6)]
|
||
).sqrt()
|
||
travel_norm = pl.sum_horizontal(
|
||
[pl.col(pred_cols[k]) ** 2 for k in range(6, 9)]
|
||
).sqrt()
|
||
raw_log_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_log_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_log_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(collection: SampleCollection):
|
||
"""Histograms backing `constraint_report`: direction norms and sign of the log-dims."""
|
||
gen = collection.gen_raw
|
||
post_norm = np.linalg.norm(gen[:, 3:6], axis=1)
|
||
travel_norm = np.linalg.norm(gen[:, 6:9], axis=1)
|
||
|
||
n_panels = 2 + _N_SCALAR_DIMS
|
||
fig, axes = plt.subplots(1, n_panels, figsize=(4 * n_panels, 3.5))
|
||
for ax, norm, title in [
|
||
(axes[0], post_norm, "||post_dir||"),
|
||
(axes[1], travel_norm, "||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)
|
||
for k, name in enumerate(RAW_TARGET_NAMES[:_N_SCALAR_DIMS]):
|
||
ax = axes[2 + k]
|
||
ax.hist(gen[:, k], bins=_hist_edges(gen[:, k], bins=50), histtype="step")
|
||
ax.set_yscale("log")
|
||
ax.axvline(0.0, color="k", linestyle="--", linewidth=1)
|
||
ax.set_title(f"generated {name}")
|
||
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 into a SampleCollection: per-event sums
|
||
# (total deposited energy, etc.) would be silently corrupted by the row
|
||
# subsampling `load_predicted_local(sample_frac=...)` uses to keep large files
|
||
# in memory, and the real files here 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
|
||
|
||
|
||
def _entry_axis_and_bin_edges(
|
||
lf: pl.LazyFrame, depth_bins: int, transverse_bins: int
|
||
) -> tuple[pl.DataFrame, np.ndarray, np.ndarray]:
|
||
"""Per-event shower axis (highest-pre_E row) plus depth/transverse bin edges.
|
||
|
||
Bin edges are sized from `pre_pos` alone (no post_pos reconstruction
|
||
needed) — pre-step positions already trace the shower's extent closely
|
||
enough to pick a sensible range, which avoids a second full streaming pass
|
||
just to size the bins.
|
||
"""
|
||
narrow = lf.select(["event_id", *_PRE_COLS])
|
||
entry = narrow.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"),
|
||
)
|
||
|
||
joined = narrow.join(entry, on="event_id")
|
||
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()
|
||
|
||
stats = (
|
||
joined.select(depth.alias("depth_proxy"), transverse.alias("transverse_proxy"))
|
||
.select(
|
||
pl.col("depth_proxy").quantile(0.001).alias("depth_lo"),
|
||
pl.col("depth_proxy").quantile(0.999).alias("depth_hi"),
|
||
pl.col("transverse_proxy").quantile(0.999).alias("transverse_hi"),
|
||
)
|
||
.collect(engine="streaming")
|
||
.row(0, named=True)
|
||
)
|
||
|
||
entry_df = entry.collect(engine="streaming").sort("event_id")
|
||
|
||
depth_lo, depth_hi = stats["depth_lo"], stats["depth_hi"]
|
||
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(stats["transverse_hi"], 1e-6)
|
||
transverse_edges = np.linspace(0.0, transverse_hi, transverse_bins + 1)
|
||
return entry_df, depth_edges, transverse_edges
|
||
|
||
|
||
def _iter_predicted_local_batches(
|
||
source: str | Path | pl.LazyFrame, columns: list[str], batch_size: int
|
||
):
|
||
"""Yield the needed columns in bounded-memory chunks.
|
||
|
||
A path is streamed via pyarrow's row-batch reader so the file is never
|
||
fully materialized; a `pl.LazyFrame` (the in-memory test-fixture case) is
|
||
just collected once, since that data is small by construction.
|
||
"""
|
||
if isinstance(source, pl.LazyFrame):
|
||
yield source.select(columns).collect()
|
||
return
|
||
pf = pq.ParquetFile(Path(source))
|
||
for batch in pf.iter_batches(batch_size=batch_size, columns=columns):
|
||
yield pl.from_arrow(batch)
|
||
|
||
|
||
def compute_event_observables_pl(
|
||
source: str | Path | pl.LazyFrame,
|
||
depth_bins: int = 20,
|
||
transverse_bins: int = 20,
|
||
batch_size: int = 1_000_000,
|
||
) -> EventObservables:
|
||
"""Stream a `giant predict --coord local` parquet file into event-level observables.
|
||
|
||
For each event, the highest-`pre_E` row is taken as the primary's entry
|
||
step (secondaries always carry less energy than their parent), fixing a
|
||
shower axis/entry point shared by both real and generated rows (both are
|
||
conditioned on the same real pre-step state). Every row's `post_pos`/`edep`
|
||
is reconstructed into the world frame in physical units (mm, MeV) via
|
||
`giant.data.transforms.reconstruct_post_pos`/`inv_log_transform` — the same
|
||
functions `giant predict --coord global` uses — and projected onto
|
||
depth-along-axis / transverse-distance-from-axis.
|
||
|
||
`event_table` also carries `real_total_length`/`gen_total_length` —
|
||
`sum(step_length)` per event, the total path length traveled by every
|
||
track in the shower (not the same as the depth of any single point).
|
||
|
||
Runs in two passes: a cheap pure-polars pass over `pre_*` columns only
|
||
(shower axis + bin-edge sizing), then one streaming pass over the full
|
||
file accumulating per-event and per-bin sums in numpy. Never materializes
|
||
the file as a `SampleCollection` — see the module-level note above.
|
||
"""
|
||
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()
|
||
n_events = len(event_ids)
|
||
entry_pos = (
|
||
entry_df.select(["entry_x", "entry_y", "entry_z"]).to_numpy().astype(np.float32)
|
||
)
|
||
axis_dir = (
|
||
entry_df.select(["axis_x", "axis_y", "axis_z"]).to_numpy().astype(np.float32)
|
||
)
|
||
|
||
n_steps = np.zeros(n_events, dtype=np.int64)
|
||
real_total_edep = np.zeros(n_events, dtype=np.float64)
|
||
gen_total_edep = np.zeros(n_events, dtype=np.float64)
|
||
real_total_length = np.zeros(n_events, dtype=np.float64)
|
||
gen_total_length = np.zeros(n_events, dtype=np.float64)
|
||
real_sum_edep_depth = np.zeros(n_events, dtype=np.float64)
|
||
gen_sum_edep_depth = np.zeros(n_events, dtype=np.float64)
|
||
real_sum_edep_transverse2 = np.zeros(n_events, dtype=np.float64)
|
||
gen_sum_edep_transverse2 = np.zeros(n_events, dtype=np.float64)
|
||
real_depth_bin_edep = np.zeros((n_events, depth_bins), dtype=np.float64)
|
||
gen_depth_bin_edep = np.zeros((n_events, depth_bins), dtype=np.float64)
|
||
real_transverse_bin_edep = np.zeros((n_events, transverse_bins), dtype=np.float64)
|
||
gen_transverse_bin_edep = np.zeros((n_events, transverse_bins), dtype=np.float64)
|
||
|
||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||
true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
|
||
needed_cols = [
|
||
"event_id",
|
||
"pre_x",
|
||
"pre_y",
|
||
"pre_z",
|
||
"pre_E", # needed to decode the energy simplex into physical edep
|
||
"pre_dx",
|
||
"pre_dy",
|
||
"pre_dz",
|
||
*pred_cols,
|
||
*true_cols,
|
||
]
|
||
|
||
for batch_df in _iter_predicted_local_batches(source, needed_cols, batch_size):
|
||
eid = batch_df["event_id"].to_numpy()
|
||
idx = np.searchsorted(event_ids, eid)
|
||
|
||
pre_pos = (
|
||
batch_df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32)
|
||
)
|
||
pre_dir = (
|
||
batch_df.select(["pre_dx", "pre_dy", "pre_dz"])
|
||
.to_numpy()
|
||
.astype(np.float32)
|
||
)
|
||
|
||
pre_E = batch_df["pre_E"].to_numpy().astype(np.float32)
|
||
|
||
def _reconstruct(cols: list[str]) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||
raw = batch_df.select(cols).to_numpy().astype(np.float32)
|
||
step_length = inv_log_transform(raw[:, 0])
|
||
edep, _e_sec, _post_E, _delta_e = energy_simplex_decode(raw[:, 1:3], pre_E)
|
||
travel_dir_local = raw[:, 6:9]
|
||
post_pos = reconstruct_post_pos(
|
||
pre_pos, pre_dir, step_length, travel_dir_local
|
||
)
|
||
return post_pos, edep, step_length
|
||
|
||
real_post_pos, real_edep, real_step_length = _reconstruct(true_cols)
|
||
gen_post_pos, gen_edep, gen_step_length = _reconstruct(pred_cols)
|
||
|
||
e_pos = entry_pos[idx]
|
||
a_dir = axis_dir[idx]
|
||
|
||
def _depth_transverse(post_pos: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||
disp = post_pos - e_pos
|
||
depth = np.sum(disp * a_dir, axis=1)
|
||
perp = disp - depth[:, None] * a_dir
|
||
transverse = np.linalg.norm(perp, axis=1)
|
||
return depth, transverse
|
||
|
||
real_depth, real_transverse = _depth_transverse(real_post_pos)
|
||
gen_depth, gen_transverse = _depth_transverse(gen_post_pos)
|
||
|
||
real_depth_bin = np.digitize(real_depth, depth_edges[1:-1])
|
||
gen_depth_bin = np.digitize(gen_depth, depth_edges[1:-1])
|
||
real_transverse_bin = np.digitize(real_transverse, transverse_edges[1:-1])
|
||
gen_transverse_bin = np.digitize(gen_transverse, transverse_edges[1:-1])
|
||
|
||
# np.add.at is an unbuffered ufunc method — a well-known numpy slow path
|
||
# for scatter-add (44% of this function's runtime on a 114M-row profile).
|
||
# np.bincount does the same accumulation with a single optimized pass;
|
||
# the 2D (per-event, per-bin) accumulators flatten (idx, bin) into one
|
||
# bincount index and reshape back, since bincount only scatters into 1D.
|
||
n_steps += np.bincount(idx, minlength=n_events)
|
||
real_total_edep += np.bincount(idx, weights=real_edep, minlength=n_events)
|
||
gen_total_edep += np.bincount(idx, weights=gen_edep, minlength=n_events)
|
||
real_total_length += np.bincount(
|
||
idx, weights=real_step_length, minlength=n_events
|
||
)
|
||
gen_total_length += np.bincount(
|
||
idx, weights=gen_step_length, minlength=n_events
|
||
)
|
||
real_sum_edep_depth += np.bincount(
|
||
idx, weights=real_edep * real_depth, minlength=n_events
|
||
)
|
||
gen_sum_edep_depth += np.bincount(
|
||
idx, weights=gen_edep * gen_depth, minlength=n_events
|
||
)
|
||
real_sum_edep_transverse2 += np.bincount(
|
||
idx, weights=real_edep * real_transverse**2, minlength=n_events
|
||
)
|
||
gen_sum_edep_transverse2 += np.bincount(
|
||
idx, weights=gen_edep * gen_transverse**2, minlength=n_events
|
||
)
|
||
real_depth_bin_edep += np.bincount(
|
||
idx * depth_bins + real_depth_bin,
|
||
weights=real_edep,
|
||
minlength=n_events * depth_bins,
|
||
).reshape(n_events, depth_bins)
|
||
gen_depth_bin_edep += np.bincount(
|
||
idx * depth_bins + gen_depth_bin,
|
||
weights=gen_edep,
|
||
minlength=n_events * depth_bins,
|
||
).reshape(n_events, depth_bins)
|
||
real_transverse_bin_edep += np.bincount(
|
||
idx * transverse_bins + real_transverse_bin,
|
||
weights=real_edep,
|
||
minlength=n_events * transverse_bins,
|
||
).reshape(n_events, transverse_bins)
|
||
gen_transverse_bin_edep += np.bincount(
|
||
idx * transverse_bins + gen_transverse_bin,
|
||
weights=gen_edep,
|
||
minlength=n_events * transverse_bins,
|
||
).reshape(n_events, transverse_bins)
|
||
|
||
safe_real_total = np.where(real_total_edep > 0, real_total_edep, 1.0)
|
||
safe_gen_total = np.where(gen_total_edep > 0, gen_total_edep, 1.0)
|
||
real_centroid_depth = real_sum_edep_depth / safe_real_total
|
||
gen_centroid_depth = gen_sum_edep_depth / safe_gen_total
|
||
real_transverse_rms = np.sqrt(real_sum_edep_transverse2 / safe_real_total)
|
||
gen_transverse_rms = np.sqrt(gen_sum_edep_transverse2 / safe_gen_total)
|
||
|
||
depth_centers = 0.5 * (depth_edges[:-1] + depth_edges[1:])
|
||
real_max_depth = depth_centers[np.argmax(real_depth_bin_edep, axis=1)]
|
||
gen_max_depth = depth_centers[np.argmax(gen_depth_bin_edep, axis=1)]
|
||
|
||
safe_n_steps = np.where(n_steps > 0, n_steps, 1)
|
||
real_mean_edep = real_total_edep / safe_n_steps
|
||
gen_mean_edep = gen_total_edep / safe_n_steps
|
||
real_mean_length = real_total_length / safe_n_steps
|
||
gen_mean_length = gen_total_length / safe_n_steps
|
||
|
||
# The inverse transform must be applied *before* the median, not after:
|
||
# median commutes with inv_log_transform only for odd-length groups. For an
|
||
# even number of steps polars' .median() averages the two central order
|
||
# statistics, and that averaging does not commute with the nonlinear exp
|
||
# (it would yield their geometric mean rather than the true median). So take
|
||
# the median on the raw (exp-transformed) values directly. No post_pos
|
||
# reconstruction is needed — only the magnitudes matter.
|
||
medians = (
|
||
lf.select(
|
||
"event_id",
|
||
"true_edep_logit",
|
||
"true_sec_logit",
|
||
"pred_edep_logit",
|
||
"pred_sec_logit",
|
||
"pre_E",
|
||
"true_log_step_length",
|
||
"pred_log_step_length",
|
||
)
|
||
.group_by("event_id")
|
||
.agg(
|
||
_edep_pl("true").median().alias("real_median_edep"),
|
||
_edep_pl("pred").median().alias("gen_median_edep"),
|
||
(pl.col("true_log_step_length").exp() - _LOG_EPS)
|
||
.median()
|
||
.alias("real_median_length"),
|
||
(pl.col("pred_log_step_length").exp() - _LOG_EPS)
|
||
.median()
|
||
.alias("gen_median_length"),
|
||
)
|
||
.sort("event_id")
|
||
.collect(engine="streaming")
|
||
)
|
||
real_median_edep = medians["real_median_edep"].to_numpy()
|
||
gen_median_edep = medians["gen_median_edep"].to_numpy()
|
||
real_median_length = medians["real_median_length"].to_numpy()
|
||
gen_median_length = medians["gen_median_length"].to_numpy()
|
||
|
||
event_table = pl.DataFrame(
|
||
{
|
||
"event_id": event_ids,
|
||
"n_steps": 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": real_median_edep,
|
||
"gen_median_edep": gen_median_edep,
|
||
"real_median_length": real_median_length,
|
||
"gen_median_length": 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_bin_edep.mean(axis=0),
|
||
gen_depth_profile=gen_depth_bin_edep.mean(axis=0),
|
||
real_depth_profile_std=real_depth_bin_edep.std(axis=0),
|
||
gen_depth_profile_std=gen_depth_bin_edep.std(axis=0),
|
||
real_transverse_profile=real_transverse_bin_edep.mean(axis=0),
|
||
gen_transverse_profile=gen_transverse_bin_edep.mean(axis=0),
|
||
real_transverse_profile_std=real_transverse_bin_edep.std(axis=0),
|
||
gen_transverse_profile_std=gen_transverse_bin_edep.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 (rather than overlaying on one axis) keeps each comparison
|
||
legible; if a tail-compression bias is the explanation for a low generated
|
||
mean, the median panel should overlap far more closely than the mean panel.
|
||
|
||
Each panel gets its own bin edges (`bins` for the mean, `median_bins` for
|
||
the median): the median is far less spread than the mean, so reusing the
|
||
mean's range would crush it into a few bins — `median_bins` is larger to
|
||
resolve that tighter distribution.
|
||
"""
|
||
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. Mean and median
|
||
are drawn in separate panels, each with its own binning (`bins` /
|
||
`median_bins`); see `_plot_mean_median_per_step`.
|
||
"""
|
||
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. Mean and median
|
||
are drawn in separate panels, each with its own binning (`bins` /
|
||
`median_bins`); see `_plot_mean_median_per_step`.
|
||
"""
|
||
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. Unlike the
|
||
# event-level checks above, this 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,
|
||
)
|