lars
81ec225b86
Merged in every non-analysis change from the MoE-prototype branch (routing, training, data pipeline, streaming rollout output), keeping this branch's lean streaming giant/analysis.py and rebuilding the rollout-vs-truth feature natively on it instead of resurrecting the old numpy SampleCollection path. - Add RolloutVsTruth, accepted anywhere Tier 1-3 functions take a predict-parquet source: decodes a giant rollout file and a held-out truth file into RAW_TARGET_NAMES space via a polars port of the forward local-frame rotation, fully streaming (no SampleCollection, no eager materialization). - Add compute_rollout_vs_truth_observables_pl for Tier 4, reusing EventObservables (now backed by independent real_table/gen_table to support unequal rollout/truth event counts) so every existing shower-observable plot function works unchanged for both one-step and full-rollout comparisons. - Update analysis/rollout_validation.ipynb to the new API and CLAUDE.md's architecture description; add test coverage for the new source type. - Fix a pre-existing return-type mismatch in giant.rollout.rollout() (found by `ty check`): the on_chunk summary-dict branch didn't match the declared dict[str, np.ndarray] return type, now expressed as a RolloutSummary TypedDict. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
321 lines
11 KiB
Plaintext
321 lines
11 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "f91460f3",
|
|
"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": "0f10da93",
|
|
"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 — 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 — 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 `RolloutVsTruth`, which treats the rollout file as \"generated\" and the truth file as \"real\". Unlike the paired predict-parquet `source` (`pred_*`/`true_*` columns of the same row), the two files here are **independent, unpaired datasets** — 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, and every check still streams (no `SampleCollection`, no full-file materialization) — see `giant.analysis`'s module docstring for the `RolloutVsTruth` mechanics.\n",
|
|
"\n",
|
|
"Same four tiers as `validation.ipynb`, all built on the same functions — pass a `RolloutVsTruth` in place of the predict-parquet path/LazyFrame everywhere:\n",
|
|
"\n",
|
|
"1. **stratified marginals** — per-dimension real-vs-generated, sliced by pdg/material/energy\n",
|
|
"2. **joint structure** — correlation matrices, physically-coupled pairwise plots, direction alignment\n",
|
|
"3. **physical constraints** — unit-norm directions, non-negative step_length/delta_e/edep (checked on the rollout's own output — 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** — total/mean/median energy and length per event, longitudinal/transverse profiles, shower-max depth, computed directly from the rollout shower against the truth file's own events (`compute_rollout_vs_truth_observables_pl`, the Tier 4 counterpart to `RolloutVsTruth`)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "251dc1f7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from giant.analysis import RolloutVsTruth, plot_kl_bars_pl\n",
|
|
"\n",
|
|
"# `giant rollout` output for the shower(s) under test.\n",
|
|
"ROLLOUT_FILE = (\n",
|
|
" \"/ceph/lbogner/geant_steps/predictions/9e76bc2c-f4ef-4488-9f62-b6d14e1f298e.parquet\"\n",
|
|
")\n",
|
|
"# Any held-out file sharing giant train's input schema (real miniCaloSim\n",
|
|
"# steps) — e.g. the val split the rollout's seed events were drawn from.\n",
|
|
"TRUTH_FILE = (\n",
|
|
" \"/ceph/lbogner/geant_steps/processed/steps/gen3/schema2/pbwo4/shard-009.parquet\"\n",
|
|
")\n",
|
|
"\n",
|
|
"# sample_frac subsamples each side of the Tier 1-3 checks independently\n",
|
|
"# (kept memory-bounded for large files); defaults to every row. Tier 4\n",
|
|
"# (compute_rollout_vs_truth_observables_pl, below) always streams every row\n",
|
|
"# regardless — per-event sums would be silently corrupted by row subsampling.\n",
|
|
"SOURCE = RolloutVsTruth(rollout=ROLLOUT_FILE, truth=TRUTH_FILE)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "df4bf24b",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Tier 1: stratified marginals\n",
|
|
"\n",
|
|
"KL(real || generated) per target dimension, streamed straight from both files."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "13cd0838",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"for grouping in [None, \"energy\", \"pdg\", \"material\"]:\n",
|
|
" fig = plot_kl_bars_pl(SOURCE, group_by=grouping)\n",
|
|
" fig.show()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "aba92bf9",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Detailed marginals (Tier 1, overlaid histograms)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "7bab2e35",
|
|
"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(SOURCE)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "6c2958e7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_marginals(SOURCE, group_by=\"energy\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "0fe70835",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_marginals(SOURCE, group_by=\"pdg\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "038dda3b",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Tier 2: joint structure"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3c463d50",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Real vs. generated Pearson correlation matrices (+ their difference) over\n",
|
|
"# the 9 raw target dims — catches a model that decorrelates targets that are\n",
|
|
"# physically coupled even when every individual marginal looks clean.\n",
|
|
"_ = plot_correlation_matrices(SOURCE)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "393c6845",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Scatter for physically-coupled pairs (step_length/delta_e/edep) — the\n",
|
|
"# joint-structure check correlation matrices alone can't fully capture.\n",
|
|
"_ = plot_pairwise(SOURCE, n_sample=10000)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1dab5305",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# cos(angle) between post_dir and travel_dir — coupled through the\n",
|
|
"# scattering physics, so this is another joint-structure check.\n",
|
|
"_ = plot_direction_alignment(SOURCE)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "7232d3b4",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Tier 3: physical constraints\n",
|
|
"\n",
|
|
"Unit-norm direction vectors, non-negative step_length/delta_e/edep. `constraint_report_pl`/`plot_constraint_violations` only ever check the *generated* side (here the rollout output) — 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": "86be25fa",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_constraint_violations(SOURCE)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ed4d3037",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Tier 4: event-level (shower) observables\n",
|
|
"\n",
|
|
"Built on `compute_rollout_vs_truth_observables_pl`, not `compute_event_observables_pl` — the rollout file carries its own `track_id`/`termination_reason` columns the event-level aggregation needs, and the shower here already *is* a full autoregressive rollout rather than one-step generations re-aggregated by event. Entry axis/point and per-event totals are computed separately per side (rollout and truth events are unrelated), but depth/transverse bin edges are shared across both so the profiles below overlay on one binning.\n",
|
|
"\n",
|
|
"Returns the same `EventObservables` `compute_event_observables_pl` does, so every plot function from `validation.ipynb` works unchanged here too."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9175876a",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"from giant.analysis import compute_rollout_vs_truth_observables_pl\n",
|
|
"from giant.analysis import plot_total_energy, plot_total_length\n",
|
|
"from giant.analysis import plot_mean_energy_per_step, plot_mean_length_per_step\n",
|
|
"from giant.analysis import plot_longitudinal_profile, plot_transverse_profile\n",
|
|
"from giant.analysis import plot_shower_max_depth\n",
|
|
"\n",
|
|
"obs = compute_rollout_vs_truth_observables_pl(ROLLOUT_FILE, TRUTH_FILE)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "83e6dd0e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_total_energy(obs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "1e35a387",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_total_length(obs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "9221e682",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_mean_energy_per_step(obs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4536acb0",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_mean_length_per_step(obs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "3a03c04f",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_longitudinal_profile(obs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "eae0536e",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_transverse_profile(obs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "98c221e7",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"_ = plot_shower_max_depth(obs)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "208ca6e4",
|
|
"metadata": {},
|
|
"source": [
|
|
"---\n",
|
|
"\n",
|
|
"For the dataset-wide breakdown of which particle species contributed how much of the total energy/length (`pdg_contribution_table_pl`), see `validation.ipynb` — it needs the paired predict schema, which this rollout-vs-truth comparison doesn't have."
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "giant (3.12.13.final.0)",
|
|
"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
|
|
}
|