From 453f9f9e208d981d2d2c01d86183447d5643747d Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 22 Jun 2026 16:47:51 +0200 Subject: [PATCH] Add total length traveled per event to event observables sum(step_length) per event_id, alongside the existing total deposited energy, since path length and energy deposit aren't interchangeable once tracks scatter. Adds plot_total_length and a matching notebook cell. Co-Authored-By: Claude Sonnet 4.6 --- analysis/validation.ipynb | 61 +++++++++++++++++++++++++------------- giant/analysis.py | 62 ++++++++++++++++++++++++++++++++------- tests/test_analysis.py | 48 ++++++++++++++++++++++-------- 3 files changed, 129 insertions(+), 42 deletions(-) diff --git a/analysis/validation.ipynb b/analysis/validation.ipynb index ca85a71..a3f975a 100644 --- a/analysis/validation.ipynb +++ b/analysis/validation.ipynb @@ -7,12 +7,12 @@ "source": [ "# GIANT validation notebook\n", "\n", - "Diagnostics for a trained checkpoint's sample quality, run against `giant predict --coord local` output (`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled — see `giant.analysis`'s module docstring). Four tiers, each building on the last:\n", + "Diagnostics for a trained checkpoint's sample quality, run against `giant predict --coord local` output (`pred_*`/`true_*` columns, denormalized but still local-frame/log-scaled \u2014 see `giant.analysis`'s module docstring). Four tiers, each building on the last:\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\n", - "4. **event-level (shower) observables** — total energy, longitudinal/transverse profiles, shower-max depth, in world-frame physical units (mm, MeV)\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\n", + "4. **event-level (shower) observables** \u2014 total energy, longitudinal/transverse profiles, shower-max depth, in world-frame physical units (mm, MeV)\n" ] }, { @@ -115,7 +115,7 @@ "id": "22c67dc4", "metadata": {}, "source": [ - "## Detailed marginals, correlation & constraints (Tiers 1–3, in-memory sample)\n", + "## Detailed marginals, correlation & constraints (Tiers 1\u20133, in-memory sample)\n", "\n", "The richer per-row diagnostics below (overlaid histograms, correlation matrices, pairwise scatter, direction alignment, constraint violations) need `real_raw`/`gen_raw` materialized as numpy arrays, so they run on a `SampleCollection` built from a 50% row sample rather than the lazy, full-file path used above." ] @@ -152,7 +152,7 @@ } ], "source": [ - "# sample_frac=0.5 keeps this a manageable in-memory size — fine for these\n", + "# sample_frac=0.5 keeps this a manageable in-memory size \u2014 fine for these\n", "# per-row diagnostics, unlike the event-level checks further down, which\n", "# need every row of an event present to sum correctly.\n", "samples = load_predicted_local(FILE, sample_frac=0.5)\n", @@ -228,7 +228,7 @@ ], "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", + "# 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)" ] @@ -251,7 +251,7 @@ } ], "source": [ - "# Scatter for physically-coupled pairs (step_length/delta_e/edep) — the\n", + "# 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=len(samples.gen_raw))" ] @@ -274,7 +274,7 @@ } ], "source": [ - "# cos(angle) between post_dir and travel_dir — coupled through the\n", + "# 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)" ] @@ -305,7 +305,7 @@ } ], "source": [ - "# Unit-norm direction vectors, non-negative step_length/delta_e/edep — the\n", + "# Unit-norm direction vectors, non-negative step_length/delta_e/edep \u2014 the\n", "# unconstrained MLP has nothing enforcing these, so any violation here is a\n", "# pure generation artifact rather than a real-data property.\n", "_ = plot_constraint_violations(samples)" @@ -318,11 +318,11 @@ "source": [ "## Tier 4: event-level (shower) observables\n", "\n", - "Everything above is a **step-level** check: one row in, one row out, compared in the local frame (`pre_dir = ẑ`). This section aggregates those same rows **per `event_id`**, reconstructed into world-frame physical units (mm, MeV), to check the shower-level quantities that actually matter physically: total deposited energy, longitudinal/transverse shower profiles, and shower-max depth (see `diffusion-model-tutorial.md` §7.2).\n", + "Everything above is a **step-level** check: one row in, one row out, compared in the local frame (`pre_dir = \u1e91`). This section aggregates those same rows **per `event_id`**, reconstructed into world-frame physical units (mm, MeV), to check the shower-level quantities that actually matter physically: total deposited energy, longitudinal/transverse shower profiles, and shower-max depth (see `diffusion-model-tutorial.md` \u00a77.2).\n", "\n", - "**Caveat:** this re-aggregates one-step-ahead generations — each row is generated conditioned on the *real* preceding state, then grouped by event — not a full autoregressive shower rollout. 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.\n", + "**Caveat:** this re-aggregates one-step-ahead generations \u2014 each row is generated conditioned on the *real* preceding state, then grouped by event \u2014 not a full autoregressive shower rollout. 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.\n", "\n", - "`compute_event_observables_pl` streams the full file directly (two polars passes, no `SampleCollection`) rather than reusing `samples` above — per-event sums would be silently corrupted by `sample_frac`-style row subsampling, since a partially-sampled event no longer sums to the true per-event total." + "`compute_event_observables_pl` streams the full file directly (two polars passes, no `SampleCollection`) rather than reusing `samples` above \u2014 per-event sums would be silently corrupted by `sample_frac`-style row subsampling, since a partially-sampled event no longer sums to the true per-event total." ] }, { @@ -333,10 +333,11 @@ "outputs": [], "source": [ "from giant.analysis import compute_event_observables_pl\n", - "from giant.analysis import plot_total_energy, plot_longitudinal_profile\n", + "from giant.analysis import plot_total_energy, plot_total_length\n", + "from giant.analysis import plot_longitudinal_profile\n", "from giant.analysis import plot_transverse_profile, plot_shower_max_depth\n", "\n", - "# Full file, not `samples` — see the markdown cell above for why.\n", + "# Full file, not `samples` \u2014 see the markdown cell above for why.\n", "obs = compute_event_observables_pl(FILE)" ] }, @@ -347,7 +348,7 @@ "source": [ "### Total deposited energy per event\n", "\n", - "`sum(edep)` grouped by `event_id`, real vs. generated, with the resolution (σ/μ) for each annotated in the legend." + "`sum(edep)` grouped by `event_id`, real vs. generated, with the resolution (\u03c3/\u03bc) for each annotated in the legend." ] }, { @@ -371,6 +372,26 @@ "_ = plot_total_energy(obs)" ] }, + { + "cell_type": "markdown", + "id": "8c63c999", + "metadata": {}, + "source": [ + "### Total length traveled per event\n", + "\n", + "`sum(step_length)` grouped by `event_id` \u2014 total path length traveled by every track in the shower, real vs. generated (not the same as the depth of any single point, since tracks scatter)." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d6561b73", + "metadata": {}, + "outputs": [], + "source": [ + "_ = plot_total_length(obs)" + ] + }, { "cell_type": "markdown", "id": "cf3ea615", @@ -378,7 +399,7 @@ "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 — the classic `E_dep(depth)` profile." + "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 \u2014 the classic `E_dep(depth)` profile." ] }, { @@ -409,7 +430,7 @@ "source": [ "### Transverse profile\n", "\n", - "Same idea, binned by perpendicular distance from the shower axis instead of depth — a Molière-radius-style lateral containment check." + "Same idea, binned by perpendicular distance from the shower axis instead of depth \u2014 a Moli\u00e8re-radius-style lateral containment check." ] }, { @@ -440,7 +461,7 @@ "source": [ "### Shower-maximum depth\n", "\n", - "Per event, the depth bin where that event's longitudinal profile peaks — compares the real vs. generated distribution of shower-max depth across events, rather than the pooled profile above." + "Per event, the depth bin where that event's longitudinal profile peaks \u2014 compares the real vs. generated distribution of shower-max depth across events, rather than the pooled profile above." ] }, { diff --git a/giant/analysis.py b/giant/analysis.py index 39393ec..09391dd 100644 --- a/giant/analysis.py +++ b/giant/analysis.py @@ -33,11 +33,13 @@ output, see `compute_event_observables_pl` — it streams the file directly corrupted by partial events:: from giant.analysis import compute_event_observables_pl - from giant.analysis import plot_total_energy, plot_longitudinal_profile - from giant.analysis import plot_transverse_profile, plot_shower_max_depth + 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_longitudinal_profile(obs) plot_transverse_profile(obs) plot_shower_max_depth(obs) @@ -54,10 +56,11 @@ Four tiers of checks, building on the aggregate marginal/KL check in 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, 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) +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 @@ -1069,6 +1072,10 @@ def compute_event_observables_pl( 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 @@ -1091,6 +1098,8 @@ def compute_event_observables_pl( 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) @@ -1127,7 +1136,7 @@ def compute_event_observables_pl( .astype(np.float32) ) - def _reconstruct(cols: list[str]) -> tuple[np.ndarray, np.ndarray]: + 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]) @@ -1135,10 +1144,10 @@ def compute_event_observables_pl( post_pos = reconstruct_post_pos( pre_pos, pre_dir, step_length, travel_dir_local ) - return post_pos, edep + return post_pos, edep, step_length - real_post_pos, real_edep = _reconstruct(true_cols) - gen_post_pos, gen_edep = _reconstruct(pred_cols) + 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] @@ -1161,6 +1170,8 @@ def compute_event_observables_pl( 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) @@ -1187,6 +1198,8 @@ def compute_event_observables_pl( "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_centroid_depth": real_centroid_depth, "gen_centroid_depth": gen_centroid_depth, "real_transverse_rms": real_transverse_rms, @@ -1240,6 +1253,35 @@ def plot_total_energy(observables: EventObservables, bins: int = 50): 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_profile( centers: np.ndarray, real_mean: np.ndarray, diff --git a/tests/test_analysis.py b/tests/test_analysis.py index 5525699..a5ee5f8 100644 --- a/tests/test_analysis.py +++ b/tests/test_analysis.py @@ -29,6 +29,7 @@ from giant.analysis import ( plot_pairwise, plot_shower_max_depth, plot_total_energy, + plot_total_length, plot_transverse_profile, ) from giant.constants import ( @@ -454,16 +455,21 @@ def _expected_event_table( disp = post_pos - entry_pos depth = disp @ axis_dir transverse = np.linalg.norm(disp - depth[:, None] * axis_dir, axis=1) - total = float(edep.sum()) - centroid = float((edep * depth).sum() / total) - rms = float(np.sqrt((edep * transverse**2).sum() / total)) - return total, centroid, rms + total_edep = float(edep.sum()) + total_length = float(step_length.sum()) + centroid = float((edep * depth).sum() / total_edep) + rms = float(np.sqrt((edep * transverse**2).sum() / total_edep)) + return total_edep, total_length, centroid, rms - real_total, real_centroid, real_rms = agg(true_log_local) - gen_total, gen_centroid, gen_rms = agg(pred_log_local) + real_total_edep, real_total_length, real_centroid, real_rms = agg( + true_log_local + ) + gen_total_edep, gen_total_length, gen_centroid, gen_rms = agg(pred_log_local) expected[e] = ( - real_total, - gen_total, + real_total_edep, + gen_total_edep, + real_total_length, + gen_total_length, real_centroid, gen_centroid, real_rms, @@ -485,11 +491,28 @@ def test_compute_event_observables_pl_matches_manual_reconstruction(tmp_path): table = obs.event_table.sort("event_id") for i, eid in enumerate(table["event_id"].to_list()): - real_total, gen_total, real_centroid, gen_centroid, real_rms, gen_rms = ( - expected[eid] + ( + real_total_edep, + gen_total_edep, + real_total_length, + gen_total_length, + real_centroid, + gen_centroid, + real_rms, + gen_rms, + ) = expected[eid] + np.testing.assert_allclose( + table["real_total_edep"][i], real_total_edep, rtol=1e-4 + ) + np.testing.assert_allclose( + table["gen_total_edep"][i], gen_total_edep, rtol=1e-4 + ) + np.testing.assert_allclose( + table["real_total_length"][i], real_total_length, rtol=1e-4 + ) + np.testing.assert_allclose( + table["gen_total_length"][i], gen_total_length, rtol=1e-4 ) - np.testing.assert_allclose(table["real_total_edep"][i], real_total, rtol=1e-4) - np.testing.assert_allclose(table["gen_total_edep"][i], gen_total, rtol=1e-4) np.testing.assert_allclose( table["real_centroid_depth"][i], real_centroid, rtol=1e-3, atol=1e-4 ) @@ -539,6 +562,7 @@ def test_event_level_plots_run_without_error(tmp_path): obs = compute_event_observables_pl(path, depth_bins=5, transverse_bins=5) assert plot_total_energy(obs) is not None + assert plot_total_length(obs) is not None assert plot_longitudinal_profile(obs) is not None assert plot_transverse_profile(obs) is not None assert plot_shower_max_depth(obs) is not None