{
 "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 \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 `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** \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, and every check still streams (no `SampleCollection`, no full-file materialization) \u2014 see `giant.analysis`'s module docstring for the `RolloutVsTruth` mechanics.\n",
    "\n",
    "Same four tiers as `validation.ipynb`, all built on the same functions \u2014 pass a `RolloutVsTruth` in place of the predict-parquet path/LazyFrame everywhere:\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/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) \u2014 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 \u2014 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 \u2014 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) \u2014 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 \u2014 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) \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": "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` \u2014 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` \u2014 it needs the paired predict schema, which this rollout-vs-truth comparison doesn't have."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Router gating showcase (MoE)\n",
    "\n",
    "Every other section above is file-only \u2014 it reads `ROLLOUT_FILE` and never touches\n",
    "a checkpoint (see `giant.analysis`'s module docstring). This section is the one\n",
    "deliberate exception: soft gate weights only exist inside the trained `Router`,\n",
    "not in the rollout parquet, so this loads the checkpoint that produced\n",
    "`ROLLOUT_FILE` and calls `model.router.gate(...)` directly on that shower's\n",
    "pre-step conditioning.\n",
    "\n",
    "`model.router` is Stage 1's router; Stage 2 (`sec_decoder.router`) is a separate,\n",
    "independently trained `Router` instance over the same axis (see\n",
    "`giant.model.network.build_models`) and isn't shown here.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import polars as pl\n",
    "import torch\n",
    "\n",
    "from giant.analysis import plot_router_gating\n",
    "from giant.data.transforms import Normalizer, build_cond_features\n",
    "from giant.model.network import build_models\n",
    "\n",
    "# Checkpoint that produced ROLLOUT_FILE (needs `model.router` enabled at\n",
    "# train time, i.e. trained with `--router` / `model.router.enabled = true`).\n",
    "CHECKPOINT = \"/ceph/lbogner/geant_steps/checkpoints/REPLACE_ME/best.pt\"\n",
    "\n",
    "ckpt = torch.load(CHECKPOINT, map_location=\"cpu\", weights_only=False)\n",
    "model_cfg = ckpt[\"model_config\"]\n",
    "conditioning = model_cfg.get(\"conditioning\", \"embedding\")\n",
    "pdg_map = {int(k): v for k, v in ckpt[\"pdg_map\"].items()}\n",
    "mat_map = {str(k): v for k, v in ckpt[\"mat_map\"].items()}\n",
    "cond_norm = Normalizer.from_dict(ckpt[\"normalizer\"][\"cond\"])\n",
    "\n",
    "model, _sec_decoder = build_models(model_cfg)\n",
    "model.load_state_dict(ckpt[\"model\"])\n",
    "model.eval()\n",
    "\n",
    "if not hasattr(model, \"router\"):\n",
    "    raise RuntimeError(\n",
    "        f\"{CHECKPOINT} has no router \u2014 it was trained with model.router.enabled=False\"\n",
    "    )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pre-step conditioning for every row of the rollout shower, reconstructed\n",
    "# the same way `giant predict`/`giant rollout` do (giant.data.transforms).\n",
    "cols = [\"pdg\", \"pre_x\", \"pre_y\", \"pre_z\", \"pre_E\", \"pre_dx\", \"pre_dy\", \"pre_dz\",\n",
    "        \"material\", \"layer_id\"]\n",
    "df = pl.read_parquet(ROLLOUT_FILE, columns=cols)\n",
    "\n",
    "# Rows whose pdg/material fell outside the training vocab can't be encoded\n",
    "# (mirrors the pdg_mask filtering in `giant predict`'s CLI path).\n",
    "known = df[\"pdg\"].map_elements(lambda p: int(p) in pdg_map, return_dtype=pl.Boolean) & df[\n",
    "    \"material\"\n",
    "].map_elements(lambda m: str(m) in mat_map, return_dtype=pl.Boolean)\n",
    "n_dropped = (~known).sum()\n",
    "if n_dropped:\n",
    "    print(f\"dropping {n_dropped}/{len(df)} rows with unknown pdg/material\")\n",
    "df = df.filter(known)\n",
    "\n",
    "data = {\n",
    "    \"pre_pos\": df.select(\"pre_x\", \"pre_y\", \"pre_z\").to_numpy().astype(np.float32),\n",
    "    \"pre_E\": df[\"pre_E\"].to_numpy().astype(np.float32),\n",
    "    \"pre_dir\": df.select(\"pre_dx\", \"pre_dy\", \"pre_dz\").to_numpy().astype(np.float32),\n",
    "    \"layer_id\": df[\"layer_id\"].to_numpy(),\n",
    "    \"pdg\": df[\"pdg\"].to_numpy(),\n",
    "    \"material\": df[\"material\"].to_numpy(),\n",
    "}\n",
    "cond_cont, cond_cat = build_cond_features(\n",
    "    data, pdg_map, mat_map, cond_norm, conditioning=conditioning\n",
    ")\n",
    "cc = torch.from_numpy(cond_cont).float()\n",
    "ck = torch.from_numpy(cond_cat).long()\n",
    "\n",
    "with torch.no_grad():\n",
    "    gate_weights = model.router.gate(cc, ck).numpy()  # (N, n_experts), rows sum to 1\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# EnergyRouter gates on pre-step energy, so that's the natural x-axis here \u2014\n",
    "# swap for a categorical plot if this checkpoint used a different router type.\n",
    "_ = plot_router_gating(data[\"pre_E\"], gate_weights, x_label=\"pre_E\")\n"
   ]
  }
 ],
 "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
}