From 36bdbf7cc174df12977a458e3a07580ec05bbc21 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 13 Jul 2026 10:58:42 +0200 Subject: [PATCH] Add load_rollout_vs_truth to compare rollouts against held-out truth data 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 --- analysis/rollout_validation.ipynb | 305 ++++++++++++++++++++++++++++ giant/analysis.py | 327 +++++++++++++++++++++++++++--- tests/test_analysis.py | 178 ++++++++++++++++ 3 files changed, 784 insertions(+), 26 deletions(-) create mode 100644 analysis/rollout_validation.ipynb diff --git a/analysis/rollout_validation.ipynb b/analysis/rollout_validation.ipynb new file mode 100644 index 0000000..70af3f2 --- /dev/null +++ b/analysis/rollout_validation.ipynb @@ -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 +} diff --git a/giant/analysis.py b/giant/analysis.py index 45ad39c..d491fb9 100644 --- a/giant/analysis.py +++ b/giant/analysis.py @@ -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, + ) diff --git a/tests/test_analysis.py b/tests/test_analysis.py index a2b6a90..b46e556 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -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)