lars
b4ce04e772
Move ipykernel into the analysis extra instead of a separate dependency group, since it's needed wherever analysis plotting runs. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1633 lines
63 KiB
Python
1633 lines
63 KiB
Python
"""Notebook-friendly diagnostics for a trained model's sample quality.
|
||
|
||
Typical use from a Jupyter notebook::
|
||
|
||
from giant.analysis import load_model_bundle, make_val_loader, collect_samples
|
||
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
|
||
|
||
bundle = load_model_bundle("runs/my_run/best.pt")
|
||
val_loader = make_val_loader(bundle, "path/to/steps.parquet")
|
||
samples = collect_samples(bundle, val_loader)
|
||
|
||
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)
|
||
|
||
If predictions were already generated offline via `giant predict --coord local`,
|
||
skip the checkpoint/model entirely and load the parquet directly::
|
||
|
||
from giant.analysis import load_predicted_local
|
||
samples = load_predicted_local("path/to/steps_predicted_local.parquet")
|
||
|
||
(`--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.
|
||
|
||
`collect_samples` takes a `steps` argument (forwarded to the flow ODE
|
||
integrator or, in ddim mode, the DDIM substep count) so a later
|
||
sampler-step-count ablation can sweep it by calling this function repeatedly
|
||
without any new plumbing.
|
||
"""
|
||
|
||
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 torch
|
||
from torch.utils.data import DataLoader
|
||
|
||
import pyarrow.parquet as pq
|
||
|
||
from giant.config import warn_if_checkpoint_config_mismatch
|
||
from giant.constants import (
|
||
LOCAL_TARGET_NAMES,
|
||
PREDICT_COORD_METADATA_KEY,
|
||
PREDICT_SCHEMA_VERSION,
|
||
PREDICT_SCHEMA_VERSION_KEY,
|
||
)
|
||
from giant.data.dataset import train_val_split
|
||
from giant.data.loader import find_parquet_files, load_steps
|
||
from giant.data.transforms import (
|
||
Normalizer,
|
||
build_features,
|
||
inv_log_transform,
|
||
reconstruct_post_pos,
|
||
)
|
||
from giant.model.network import DenoisingMLP
|
||
from giant.model.schedule import CosineSchedule
|
||
from giant.sample import sample_ddim, sample_ddpm, sample_flow
|
||
from giant.validate import _histogram_kl
|
||
|
||
_N_LOG_DIMS = 3 # log_step_length, log_delta_e, log_edep are the first 3 target dims
|
||
RAW_TARGET_NAMES = [n.removeprefix("log_") for n in LOCAL_TARGET_NAMES]
|
||
_LOG_EPS = (
|
||
1e-8 # mirrors giant.data.transforms._EPS, duplicated for use in polars exprs
|
||
)
|
||
|
||
|
||
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 ModelBundle:
|
||
model: torch.nn.Module
|
||
cond_normalizer: Normalizer
|
||
target_normalizer: Normalizer
|
||
pdg_map: dict[int, int]
|
||
mat_map: dict[str, int]
|
||
model_config: dict
|
||
mode: str
|
||
schedule: CosineSchedule | None
|
||
device: torch.device
|
||
|
||
@property
|
||
def idx_to_pdg(self) -> dict[int, int]:
|
||
return {v: k for k, v in self.pdg_map.items()}
|
||
|
||
@property
|
||
def idx_to_mat(self) -> dict[int, str]:
|
||
return {v: k for k, v in self.mat_map.items()}
|
||
|
||
|
||
@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
|
||
# Normalized-space (model's native training space) targets — only available
|
||
# when collected live with the normalizer (collect_samples). A predict
|
||
# parquet has already been denormalized on disk with no normalizer
|
||
# attached, so loaders built from one (e.g. load_predicted_local) leave
|
||
# these as None rather than fabricate a value.
|
||
real_norm: np.ndarray | None = None
|
||
gen_norm: np.ndarray | None = None
|
||
|
||
|
||
def load_model_bundle(
|
||
ckpt_path: str | Path,
|
||
mode: str = "flow",
|
||
device: torch.device | None = None,
|
||
) -> ModelBundle:
|
||
"""Reconstruct a trained model and its normalizers/vocab from a training checkpoint."""
|
||
warn_if_checkpoint_config_mismatch(ckpt_path)
|
||
device = device or torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
|
||
|
||
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
|
||
mat_map = dict(ckpt["mat_map"])
|
||
|
||
model = DenoisingMLP(**ckpt["model_config"])
|
||
model.load_state_dict(ckpt["model"])
|
||
model.to(device).eval()
|
||
|
||
cond_normalizer = Normalizer.from_dict(ckpt["normalizer"]["cond"])
|
||
target_normalizer = Normalizer.from_dict(ckpt["normalizer"]["target"])
|
||
schedule = CosineSchedule().to(device) if mode != "flow" else None
|
||
|
||
return ModelBundle(
|
||
model=model,
|
||
cond_normalizer=cond_normalizer,
|
||
target_normalizer=target_normalizer,
|
||
pdg_map=pdg_map,
|
||
mat_map=mat_map,
|
||
model_config=ckpt["model_config"],
|
||
mode=mode,
|
||
schedule=schedule,
|
||
device=device,
|
||
)
|
||
|
||
|
||
def make_val_loader(
|
||
bundle: ModelBundle,
|
||
data: str | Path,
|
||
val_fraction: float = 0.1,
|
||
seed: int = 42,
|
||
batch_size: int = 4096,
|
||
) -> DataLoader:
|
||
"""Build a DataLoader over the val split, normalized with the bundle's fitted stats.
|
||
|
||
Loads the whole file into memory — fine for typical validation-set sizes; for
|
||
very large datasets, build a StreamingStepsDataset directly (see giant.pipeline).
|
||
"""
|
||
files = find_parquet_files(data)
|
||
chunks = [load_steps(f) for f in files]
|
||
full = {k: np.concatenate([c[k] for c in chunks]) for k in chunks[0]}
|
||
|
||
cond_cont, cond_cat, target, _, _ = build_features(
|
||
full,
|
||
bundle.pdg_map,
|
||
bundle.mat_map,
|
||
cond_normalizer=bundle.cond_normalizer,
|
||
target_normalizer=bundle.target_normalizer,
|
||
)
|
||
_, val_ds = train_val_split(
|
||
full, cond_cont, cond_cat, target, val_fraction=val_fraction, seed=seed
|
||
)
|
||
return DataLoader(val_ds, batch_size=batch_size, shuffle=False)
|
||
|
||
|
||
def _to_raw_targets(target_norm: np.ndarray, normalizer: Normalizer) -> np.ndarray:
|
||
raw = normalizer.inverse_transform(target_norm)
|
||
raw[:, :_N_LOG_DIMS] = inv_log_transform(raw[:, :_N_LOG_DIMS])
|
||
return raw
|
||
|
||
|
||
@torch.no_grad()
|
||
def collect_samples(
|
||
bundle: ModelBundle,
|
||
val_loader: DataLoader,
|
||
n_batches: int | None = None,
|
||
steps: int | None = None,
|
||
) -> SampleCollection:
|
||
"""Run the sampler over `val_loader`, pairing generations with real targets + conditioning.
|
||
|
||
`steps` overrides the number of sampler steps (flow ODE steps or DDIM
|
||
substeps); `None` keeps each sampler's own default. Unused in "ddpm"
|
||
mode, which always runs the full schedule.
|
||
"""
|
||
model = bundle.model
|
||
device = bundle.device
|
||
steps_kw = {} if steps is None else {"steps": steps}
|
||
|
||
cond_list, real_list, gen_list = [], [], []
|
||
for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
|
||
if n_batches is not None and i >= n_batches:
|
||
break
|
||
cond_cont = cond_cont.to(device)
|
||
cond_cat = cond_cat.to(device)
|
||
|
||
if bundle.mode == "flow":
|
||
gen = sample_flow(model, cond_cont, cond_cat, **steps_kw)
|
||
elif bundle.mode == "ddpm":
|
||
gen = sample_ddpm(model, cond_cont, cond_cat, bundle.schedule)
|
||
else:
|
||
gen = sample_ddim(model, cond_cont, cond_cat, bundle.schedule, **steps_kw)
|
||
|
||
cond_full = torch.cat([cond_cont.cpu(), cond_cat.cpu().float()], dim=-1)
|
||
cond_list.append(cond_full.numpy())
|
||
real_list.append(x1.numpy())
|
||
gen_list.append(gen.cpu().numpy())
|
||
|
||
cond_all = np.concatenate(cond_list, axis=0)
|
||
real_norm = np.concatenate(real_list, axis=0)
|
||
gen_norm = np.concatenate(gen_list, axis=0)
|
||
|
||
cond_cont_norm, cond_cat_arr = cond_all[:, :-2], cond_all[:, -2:]
|
||
cond_cont_raw = bundle.cond_normalizer.inverse_transform(cond_cont_norm)
|
||
cond_cont_raw[:, 3] = inv_log_transform(cond_cont_raw[:, 3]) # log(pre_E) -> pre_E
|
||
|
||
idx_to_pdg, idx_to_mat = bundle.idx_to_pdg, bundle.idx_to_mat
|
||
pdg = np.array([idx_to_pdg[i] for i in cond_cat_arr[:, 0].astype(np.int64)])
|
||
material = np.array([idx_to_mat[i] for i in cond_cat_arr[:, 1].astype(np.int64)])
|
||
|
||
return SampleCollection(
|
||
cond_cont_raw=cond_cont_raw,
|
||
pdg=pdg,
|
||
material=material,
|
||
real_raw=_to_raw_targets(real_norm, bundle.target_normalizer),
|
||
gen_raw=_to_raw_targets(gen_norm, bundle.target_normalizer),
|
||
real_norm=real_norm,
|
||
gen_norm=gen_norm,
|
||
)
|
||
|
||
|
||
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
|
||
) -> 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 the same log-scaled,
|
||
local-frame space `collect_samples` produces internally before raw conversion.
|
||
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 a lazy polars scan with the needed columns selected before
|
||
`.collect()`, so column projection is pushed down into the parquet reader
|
||
(e.g. `event_id` is never read) instead of materializing every column of
|
||
the file as a pandas DataFrame first.
|
||
|
||
`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`. Sampling happens
|
||
after `.collect()` since polars' row-level sampling isn't pushed down into
|
||
the lazy scan; `seed` makes the subsample reproducible.
|
||
"""
|
||
if not (0 < sample_frac <= 1):
|
||
raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
|
||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||
true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
|
||
df = (
|
||
_scan_predicted_local(path)
|
||
.select(pred_cols + true_cols + _COND_CONT_COLS + ["pdg", "material"])
|
||
.collect()
|
||
)
|
||
if sample_frac < 1.0:
|
||
df = df.sample(fraction=sample_frac, seed=seed)
|
||
|
||
gen_log_local = df.select(pred_cols).to_numpy().astype(np.float32)
|
||
real_log_local = df.select(true_cols).to_numpy().astype(np.float32)
|
||
|
||
def to_raw(log_local: np.ndarray) -> np.ndarray:
|
||
raw = log_local.copy()
|
||
raw[:, :_N_LOG_DIMS] = inv_log_transform(raw[:, :_N_LOG_DIMS])
|
||
return raw
|
||
|
||
cond_cont_raw = df.select(_COND_CONT_COLS).to_numpy().astype(np.float32)
|
||
|
||
return SampleCollection(
|
||
cond_cont_raw=cond_cont_raw,
|
||
pdg=df["pdg"].to_numpy(),
|
||
material=df["material"].to_numpy(),
|
||
real_raw=to_raw(real_log_local),
|
||
gen_raw=to_raw(gen_log_local),
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 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
|
||
# end to end: each (group, dim) pair is filtered, projected to just the two
|
||
# columns it needs, and collected on its own, so peak memory is one column
|
||
# pair rather than the whole file — useful when `load_predicted_local` itself
|
||
# would be too large to hold in memory at once.
|
||
# ---------------------------------------------------------------------------
|
||
|
||
|
||
def _histogram_kl_pl(
|
||
p: pl.Series, q: pl.Series, bins: int = 50, eps: float = 1e-8
|
||
) -> float:
|
||
"""Polars duplicate of `giant.validate._histogram_kl`, binning via `Series.hist`."""
|
||
lo, hi = min(p.min(), q.min()), max(p.max(), q.max())
|
||
if hi <= lo:
|
||
return 0.0
|
||
edges = np.linspace(lo, hi, bins + 1).tolist()
|
||
p_hist = p.hist(bins=edges)["count"].to_numpy().astype(np.float64) + eps
|
||
q_hist = q.hist(bins=edges)["count"].to_numpy().astype(np.float64) + eps
|
||
p_hist /= p_hist.sum()
|
||
q_hist /= q_hist.sum()
|
||
return float(np.sum(p_hist * np.log(p_hist / q_hist)))
|
||
|
||
|
||
def _raw_expr(col: str, j: int) -> pl.Expr:
|
||
"""`pred_*`/`true_*` columns are log-scaled for the first `_N_LOG_DIMS` dims."""
|
||
e = pl.col(col)
|
||
return (e.exp() - _LOG_EPS) if j < _N_LOG_DIMS else e
|
||
|
||
|
||
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 _group_filters_pl(
|
||
lf: pl.LazyFrame,
|
||
group_by: str | None,
|
||
n_energy_bins: int,
|
||
) -> list[tuple[str, pl.Expr]]:
|
||
if group_by is None:
|
||
return [("all", pl.lit(True))]
|
||
if group_by == "pdg":
|
||
vals = lf.select("pdg").unique().collect()["pdg"].sort().to_list()
|
||
return [(f"pdg={v}", pl.col("pdg") == v) for v in vals]
|
||
if group_by == "material":
|
||
vals = lf.select("material").unique().collect()["material"].sort().to_list()
|
||
return [(f"material={v}", pl.col("material") == v) for v in vals]
|
||
if group_by == "energy":
|
||
pre_E = lf.select("pre_E").collect().to_series().to_numpy()
|
||
edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
|
||
edges[-1] += 1e-6
|
||
return [
|
||
(
|
||
f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})",
|
||
(pl.col("pre_E") >= edges[i]) & (pl.col("pre_E") < edges[i + 1]),
|
||
)
|
||
for i in range(n_energy_bins)
|
||
]
|
||
raise ValueError(f"unknown group_by={group_by!r}")
|
||
|
||
|
||
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.
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
pred_cols = [f"pred_{name}" for name in LOCAL_TARGET_NAMES]
|
||
true_cols = [f"true_{name}" for name in LOCAL_TARGET_NAMES]
|
||
|
||
rows = []
|
||
for label, cond in _group_filters_pl(lf, group_by, n_energy_bins):
|
||
glf = lf.filter(cond)
|
||
n = glf.select(pl.len()).collect().item()
|
||
if n < 2:
|
||
continue
|
||
for j, name in enumerate(RAW_TARGET_NAMES):
|
||
pair = glf.select(
|
||
[
|
||
_raw_expr(true_cols[j], j).alias("real"),
|
||
_raw_expr(pred_cols[j], j).alias("gen"),
|
||
]
|
||
).collect()
|
||
real_s, gen_s = pair["real"], pair["gen"]
|
||
rows.append(
|
||
{
|
||
"group": label,
|
||
"dim": name,
|
||
"n": n,
|
||
"real_mean": real_s.mean(),
|
||
"gen_mean": gen_s.mean(),
|
||
# ddof=0 to match numpy's (population-std) default used by marginal_table
|
||
"real_std": real_s.std(ddof=0),
|
||
"gen_std": gen_s.std(ddof=0),
|
||
"kl_real_gen": _histogram_kl_pl(real_s, gen_s, bins=bins),
|
||
}
|
||
)
|
||
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)
|
||
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_LOG_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 = sum(pl.col(pred_cols[k]) ** 2 for k in range(3, 6)).sqrt()
|
||
travel_norm = sum(pl.col(pred_cols[k]) ** 2 for k in range(6, 9)).sqrt()
|
||
raw_log_dims = [_raw_expr(pred_cols[j], j) for j in range(_N_LOG_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_LOG_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_LOG_DIMS])
|
||
],
|
||
]
|
||
)
|
||
.collect()
|
||
.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_LOG_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_LOG_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_LOG_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()
|
||
.row(0, named=True)
|
||
)
|
||
|
||
entry_df = entry.collect().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_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)
|
||
)
|
||
|
||
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 = inv_log_transform(raw[:, 2])
|
||
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(n_steps, idx, 1)
|
||
np.add.at(real_total_edep, idx, real_edep)
|
||
np.add.at(gen_total_edep, idx, gen_edep)
|
||
np.add.at(real_total_length, idx, real_step_length)
|
||
np.add.at(gen_total_length, idx, gen_step_length)
|
||
np.add.at(real_sum_edep_depth, idx, real_edep * real_depth)
|
||
np.add.at(gen_sum_edep_depth, idx, gen_edep * gen_depth)
|
||
np.add.at(real_sum_edep_transverse2, idx, real_edep * real_transverse**2)
|
||
np.add.at(gen_sum_edep_transverse2, idx, gen_edep * gen_transverse**2)
|
||
np.add.at(real_depth_bin_edep, (idx, real_depth_bin), real_edep)
|
||
np.add.at(gen_depth_bin_edep, (idx, gen_depth_bin), gen_edep)
|
||
np.add.at(real_transverse_bin_edep, (idx, real_transverse_bin), real_edep)
|
||
np.add.at(gen_transverse_bin_edep, (idx, gen_transverse_bin), gen_edep)
|
||
|
||
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_log_edep",
|
||
"pred_log_edep",
|
||
"true_log_step_length",
|
||
"pred_log_step_length",
|
||
)
|
||
.group_by("event_id")
|
||
.agg(
|
||
(pl.col("true_log_edep").exp() - _LOG_EPS)
|
||
.median()
|
||
.alias("real_median_edep"),
|
||
(pl.col("pred_log_edep").exp() - _LOG_EPS)
|
||
.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()
|
||
)
|
||
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)
|
||
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)
|
||
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)
|
||
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, so this only needs the same `exp(...) - eps` de-log
|
||
transform `inv_log_transform` does, expressed directly as a polars expr.
|
||
"""
|
||
lf = _scan_predicted_local(source)
|
||
|
||
def _delog(col: str) -> pl.Expr:
|
||
return pl.col(col).exp() - _LOG_EPS
|
||
|
||
return (
|
||
lf.group_by("pdg")
|
||
.agg(
|
||
_delog("true_log_edep").sum().alias("real_total_edep"),
|
||
_delog("pred_log_edep").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()
|
||
.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,
|
||
)
|