{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bb22f42a",
   "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": "f65c3ad9",
   "metadata": {},
   "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",
    "\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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1b65ca80",
   "metadata": {},
   "outputs": [],
   "source": [
    "from giant.analysis import plot_kl_bars_pl\n",
    "\n",
    "# Predict parquet produced by `giant predict --coord local --checkpoint ...`,\n",
    "# carrying both pred_*/true_* columns so real vs. generated can be compared.\n",
    "FILE = \"/home/lars/Programming/giant/0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "82142e6b",
   "metadata": {},
   "source": [
    "## Tier 1: stratified marginals\n",
    "\n",
    "KL(real || generated) per target dimension, computed lazily straight from the parquet over every row in the file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f2752229",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_1223195/3415847466.py:7: UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown\n",
      "  fig.show()\n",
      "/tmp/ipykernel_1223195/3415847466.py:7: UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown\n",
      "  fig.show()\n",
      "/tmp/ipykernel_1223195/3415847466.py:7: UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown\n",
      "  fig.show()\n",
      "/tmp/ipykernel_1223195/3415847466.py:7: UserWarning: FigureCanvasAgg is non-interactive, and thus cannot be shown\n",
      "  fig.show()\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 27448 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 36172 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 31828 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 44792 characters]",
      "text/plain": [
       "<Figure size 800x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# KL(real || generated) per target dimension, read lazily straight from the\n",
    "# parquet (no SampleCollection, no row subsampling) so this covers every row\n",
    "# in the file regardless of size. \"energy\"/\"pdg\"/\"material\" stratify the\n",
    "# aggregate check so a failure hidden by the overall KL doesn't go unnoticed.\n",
    "for grouping in [None, \"energy\", \"pdg\", \"material\"]:\n",
    "    fig = plot_kl_bars_pl(FILE, group_by=grouping)\n",
    "    fig.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22c67dc4",
   "metadata": {},
   "source": [
    "## Detailed marginals, correlation & constraints (Tiers 1–3, 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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6f0903ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Tiers 1-3 detailed plots. Every one streams straight from the predict\n",
    "# parquet path (bounded memory, no in-memory SampleCollection) — pass FILE\n",
    "# directly, exactly like the lazy KL bars above and the event-level checks below.\n",
    "from giant.analysis import plot_marginals, plot_correlation_matrices, plot_pairwise\n",
    "from giant.analysis import plot_direction_alignment, plot_constraint_violations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "6d292c9c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 95200 characters]",
      "text/plain": [
       "<Figure size 3150x280 with 9 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_marginals(FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "79e066b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 440628 characters]",
      "text/plain": [
       "<Figure size 3150x1120 with 36 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_marginals(FILE, group_by=\"energy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d5a0bdb6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 579200 characters]",
      "text/plain": [
       "<Figure size 3150x1680 with 54 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_marginals(FILE, group_by=\"pdg\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "87ed1985",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 650388 characters]",
      "text/plain": [
       "<Figure size 3150x1680 with 54 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_marginals(FILE, group_by=\"material\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3585ead",
   "metadata": {},
   "source": [
    "## Tier 2: joint structure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8cdb9c2a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 93520 characters]",
      "text/plain": [
       "<Figure size 1300x400 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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(FILE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "e88c8bd9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 87228 characters]",
      "text/plain": [
       "<Figure size 1200x700 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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(FILE, n_sample=10000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "8c6e5d81",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 22728 characters]",
      "text/plain": [
       "<Figure size 500x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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(FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6aed1234",
   "metadata": {},
   "source": [
    "## Tier 3: physical constraints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "2fb6c511",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 84080 characters]",
      "text/plain": [
       "<Figure size 2000x350 with 5 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "# Unit-norm direction vectors, non-negative step_length/delta_e/edep — 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(FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0a0446d6",
   "metadata": {},
   "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",
    "\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",
    "\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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "d9a6c47b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from giant.analysis import compute_event_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\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",
    "obs = compute_event_observables_pl(FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "078d4844",
   "metadata": {},
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "0e5d3f74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 26436 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = 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` — 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": 14,
   "id": "d6561b73",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 23628 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_total_length(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab935fd",
   "metadata": {},
   "source": [
    "### Mean deposited energy per step, per event\n",
    "\n",
    "`total_edep / n_steps` grouped by `event_id` — the per-event mean, as opposed to the per-event total above, so a difference here points at the typical step's energy deposit rather than at how many steps the event happened to have."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e761bdc2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 52556 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_mean_energy_per_step(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6b6ae02",
   "metadata": {},
   "source": [
    "### Mean step length per step, per event\n",
    "\n",
    "`total_length / n_steps` grouped by `event_id` — the per-event mean step length, as opposed to the per-event total above."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "79b08c5a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50136 characters]",
      "text/plain": [
       "<Figure size 1200x400 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_mean_length_per_step(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf3ea615",
   "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 — the classic `E_dep(depth)` profile."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "a94c6034",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 49264 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_longitudinal_profile(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adbbd0ef",
   "metadata": {},
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "f3652b29",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 41104 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_transverse_profile(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de4713d9",
   "metadata": {},
   "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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "78bf19d5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 20052 characters]",
      "text/plain": [
       "<Figure size 600x400 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_shower_max_depth(obs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8dabe1c",
   "metadata": {},
   "source": [
    "## Particle-species (pdg) contribution shares\n",
    "\n",
    "Dataset-wide (not per-event) breakdown of which pdg species contributed how much of the total deposited energy / total length traveled, real vs. generated. Pure lazy `group_by(\"pdg\")` over the whole file — no post_pos reconstruction needed since these are scalar sums."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "d4ef0559",
   "metadata": {},
   "outputs": [],
   "source": [
    "from giant.analysis import pdg_contribution_table_pl\n",
    "from giant.analysis import plot_pdg_energy_share, plot_pdg_length_share\n",
    "\n",
    "pdg_table = pdg_contribution_table_pl(FILE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44aa3ad3",
   "metadata": {},
   "source": [
    "### Energy share by particle type\n",
    "\n",
    "Two pies side by side — real and generated — so the breakdown of which species deposits the energy can be compared directly. Small contributors are lumped into \"other\" (see `max_slices`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "06952b61",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54312 characters]",
      "text/plain": [
       "<Figure size 900x450 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_pdg_energy_share(pdg_table)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b1ee5f0",
   "metadata": {},
   "source": [
    "### Length-traveled share by particle type\n",
    "\n",
    "Same breakdown, but by total path length traveled (`sum(step_length)`) rather than deposited energy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "d46f1d96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 50804 characters]",
      "text/plain": [
       "<Figure size 900x450 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "_ = plot_pdg_length_share(pdg_table)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c61eeb68",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "giant",
   "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
}
