{
 "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
}