lars/giant
1
0
Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

Add load_rollout_vs_truth to compare rollouts against held-out truth data

Browse Source 
Extends the Tier 1-3 SampleCollection diagnostics (marginals, correlations,
pairwise, direction alignment, constraints) to work on a full autoregressive
giant rollout shower checked against an independent ground-truth steps file,
rather than only paired giant predict --coord local output. The two files
are unpaired (different lengths, own conditioning), so SampleCollection
gains optional *_gen fields and _group_labels/marginal_table/plot_marginals/
plot_pairwise build independent real/gen masks instead of assuming one.

Adds analysis/rollout_validation.ipynb, a sibling of validation.ipynb built
around this workflow.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
This commit is contained in:
Lars Bogner
2026-07-13 10:58:42 +02:00
parent 740ebdf6b6
commit 36bdbf7cc1
3 changed files with 784 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
analysis/rollout_validation.ipynb
+305
View File
@@ -0,0 +1,305 @@
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "d6c26bca",
"metadata": {},
"outputs": [],
"source": [
"# Auto-reload edited modules (e.g. giant.analysis) without restarting the kernel.\n",
"%load_ext autoreload\n",
"%autoreload 2"
]
},
{
"cell_type": "markdown",
"id": "076c43a2",
"metadata": {},
"source": [
"# GIANT rollout-vs-truth validation notebook\n",
"\n",
"Diagnostics for a full autoregressive `giant rollout` shower, compared against a held-out ground-truth steps file (the same schema `giant train` consumes \u2014 see `giant.data.loader.load_steps`) rather than one-step-ahead `giant predict` output.\n",
"\n",
"This is the sibling of `validation.ipynb`: that notebook checks whether one-step generation (conditioned on the *real* preceding state, every row) reproduces real marginals/correlations/shower observables. This one checks the thing that actually matters for deployment \u2014 whether a shower **rolled out autoregressively from the model's own outputs** still looks physical, which is where covariate shift (small per-step errors compounding across a track) would show up.\n",
"\n",
"Built on `load_rollout_vs_truth`, which treats the rollout file as \"generated\" and the truth file as \"real\". Unlike `load_predicted_local`, the two files are **independent, unpaired datasets** \u2014 a rollout doesn't replay real events row-for-row, so real/generated may have different lengths and there's no per-row correspondence. Everything below only ever compares real-vs-generated *distributions*, never individual paired rows, so this is transparent to the checks themselves; see `giant.analysis`'s module docstring for the `SampleCollection.*_gen` mechanics.\n",
"\n",
"Same three tiers as `validation.ipynb` for the step-level checks (stratified marginals, joint structure, physical constraints), plus a rollout-only event-level tier built on `compute_rollout_observables` instead of `compute_event_observables_pl`:\n",
"\n",
"1. **stratified marginals** \u2014 per-dimension real-vs-generated, sliced by pdg/material/energy\n",
"2. **joint structure** \u2014 correlation matrices, physically-coupled pairwise plots, direction alignment\n",
"3. **physical constraints** \u2014 unit-norm directions, non-negative step_length/delta_e/edep (checked on the rollout's own output \u2014 with autoregression, a constraint violation early in a track can compound into later steps, unlike one-step-ahead validation)\n",
"4. **event-level (shower) observables** \u2014 total energy, longitudinal/transverse profiles, computed from the rollout shower itself; optionally overlaid against a real reference computed from a *paired* `giant predict --coord local` file, if one exists for the same held-out events (see the markdown note in that section \u2014 the raw truth-schema file used above doesn't carry the columns `compute_event_observables_pl` needs)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f9741197",
"metadata": {},
"outputs": [],
"source": [
"from giant.analysis import load_rollout_vs_truth, plot_kl_bars\n",
"\n",
"# `giant rollout` output for the shower(s) under test.\n",
"ROLLOUT_FILE = \"/home/lars/Programming/giant/rollout.parquet\"\n",
"# Any held-out file sharing giant train's input schema (real miniCaloSim\n",
"# steps) \u2014 e.g. the val split the rollout's seed events were drawn from.\n",
"TRUTH_FILE = \"/home/lars/Programming/giant/val.parquet\"\n",
"\n",
"# sample_frac subsamples each file independently (kept memory-bounded for\n",
"# large files); both default to every row when omitted.\n",
"samples = load_rollout_vs_truth(ROLLOUT_FILE, TRUTH_FILE)"
]
},
{
"cell_type": "markdown",
"id": "1297bd12",
"metadata": {},
"source": [
"## Tier 1: stratified marginals\n",
"\n",
"KL(real || generated) per target dimension. Unlike `validation.ipynb`'s first cell, there's no lazy full-file `_pl` path for this unpaired comparison \u2014 `load_rollout_vs_truth` always materializes both sides as numpy arrays (see `sample_frac` above for large files)."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "df53a498",
"metadata": {},
"outputs": [],
"source": [
"for grouping in [None, \"energy\", \"pdg\", \"material\"]:\n",
" fig = plot_kl_bars(samples, group_by=grouping)\n",
" fig.show()"
]
},
{
"cell_type": "markdown",
"id": "93d2446e",
"metadata": {},
"source": [
"## Detailed marginals (Tier 1, overlaid histograms)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2f717b07",
"metadata": {},
"outputs": [],
"source": [
"from giant.analysis import plot_marginals, plot_correlation_matrices, plot_pairwise\n",
"from giant.analysis import plot_direction_alignment, plot_constraint_violations\n",
"\n",
"_ = plot_marginals(samples)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "7c9e7b98",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_marginals(samples, group_by=\"energy\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "8202f2f2",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_marginals(samples, group_by=\"pdg\")"
]
},
{
"cell_type": "markdown",
"id": "a232b82e",
"metadata": {},
"source": [
"## Tier 2: joint structure"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e0618b09",
"metadata": {},
"outputs": [],
"source": [
"# Real vs. generated Pearson correlation matrices (+ their difference) over\n",
"# the 9 raw target dims \u2014 catches a model that decorrelates targets that are\n",
"# physically coupled even when every individual marginal looks clean.\n",
"_ = plot_correlation_matrices(samples)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "75a6624b",
"metadata": {},
"outputs": [],
"source": [
"# Scatter for physically-coupled pairs (step_length/delta_e/edep) \u2014 the\n",
"# joint-structure check correlation matrices alone can't fully capture.\n",
"_ = plot_pairwise(samples, n_sample=10000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "69ca0042",
"metadata": {},
"outputs": [],
"source": [
"# cos(angle) between post_dir and travel_dir \u2014 coupled through the\n",
"# scattering physics, so this is another joint-structure check.\n",
"_ = plot_direction_alignment(samples)"
]
},
{
"cell_type": "markdown",
"id": "07f324fd",
"metadata": {},
"source": [
"## Tier 3: physical constraints\n",
"\n",
"Unit-norm direction vectors, non-negative step_length/delta_e/edep. `constraint_report`/`plot_constraint_violations` only ever check the *generated* side (`samples.gen_raw`, here the rollout output) \u2014 under autoregression a violation isn't just a one-off artifact, it can feed the next step's conditioning, so this is worth watching more closely here than in one-step-ahead validation."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "db94a36f",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_constraint_violations(samples)"
]
},
{
"cell_type": "markdown",
"id": "1a3a3519",
"metadata": {},
"source": [
"## Tier 4: event-level (shower) observables\n",
"\n",
"Built on `compute_rollout_observables`, not `compute_event_observables_pl` \u2014 the rollout file carries its own `track_id`/`termination_reason` columns that the event-level aggregation needs, and the shower here already *is* a full autoregressive rollout rather than one-step generations re-aggregated by event.\n",
"\n",
"To overlay a real reference profile, pass a *paired* `giant predict --coord local` file for the same held-out events as `reference_path` below (see `analysis/export_rollout_observables.py`) \u2014 `TRUTH_FILE` above can't serve as that reference directly, since it's the raw training-input schema, not predict output. Leave `reference_path = None` to skip the overlay."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a717f38e",
"metadata": {},
"outputs": [],
"source": [
"from giant.analysis import compute_rollout_observables, compute_event_observables_pl\n",
"from giant.analysis import plot_rollout_longitudinal, plot_rollout_transverse\n",
"from giant.analysis import plot_rollout_total_energy\n",
"\n",
"obs = compute_rollout_observables(ROLLOUT_FILE)\n",
"\n",
"reference_path = None # optional: a `giant predict --coord local` file, see above\n",
"reference = compute_event_observables_pl(reference_path) if reference_path else None"
]
},
{
"cell_type": "markdown",
"id": "553c4353",
"metadata": {},
"source": [
"### Total deposited energy per event\n",
"\n",
"`sum(edep)` per event, rollout vs. (optionally) real reference."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2372ff43",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_rollout_total_energy(obs, reference=reference)"
]
},
{
"cell_type": "markdown",
"id": "c98da271",
"metadata": {},
"source": [
"### Longitudinal profile\n",
"\n",
"Mean deposited energy per event, binned by depth along the shower axis (the `pre_dir` of each event's highest-`pre_E` row), with the event-to-event RMS as error bars."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a3b34ad9",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_rollout_longitudinal(obs, reference=reference)"
]
},
{
"cell_type": "markdown",
"id": "14f83902",
"metadata": {},
"source": [
"### Transverse profile\n",
"\n",
"Same idea, binned by perpendicular distance from the shower axis instead of depth \u2014 a Moli\u00e8re-radius-style lateral containment check."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "25971f55",
"metadata": {},
"outputs": [],
"source": [
"_ = plot_rollout_transverse(obs, reference=reference)"
]
},
{
"cell_type": "markdown",
"id": "c5188a68",
"metadata": {},
"source": [
"---\n",
"\n",
"Not covered here (both need the paired predict schema, see `validation.ipynb` instead): mean deposited-energy/step-length per step, shower-maximum depth, and the dataset-wide pdg energy/length contribution shares (`pdg_contribution_table_pl`)."
]
}
],
"metadata": {
"kernelspec": {
"display_name": "giant (3.12.13)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.13"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
giant/analysis.py
+301 -26
View File
@@ -51,6 +51,21 @@ how much of the total energy/length, see `pdg_contribution_table_pl`::
plot_pdg_energy_share(table)
plot_pdg_length_share(table)
To run the tiers 1-3 checks above against a full `giant rollout` shower
instead of one-step-ahead predict output, use `load_rollout_vs_truth` in place
of `load_predicted_local` — it builds the same `SampleCollection` from a
`giant rollout` output file ("generated") and any held-out file sharing
`giant train`'s input schema ("real"); the two are independent, unpaired
files (a rollout doesn't replay real events row-for-row), unlike the other
loader's paired pred_*/true_* columns::
from giant.analysis import load_rollout_vs_truth
samples = load_rollout_vs_truth("path/to/rollout.parquet", "path/to/val.parquet")
plot_marginals(samples, group_by="energy")
# ... same plot_kl_bars / plot_correlation_matrices / plot_pairwise /
# plot_direction_alignment / plot_constraint_violations as above.
Four tiers of checks, building on the aggregate marginal/KL check in
`giant.validate.validate_marginals`:
@@ -97,7 +112,9 @@ from giant.constants import (
from giant.data.transforms import (
energy_simplex_decode,
inv_log_transform,
local_frame_rotation,
reconstruct_post_pos,
travel_direction,
)
from giant.validate import _histogram_kl
@@ -179,7 +196,16 @@ class SampleCollection:
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
gen_raw: np.ndarray # (M, 9) denormalized + delogged generated targets
# Set these three when real_raw/gen_raw are *unpaired* — independent files
# with their own row counts and conditioning (e.g. `load_rollout_vs_truth`,
# comparing a `giant rollout` shower against a held-out truth file) — rather
# than the row-for-row pred_*/true_* pairing `load_predicted_local` produces.
# None (the default) means "same as the real-side field above", which
# reproduces the original paired behavior exactly.
cond_cont_raw_gen: np.ndarray | None = None
pdg_gen: np.ndarray | None = None
material_gen: np.ndarray | None = None
def _check_predict_metadata(path: Path) -> None:
@@ -294,28 +320,66 @@ def load_predicted_local(
# ---------------------------------------------------------------------------
def _gen_pdg(collection: SampleCollection) -> np.ndarray:
return collection.pdg if collection.pdg_gen is None else collection.pdg_gen
def _gen_material(collection: SampleCollection) -> np.ndarray:
return (
collection.material
if collection.material_gen is None
else collection.material_gen
)
def _gen_cond(collection: SampleCollection) -> np.ndarray:
return (
collection.cond_cont_raw
if collection.cond_cont_raw_gen is None
else collection.cond_cont_raw_gen
)
def _group_labels(
collection: SampleCollection,
group_by: str | None,
n_energy_bins: int,
) -> list[tuple[str, np.ndarray]]:
n = len(collection.pdg)
) -> list[tuple[str, np.ndarray, np.ndarray]]:
"""Per-group `(label, mask_real, mask_gen)` triples.
`mask_real` indexes `collection.real_raw` (via `pdg`/`material`/
`cond_cont_raw`); `mask_gen` indexes `collection.gen_raw` via the `*_gen`
fields when set (unpaired real/gen — see `SampleCollection`), or the same
real-side arrays otherwise, which collapses to a single shared mask — the
original paired behavior (real/gen same length, row-for-row).
"""
gen_pdg, gen_material, gen_cond = (
_gen_pdg(collection),
_gen_material(collection),
_gen_cond(collection),
)
n_real, n_gen = len(collection.pdg), len(gen_pdg)
if group_by is None:
return [("all", np.ones(n, dtype=bool))]
return [("all", np.ones(n_real, dtype=bool), np.ones(n_gen, dtype=bool))]
if group_by == "pdg":
return [(f"pdg={v}", collection.pdg == v) for v in np.unique(collection.pdg)]
values = np.unique(np.concatenate([collection.pdg, gen_pdg]))
return [(f"pdg={v}", collection.pdg == v, gen_pdg == v) for v in values]
if group_by == "material":
values = np.unique(np.concatenate([collection.material, gen_material]))
return [
(f"material={v}", collection.material == v)
for v in np.unique(collection.material)
(f"material={v}", collection.material == v, gen_material == v)
for v in values
]
if group_by == "energy":
pre_E = collection.cond_cont_raw[:, 3]
edges = np.quantile(pre_E, np.linspace(0, 1, n_energy_bins + 1))
real_E, gen_E = collection.cond_cont_raw[:, 3], gen_cond[:, 3]
edges = np.quantile(
np.concatenate([real_E, gen_E]), np.linspace(0, 1, n_energy_bins + 1)
)
edges[-1] += 1e-6
bin_idx = np.digitize(pre_E, edges[1:-1])
real_bin = np.digitize(real_E, edges[1:-1])
gen_bin = np.digitize(gen_E, edges[1:-1])
return [
(f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", bin_idx == i)
(f"E∈[{edges[i]:.3g},{edges[i + 1]:.3g})", real_bin == i, gen_bin == i)
for i in range(n_energy_bins)
]
raise ValueError(f"unknown group_by={group_by!r}")
@@ -334,16 +398,19 @@ def marginal_table(
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:
for label, mask_real, mask_gen in _group_labels(
collection, group_by, n_energy_bins
):
if mask_real.sum() < 2 or mask_gen.sum() < 2:
continue
real, gen = collection.real_raw[mask], collection.gen_raw[mask]
real, gen = collection.real_raw[mask_real], collection.gen_raw[mask_gen]
for j, name in enumerate(RAW_TARGET_NAMES):
rows.append(
{
"group": label,
"dim": name,
"n": int(mask.sum()),
"n": int(mask_real.sum()),
"n_gen": int(mask_gen.sum()),
"real_mean": real[:, j].mean(),
"gen_mean": gen[:, j].mean(),
"real_std": real[:, j].std(),
@@ -649,8 +716,8 @@ def plot_marginals(
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 row, (label, mask_real, mask_gen) in enumerate(groups):
real, gen = collection.real_raw[mask_real], collection.gen_raw[mask_gen]
for col, j in enumerate(dim_idx):
ax = axes[row][col]
edges = _hist_edges(real[:, j], gen[:, j], bins=bins)
@@ -818,8 +885,6 @@ def plot_pairwise(
"""
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):
@@ -827,6 +892,8 @@ def plot_pairwise(
for row, (data, title) in enumerate(
[(collection.real_raw, "real"), (collection.gen_raw, "generated")]
):
n = len(data)
idx = rng.choice(n, size=min(n_sample, n), replace=False)
ax = axes[row][col]
ax.scatter(data[idx, ia], data[idx, ib], s=3, alpha=0.3)
ax.set_xlabel(a)
@@ -1723,6 +1790,22 @@ _ROLLOUT_COLS = [
]
def _check_rollout_metadata(path: Path) -> None:
"""Raise if `path` carries coord metadata that isn't `ROLLOUT_COORD_VALUE`.
A missing tag (older rollout output, predating tagging) is let through
silently, matching `giant predict`/`giant rollout`'s own leniency; a tag
that's present but wrong is a real mismatch.
"""
metadata = pq.read_schema(path).metadata or {}
coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
raise ValueError(
f"{path} is not a rollout file (coord={coord.decode()!r}); "
"expected a `giant rollout` output"
)
@dataclass
class RolloutObservables:
event_table: pd.DataFrame # one row per event_id (mm/MeV)
@@ -1746,13 +1829,7 @@ def compute_rollout_observables(
depth-along-axis and transverse-distance-from-axis, then binned. Returns
per-event totals plus dataset-mean longitudinal/transverse profiles.
"""
pf = pq.ParquetFile(Path(path))
coord = (pf.schema_arrow.metadata or {}).get(PREDICT_COORD_METADATA_KEY.encode())
if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
raise ValueError(
f"{path} is not a rollout file (coord={coord.decode()!r}); "
"expected a `giant rollout` output"
)
_check_rollout_metadata(Path(path))
df = pd.read_parquet(path, columns=_ROLLOUT_COLS)
@@ -1899,3 +1976,201 @@ def plot_rollout_total_energy(obs: RolloutObservables, bins: int = 50, reference
ax.legend()
fig.tight_layout()
return fig
# ── Rollout vs. held-out truth (Tier 1-3, unpaired) ────────────────────────────
#
# `compute_rollout_observables` above covers Tier 4 (event-level) comparisons.
# This builds a `SampleCollection` instead, so the Tier 1-3 diagnostics
# (plot_marginals, plot_kl_bars, correlation_matrices, plot_pairwise,
# direction_alignment, constraint_report) also work on a rollout: `giant
# rollout` output and a training-input-schema truth file (see
# `giant.data.loader.load_steps`) both carry pre_*/post_*/edep/step_length in
# the same physical, world-frame units, so both sides decode into
# RAW_TARGET_NAMES space via the same local_frame_rotation/travel_direction
# construction `giant.data.transforms.build_features` uses for `target_s1` —
# just skipping the log/ALR encode step, since neither file needs it decoded.
#
# Unlike `load_predicted_local`, the two files are independent rather than
# row-for-row paired (a rollout doesn't replay real events step-by-step), so
# real_raw/gen_raw may have different lengths; `SampleCollection`'s `*_gen`
# fields carry the rollout side's own pdg/material/conditioning for grouping.
_WORLD_FRAME_STEP_COLS = [
"pdg",
"material",
"layer_id",
"pre_x",
"pre_y",
"pre_z",
"pre_E",
"pre_dx",
"pre_dy",
"pre_dz",
"post_x",
"post_y",
"post_z",
"post_E",
"post_dx",
"post_dy",
"post_dz",
"edep",
"step_length",
]
def _world_frame_raw_targets(batch: pl.DataFrame) -> np.ndarray:
"""(N, 9) RAW_TARGET_NAMES array from a world-frame steps batch.
See the module note above `_WORLD_FRAME_STEP_COLS` — both `load_rollout_vs_truth`
inputs share this schema.
"""
pre_pos = batch.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32)
pre_dir = batch.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32)
post_pos = (
batch.select(["post_x", "post_y", "post_z"]).to_numpy().astype(np.float32)
)
post_dir = (
batch.select(["post_dx", "post_dy", "post_dz"]).to_numpy().astype(np.float32)
)
pre_E = batch["pre_E"].to_numpy().astype(np.float32)
post_E = batch["post_E"].to_numpy().astype(np.float32)
post_dir_local = local_frame_rotation(pre_dir, post_dir)
travel_dir_local = local_frame_rotation(
pre_dir, travel_direction(pre_pos, post_pos)
)
return np.column_stack(
[
batch["step_length"].to_numpy().astype(np.float32),
pre_E - post_E,
batch["edep"].to_numpy().astype(np.float32),
post_dir_local,
travel_dir_local,
]
).astype(np.float32)
def _world_frame_cond(batch: pl.DataFrame) -> np.ndarray:
"""(N, 9) `_COND_CONT_COLS`-layout array from a world-frame steps batch.
`n_sec` comes from `child_track_ids` (truth schema) or `n_sec_pred`
(rollout schema) — whichever the batch has; kept only for shape parity
with `load_predicted_local`'s `cond_cont_raw` — nothing downstream in this
module groups by it, only `pre_E` (index 3, for `group_by="energy"`).
"""
if "n_sec_pred" in batch.columns:
n_sec = batch["n_sec_pred"].to_numpy().astype(np.float32)
elif "child_track_ids" in batch.columns:
n_sec = batch["child_track_ids"].list.len().to_numpy().astype(np.float32)
else:
n_sec = np.zeros(batch.height, dtype=np.float32)
return np.column_stack(
[
batch.select(["pre_x", "pre_y", "pre_z"]).to_numpy(),
batch["pre_E"].to_numpy(),
batch.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy(),
batch["layer_id"].to_numpy().astype(np.float32),
n_sec,
]
).astype(np.float32)
def _existing_columns(
source: str | Path | pl.LazyFrame, wanted: list[str]
) -> list[str]:
if isinstance(source, pl.LazyFrame):
available = set(source.collect_schema().names())
else:
available = set(pq.ParquetFile(Path(source)).schema_arrow.names)
return [c for c in wanted if c in available]
def _load_world_frame_side(
source: str | Path | pl.LazyFrame,
sample_frac: float,
seed: int,
batch_size: int,
extra_cols: list[str],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
"""Stream one world-frame steps file/LazyFrame into (raw9, cond9, pdg, material).
`extra_cols` (e.g. `["n_sec_pred"]` or `["child_track_ids"]`) are included
only when present, so the same helper serves both the rollout and truth
schemas without either needing the other's columns.
"""
columns = _WORLD_FRAME_STEP_COLS + _existing_columns(source, extra_cols)
threshold = int(sample_frac * 2**32) if sample_frac < 1.0 else None
raw_parts, cond_parts, pdg_parts, mat_parts = [], [], [], []
offset = 0
for batch in _iter_predicted_local_batches(source, columns, batch_size):
n = batch.height
if threshold is not None:
row_idx = pl.arange(offset, offset + n, eager=True).cast(pl.UInt32)
batch = batch.filter((row_idx.hash(seed=seed) % 2**32) < threshold)
offset += n
if batch.height == 0:
continue
raw_parts.append(_world_frame_raw_targets(batch))
cond_parts.append(_world_frame_cond(batch))
pdg_parts.append(batch["pdg"].to_numpy())
mat_parts.append(batch["material"].to_numpy())
if not raw_parts:
raise ValueError(f"{source}: no rows survived (sample_frac={sample_frac})")
return (
np.concatenate(raw_parts, axis=0),
np.concatenate(cond_parts, axis=0),
np.concatenate(pdg_parts, axis=0),
np.concatenate(mat_parts, axis=0),
)
def load_rollout_vs_truth(
rollout_path: str | Path | pl.LazyFrame,
truth_path: str | Path | pl.LazyFrame,
sample_frac: float = 1.0,
seed: int = 0,
batch_size: int = 1_000_000,
) -> SampleCollection:
"""Build a `SampleCollection` comparing a `giant rollout` shower to a truth file.
`truth_path` — any file sharing `giant train`'s input schema (real
miniCaloSim steps, e.g. a held-out/val parquet) — is treated as "real";
`rollout_path` (`giant rollout` output) is treated as "generated". The two
are independent files (not row-for-row paired, since a rollout doesn't
replay real events step-by-step): `real_raw`/`gen_raw` may have different
lengths, and `pdg`/`material`/`cond_cont_raw` are computed separately per
side (`SampleCollection`'s `*_gen` fields) — `_group_labels` (used by
`marginal_table`/`plot_marginals`/`plot_kl_bars`) builds independent masks
for each. `correlation_matrices`, `direction_alignment`, and
`constraint_report` never paired real/gen row-for-row to begin with, so
they need no special handling here.
Every raw target dim is decoded from the world-frame `pre_*`/`post_*`/
`edep`/`step_length` columns both files share — see the module note above
`_WORLD_FRAME_STEP_COLS`. `sample_frac`/`seed`/`batch_size` behave as in
`load_predicted_local`, applied independently to each file.
"""
if not (0 < sample_frac <= 1):
raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}")
if not isinstance(rollout_path, pl.LazyFrame):
_check_rollout_metadata(Path(rollout_path))
gen_raw, gen_cond, gen_pdg, gen_material = _load_world_frame_side(
rollout_path, sample_frac, seed, batch_size, extra_cols=["n_sec_pred"]
)
real_raw, real_cond, real_pdg, real_material = _load_world_frame_side(
truth_path, sample_frac, seed, batch_size, extra_cols=["child_track_ids"]
)
return SampleCollection(
cond_cont_raw=real_cond,
pdg=real_pdg,
material=real_material,
real_raw=real_raw,
gen_raw=gen_raw,
cond_cont_raw_gen=gen_cond,
pdg_gen=gen_pdg,
material_gen=gen_material,
)
tests/test_analysis.py
+178
View File
@@ -17,6 +17,7 @@ from giant.analysis import (
correlation_matrices,
direction_alignment,
load_predicted_local,
load_rollout_vs_truth,
marginal_table,
marginal_table_pl,
pdg_contribution_table_pl,
@@ -40,12 +41,15 @@ from giant.constants import (
PREDICT_COORD_METADATA_KEY,
PREDICT_SCHEMA_VERSION,
PREDICT_SCHEMA_VERSION_KEY,
ROLLOUT_COORD_VALUE,
)
from giant.data.transforms import (
energy_simplex_decode,
inv_log_transform,
local_frame_rotation,
log_transform,
reconstruct_post_pos,
travel_direction,
)
@@ -660,3 +664,177 @@ def test_plot_pdg_energy_share_caps_slices():
fig = plot_pdg_energy_share(table, max_slices=4)
for ax in fig.axes:
assert len(ax.patches) == 4
# ---------------------------------------------------------------------------
# load_rollout_vs_truth: unpaired rollout-vs-truth SampleCollection
# ---------------------------------------------------------------------------
def _make_world_frame_physical(rng, n):
"""Random-but-physical pre/post step fields shared by the truth/rollout schemas."""
pre_pos = rng.uniform(-5.0, 5.0, (n, 3)).astype(np.float32)
pre_dir = _unit_vectors(rng, n)
pre_E = rng.uniform(1.0, 100.0, n).astype(np.float32)
step_length = rng.uniform(0.1, 5.0, n).astype(np.float32)
travel_dir_world = _unit_vectors(rng, n)
post_pos = pre_pos + step_length[:, None] * travel_dir_world
post_dir_world = _unit_vectors(rng, n)
delta_e = (rng.uniform(0.0, 1.0, n) * pre_E).astype(np.float32)
post_E = pre_E - delta_e
edep = (delta_e * rng.uniform(0.0, 1.0, n)).astype(np.float32)
return pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep
def _expected_raw9(
pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep
):
post_dir_local = local_frame_rotation(pre_dir, post_dir_world)
travel_dir_local = local_frame_rotation(
pre_dir, travel_direction(pre_pos, post_pos)
)
return np.column_stack(
[step_length, pre_E - post_E, edep, post_dir_local, travel_dir_local]
).astype(np.float32)
def _write_truth_parquet(path, n=200, seed=0):
rng = np.random.default_rng(seed)
fields = _make_world_frame_physical(rng, n)
pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep = (
fields
)
table = pa.table(
{
"event_id": rng.integers(0, 20, n),
"pdg": rng.choice([11, -11, 22], n),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"child_track_ids": [list(range(int(k))) for k in rng.integers(0, 3, n)],
"e_sec": rng.uniform(0.0, 1.0, n).astype(np.float32),
"step_length": step_length,
"post_E": post_E,
"edep": edep,
"post_dx": post_dir_world[:, 0],
"post_dy": post_dir_world[:, 1],
"post_dz": post_dir_world[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
}
)
pq.write_table(table, path)
return _expected_raw9(*fields)
def _write_rollout_parquet(path, n=150, seed=1, coord=ROLLOUT_COORD_VALUE):
rng = np.random.default_rng(seed)
fields = _make_world_frame_physical(rng, n)
pre_pos, pre_dir, pre_E, step_length, post_pos, post_dir_world, post_E, edep = (
fields
)
table = pa.table(
{
"event_id": rng.integers(0, 20, n),
"track_id": rng.integers(0, 3, n),
"parent_id": np.full(n, -1, dtype=np.int64),
"generation": np.zeros(n, dtype=np.int64),
"step_no": np.zeros(n, dtype=np.int64),
"pdg": rng.choice([11, -11, 22], n),
"pre_x": pre_pos[:, 0],
"pre_y": pre_pos[:, 1],
"pre_z": pre_pos[:, 2],
"pre_E": pre_E,
"pre_dx": pre_dir[:, 0],
"pre_dy": pre_dir[:, 1],
"pre_dz": pre_dir[:, 2],
"post_x": post_pos[:, 0],
"post_y": post_pos[:, 1],
"post_z": post_pos[:, 2],
"post_E": post_E,
"post_dx": post_dir_world[:, 0],
"post_dy": post_dir_world[:, 1],
"post_dz": post_dir_world[:, 2],
"edep": edep,
"step_length": step_length,
"material": rng.choice(["W", "Pb"], n),
"layer_id": rng.integers(0, 10, n).astype(np.int32),
"n_sec_pred": rng.integers(0, 3, n).astype(np.int32),
"termination_reason": rng.choice(["natural_end", "energy_cutoff"], n),
}
)
if coord is not None:
table = table.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: coord})
pq.write_table(table, path)
return _expected_raw9(*fields)
def test_load_rollout_vs_truth_decodes_raw_targets_correctly(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
expected_real = _write_truth_parquet(truth_path, n=200, seed=0)
expected_gen = _write_rollout_parquet(rollout_path, n=150, seed=1)
samples = load_rollout_vs_truth(rollout_path, truth_path)
np.testing.assert_allclose(samples.real_raw, expected_real, atol=1e-4)
np.testing.assert_allclose(samples.gen_raw, expected_gen, atol=1e-4)
def test_load_rollout_vs_truth_allows_unpaired_lengths(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path, n=200)
_write_rollout_parquet(rollout_path, n=150)
samples = load_rollout_vs_truth(rollout_path, truth_path)
assert samples.real_raw.shape == (200, 9)
assert samples.gen_raw.shape == (150, 9)
assert samples.pdg.shape == (200,)
assert samples.pdg_gen.shape == (150,)
@pytest.mark.parametrize("group_by", [None, "pdg", "material", "energy"])
def test_load_rollout_vs_truth_downstream_plots_run_without_error(tmp_path, group_by):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path, n=200)
_write_rollout_parquet(rollout_path, n=150)
samples = load_rollout_vs_truth(rollout_path, truth_path)
table = marginal_table(samples, group_by=group_by)
assert set(table["dim"]) == set(RAW_TARGET_NAMES)
assert plot_marginals(samples, group_by=group_by) is not None
assert plot_kl_bars(samples, group_by=group_by) is not None
def test_load_rollout_vs_truth_joint_and_constraint_checks_run(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path, n=200)
_write_rollout_parquet(rollout_path, n=150)
samples = load_rollout_vs_truth(rollout_path, truth_path)
assert plot_correlation_matrices(samples) is not None
assert plot_pairwise(samples) is not None
assert plot_direction_alignment(samples) is not None
assert plot_constraint_violations(samples) is not None
assert constraint_report(samples) is not None
def test_load_rollout_vs_truth_rejects_wrong_coord_metadata(tmp_path):
truth_path = tmp_path / "truth.parquet"
rollout_path = tmp_path / "rollout.parquet"
_write_truth_parquet(truth_path)
_write_rollout_parquet(rollout_path, coord="local")
with pytest.raises(ValueError, match="not a rollout file"):
load_rollout_vs_truth(rollout_path, truth_path)