chore: release v0.3.22

Merge pull request 'feat(analyze): add paired truth/pred plots from giant predict' (#98 ) from feature/predict-analysis-plots into master
Reviewed-on: #98
2026-09-07 10:02:33 +00:00 · 2026-09-07 11:57:34 +02:00 · 2026-09-07 11:52:45 +02:00 · 2026-09-07 11:35:34 +02:00 · 2026-09-07 11:19:19 +02:00 · 2026-09-07 07:05:16 +00:00
78 changed files with 4801 additions and 1817 deletions
@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.16"
+current_version = "0.3.22"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"
@@ -0,0 +1,20 @@
+name: Setup uv + sync deps
+description: >-
+  Install uv, point its cache at the runner-local mount, and sync the
+  project with the cpu + dev extras. Every CI job does this identically;
+  the caller must still mount /srv/act-runner-cache/uv:/uv-cache on its
+  own container (a composite action can't set that) and check out the
+  repo before this runs.
+
+runs:
+  using: composite
+  steps:
+    - uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: false
+    - shell: bash
+      run: |
+        echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+        echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+    - shell: bash
+      run: uv sync --extra cpu --extra dev
@@ -0,0 +1,63 @@
+#!/usr/bin/env bash
+# Build one complete release commit: version bump (or sync to a given
+# version, or no bump at all), changelog entry, and refreshed README
+# badges — all in a single commit, tagged at the end. Used by both the
+# merge-to-master release job and the hand-pushed-tag sync job, so "the
+# tag == a complete release commit" holds either way a release gets made.
+#
+# Usage: release-commit.sh [--no-bump|VERSION]
+#   (no argument)  bump the current patch version (normal merge path).
+#   --no-bump      don't touch the version — the branch already bumped it
+#                   (e.g. a manual minor/major bump); just build the
+#                   changelog/badge commit for whatever version is current.
+#   VERSION        sync the project to this exact version (tag-sync path).
+#
+# Preconditions: repo is checked out with full history (fetch-depth: 0),
+# `uv` is available and synced, and git user.name/user.email are configured.
+# Idempotent: if the resulting tag already exists, this is a no-op.
+set -euo pipefail
+
+VERSION_ARG="${1:-}"
+
+if [ "$VERSION_ARG" = "--no-bump" ]; then
+  echo "Version already bumped by this branch; using current version as-is"
+elif [ -n "$VERSION_ARG" ]; then
+  echo "Syncing project version to $VERSION_ARG"
+  uv version "$VERSION_ARG" --no-sync
+  uv lock
+else
+  CURRENT_VERSION=$(uv version --short)
+  echo "Bumping patch version from $CURRENT_VERSION"
+  uv run bump-my-version bump patch --current-version "$CURRENT_VERSION" --no-commit --no-tag
+  # bump-my-version's pre_commit_hooks (uv lock + git add uv.lock) only run
+  # on its own commit path, which we skipped with --no-commit — so do it here.
+  uv lock
+fi
+
+VERSION=$(uv version --short)
+TAG="v$VERSION"
+
+if git rev-parse "$TAG" >/dev/null 2>&1; then
+  echo "Tag $TAG already exists; nothing to do"
+  exit 0
+fi
+
+uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
+
+TEST_COUNT=$(uv run pytest --collect-only -q 2>/dev/null | grep -oE '^[0-9]+ tests? collected' | grep -oE '^[0-9]+')
+sed -i -E "s|badge/version-[^-]+-informational|badge/version-${VERSION}-informational|" README.md
+sed -i -E "s|badge/tests-[0-9]+%20passing-brightgreen|badge/tests-${TEST_COUNT}%20passing-brightgreen|" README.md
+
+# .bumpversion.toml stores its own current_version, which bump-my-version
+# rewrites even with --no-commit — must be staged or the next run sees a
+# dirty tree and bump-my-version refuses (allow_dirty = false).
+git add pyproject.toml .bumpversion.toml uv.lock CHANGELOG.md README.md
+if git diff --cached --quiet; then
+  echo "Nothing changed; skipping release commit"
+  exit 0
+fi
+
+git commit -m "chore: release $TAG"
+git tag -a "$TAG" -m "$TAG"
+
+echo "Created release commit and tag $TAG"
@@ -12,6 +12,7 @@ env:
 jobs:
  ruff-check:
    name: Lint (ruff check)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -19,17 +20,12 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
      - run: uv run ruff check .

  ruff-format:
    name: Format (ruff format)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -37,17 +33,12 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
      - run: uv run ruff format --check .

  type-check:
    name: Type check (ty)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -55,18 +46,13 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
      - run: uv run ty check .

  test:
    name: Tests
-    needs: [ruff-check, type-check]
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
+    needs: [ruff-check, ruff-format, type-check]
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -74,21 +60,15 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
      - run: uv run pytest --cov --cov-report=term-missing --cov-report=xml
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
        with:
          name: coverage-report
          path: coverage.xml

-  bump-version:
-    name: Bump version, tag, and update changelog on merge to master
+  release:
+    name: Release (bump, changelog, badges, tag) on merge to master
    needs: [ruff-check, ruff-format, type-check, test]
    if: github.ref == 'refs/heads/master' && github.event_name == 'push'
    runs-on: ubuntu-latest
@@ -98,7 +78,7 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      # CI_TOKEN needs write:repository scope (not just read) — this job
-      # pushes commits and tags to master, unlike ruff-check/ruff-format/
+      # pushes a commit and a tag to master, unlike ruff-check/ruff-format/
      # type-check/test above, which only need to check out the repo.
      - uses: actions/checkout@v4
        with:
@@ -114,22 +94,14 @@ jobs:
          else
            echo "is_merge=false" >> "$GITHUB_OUTPUT"
          fi
-      - uses: astral-sh/setup-uv@v5
-        if: steps.merge_check.outputs.is_merge == 'true'
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-        if: steps.merge_check.outputs.is_merge == 'true'
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
        if: steps.merge_check.outputs.is_merge == 'true'
      - name: Configure git identity
        if: steps.merge_check.outputs.is_merge == 'true'
        run: |
          git config user.name "gitea-actions"
          git config user.email "actions@git.larsbogner.de"
-      - name: Bump patch version if this merge didn't already bump it
+      - name: Build the release commit
        if: steps.merge_check.outputs.is_merge == 'true'
        run: |
          OLD_VERSION=$(git show "${{ github.event.before }}:pyproject.toml" 2>/dev/null | grep -m1 '^version = ' | sed -E 's/version = "(.*)"/\1/')
@@ -140,70 +112,64 @@ jobs:
          fi
          if [ "$OLD_VERSION" = "$CURRENT_VERSION" ]; then
            echo "Version unchanged by this merge ($CURRENT_VERSION); bumping patch"
-            uv run bump-my-version bump patch --current-version "$CURRENT_VERSION"
+            .gitea/scripts/release-commit.sh
          else
-            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); skipping auto-bump"
+            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); building release commit without bumping"
+            .gitea/scripts/release-commit.sh --no-bump
          fi
-      - name: Update changelog for the current version if not already tagged
+      - name: Push the release commit and its tag together
        if: steps.merge_check.outputs.is_merge == 'true'
        run: |
          VERSION=$(uv version --short)
          TAG="v$VERSION"
-          if git rev-parse "$TAG" >/dev/null 2>&1; then
-            echo "Tag $TAG already exists; skipping changelog update"
+          if git rev-parse "$TAG" >/dev/null 2>&1 && [ "$(git rev-parse "$TAG^{commit}")" = "$(git rev-parse HEAD)" ]; then
+            git push --atomic origin HEAD:master "refs/tags/$TAG"
          else
-            uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
-            git add CHANGELOG.md
-            if ! git diff --cached --quiet -- CHANGELOG.md; then
-              git commit -m "chore: update changelog for $TAG"
-            else
-              git restore --staged CHANGELOG.md
-            fi
-          fi
-      - name: Push commits and tag the current version
-        if: steps.merge_check.outputs.is_merge == 'true'
-        run: |
-          git push origin HEAD:master
-          VERSION=$(uv version --short)
-          TAG="v$VERSION"
-          if git rev-parse "$TAG" >/dev/null 2>&1; then
-            echo "Tag $TAG already exists"
-          else
-            git tag -a "$TAG" -m "$TAG"
-            git push origin "refs/tags/$TAG"
+            echo "No new release commit/tag to push (already released, or nothing changed)"
+            git push origin HEAD:master
          fi

  sync-version-on-tag:
-    name: Sync project version with tag
+    name: Sync project version with tag (hand-pushed tags only)
    if: startsWith(github.ref, 'refs/tags/')
    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
        with:
          token: ${{ secrets.CI_TOKEN }}
-      - uses: astral-sh/setup-uv@v5
-      - name: Check tag against project version, update if they differ
+          fetch-depth: 0
+      - name: Require the tagged commit to already be on master
+        run: |
+          git fetch origin master
+          if ! git merge-base --is-ancestor "${{ github.sha }}" origin/master; then
+            echo "::error::Tag ${GITHUB_REF_NAME} points at a commit not on master; refusing to publish an unreviewed tree. Push the commit to master first, or delete and re-push the tag once it is."
+            exit 1
+          fi
+      - uses: ./.gitea/actions/setup
+      - name: Check tag against project version, build a release commit if they differ
        run: |
          TAG_VERSION="${GITHUB_REF_NAME#v}"
          CURRENT_VERSION=$(uv version --short)
-          if [ "$TAG_VERSION" != "$CURRENT_VERSION" ]; then
-            echo "Tag version ($TAG_VERSION) != project version ($CURRENT_VERSION); updating pyproject.toml"
-            uv version "$TAG_VERSION" --no-sync
+          if [ "$TAG_VERSION" = "$CURRENT_VERSION" ]; then
+            echo "Tag version matches project version ($CURRENT_VERSION)"
+          else
+            echo "Tag version ($TAG_VERSION) != project version ($CURRENT_VERSION); building a release commit"
            git config user.name "gitea-actions"
            git config user.email "actions@git.larsbogner.de"
-            git add pyproject.toml uv.lock
-            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME}"
+            .gitea/scripts/release-commit.sh "$TAG_VERSION"
            git push origin HEAD:master
            git push origin ":refs/tags/${GITHUB_REF_NAME}"
            git tag -f "${GITHUB_REF_NAME}" HEAD
            git push origin "refs/tags/${GITHUB_REF_NAME}"
-          else
-            echo "Tag version matches project version ($CURRENT_VERSION)"
          fi

  publish-package:
    name: Publish package to Gitea package registry
-    needs: [ruff-check, ruff-format, type-check, test, sync-version-on-tag]
+    needs: [sync-version-on-tag]
    if: startsWith(github.ref, 'refs/tags/')
    runs-on: ubuntu-latest
    container:
@@ -216,12 +182,7 @@ jobs:
      - uses: actions/checkout@v4
        with:
          ref: ${{ github.ref_name }}
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+      - uses: ./.gitea/actions/setup
      - run: uv build
      # CI_TOKEN needs write:package scope (in addition to write:repository,
      # used elsewhere) for this upload to authenticate.
@@ -1,5 +1,53 @@
 # Changelog

+## [0.3.22] - 2026-09-07
+
+### Changed
+
+- Feat(analyze): add paired truth/pred plots from giant predict
+
+- Fix(tests): make predict --truth flag test robust to terminal rendering
+
+- Feat(predict): enrich YAML sidecar with provenance and timing
+
+## [0.3.21] - 2026-09-07
+
+### Changed
+
+- Chore: bump uv.lock and fix ruff 0.16 default-rule lint findings
+
+- Chore: raise pyarrow ceiling to <26, bump to 25.0.1
+
+## [0.3.20] - 2026-09-04
+
+### Changed
+
+- Fix: dereference annotated tag to its commit before the push comparison
+
+## [0.3.19] - 2026-09-04
+
+### Changed
+
+- Chore: update README badges (version 0.3.18, 1138 tests)
+
+- Ci: restructure release pipeline into a single atomic release commit
+
+## [0.3.18] - 2026-09-03
+
+### Changed
+
+- Docs: rewrite README, keep version/test badges live via Gitea Actions
+
+- Perf: compact Stage-2 AR inference loop to active rows only
+
+## [0.3.17] - 2026-09-02
+
+### Changed
+
+- Feat: add WGAN + AR stop-token config variant
+
+- Perf: replace pandas with polars in the setup-stage scan
+
 ## [0.3.16] - 2026-08-31

 ### Changed
@@ -18,7 +18,7 @@ giant train path/to/steps.parquet --router --router-type energy  # MoE routing t
 giant model summary --config config.toml         # build-only: parameter counts + which config keys actually bite
 giant predict path/to/steps.parquet --checkpoint ckpt/best.pt   # per-step predictions
 giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl  # full showers
-giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs-reference analysis on HTCondor
+giant analyze submit rollout.yaml --prediction pred.yaml --accounting-group cms  # + paired truth/pred plots
 giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep/submit)
 giant analyze metrics <train_run_dir>                      # training-progress plots from metrics.csv
 dwarf --help                                     # dataset/tooling CLI: convert, migrate, bump-gen,
@@ -101,6 +101,8 @@ Secondary energies are a **stick-breaking partition of the `e_sec` budget** from

 **Input is one or more `giant rollout` YAML sidecars** (`condor.py:load_rollout_yamls`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML). `prep` creates a **run directory** (`<cwd>/analysis_runs/analysis_<id>/` by default, `--run-dir` to override) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit a.yaml [b.yaml ...] --chunks N` runs `prep` (recording `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` splits into a `compute_partial`/`finalize` pair so chunks can be summed/concatenated back per rollout (`chunkable=False` specs — the checkpoint-bound diagnostics, already bounded/subsampled — always run as a single chunk). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`; `merge-one` does a single plot for debugging), then turns those into the styled PDF/gallery tree. `giant analyze metrics <train_run_dir>` is a separate, unrelated entry point: training-progress plots straight from a run's `metrics.csv`.

+**`prediction` family (paired truth/pred, `giant/analysis/prediction.py`):** an optional add-on to the rollout comparison, driven by `--prediction`/`--prediction-label` on `analyze prep`/`submit` (repeatable, same convention as `--label`/rollout YAMLs; series name defaults to the YAML stem for N>1 or `"prediction"` for one). Unlike a rollout (freely generated, no row-level correspondence to truth), a `giant predict` output has a matching truth row for every prediction — a paired, not distributional, comparison. `giant predict --coord global` (schema v3, `--truth` on by default) writes both `pred_*` and `true_*` physical columns plus truth/predicted secondary lists; `--coord local` is the older, always-paired 9D model-space output (`pred_{name}`/`true_{name}` for `LOCAL_TARGET_NAMES`, no secondaries — stage 2 doesn't run there). `paired_frame()` normalizes either coord into one canonical `true_<var>`/`pred_<var>` frame over `PAIRED_VARS` (`step_length`, `edep`, `delta_e`, `post_E`, `cos_scatter`, `cos_travel`), decoding local coord's ALR energy logits the same way `energy_simplex_decode` does. Every prediction in one run must share one `--coord` and the rollouts' `dataset` (`condor.load_prediction_yamls`). The catalog's `prediction` family (`catalog.py`, ids prefixed `pred_`) covers per-variable marginals (new `paired_hist` kind: true dashed / pred solid) and truth-vs-pred 2D histograms (new `heatmap2d` kind, with a y=x guide), residuals/relative-residuals/residual-vs-truth profiles, KS/bias/RMSE scorecards (reusing `heatmap`), `n_sec` and secondary-species confusion matrices (row-normalised `heatmap`), direction-alignment and physical-constraint-violation checks, and a pred/true correlation-matrix delta. Every spec degrades to `kind="unavailable"` when no `--prediction` was given, so a rollout-only run is unaffected. `giant predict` also writes a YAML sidecar next to the checkpoint (`cli.py:_write_prediction_ref`, mirroring `giant rollout`'s) carrying the run's provenance and timing — coord/weights/steps/batch size, row/skip/unknown-PDG counts, a `timing` block, and the checkpoint's `model_config`/`training_epoch`/`training_config` — which `--prediction` consumes the same way `--label` rollout YAMLs are consumed, surfacing those keys into each plot's gallery `metadata.yaml` (`condor.py:_PLOT_META_KEYS`).
+
 **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower, advancing tracks breadth-first (every sweep steps all active tracks once, in `batch_size` chunks, so many tracks share each forward pass). Each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on one of the `TERM_*` reasons in `constants.py` (energy cutoff, max steps, escape, natural end, unknown pdg, max tracks); energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. `giant/checkpoint_io.py` is the shared checkpoint → ready-to-run-models path used by both `predict` and `rollout`.

 ## Roadmap
@@ -1,203 +1,295 @@
-# giant
+<div align="center">

-**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
+# GIANT

-A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.
+### **G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate
+
+A conditional generative model that replaces the Geant4 step function — sample a
+post-step outcome instead of simulating one, then roll that out into full
+calorimeter showers.
+
+[![python](https://img.shields.io/badge/python-3.12%2B-3776AB?logo=python&logoColor=white)](pyproject.toml)
+[![torch](https://img.shields.io/badge/torch-2.3.x-EE4C2C?logo=pytorch&logoColor=white)](pyproject.toml)
+[![version](https://img.shields.io/badge/version-0.3.22-informational)](CHANGELOG.md)
+[![tests](https://img.shields.io/badge/tests-1167%20passing-brightgreen)](tests/)
+[![CI](https://git.larsbogner.de/lars/giant/actions/workflows/ci.yml/badge.svg?branch=master)](https://git.larsbogner.de/lars/giant/actions)
+[![license](https://img.shields.io/badge/license-unlicensed-lightgrey)](#license)
+
+</div>
+
+---
+
+## The idea
+
+Geant4's step function is the innermost loop of detector simulation — for every
+particle, at every step, it stochastically samples where the particle goes next,
+how much energy it deposits, and what secondaries it spawns. GIANT learns that
+function instead of running it: given a pre-step particle state (position,
+energy, direction, particle species, material), a two-stage model samples a
+post-step outcome — including the variable-length list of secondaries — and
+autoregressively rolls that out into whole showers. It trains entirely from
+parquet dumps of a Geant4 steps tree; nothing downstream needs a Geant4 runtime.
+
+Two guarantees are architectural, not learned:
+
+- **Energy is conserved by construction.** Stage 1 decodes deposit / secondary /
+  post-step energy through a softmax simplex that sums to the pre-step energy
+  exactly; Stage 2's secondaries stick-break that same energy budget.
+- **No shower leaks across the train/val split.** Steps are split by `event_id`,
+  never by row, so correlated steps from the same shower can't appear on both
+  sides.
+
+## How a step becomes a shower
+
+```mermaid
+flowchart LR
+    A["pre-step state\nposition · energy · direction\nspecies · material"] --> B["ConditionEncoder\nphysical / embedding / onehot"]
+    B --> C["Stage 1\n9D post-step outcome"]
+    C --> D["Stage 2 (autoregressive)\nsecondaries, descending energy"]
+    D --> E["rollout step"]
+    E -->|"primary continues"| F["GeometryOracle\nposition to material, layer"]
+    E -->|"secondaries pushed"| G["track queue"]
+    F --> A
+    G --> A
+    E -->|"terminated"| H["deposited shower"]
+```

 ## Quick start

 ```bash
-uv sync --extra cpu               # install deps (CPU torch; use --extra cuda for GPU)
-
-giant new-run --hidden-dim 512 --lr 3e-4          # scaffold config.toml + run dir
-giant model summary --config config.toml          # parameter counts + which config keys actually bite
-giant train path/to/steps.parquet                 # train (flow + wgan by default)
-giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
+uv sync --extra cpu                                  # install deps (CPU torch; --extra cuda for GPU)

+giant new-run --hidden-dim 512 --lr 3e-4              # scaffold config.toml + run dir
+giant train path/to/steps.parquet                     # train (flow stage 1 + wgan stage 2, by default)
 dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # needed for rollout
 giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
+
+giant analyze prep rollout.yaml && giant analyze render <run_dir> --gallery   # rollout-vs-Geant4 diagnostics
 ```

-Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.
+Every command takes `--help` for its full flag list, and `--config config.toml`
+for anything not exposed as a flag.

-## Architecture
+---

-A **two-stage model**, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (`--stage1-generator`/`--stage2-generator`, or `--mode` to set both at once): `flow` (conditional flow matching, ODE-sampled in ~10 steps), `ddpm` (denoising diffusion), or `wgan` (single-pass WGAN-GP generator/critic).
+<details>
+<summary><h2 style="display:inline">Architecture</h2></summary>

-**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
+**Stage 1 — primary step.** Predicts the 9D post-step outcome
+(`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:

 | Index | Variable | Encoding |
 |-------|----------|----------|
 | 0 | `step_length` [mm] | log |
-| 1–2 | `edep_logit`, `sec_logit` | ALR coords of the deposit/secondary/post-energy simplex |
-| 3–5 | `post_dir` in local frame | unit vector |
-| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
+| 1–2 | `edep_logit`, `sec_logit` | ALR coordinates of the deposit / secondary / post-energy simplex |
+| 3–5 | `post_dir` | unit vector, local frame (`pre_dir = ẑ`) |
+| 6–8 | `travel_dir` (`post_pos − pre_pos`) | unit vector, local frame |

- Energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — conservation is architectural, not learned.
- `post_dir`/`travel_dir` live in the frame where `pre_dir = ẑ`. `post_pos` isn't a target — it's reconstructed as `pre_pos + step_length * world_frame(travel_dir)`.
+Energy logits decode via `softmax([edep_logit, sec_logit, 0]) × pre_E`, so
+`edep + e_sec + post_E == pre_E` holds exactly. `post_pos` is not itself a
+target — it's reconstructed as `pre_pos + step_length · world_frame(travel_dir)`,
+since duplicating that magnitude in a second target would let the two drift out
+of sync.

-**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (`--stage2-decoder`):
+**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's
+outcome, it generates the variable-length secondary list, one token at a time in
+descending-energy order (`autoregressive`, default) or all `K_MAX` slots in one
+masked pass (`one_shot`). Autoregressive tokens condition on a running history —
+`markov` (previous token only) or `attention` (causal self-attention, KV-cached
+at inference). Either way, secondary energies stick-break the `e_sec` budget
+handed down from Stage 1, so the whole chain conserves energy. A secondary's
+species is represented `onehot` (categorical, top-N PDG codes + "other"),
+`physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour
+lookup).

- `autoregressive` — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (`markov`: previous token only, or `attention`: causal self-attention, KV-cached at inference)
- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
+**Conditioning (15D).** Pre-step position / energy / direction / layer, plus
+particle mass/charge and material Z_eff/A_eff/density/X0/λ_int — encoded the
+same three ways as secondary species above, configured *independently* per axis
+(`conditioning.particle.type` / `conditioning.material.type`). `physical`
+computes rather than looks up, so it generalizes to species and materials
+outside the training menu; that's the default. `n_sec`/`e_sec` are always model
+outputs, never conditioning inputs.

-Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).
+**Composable by design** — every stage assembles from small registries, so
+swapping one axis doesn't touch the others:

-**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above. The particle and material axes are configured independently (`conditioning.particle.type` / `conditioning.material.type`; `--conditioning` sets both at once) and may mix — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
+| Registry | Choices |
+|---|---|
+| Objective | `flow` (matching, ~10-step ODE sample) · `ddpm` (denoising diffusion) · `wgan` (single-pass GAN) |
+| Trunk | `resmlp` · `none`, optionally MoE-routed (`RoutedTrunk`) |
+| Router | `energy` · `pdg` · `process` · `composed` · `none` — soft-mixed at train time, **top-1 dispatched at eval time**, which is the actual inference-speed win |
+| History (stage 2 AR) | `markov` · `attention` · `none` |

-**MoE routing** (`--router`): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk. The CLI flags configure Stage 1's router; Stage 2 has its own `stage2_model.router` block, config-file only.
+</details>

-## Data
+<details>
+<summary><h2 style="display:inline">Configuration</h2></summary>

- Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from ROOT via `dwarf convert`. One row = one Geant4 step.
- **Conditioning (pre-step) columns:** `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz` (direction), `material`, `layer_id`.
- **Primary outcome (post-step) columns:** `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep` (energy deposited in this step), `e_sec` (total energy carried off by secondaries), `child_track_ids` (its length gives `n_sec`).
- **Secondary columns**, one variable-length list per step: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX` (15) slots on load, ordered by descending energy.
- **Optional:** `process` — the physics process that produced the step (e.g. `compt`, `phot`, `eBrem`); a post-step label used only as classifier supervision (`ProcessRouter`), never as conditioning.
- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
- Loading a directory or `.manifest` of multiple parquet files (each one Geant4 job, `event_id` restarting from 0) offsets each file's `event_id`s by a fixed per-file stride so ids stay globally unique across files.
+Every default lives in one place: frozen dataclasses in `giant/config.py`,
+composed into `GiantConfig` (`conditioning` / `stage1_model` / `stage2_model` /
+`train`). `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather
+than hand-maintained, so the dataclasses can't drift from what actually gets
+merged. TOML config keys are validated against that shape — an unknown key is
+rejected with a did-you-mean suggestion. Precedence: CLI flag > `--config` file
+> default.

-## Project structure
+```toml
+# config.toml — resolved shape of the four blocks
+[conditioning]
+particle.type = "physical"
+material.type = "physical"

-```
-giant/
-├── giant/
-│   ├── data/
-│   │   ├── loader.py       # parquet → numpy arrays (incl. streaming/chunked reads)
-│   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
-│   │   ├── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
-│   │   └── setup_cache.py  # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
-│   ├── model/
-│   │   ├── models.py       # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
-│   │   ├── builders.py     # build_models / build_critics — config dict → assembled stage models
-│   │   ├── encoders.py     # ConditionEncoder (physical / embedding / onehot, per axis)
-│   │   ├── layers.py       # ResBlock/AdaLNResBlock registry, SinusoidalEmbedding, MLP heads
-│   │   ├── trunks.py       # trunk registry (resmlp, none) + RoutedTrunk (MoE expert bodies)
-│   │   ├── routers.py      # Router registry: energy / pdg / process / composed / none
-│   │   ├── history.py      # stage-2 AR history encoders: markov / attention (KV-cached) / none
-│   │   ├── objectives.py   # flow / ddpm / wgan objective registry
-│   │   ├── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
-│   │   ├── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
-│   │   ├── summary.py      # build-only introspection behind `giant model summary`
-│   │   ├── _legacy.py      # v0.2 checkpoint model_config/state-dict migration
-│   │   └── network.py      # re-export shim over all of the above
-│   ├── constants.py        # output/conditioning dims, K_MAX, secondary slot layout, schema keys
-│   ├── cond_layout.py      # single source of truth for the cond_cont/cond_cat column layout
-│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
-│   ├── materials.py        # material name → (Z_eff, A_eff, density, X0, λ_int)
-│   ├── config.py           # default hyperparameters, TOML config merging, device autodetect
-│   ├── pipeline.py         # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
-│   ├── training/           # two-stage training: loop, per-stage trainers, metrics, checkpointing
-│   │   ├── loop.py         #   epoch loop, graceful shutdown, best-checkpoint selection
-│   │   ├── trainers.py     #   StageSpec + flow/ddpm and WGAN-GP per-stage trainers
-│   │   ├── stage2_inputs.py#   ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
-│   │   ├── metrics.py      #   MetricsCollector: metrics.csv columns, W&B logging, progress/summary
-│   │   ├── amp.py          #   bf16 autocast (`train.precision`)
-│   │   ├── plots.py        #   training-progress plots (`giant analyze metrics`)
-│   │   └── checkpoint.py   #   checkpoint assembly/restore (format unchanged since v0.2)
-│   ├── sample.py           # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
-│   ├── checkpoint_io.py    # checkpoint → ready-to-run models/normalizers (predict + rollout)
-│   ├── geometry.py         # GeometryOracle: position → (material, layer_id, escaped) for rollout
-│   ├── rollout.py          # autoregressive shower rollout driver
-│   ├── validate.py         # step-level marginal + KL-divergence validation
-│   ├── _migration.py       # shared v0.2 → v0.3 facts used by both migration surfaces
-│   ├── analysis/           # rollout-vs-reference analysis pipeline (see `giant analyze` below)
-│   │   ├── sources.py      #   canonical LazyFrames + secondary view
-│   │   ├── variables.py    #   per-step value expressions shared by range sizing and the catalog
-│   │   ├── reduce.py       #   streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
-│   │   ├── grouping.py     #   fixed bin edges + energy/pdg/material group sets
-│   │   ├── context.py      #   resolves grouping into `shared.json` once per run
-│   │   ├── reduced.py      #   Partial/Reduced — the compact JSON a compute job emits
-│   │   ├── catalog.py      #   declarative PlotSpec registry (`giant analyze list`)
-│   │   ├── router_gating.py / type_embedding_distance.py   # checkpoint-bound diagnostics
-│   │   ├── runtime_estimate.py  #   per-(plot, chunk) walltime estimates for submit
-│   │   ├── condor.py       #   prep / compute-one / merge / submit-description plumbing
-│   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
-│   └── cli.py              # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze`
-├── giant/tools/                        # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
-│   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
-│   │                                    #   update-manifest, create-manifest, make-root,
-│   │                                    #   build-geometry-oracle, warm-cache, hparam-scan
-│   ├── steps_to_parquet.py             # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
-│   ├── steps_to_parquet_parallel.py    # fan out conversion over several ROOT files — `dwarf convert --jobs N`
-│   ├── migrate_geant_steps.py          # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
-│   ├── bump_dataset_version.py         # cut a new raw gen or parquet schema, with a logged reason —
-│   │                                    #   `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
-│   ├── create_root_files.py            # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
-│   ├── geometry_oracle.py              # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
-│   ├── warm_setup_cache.py             # precompute `giant train`'s setup-stage sidecar — `dwarf warm-cache`
-│   ├── hparam_scan.py                  # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
-│   └── profile_analysis_costs.py       # profiling helper for the `giant analyze` reduction pipeline
-└── tests/
+[stage1_model]
+generator = "flow"
+
+[stage2_model]
+generator = "wgan"
+decoder   = "autoregressive"
+
+[train]
+epochs     = 100
+batch_size = 4096
+lr         = 3e-4
 ```

-## Setup
+Some knobs only exist in the config file, with no CLI flag:
+`stage2_model.autoregressive.teacher_forcing`/`.history`,
+`stage2_model.particle_type.target`/`.class_weighting`,
+`stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`,
+`stage2_model.router.*`, and the finer `router` knobs (`lambda_balance`,
+`gumbel`, `learn_width`, …).
+
+`configs/` holds kept reference configs — `baseline.toml` is the fixed
+comparison point every experimental variant (routed trunk, WGAN, attention
+history, embedding conditioning) is a single edit away from. v0.2 flat-schema
+configs and checkpoints load and auto-migrate.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Data</h2></summary>
+
+Input is parquet — one row per Geant4 step — from
+[miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from
+ROOT via `dwarf convert`.
+
+| Group | Columns |
+|---|---|
+| **Conditioning (pre-step)** | `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz`, `material`, `layer_id` |
+| **Primary outcome (post-step)** | `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep`, `e_sec`, `child_track_ids` (length → `n_sec`) |
+| **Secondaries** (variable-length lists) | `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX = 15` slots, descending energy |
+| **Optional** | `process` — physics-process label, classifier supervision only (`ProcessRouter`), never conditioning |
+
+Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so a
+shower's correlated steps never straddle the split. Loading a directory or
+`.manifest` of several parquet files offsets each file's `event_id`s by a
+per-file stride so ids stay globally unique. The pre-epoch setup scan (vocab
+maps, event split, normalizer stats) persists to a sidecar cache
+(`--cache-setup`/`--rebuild-setup-cache`), precomputable ahead of time via
+`dwarf warm-cache`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">CLI reference</h2></summary>
+
+**`giant`** — train, run, and analyze the surrogate:
+
+| Command | Does |
+|---|---|
+| `new-run` | scaffold a `config.toml` + run directory from flags |
+| `train DATA` | train the two-stage model |
+| `model summary` | build-only parameter counts, without training |
+| `predict DATA --checkpoint …` | per-step predictions from a checkpoint |
+| `rollout DATA --checkpoint … --geometry …` | full autoregressive shower rollout |
+| `analyze prep/submit` | build a run dir; `submit` also queues HTCondor compute jobs |
+| `analyze compute-one` / `merge-one` | one plot × chunk reduction / merge (what a condor job runs) |
+| `analyze render <run_dir> --gallery` | merge chunks → styled PDFs + HTML gallery (local, needs LaTeX) |
+| `analyze metrics <train_run_dir>` | training-progress plots from `metrics.csv` |
+| `analyze list` | every catalog plot id |
+
+**`dwarf`** — dataset/tooling CLI:
+
+| Command | Does |
+|---|---|
+| `convert` | ROOT Steps tree → parquet (`--jobs N` fans out) |
+| `migrate` | one-time move into the raw/processed/pools/derived layout |
+| `bump-gen` / `bump-schema` / `status` | dataset versioning |
+| `update-manifest` / `create-manifest` | point/build a manifest of parquet files |
+| `make-root` | generate new ROOT shards via a minicalosim executable |
+| `build-geometry-oracle` | fit position → (material, layer_id) for rollout |
+| `warm-cache` | precompute `giant train`'s setup-stage sidecar |
+| `hparam-scan` | grid-scan dropout × n_blocks × hidden_dim |
+
+Worth knowing on `giant train` (full surface behind `--help`):
+`--mode {flow,ddpm,wgan}` / `--stage1-generator` / `--stage2-generator`,
+`--stage2-decoder {autoregressive,one_shot}`, `--conditioning
+{physical,embedding,onehot}`, `--router` / `--router-type` / `--n-experts` /
+`--router-axis`, `--stage{1,2}-init-from` + `--stage{1,2}-freeze` (retrain one
+stage against a fixed other one), `--precision {fp32,bf16}`, `--wandb`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Rollout &amp; analysis</h2></summary>
+
+`giant rollout` seeds showers from each event's highest-energy entry step, then
+autoregressively steps the model to completion — advancing all active tracks
+breadth-first, batched — pushing secondaries as new tracks and looking up
+`material`/`layer_id` from the geometry oracle each step. Tracks terminate on
+one of six reasons (energy cutoff, max steps, detector escape, natural end,
+unknown pdg, max tracks); every reason but escape deposits the remaining energy
+locally, so showers conserve energy by construction — only `escaped` counts as
+leakage.
+
+`giant analyze` compares one or more rollouts against a single held-out
+reference: `prep` resolves shared bin edges/groups once, `submit`/`compute-one`
+run each (plot, `event_id`-disjoint chunk) pair as a polars/numpy-only HTCondor
+job, `render` merges the chunks and produces the styled PDFs + HTML gallery
+locally (the only step that needs LaTeX). Each rollout gets its own colored
+series against one shared reference line. `giant analyze metrics` is a separate
+entry point — training-progress plots straight from a run's `metrics.csv`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Install</h2></summary>
+
+| Extra | Adds | For |
+|---|---|---|
+| `cpu` **or** `cuda` | torch 2.3.x | required — mutually exclusive, pick one |
+| `geometry` | scikit-learn | `dwarf build-geometry-oracle`, rollout |
+| `analysis` | matplotlib, plotstyle | `giant analyze render` |
+| `convert` | uproot, awkward | `dwarf convert` |
+| `wandb` | wandb | `giant train --wandb` |
+| `dev` | pytest, ruff, ty, + all of the above | development |

 ```bash
-uv sync --extra cpu                   # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
-uv sync --extra cpu --extra dev       # add dev tools (pytest, ruff, ty)
-uv sync --extra cpu --extra geometry  # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
-uv sync --extra cpu --extra analysis  # matplotlib/polars/plotstyle, for `giant analyze render`
-uv sync --extra cpu --extra convert   # uproot/awkward/polars, for `dwarf convert`
-uv sync --extra cpu --extra wandb     # W&B logging (`giant train --wandb`)
+uv sync --extra cpu --extra dev     # everything needed to develop
 ```

-The `dev` extra pulls in `convert`, `analysis`, `geometry` and `wandb` as well.
+Plain `uv sync` with no extra installs **no torch at all** — always include
+`--extra cpu` or `--extra cuda`.

-`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
+</details>

-## Training, prediction, rollout
+<details>
+<summary><h2 style="display:inline">Development</h2></summary>

 ```bash
-giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir
-giant train path/to/steps.parquet                          # train (flow stage 1 + wgan stage 2, default)
-giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
-
-dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # position → material/layer_id
-giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
-```
-
-Useful flags on `giant train`:
-
- `--mode {flow,ddpm,wgan}` sets both stages' objective at once; `--stage1-generator`/`--stage2-generator` override per stage
- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
- `--conditioning {physical,embedding,onehot}` — conditioning representation
- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
- `--stage2-stage1-context {truth,sampled}` — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via `stage2_model.ctx_p_start`/`ctx_p_end`)
- `--precision {fp32,bf16}` — bf16 autocast in the training loop
- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
- `--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
-
-Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`/`.class_weighting`, `stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`, `stage*_model.trunk.*` and the finer `router` knobs (`lambda_balance`, `gumbel`, `learn_width`, …). `configs/` holds kept reference configs. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
-
-`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.
-
-## Validation and analysis
-
- `giant.validate.validate_marginals` — step-level marginal + KL-divergence checks during training (`--validate-every`)
- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
-
-```bash
-giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot × chunk (compute only)
-giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan  # N rollouts vs one shared reference
-giant analyze render <run_dir> --gallery                   # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
-
-giant analyze list                                         # every catalog plot id
-giant analyze prep rollout.yaml --chunks 8                 # just the run directory, no submission
-giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0   # what a condor job runs
-giant analyze merge-one --id marginal_edep --run-dir <run_dir>               # merge one plot's chunks (debugging)
-```
-
-`<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
-
-Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.
-
-## Development
-
-```bash
-uv run pytest            # run tests
+uv run pytest            # 964 tests
 uv run ruff check .      # lint
 uv run ruff format .     # format
 uv run ty check .        # type check
 ```
+
+Gitea Actions (`.gitea/workflows/ci.yml`) runs lint + format-check + type-check
+ tests on every push and PR; merges to `master` auto-bump the patch version
+and regenerate `CHANGELOG.md` — don't hand-edit either.
+
+</details>
+
+## License
+
+Not yet decided — treat this repository as all-rights-reserved until a
+`LICENSE` file is added.
@@ -37,7 +37,7 @@ commit_preprocessors = [
 protect_breaking_commits = false
 commit_parsers = [
    { message = "^Merge ", skip = true },
-    { message = "^chore: (bump version|update changelog|sync project version)", skip = true },
+    { message = "^chore: (release|bump version|update changelog|sync project version)", skip = true },
    { message = "^Add", group = "<!-- 0 -->Added" },
    { message = "^(Fix|Clamp|Clip)", group = "<!-- 1 -->Fixed" },
    { message = "^(Remove|Drop|Deprecate)", group = "<!-- 2 -->Removed" },
@@ -84,15 +84,30 @@ dropout      = 0.0
 # secondary-species failure. Flow (not the schema default wgan) so the
 # baseline varies only the decoder relative to the best v0.2 result.
 #
-# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
-# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
-# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
-# k_max * steps = 150 stage-2 calls per physics step. That makes this block
-# the dominant cost on both sides:
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated — but see
+# the row-compaction note below, which changes the INFERENCE side of this:
 #   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
 #   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
 # Accepted deliberately: one-shot is the configuration whose secondary
 # species distribution failed, and that failure is what v0.3 exists to fix.
+#
+# Row compaction (landed after the above measurement): at inference,
+# sample.sample_secondaries_ar used to loop `for k in range(k_max)`
+# unconditionally — all 15 slots regardless of predicted n_sec — so a flow
+# AR token cost k_max * steps = 150 stage-2 calls per physics step. It now
+# drops a row from the batch the moment its own secondary count is
+# exhausted, so the real inference cost is ~n_sec * steps stage-2 calls
+# (this checkpoint's own rollout measured 0.382 secondaries/step — see
+# giant-baseline-flow-ar-rollout-validation.md), not k_max * steps. A CPU
+# micro-benchmark at that multiplicity (giant/model/history.py's
+# hidden_dim=512/6-block shape, k_max=15, batch 512) measured 17.6-22.9x
+# fewer wall-clock seconds for the AR loop alone (markov/attention history
+# respectively) — directional only (CPU, synthetic n_sec distribution, not
+# an end-to-end rollout); the 8.1x inference ratio above is now stale and
+# should be re-measured on GPU via a real rollout + `eval_cost_per_step`
+# once one is run against this checkpoint. Training cost (the 6.5x/29.5k
+# figures) is untouched by this: teacher_forcing = "always" here never
+# calls the AR sampler at train time (see [stage2_model.autoregressive]).
 decoder      = "autoregressive"
 generator    = "flow"
 hidden_dim   = 512
@@ -13,12 +13,15 @@ re-exported here is plotstyle-free so it runs on a compute worker. Import

 from giant.analysis.catalog import build_catalog, catalog_ids, get_spec
 from giant.analysis.condor import (
+    LoadedPrediction,
    LoadedRollout,
    RunMeta,
    SubmitConfig,
    compute_one,
    compute_reduced,
    derive_run_dir,
+    load_prediction_yaml,
+    load_prediction_yamls,
    load_rollout_yaml,
    load_rollout_yamls,
    merge_all,
@@ -27,32 +30,37 @@ from giant.analysis.condor import (
    write_submit,
 )
 from giant.analysis.context import Context, build_context
+from giant.analysis.prediction import PredictionSpec
 from giant.analysis.reduced import Partial, Reduced
 from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
 from giant.analysis.sources import RolloutSpec, Side

 __all__ = [
-    "build_catalog",
-    "catalog_ids",
-    "get_spec",
+    "RUNTIME_SAFETY_MARGIN",
+    "Context",
+    "LoadedPrediction",
    "LoadedRollout",
+    "Partial",
+    "PredictionSpec",
+    "Reduced",
+    "RolloutSpec",
    "RunMeta",
+    "Side",
    "SubmitConfig",
+    "build_catalog",
+    "build_context",
+    "catalog_ids",
    "compute_one",
    "compute_reduced",
    "derive_run_dir",
+    "estimate_runtime_s",
+    "get_spec",
+    "load_prediction_yaml",
+    "load_prediction_yamls",
    "load_rollout_yaml",
    "load_rollout_yamls",
    "merge_all",
    "merge_one",
    "prep",
    "write_submit",
-    "Context",
-    "build_context",
-    "Partial",
-    "Reduced",
-    "RolloutSpec",
-    "Side",
-    "RUNTIME_SAFETY_MARGIN",
-    "estimate_runtime_s",
 ]
@@ -56,11 +56,14 @@ import yaml

 from giant.analysis.catalog import Bundle, catalog_ids, get_spec
 from giant.analysis.context import Context, build_context
+from giant.analysis.prediction import PredictionSpec, open_prediction
 from giant.analysis.reduced import Partial
 from giant.analysis.runtime_estimate import estimate_runtime_s
 from giant.analysis.sources import RolloutSpec, Side, open_side

-# Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
+# Keys copied verbatim from a rollout or prediction YAML into each plot's
+# gallery metadata. Rollout-only and predict-only keys both live here —
+# `_plot_meta` copies only whichever of these are present in a given YAML.
 _PLOT_META_KEYS = (
    "prediction_id",
    "checkpoint",
@@ -84,10 +87,20 @@ _PLOT_META_KEYS = (
    "termination_reason_counts",
    "timing",
    "model_config",
+    "config_overrides",
    "training_epoch",
    "best_val_loss",
    "training_config",
    "training_meta",
+    # giant predict only (giant/cli.py's predict command).
+    "coord",
+    "has_truth",
+    "schema_version",
+    "n_input_rows",
+    "n_files",
+    "n_skipped_rows",
+    "unknown_pdg_counts",
+    "batch_size_auto",
    # Diagnostic — only present when giant rollout ran under
    # stage2_model.particle_type.target="embedding" (see giant/cli.py's
    # rollout command and giant.rollout.L1DistCollector); absent otherwise,
@@ -163,6 +176,72 @@ def load_rollout_yamls(
    return [LoadedRollout(name=n, yaml=y) for n, y in zip(names, yamls)], yamls[0]["dataset"]


+def load_prediction_yaml(path: str | Path) -> dict:
+    """Load a `giant predict` YAML sidecar, requiring the two file paths."""
+    d = yaml.safe_load(Path(path).read_text())
+    for key in ("output", "dataset"):
+        if key not in d:
+            raise ValueError(
+                f"{path} is not a prediction YAML (missing {key!r}); expected the "
+                "sidecar `giant predict` writes next to the checkpoint"
+            )
+    if d.get("kind") not in (None, "prediction"):
+        raise ValueError(f"{path} has kind={d.get('kind')!r}, not a prediction YAML")
+    return d
+
+
+@dataclass
+class LoadedPrediction:
+    """One prediction YAML plus its resolved series ``name`` and predict ``coord``."""
+
+    name: str
+    yaml: dict
+    coord: str
+
+
+def load_prediction_yamls(
+    paths: Sequence[str | Path], reference: str, labels: Sequence[str] | None = None
+) -> list[LoadedPrediction]:
+    """Load every prediction YAML, resolve each one's series name, and verify
+    they're seeded from the same ``reference`` as the rollout(s) and all share
+    one predict ``--coord`` (direction components mean different things in
+    the two coords — see ``giant.analysis.prediction``'s module docstring).
+
+    Names follow the same convention as ``load_rollout_yamls``: an explicit
+    ``labels[i]`` if given, else the YAML stem for N>1, or ``"prediction"``
+    for the single-YAML case.
+    """
+    if labels and len(labels) != len(paths):
+        raise ValueError(
+            f"--prediction-label given {len(labels)} time(s) but {len(paths)} --prediction YAML(s) were passed"
+        )
+    yamls = [load_prediction_yaml(p) for p in paths]
+    if labels:
+        names = list(labels)
+    elif len(paths) == 1:
+        names = ["prediction"]
+    else:
+        names = [Path(p).stem for p in paths]
+    if len(set(names)) != len(names):
+        dupes = sorted({n for n in names if names.count(n) > 1})
+        raise ValueError(f"prediction series names collide: {dupes} — pass --prediction-label to disambiguate")
+
+    bad_ref = [(p, y) for p, y in zip(paths, yamls) if str(y["dataset"]) != str(reference)]
+    if bad_ref:
+        detail = "\n".join(f"  {p}: dataset={y['dataset']!r}" for p, y in bad_ref)
+        raise ValueError(
+            f"every --prediction must be seeded from the same reference as the rollout(s) "
+            f"({reference!r}) — mismatched:\n{detail}"
+        )
+
+    coords = {str(p): open_prediction(y["output"]).coord for p, y in zip(paths, yamls)}
+    if len(set(coords.values())) > 1:
+        detail = "\n".join(f"  {p}: coord={c!r}" for p, c in coords.items())
+        raise ValueError(f"every --prediction in one run must share one --coord — got:\n{detail}")
+
+    return [LoadedPrediction(name=n, yaml=y, coord=coords[str(p)]) for n, y, p in zip(names, yamls, paths)]
+
+
 def _run_tag(y: dict) -> str:
    rollout = Path(y["output"])
    return str(y.get("prediction_id") or rollout.stem)[:8]
@@ -223,17 +302,26 @@ class RunMeta:
    # Empty/0 on run directories written before this field existed.
    rows_per_chunk: list[int] = field(default_factory=list)
    total_rows: int = 0
+    # `giant predict` inputs (the paired-truth "prediction" family) — same
+    # shape as `rollouts`. Empty on a run with no --prediction, so old
+    # run_meta.json files still load.
+    predictions: list[dict] = field(default_factory=list)

    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps(self.__dict__, indent=2))

    @classmethod
-    def load(cls, path: str | Path) -> "RunMeta":
+    def load(cls, path: str | Path) -> RunMeta:
        return cls(**json.loads(Path(path).read_text()))


-def _rows_per_chunk(rollouts: list[str | Path], reference: str | Path, n_chunks: int) -> list[int]:
-    """Combined rollout+reference row count of each ``event_id % n_chunks`` chunk.
+def _rows_per_chunk(
+    rollouts: list[str | Path],
+    reference: str | Path,
+    n_chunks: int,
+    predictions: Sequence[str | Path] = (),
+) -> list[int]:
+    """Combined rollout+reference+prediction row count of each ``event_id % n_chunks`` chunk.

    One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
    the sizing input every job's estimated walltime
@@ -249,7 +337,11 @@ def _rows_per_chunk(rollouts: list[str | Path], reference: str | Path, n_chunks:
        )

    out = [0] * n_chunks
-    sides = [open_side(reference, Side.reference)] + [open_side(r, Side.rollout) for r in rollouts]
+    sides = (
+        [open_side(reference, Side.reference)]
+        + [open_side(r, Side.rollout) for r in rollouts]
+        + [open_prediction(p).lf for p in predictions]
+    )
    for lf in sides:
        df = counts(lf)
        for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
@@ -263,26 +355,33 @@ def prep(
    n_chunks: int = 1,
    default_base: str | Path | None = None,
    labels: Sequence[str] | None = None,
+    prediction_yamls: Sequence[str | Path] = (),
+    prediction_labels: Sequence[str] | None = None,
    **ctx_kwargs,
 ) -> Path:
-    """Read the rollout YAML(s), build the shared context, and lay out the run dir.
+    """Read the rollout (+ optional prediction) YAML(s), build the shared
+    context, and lay out the run dir.

    Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
    ``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one``
    job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is
    resolved once here rather than re-passed (and risking disagreement) at every
    later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
-    resolve the actual directory, and ``load_rollout_yamls`` for how
-    ``labels``/YAML stems resolve each rollout's series name.
+    resolve the actual directory, ``load_rollout_yamls`` for how
+    ``labels``/YAML stems resolve each rollout's series name, and
+    ``load_prediction_yamls`` for the same on ``prediction_yamls`` (which,
+    unlike rollouts, is optional — the ``prediction`` plot family degrades to
+    ``kind="unavailable"`` when it's empty).

    Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
    this same ``run_dir``: partial files carry no record of what context
    (``n_chunks``, bin edges, group sets) they were computed under, so
    re-prepping with a different ``n_chunks``/``**ctx_kwargs`` (or after the
-    rollout/reference files changed) would otherwise let ``merge_one`` silently
-    merge stale partials against the new ``shared.json``.
+    rollout/reference/prediction files changed) would otherwise let
+    ``merge_one`` silently merge stale partials against the new ``shared.json``.
    """
    loaded, reference = load_rollout_yamls(list(rollout_yamls), labels)
+    loaded_preds = load_prediction_yamls(list(prediction_yamls), reference, prediction_labels)
    run_path = derive_run_dir([lr.yaml for lr in loaded], run_dir, default_base=default_base)
    run_path.mkdir(parents=True, exist_ok=True)

@@ -292,14 +391,20 @@ def prep(
            shutil.rmtree(stale_dir)

    rollout_specs = [RolloutSpec(name=lr.name, source=lr.yaml["output"]) for lr in loaded]
-    ctx = build_context(rollout_specs, reference, **ctx_kwargs)
+    pred_specs = [PredictionSpec(name=lp.name, source=lp.yaml["output"]) for lp in loaded_preds]
+    ctx = build_context(rollout_specs, reference, predictions=pred_specs, **ctx_kwargs)
    ctx.save(run_path / "shared.json")

-    rows_per_chunk = _rows_per_chunk([lr.yaml["output"] for lr in loaded], reference, n_chunks)
+    rows_per_chunk = _rows_per_chunk(
+        [lr.yaml["output"] for lr in loaded], reference, n_chunks, [lp.yaml["output"] for lp in loaded_preds]
+    )

    rollouts_meta = [
        {"name": lr.name, "path": str(lr.yaml["output"]), "plot_meta": _plot_meta(lr.yaml)} for lr in loaded
    ]
+    predictions_meta = [
+        {"name": lp.name, "path": str(lp.yaml["output"]), "plot_meta": _plot_meta(lp.yaml)} for lp in loaded_preds
+    ]
    ckpts = ", ".join(Path(lr.yaml.get("checkpoint", "")).name or "rollout" for lr in loaded)

    RunMeta(
@@ -310,6 +415,7 @@ def prep(
        n_chunks=n_chunks,
        rows_per_chunk=rows_per_chunk,
        total_rows=sum(rows_per_chunk),
+        predictions=predictions_meta,
    ).save(run_path / "run_meta.json")
    return run_path

@@ -327,12 +433,15 @@ def compute_reduced(
    out: str | Path,
    chunk_index: int = 0,
    n_chunks: int = 1,
+    predictions: Sequence[dict] = (),
 ) -> Path:
    """Core: run one (plot, chunk)'s partial reduction against explicit paths.

    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?,
    "timing"?}, ...]``, one per rollout series (insertion order preserved
-    through to every plot's ``Reduced.payload["series"]``).
+    through to every plot's ``Reduced.payload["series"]``). ``predictions``:
+    ``[{"name", "path"}, ...]``, one per ``giant predict`` series (the
+    ``prediction`` family; empty on a run with no ``--prediction``).

    Writes a ``Partial`` JSON — the raw, not-yet-merged output of
    ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one``
@@ -357,7 +466,8 @@ def compute_reduced(
        )
        for r in rollouts
    ]
-    bundle = Bundle.open(rollout_specs, reference, ctx, chunk=(chunk_index, effective_n))
+    pred_specs = [PredictionSpec(name=p["name"], source=p["path"]) for p in predictions]
+    bundle = Bundle.open(rollout_specs, reference, ctx, chunk=(chunk_index, effective_n), predictions=pred_specs)
    partial = Partial(
        id=spec_id,
        family=spec.family,
@@ -383,6 +493,7 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
        }
        for ro in meta.rollouts
    ]
+    predictions = [{"name": p["name"], "path": p["path"]} for p in meta.predictions]
    return compute_reduced(
        spec_id,
        rollouts,
@@ -391,6 +502,7 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
        run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json",
        chunk_index=chunk_index,
        n_chunks=meta.n_chunks,
+        predictions=predictions,
    )


@@ -20,6 +20,7 @@ import numpy as np
 import polars as pl

 from giant.analysis.grouping import energy_bin_edges
+from giant.analysis.prediction import PredictionSpec, open_prediction, paired_vars_for_coord, prediction_secondaries
 from giant.analysis.reduce import (
    attach_entry_axis,
    depth_expr,
@@ -44,16 +45,26 @@ class Context:
    sec_energy_range: tuple[float, float]
    n_sec_bins: int
    n_events: dict[str, int] = field(default_factory=dict)
+    # -- giant predict (paired truth/pred comparison) — empty when no
+    # --prediction was given to `prep`, so old shared.json files still load.
+    pred_var_ranges: dict[str, tuple[float, float]] = field(default_factory=dict)
+    pred_residual_ranges: dict[str, tuple[float, float]] = field(default_factory=dict)
+    pred_n_sec_cap: int = 10
+    pred_top_sec_pdgs: list[int] = field(default_factory=list)

    # -- (de)serialization -------------------------------------------------
    def save(self, path: str | Path) -> None:
        Path(path).write_text(json.dumps(asdict(self), indent=2))

    @classmethod
-    def load(cls, path: str | Path) -> "Context":
+    def load(cls, path: str | Path) -> Context:
        d = json.loads(Path(path).read_text())
        d["var_ranges"] = {k: tuple(v) for k, v in d["var_ranges"].items()}
        d["sec_energy_range"] = tuple(d["sec_energy_range"])
+        if "pred_var_ranges" in d:
+            d["pred_var_ranges"] = {k: tuple(v) for k, v in d["pred_var_ranges"].items()}
+        if "pred_residual_ranges" in d:
+            d["pred_residual_ranges"] = {k: tuple(v) for k, v in d["pred_residual_ranges"].items()}
        return cls(**d)

    # -- convenience -------------------------------------------------------
@@ -61,6 +72,14 @@ class Context:
        lo, hi = self.var_ranges[var]
        return np.linspace(lo, hi, self.n_marginal_bins + 1)

+    def pred_marginal_edges(self, var: str) -> np.ndarray:
+        lo, hi = self.pred_var_ranges[var]
+        return np.linspace(lo, hi, self.n_marginal_bins + 1)
+
+    def pred_residual_edges(self, var: str) -> np.ndarray:
+        lo, hi = self.pred_residual_ranges[var]
+        return np.linspace(lo, hi, self.n_marginal_bins + 1)
+

 _LO_Q, _HI_Q = 0.001, 0.999

@@ -87,10 +106,13 @@ def build_context(
    rollouts: list[RolloutSpec],
    reference: str | Path | pl.LazyFrame,
    *,
+    predictions: list[PredictionSpec] | None = None,
    n_energy_bins: int = 4,
    n_marginal_bins: int = 50,
    n_sec_bins: int = 40,
    top_k_pdg: int = 6,
+    pred_n_sec_cap: int = 10,
+    top_k_sec_pdg: int = 8,
    sample_rows: int = 1_000_000,
    seed: int = 0,
 ) -> Context:
@@ -165,6 +187,44 @@ def build_context(
    }
    sec_energy_range = _combined_quantiles([t_se, *r_se.values()], _LO_Q, _HI_Q)

+    # giant predict: paired truth/pred ranges + residual ranges + secondary
+    # species vocab, all over the union of every prediction's `paired` frame.
+    pred_var_ranges: dict[str, tuple[float, float]] = {}
+    pred_residual_ranges: dict[str, tuple[float, float]] = {}
+    top_sec_pdgs: list[int] = []
+    if predictions:
+        sides = {ps.name: open_prediction(ps.source) for ps in predictions}
+        present_vars = sorted(set().union(*(paired_vars_for_coord(s.coord) for s in sides.values())))
+        for var in present_vars:
+            true_samples, pred_samples, residual_samples = [], [], []
+            for s in sides.values():
+                if var not in paired_vars_for_coord(s.coord):
+                    continue
+                cols = [f"pred_{var}"] + ([f"true_{var}"] if s.has_truth else [])
+                sample = _row_subsample(s.paired.select(cols), sample_rows, seed).collect(engine="streaming")
+                pred_samples.append(sample[f"pred_{var}"].to_numpy())
+                if s.has_truth:
+                    true_samples.append(sample[f"true_{var}"].to_numpy())
+                    residual_samples.append(sample[f"pred_{var}"].to_numpy() - sample[f"true_{var}"].to_numpy())
+            pred_var_ranges[var] = _combined_quantiles([*true_samples, *pred_samples], _LO_Q, _HI_Q)
+            if residual_samples:
+                pred_residual_ranges[var] = _combined_quantiles(residual_samples, _LO_Q, _HI_Q)
+
+        sec_pdg_counts: dict[int, int] = {}
+        for s in sides.values():
+            if s.coord != "global" or not s.has_truth:
+                continue
+            for prefix in ("true", "pred"):
+                counts = (
+                    prediction_secondaries(s.lf, prefix)
+                    .group_by("pdg")
+                    .agg(pl.len().alias("n"))
+                    .collect(engine="streaming")
+                )
+                for pdg, n in zip(counts["pdg"].to_list(), counts["n"].to_list()):
+                    sec_pdg_counts[pdg] = sec_pdg_counts.get(pdg, 0) + n
+        top_sec_pdgs = [pdg for pdg, _ in sorted(sec_pdg_counts.items(), key=lambda kv: -kv[1])[:top_k_sec_pdg]]
+
    return Context(
        n_marginal_bins=n_marginal_bins,
        var_ranges=var_ranges,
@@ -173,6 +233,10 @@ def build_context(
        materials=materials,
        depth_edges=[float(x) for x in depth_edges],
        transverse_edges=[float(x) for x in transverse_edges],
+        pred_var_ranges=pred_var_ranges,
+        pred_residual_ranges=pred_residual_ranges,
+        pred_n_sec_cap=pred_n_sec_cap,
+        pred_top_sec_pdgs=top_sec_pdgs,
        sec_energy_range=sec_energy_range,
        n_sec_bins=n_sec_bins,
        n_events={
@@ -46,7 +46,7 @@ def pdg_label(code: int) -> str:

 def material_label(name: str) -> str:
    """Display label for a Geant4 material, dropping the ``G4_`` prefix."""
-    return name[3:] if name.startswith("G4_") else name
+    return name.removeprefix("G4_")


 def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
@@ -0,0 +1,315 @@
+"""Canonical paired truth/prediction LazyFrame for `giant predict` output.
+
+Unlike a `giant rollout` (an unpaired, freely-generated shower), `giant predict`
+runs the model once per real pre-step state, so every output row has a
+matching truth row — a paired comparison, not a distribution comparison. This
+module builds one canonical **paired** LazyFrame per prediction file, in
+either coord mode `giant predict` supports, so every catalog spec in the
+`prediction` family is coord-agnostic:
+
+    event_id, pdg, material, pre_E, n_sec, n_sec_pred,
+    true_<var>, pred_<var>   for var in PAIRED_VARS
+
+`--coord global` (v3+, `--truth` on) already carries physical `true_*`/`pred_*`-
+shaped columns directly. `--coord local` carries the raw 9D `true_{name}`/
+`pred_{name}` model-space target (`LOCAL_TARGET_NAMES`) instead — its two
+ALR energy logits are decoded into physical `edep`/`delta_e` with the same
+softmax-against-`pre_E` expressions `giant.data.transforms.energy_simplex_decode`
+uses, resurrected from the pre-package-rewrite `giant/analysis.py` (see
+`_edep_pl`/`_delta_e_pl`/`_raw_dim_expr` there). Direction components differ in
+*meaning* between the two coords (world vs. local frame), so a run must not mix
+them — `condor.load_prediction_yamls` enforces one coord across every
+prediction in a run.
+
+Secondaries only exist in `--coord global --truth` output (local mode never
+samples stage 2); `paired_secondaries` is `None` otherwise, and secondary-based
+specs render `kind="unavailable"` instead of raising.
+
+plotstyle-free (runs on HTCondor workers).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import polars as pl
+import pyarrow.parquet as pq
+
+from giant.constants import (
+    LOCAL_TARGET_NAMES,
+    PREDICT_COORD_METADATA_KEY,
+    PREDICT_TRUTH_METADATA_KEY,
+    ROLLOUT_COORD_VALUE,
+)
+
+# The paired scalar/direction variables every coord mode can produce, in
+# physical units (mm / MeV) regardless of source coord.
+PAIRED_SCALARS: tuple[str, ...] = ("step_length", "edep", "delta_e", "post_E")
+PAIRED_VARS: tuple[str, ...] = (*PAIRED_SCALARS, "cos_scatter", "cos_travel")
+
+_LOG_EPS = 1e-6
+
+
+@dataclass
+class PredictionSpec:
+    """One named prediction input, as fed to `build_context`/`Bundle.open`.
+
+    Mirrors `sources.RolloutSpec`: `name` is the series identity carried
+    through `payload["series"]` keys, legend labels, and color assignment.
+    """
+
+    name: str
+    source: str | Path | pl.LazyFrame
+    checkpoint: str | None = None
+
+
+@dataclass
+class PredictionSide:
+    """One prediction's opened frame + its coord/truth-availability facts."""
+
+    lf: pl.LazyFrame  # raw scan, chunk-filtered
+    paired: pl.LazyFrame  # canonical paired frame (see module docstring)
+    coord: str  # "global" | "local"
+    has_truth: bool
+    checkpoint: str | None = None
+
+
+def _check_predict_metadata(path: Path) -> tuple[str, bool]:
+    """Return `(coord, has_truth)`, raising if `path` isn't predict output.
+
+    Distinguishes a predict file from a rollout file (both are tagged with
+    `PREDICT_COORD_METADATA_KEY`, but a rollout's value is `ROLLOUT_COORD_VALUE`
+    rather than `"global"`/`"local"`).
+    """
+    metadata = pq.read_schema(path).metadata or {}
+    coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
+    if coord is None:
+        raise ValueError(
+            f"{path} has no {PREDICT_COORD_METADATA_KEY!r} parquet metadata — it wasn't "
+            "written by `giant predict` (or predates schema tagging)"
+        )
+    coord = coord.decode()
+    if coord == ROLLOUT_COORD_VALUE:
+        raise ValueError(f"{path} is a `giant rollout` file, not `giant predict` output")
+    if coord not in ("global", "local"):
+        raise ValueError(f"{path} has unrecognised predict coord {coord!r}")
+    # A v1 file predates truth tagging; only --coord local was paired then.
+    truth_raw = metadata.get(PREDICT_TRUTH_METADATA_KEY.encode())
+    has_truth = truth_raw.decode() == "1" if truth_raw is not None else coord == "local"
+    return coord, has_truth
+
+
+def _edep_pl(prefix: str) -> pl.Expr:
+    """Physical edep from `{prefix}_edep_logit`/`{prefix}_sec_logit` + `pre_E`.
+
+    Polars equivalent of `energy_simplex_decode(...)[0]` (the deposit
+    component): a softmax over `[z_edep, z_sec, 0]` times `pre_E`.
+    """
+    z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
+    m = pl.max_horizontal(z1, z2, pl.lit(0.0))
+    e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
+    return (e1 / (e1 + e2 + e3)) * pl.col("pre_E")
+
+
+def _delta_e_pl(prefix: str) -> pl.Expr:
+    """Physical delta_e (= edep + e_sec = pre_E - post_E) from the ALR logits + pre_E."""
+    z1, z2 = pl.col(f"{prefix}_edep_logit"), pl.col(f"{prefix}_sec_logit")
+    m = pl.max_horizontal(z1, z2, pl.lit(0.0))
+    e1, e2, e3 = (z1 - m).exp(), (z2 - m).exp(), (pl.lit(0.0) - m).exp()
+    return ((e1 + e2) / (e1 + e2 + e3)) * pl.col("pre_E")
+
+
+def _local_var_expr(prefix: str, var: str) -> pl.Expr:
+    """Physical value of one `PAIRED_VARS` entry from a `--coord local` file."""
+    if var == "step_length":
+        return pl.col(f"{prefix}_log_step_length").exp() - _LOG_EPS
+    if var == "edep":
+        return _edep_pl(prefix)
+    if var == "delta_e":
+        return _delta_e_pl(prefix)
+    if var == "post_E":
+        return pl.col("pre_E") - _delta_e_pl(prefix)
+    if var == "cos_scatter":
+        dot = pl.sum_horizontal([pl.col(f"{prefix}_post_d{ax}") * pl.col(f"{prefix}_travel_d{ax}") for ax in "xyz"])
+        return dot.clip(-1.0, 1.0)
+    raise ValueError(f"{var!r} has no direction-alignment meaning as a solo local-frame variable")
+
+
+def _g(prefix: str, name: str) -> str:
+    """Global-coord column name for `name` under `prefix`.
+
+    `giant predict --coord global` writes the *predicted* value under its bare
+    name (`step_length`, `edep`, `post_dx`, ...) and the truth under a
+    `true_` prefix (`true_step_length`, ...) — asymmetric, unlike the `local`
+    coord's symmetric `pred_*`/`true_*` naming.
+    """
+    return name if prefix == "pred" else f"true_{name}"
+
+
+def _global_var_expr(prefix: str, var: str) -> pl.Expr:
+    """Physical value of one `PAIRED_VARS` entry from a `--coord global` file."""
+    if var == "post_E":
+        # Not written directly for the prediction (it's implied by energy
+        # conservation: post_E = pre_E - delta_e); truth carries it as
+        # true_post_E.
+        return pl.col("pre_E") - pl.col(_g(prefix, "delta_e")) if prefix == "pred" else pl.col(_g(prefix, "post_E"))
+    if var == "cos_scatter":
+        dot = pl.sum_horizontal([pl.col(f"pre_d{ax}") * pl.col(_g(prefix, f"post_d{ax}")) for ax in "xyz"])
+        return dot.clip(-1.0, 1.0)
+    if var == "cos_travel":
+        # travel_dir isn't written by predict (only rollout reconstructs
+        # post_pos from it) — approximate with the post_pos - pre_pos
+        # direction instead, which is exactly what travel_dir encodes.
+        dx = pl.col(_g(prefix, "post_x")) - pl.col("pre_x")
+        dy = pl.col(_g(prefix, "post_y")) - pl.col("pre_y")
+        dz = pl.col(_g(prefix, "post_z")) - pl.col("pre_z")
+        norm = (dx**2 + dy**2 + dz**2).sqrt()
+        dot = (
+            pl.col("pre_dx") * dx / (norm + 1e-8)
+            + pl.col("pre_dy") * dy / (norm + 1e-8)
+            + pl.col("pre_dz") * dz / (norm + 1e-8)
+        )
+        return dot.clip(-1.0, 1.0)
+    return pl.col(_g(prefix, var))
+
+
+def _var_expr(coord: str, prefix: str, var: str) -> pl.Expr:
+    # `cos_travel` is excluded for `coord == "local"` by `paired_vars_for_coord`
+    # (predict never reconstructs post_pos/travel_dir there), so this only
+    # ever sees local-representable vars on that path.
+    if coord == "local":
+        return _local_var_expr(prefix, var)
+    return _global_var_expr(prefix, var)
+
+
+def paired_vars_for_coord(coord: str) -> tuple[str, ...]:
+    """The `PAIRED_VARS` a given coord mode can actually produce.
+
+    `cos_travel` needs a reconstructed `travel_dir`/`post_pos`, which
+    `--coord local` predict output never has (stage 2 doesn't run there) —
+    so local-coord predictions drop it rather than emit a meaningless value.
+    """
+    if coord == "local":
+        return PAIRED_SCALARS + ("cos_scatter",)
+    return PAIRED_VARS
+
+
+def dir_alignment_expr(coord: str, kind: str) -> pl.Expr:
+    """cos angle between the true and predicted direction vector (raw, not paired).
+
+    `kind="post"` compares `post_dir`; `kind="travel"` compares the
+    post_pos-implied travel direction. Reads the *raw* opened frame
+    (`PredictionSide.lf`), not `paired` — direction components aren't part of
+    `PAIRED_VARS` (only their two scattering cosines are), so this stays a
+    separate helper.
+    """
+    if coord == "local":
+        prefix_dim = "post_d" if kind == "post" else "travel_d"
+        true_v = [pl.col(f"true_{prefix_dim}{ax}") for ax in "xyz"]
+        pred_v = [pl.col(f"pred_{prefix_dim}{ax}") for ax in "xyz"]
+    elif kind == "post":
+        true_v = [pl.col(f"true_post_d{ax}") for ax in "xyz"]
+        pred_v = [pl.col(f"post_d{ax}") for ax in "xyz"]  # unprefixed: see paired_frame's _g
+    else:
+        true_v = [pl.col(f"true_post_{ax}") - pl.col(f"pre_{ax}") for ax in "xyz"]
+        pred_v = [pl.col(f"post_{ax}") - pl.col(f"pre_{ax}") for ax in "xyz"]
+    dot = pl.sum_horizontal([t * p for t, p in zip(true_v, pred_v)])
+    true_norm = pl.sum_horizontal([t**2 for t in true_v]).sqrt()
+    pred_norm = pl.sum_horizontal([p**2 for p in pred_v]).sqrt()
+    return (dot / (true_norm * pred_norm + 1e-8)).clip(-1.0, 1.0)
+
+
+def paired_frame(lf: pl.LazyFrame, coord: str, has_truth: bool) -> pl.LazyFrame:
+    """Canonical `event_id, pdg, material, pre_E, n_sec, n_sec_pred, true_*, pred_*` frame."""
+    schema = lf.collect_schema().names()
+    cols = [
+        "event_id",
+        "pdg",
+        "pre_E",
+        "material",
+        "n_sec",
+        pl.col("n_sec_pred") if "n_sec_pred" in schema else pl.lit(None, dtype=pl.Int64).alias("n_sec_pred"),
+    ]
+    for var in paired_vars_for_coord(coord):
+        cols.append(_var_expr(coord, "pred", var).alias(f"pred_{var}"))
+        if has_truth:
+            cols.append(_var_expr(coord, "true", var).alias(f"true_{var}"))
+    return lf.select(cols)
+
+
+def open_prediction(source: str | Path | pl.LazyFrame) -> PredictionSide:
+    """Lazily scan one prediction file, verifying its predict tag."""
+    if isinstance(source, pl.LazyFrame):
+        lf = source.with_columns(pl.col("pdg").cast(pl.Int64))
+        schema = lf.collect_schema().names()
+        coord = "local" if "pred_log_step_length" in schema else "global"
+        has_truth = f"true_{LOCAL_TARGET_NAMES[0]}" in schema or "true_step_length" in schema
+    else:
+        path = Path(source)
+        coord, has_truth = _check_predict_metadata(path)
+        lf = pl.scan_parquet(path).with_columns(pl.col("pdg").cast(pl.Int64))
+    return PredictionSide(
+        lf=lf,
+        paired=paired_frame(lf, coord, has_truth),
+        coord=coord,
+        has_truth=has_truth,
+    )
+
+
+# ---------------------------------------------------------------------------
+# Secondaries (global + truth only)
+# ---------------------------------------------------------------------------
+
+
+def prediction_secondaries(lf: pl.LazyFrame, prefix: str) -> pl.LazyFrame:
+    """One row per secondary from the true/predicted `sec_*_list` columns.
+
+    Canonical columns: `event_id, energy, pdg, sdx, sdy, sdz` — same shape as
+    `sources.secondaries`'s reference-side branch. `prefix` is `"true"` or
+    `"pred"`; matches `giant predict --coord global`'s asymmetric naming (see
+    `_g`) — the predicted lists are unprefixed (`sec_E_list`, ...), only the
+    truth ones carry `true_` (`true_sec_E_list`, ...).
+    """
+    col_prefix = "" if prefix == "pred" else "true_"
+    lists = [f"{col_prefix}sec_{c}_list" for c in ("E", "pdg", "dx", "dy", "dz")]
+    return (
+        lf.select("event_id", *lists)
+        .explode(lists)
+        .drop_nulls(lists[0])
+        .select(
+            "event_id",
+            pl.col(lists[0]).alias("energy"),
+            pl.col(lists[1]).cast(pl.Int64).alias("pdg"),
+            pl.col(lists[2]).alias("sdx"),
+            pl.col(lists[3]).alias("sdy"),
+            pl.col(lists[4]).alias("sdz"),
+        )
+    )
+
+
+def paired_secondaries(lf: pl.LazyFrame) -> pl.LazyFrame:
+    """True/predicted secondary PDG pairs, aligned by descending-energy rank.
+
+    Stage 2 emits secondaries in descending-energy order (`network.md`/
+    `giant/model/models.py`'s autoregressive decoder), so the natural
+    per-step alignment between the true and predicted secondary lists is
+    positional: rank `i` of one list vs. rank `i` of the other, for
+    `i < min(n_sec, n_sec_pred)`. Requires `--coord global --truth`.
+    """
+    return (
+        lf.select("true_sec_pdg_list", "sec_pdg_list")
+        .with_row_index("_row")
+        .with_columns(
+            pl.col("true_sec_pdg_list").list.len().alias("_n_true"),
+            pl.col("sec_pdg_list").list.len().alias("_n_pred"),
+        )
+        .with_columns(pl.min_horizontal("_n_true", "_n_pred").alias("_n_paired"))
+        .filter(pl.col("_n_paired") > 0)
+        .with_columns(pl.int_ranges(0, pl.col("_n_paired")).alias("_rank"))
+        .explode("_rank")
+        .select(
+            pl.col("true_sec_pdg_list").list.get(pl.col("_rank")).cast(pl.Int64).alias("true_pdg"),
+            pl.col("sec_pdg_list").list.get(pl.col("_rank")).cast(pl.Int64).alias("pred_pdg"),
+        )
+    )
@@ -70,6 +70,64 @@ def hist1d(
    return out


+def hist2d(
+    lf: pl.LazyFrame,
+    x: pl.Expr,
+    y: pl.Expr,
+    x_edges: np.ndarray,
+    y_edges: np.ndarray,
+) -> np.ndarray:
+    """Streaming 2D histogram of `(x, y)` over fixed uniform edges.
+
+    One `group_by([_bx, _by]).len()` pass; returns the full `(len(x_edges)-1,
+    len(y_edges)-1)` int64 count matrix (row = x bin, col = y bin) — small
+    enough (a truth-vs-pred scatter has at most a few thousand cells) to
+    materialize whole, unlike `hist1d`'s per-group dict.
+    """
+    x_lo, x_hi, x_n = float(x_edges[0]), float(x_edges[-1]), len(x_edges) - 1
+    y_lo, y_hi, y_n = float(y_edges[0]), float(y_edges[-1]), len(y_edges) - 1
+    res = (
+        lf.select(_bin_expr(x, x_lo, x_hi, x_n).alias("_bx"), _bin_expr(y, y_lo, y_hi, y_n).alias("_by"))
+        .drop_nulls(["_bx", "_by"])
+        .group_by("_bx", "_by")
+        .agg(pl.len().alias("_n"))
+        .collect(engine="streaming")
+    )
+    mat = np.zeros((x_n, y_n), dtype=np.int64)
+    mat[res["_bx"].to_numpy(), res["_by"].to_numpy()] = res["_n"].to_numpy()
+    return mat
+
+
+def binned_moments(
+    lf: pl.LazyFrame,
+    bin_value: pl.Expr,
+    agg_value: pl.Expr,
+    edges: np.ndarray,
+) -> dict[str, list]:
+    """Per-bin ``(n, sum, sumsq)`` of ``agg_value``, binned by ``bin_value`` over fixed edges.
+
+    One streaming `group_by` pass; sum-mergeable across chunks the same way
+    `hist1d` counts are — elementwise-summing `n`/`sum`/`sumsq` per bin across
+    chunks reconstructs the moments of the full merged data, from which
+    `finalize` derives mean/std (``mean = sum/n``,
+    ``std = sqrt(sumsq/n - mean**2)``).
+    """
+    lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1
+    res = (
+        lf.select(_bin_expr(bin_value, lo, hi, nbins).alias("_b"), agg_value.alias("_v"))
+        .drop_nulls(["_b", "_v"])
+        .group_by("_b")
+        .agg(pl.len().alias("_n"), pl.col("_v").sum().alias("_s"), (pl.col("_v") ** 2).sum().alias("_ss"))
+        .collect(engine="streaming")
+    )
+    n = np.zeros(nbins, dtype=np.int64)
+    s = np.zeros(nbins, dtype=np.float64)
+    ss = np.zeros(nbins, dtype=np.float64)
+    for b_, nn, ssum, sqsum in res.iter_rows():
+        n[b_], s[b_], ss[b_] = nn, ssum, sqsum
+    return {"n": n.tolist(), "sum": s.tolist(), "sumsq": ss.tolist()}
+
+
 def sum_merge(dicts: list[dict[str, Any]]) -> dict[str, Any]:
    """Elementwise-sum a list of sum-mergeable count/total dicts (JSON-safe keys).

@@ -27,8 +27,13 @@ from pathlib import Path
 #   "router_specialization"  max gate weight vs energy (one scalar trend line
 #                    summarizing "router_gating"), per rollout with an enabled router
 #   "heatmap"        row x col matrix + colorbar, one panel per rollout (a
-#                    distance scorecard)
-#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint)
+#                    distance scorecard) or per prediction (a confusion matrix)
+#   "paired_hist"    per-prediction true/pred density histogram over shared
+#                    edges (giant predict's paired truth, not a rollout)
+#   "heatmap2d"      numeric x/y-binned true-vs-pred count matrix + colorbar,
+#                    one panel per prediction, with a y=x diagonal guide
+#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint,
+#                    or no --prediction given)


@dataclass
@@ -46,7 +51,7 @@ class Reduced:
        Path(path).write_text(json.dumps(asdict(self)))

    @classmethod
-    def load(cls, path: str | Path) -> "Reduced":
+    def load(cls, path: str | Path) -> Reduced:
        return cls(**json.loads(Path(path).read_text()))


@@ -70,5 +75,5 @@ class Partial:
        Path(path).write_text(json.dumps(asdict(self)))

    @classmethod
-    def load(cls, path: str | Path) -> "Partial":
+    def load(cls, path: str | Path) -> Partial:
        return cls(**json.loads(Path(path).read_text()))
@@ -27,8 +27,8 @@ from pathlib import Path

 import numpy as np
 import plotstyle as ps
-from matplotlib.colors import LogNorm
 import yaml
+from matplotlib.colors import LogNorm

 from giant.analysis.reduced import Reduced

@@ -422,6 +422,78 @@ def _render_heatmap(r: Reduced, params: dict):
    return fig


+def _render_paired_hist(r: Reduced, params: dict):
+    """`giant predict`'s paired truth/pred density histogram (see
+    `giant.analysis.prediction`) — unlike `_render_overlay`, there's no single
+    shared reference: each prediction carries its own truth. A lone prediction
+    draws its truth in the reference ink so a single-series run reads exactly
+    like an `overlay_hist` figure; two-or-more predictions each get their own
+    color, pred solid / true dashed, so a same-colored pair is directly
+    comparable.
+    """
+    edges = np.asarray(r.payload["edges"])
+    series = r.payload.get("series", {})
+    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
+    solo = len(series) == 1
+    for i, (name, entry) in enumerate(series.items()):
+        color = _ref_color() if solo else ps.get_color(i)
+        if "true" in entry:
+            true_label = _REFERENCE_LABEL if solo else f"{name} (true)"
+            ax.stairs(_density(entry["true"], edges), edges, label=true_label, color=color, linestyle="--")
+        pred_color = ps.get_color(i)
+        pred_label = name if solo else f"{name} (pred)"
+        ax.stairs(_density(entry["pred"], edges), edges, label=pred_label, color=pred_color)
+    if r.payload.get("log_y"):
+        ax.set_yscale("log")
+    ax.set_xlabel(r.xlabel)
+    ax.set_ylabel("density")
+    ps.style_legend(ax, title="source")
+    return fig
+
+
+def _render_heatmap2d(r: Reduced, params: dict):
+    """Numeric truth-vs-pred 2D histogram, one panel per prediction, with an
+    optional y=x guide line — the direct analogue of `_render_heatmap` for
+    continuous (not categorical) axes."""
+    x_edges = np.asarray(r.payload["x_edges"])
+    y_edges = np.asarray(r.payload["y_edges"])
+    series = r.payload["series"]
+    names = list(series)
+    norm = LogNorm(vmin=1) if r.payload.get("log_color") else None
+    fig, axes = ps.new_figure(
+        "slide-16x9" if len(names) > 1 else "thesis-single",
+        title=r.title,
+        params=params,
+        nrows=1,
+        ncols=len(names),
+        squeeze=False,
+    )
+    flat = axes.ravel()
+    im = None
+    for ax, name in zip(flat, names):
+        mat = np.asarray(series[name], dtype=float)
+        im = ax.pcolormesh(
+            x_edges,
+            y_edges,
+            mat.T,
+            cmap=r.payload.get("cmap", "viridis"),
+            norm=norm,
+            vmin=None if norm else r.payload.get("vmin"),
+            vmax=None if norm else r.payload.get("vmax"),
+        )
+        if r.payload.get("diagonal"):
+            lo, hi = max(x_edges[0], y_edges[0]), min(x_edges[-1], y_edges[-1])
+            ax.plot([lo, hi], [lo, hi], color=_ref_color(), linestyle="--", linewidth=1, label="y = x")
+        ax.set_xlabel(r.xlabel)
+        if len(names) > 1:
+            ax.set_title(name, fontsize=8)
+    flat[0].set_ylabel(r.payload.get("ylabel", ""))
+    if r.payload.get("diagonal"):
+        ps.style_legend(flat[0], title="guide")
+    fig.colorbar(im, ax=list(flat), label=r.payload.get("cbar_label", "count"))
+    return fig
+
+
 def _render_unavailable(r: Reduced, params: dict):
    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
    ax.axis("off")
@@ -448,6 +520,8 @@ _RENDERERS = {
    "router_share": _render_router_share,
    "router_specialization": _render_router_specialization,
    "heatmap": _render_heatmap,
+    "paired_hist": _render_paired_hist,
+    "heatmap2d": _render_heatmap2d,
    "unavailable": _render_unavailable,
 }

@@ -58,10 +58,10 @@ _COLS = (

@dataclass
 class _RouterHandle:
-    router: "torch.nn.Module"
+    router: torch.nn.Module
    pdg_map: dict[int, int]
    mat_map: dict[str, int]
-    cond_normalizer: "Normalizer"
+    cond_normalizer: Normalizer
    particle_conditioning: str
    material_conditioning: str
    router_type: str
@@ -233,7 +233,7 @@ def _gating_entry(checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: p
    return {"router_type": handle.router_type, "n_experts": handle.router.n_experts, **sides}


-def compute_router_gating(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+def compute_router_gating(rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
    """`Reduced` for the router-gating figure: one panel-pair per rollout with
    an enabled MoE router, or an explanatory note if none of them have one."""
    series = {}
@@ -289,7 +289,7 @@ def _specialization_entry(
    }


-def compute_router_specialization(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+def compute_router_specialization(rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
    """`Reduced` for the router-specialization figure, one curve per rollout with
    an enabled MoE router (see `_specialization_entry`)."""
    series = {}
@@ -331,7 +331,7 @@ def _share_by_pdg_entry(


 def compute_router_share_by_pdg(
-    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int = 0
+    rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int = 0
 ) -> Reduced:
    """`Reduced` for the router expert-share-by-species figure, one panel-pair
    per rollout with an enabled MoE router."""
@@ -375,7 +375,7 @@ def _share_by_process_entry(checkpoint: str | Path | None, t_phys: pl.LazyFrame,


 def compute_router_share_by_process(
-    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0, top_k: int = _TOP_K_PROCESS
+    rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, seed: int = 0, top_k: int = _TOP_K_PROCESS
 ) -> Reduced:
    """Stacked-bar share of each physics process dispatched to each expert, one
    panel per rollout checkpoint with an enabled MoE router.
@@ -35,7 +35,7 @@ _NOTE_NOT_APPLICABLE = (
 )


-def compute_type_embedding_l1_distance(rollouts: dict[str, "RolloutSide"]) -> Reduced:
+def compute_type_embedding_l1_distance(rollouts: dict[str, RolloutSide]) -> Reduced:
    """`Reduced` for the type-embedding-distance figure: one series per rollout
    whose checkpoint populated the diagnostic, or an explanatory note if none did.

@@ -9,7 +9,7 @@ import subprocess
 import sys
 import tomllib
 from dataclasses import dataclass, field
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -80,7 +80,7 @@ class ConditioningAxisConfig:
    n_layers: int = 1

    @classmethod
-    def from_dict(cls, d: dict | None) -> "ConditioningAxisConfig":
+    def from_dict(cls, d: dict | None) -> ConditioningAxisConfig:
        d = d or {}
        return cls(
            type=d.get("type", "physical"),
@@ -106,7 +106,7 @@ class ConditioningConfig:
    material: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "ConditioningConfig":
+    def from_dict(cls, d: dict | None) -> ConditioningConfig:
        d = d or {}
        return cls(
            out_dim=d.get("out_dim", 128),
@@ -130,7 +130,7 @@ class FlowConfig:
    time_dim: int = 64

    @classmethod
-    def from_dict(cls, d: dict | None) -> "FlowConfig":
+    def from_dict(cls, d: dict | None) -> FlowConfig:
        d = d or {}
        return cls(time_dim=d.get("time_dim", 64))

@@ -144,7 +144,7 @@ class DdpmConfig:
    n_steps: int = 1000

    @classmethod
-    def from_dict(cls, d: dict | None) -> "DdpmConfig":
+    def from_dict(cls, d: dict | None) -> DdpmConfig:
        d = d or {}
        return cls(time_dim=d.get("time_dim", 64), n_steps=d.get("n_steps", 1000))

@@ -166,7 +166,7 @@ class Stage1WganConfig:
    critic_n_res_blocks: int = 0

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage1WganConfig":
+    def from_dict(cls, d: dict | None) -> Stage1WganConfig:
        d = d or {}
        return cls(
            noise_dim=d.get("noise_dim", 64),
@@ -198,7 +198,7 @@ class Stage2WganConfig(Stage1WganConfig):
    gumbel_tau_end: float = 0.1

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2WganConfig":
+    def from_dict(cls, d: dict | None) -> Stage2WganConfig:
        d = d or {}
        return cls(
            noise_dim=d.get("noise_dim", 64),
@@ -293,7 +293,7 @@ class RouterConfig:
    extra: dict = field(default_factory=dict)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "RouterConfig":
+    def from_dict(cls, d: dict | None) -> RouterConfig:
        d = d or {}
        return cls(
            enabled=d.get("enabled", False),
@@ -360,7 +360,7 @@ class TrunkConfig:
    block_conditioning: str = "add"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "TrunkConfig":
+    def from_dict(cls, d: dict | None) -> TrunkConfig:
        d = d or {}
        return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))

@@ -377,7 +377,7 @@ class Stage2RouterConfig(RouterConfig):
    tie_to_stage1: bool = False

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2RouterConfig":
+    def from_dict(cls, d: dict | None) -> Stage2RouterConfig:
        d = d or {}
        known = _ROUTER_KNOWN_KEYS | {"tie_to_stage1"}
        return cls(
@@ -447,7 +447,7 @@ class NSecConfig:
    sampling: str = "greedy"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "NSecConfig":
+    def from_dict(cls, d: dict | None) -> NSecConfig:
        d = d or {}
        return cls(
            mode=d.get("mode", "head"),
@@ -498,7 +498,7 @@ class ParticleTypeConfig:
    class_weighting: str = "none"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
+    def from_dict(cls, d: dict | None) -> ParticleTypeConfig:
        d = d or {}
        return cls(
            target=d.get("target", "onehot"),
@@ -537,7 +537,7 @@ class AutoregressiveConfig:
    attn_n_layers: int = 2

    @classmethod
-    def from_dict(cls, d: dict | None) -> "AutoregressiveConfig":
+    def from_dict(cls, d: dict | None) -> AutoregressiveConfig:
        d = d or {}
        return cls(
            order=d.get("order", "energy_desc"),
@@ -576,7 +576,7 @@ class HeadConfig:
    depth: int = 2  # matches build_mlp_head's depth

    @classmethod
-    def from_dict(cls, d: dict | None) -> "HeadConfig":
+    def from_dict(cls, d: dict | None) -> HeadConfig:
        d = d or {}
        return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2))

@@ -593,7 +593,7 @@ class Stage1HeadsConfig:
    n_sec: HeadConfig = field(default_factory=HeadConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig":
+    def from_dict(cls, d: dict | None) -> Stage1HeadsConfig:
        d = d or {}
        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")))

@@ -611,7 +611,7 @@ class Stage2HeadsConfig:
    type: HeadConfig = field(default_factory=HeadConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig":
+    def from_dict(cls, d: dict | None) -> Stage2HeadsConfig:
        d = d or {}
        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type")))

@@ -659,7 +659,7 @@ class Stage1ModelConfig:
    heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage1ModelConfig":
+    def from_dict(cls, d: dict | None) -> Stage1ModelConfig:
        d = d or {}
        return cls(
            active=d.get("active", True),
@@ -745,7 +745,7 @@ class Stage2ModelConfig:
    heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2ModelConfig":
+    def from_dict(cls, d: dict | None) -> Stage2ModelConfig:
        d = d or {}
        return cls(
            active=d.get("active", True),
@@ -839,7 +839,7 @@ class TrainConfig:
    precision: str = "fp32"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "TrainConfig":
+    def from_dict(cls, d: dict | None) -> TrainConfig:
        d = d or {}
        return cls(
            epochs=d.get("epochs", 100),
@@ -895,7 +895,7 @@ class GiantConfig:
    train: TrainConfig = field(default_factory=TrainConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "GiantConfig":
+    def from_dict(cls, d: dict | None) -> GiantConfig:
        d = d or {}
        return cls(
            conditioning=ConditioningConfig.from_dict(d.get("conditioning")),
@@ -938,7 +938,7 @@ def leaf_paths(node: dict, prefix: str = "") -> list[str]:
 def git_hash() -> str:
    try:
        return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
-    except Exception:
+    except Exception:  # noqa: BLE001 - any failure (no git, no repo, ...) degrades to "unknown"
        return "unknown"


@@ -1852,7 +1852,7 @@ def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
    name unboundedly. This name doubles as the run's W&B id (see
    giant.training), which is the reason a timestamp is always included.
    """
-    now = now or datetime.now()
+    now = now or datetime.now()  # noqa: DTZ005 - human-readable local wall-clock time for run/W&B naming, not stored
    tokens = []
    overflow = []
    for label, candidate in _OUT_DIR_NAME_CANDIDATES:
@@ -1951,7 +1951,7 @@ def build_run_meta(
        "config_version": CONFIG_VERSION,
        "git_hash": git_hash(),
        "seed": seed,
-        "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+        "timestamp_utc": datetime.now(UTC).isoformat(timespec="seconds"),
        "python_version": sys.version.split()[0],
        "torch_version": torch.__version__,
        "command": " ".join(sys.argv),
@@ -64,7 +64,12 @@ LOCAL_TARGET_NAMES = [
 # guessing from its column names.
 PREDICT_COORD_METADATA_KEY = "giant.predict.coord"
 PREDICT_SCHEMA_VERSION_KEY = "giant.predict.schema_version"
-PREDICT_SCHEMA_VERSION = "2"
+PREDICT_SCHEMA_VERSION = "3"
+
+# Whether a --coord global predict parquet also carries true_* / true_sec_*
+# columns (v3+; "1"/"0"). Lets analysis code tell a paired prediction file
+# apart from a --no-truth one without sniffing for column presence.
+PREDICT_TRUTH_METADATA_KEY = "giant.predict.has_truth"

 # Coord-metadata value tagging a `giant rollout` steps parquet (world frame,
 # autoregressive shower output). Distinct from predict's "global"/"local".
@@ -1,6 +1,7 @@
+from collections.abc import Iterator, Mapping
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Iterator, Mapping
+from typing import TYPE_CHECKING, Any

 import numpy as np
 import polars as pl
@@ -172,7 +172,7 @@ class NormalizerEntry:
        }

    @classmethod
-    def from_json(cls, d: dict) -> "NormalizerEntry":
+    def from_json(cls, d: dict) -> NormalizerEntry:
        return cls(
            cond_norm=Normalizer.from_dict(d["cond_norm"]),
            tgt_norm=Normalizer.from_dict(d["tgt_norm"]),
@@ -194,7 +194,7 @@ class SetupCache:
    """Keyed by `topn_key(axis, n_classes)`."""

    @classmethod
-    def empty(cls, files: list[Path]) -> "SetupCache":
+    def empty(cls, files: list[Path]) -> SetupCache:
        return cls(fingerprint=fingerprint_files(files))

    def to_json(self) -> dict:
@@ -222,7 +222,7 @@ class SetupCache:
        return d

    @classmethod
-    def from_json(cls, d: dict) -> "SetupCache":
+    def from_json(cls, d: dict) -> SetupCache:
        vocab = None
        if "vocab" in d:
            pdg_map = {int(k): v for k, v in d["vocab"]["pdg_map"].items()}
@@ -247,7 +247,7 @@ class SetupCache:
            topn_maps=topn_maps,
        )

-    def merge(self, other: "SetupCache") -> "SetupCache":
+    def merge(self, other: SetupCache) -> SetupCache:
        """Union of both caches; `other`'s populated fields win on a shared key.

        Used by `save` to combine freshly-computed sections with whatever a
@@ -23,9 +23,10 @@ install stays lean.

 from __future__ import annotations

+from collections.abc import Iterable
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Iterable
+from typing import Any

 import numpy as np
 import polars as pl
@@ -155,7 +156,7 @@ class GeometryOracle:
        )

    @classmethod
-    def load(cls, path: str | Path) -> "GeometryOracle":
+    def load(cls, path: str | Path) -> GeometryOracle:
        _require_sklearn()
        import joblib

@@ -353,7 +354,7 @@ def _fit_slab_lookup(
        radius_max=radius_max,
    )
    info = {
-        "n_segments": int(len(materials)),
+        "n_segments": len(materials),
        "z_range": (z_min, z_max),
        "median_z_spacing": median_spacing,
        "radius_max": radius_max,
@@ -412,7 +413,7 @@ def build_geometry_oracle(
                "method": "slab",
                "depth_axis": depth_axis,
                "n_bins": n_bins,
-                "n_reference_points": int(len(pos)),
+                "n_reference_points": len(pos),
                "escape_factor": escape_factor,
                "n_files": len(files),
                **info,
@@ -458,7 +459,7 @@ def build_geometry_oracle(
        metadata={
            "method": method,
            "k": k,
-            "n_reference_points": int(len(X)),
+            "n_reference_points": len(X),
            "median_nn_dist": median_nn,
            "escape_factor": escape_factor,
            "n_files": len(files),
@@ -1,7 +1,7 @@
 """Factories: `build_models`/`build_critics` assemble the top-level stage
 models from a config dict (issues.md Issue 8)."""

-import torch.nn as nn
+from torch import nn

 from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig
 from giant.constants import X_DIM
@@ -2,8 +2,8 @@
 identity (issues.md Issue 8)."""

 import torch
-import torch.nn as nn
 import torch.nn.functional as F
+from torch import nn

 from giant.cond_layout import CondLayout
 from giant.config import ConditioningAxisConfig
@@ -4,9 +4,10 @@ for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
 `giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""

 import inspect
+from typing import cast

 import torch
-import torch.nn as nn
+from torch import nn


 class HistoryEncoder(nn.Module):
@@ -34,6 +35,17 @@ class HistoryEncoder(nn.Module):
    def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
        return self.forward(feat, has_prev), cache

+    def select_cache(self, cache: object, idx: torch.Tensor) -> object:
+        """Row-compacts an inference cache (`init_cache`/`step`'s state) down
+        to `idx` — used by `giant.sample.sample_secondaries_ar`'s row
+        compaction to keep a shrinking active-row set's cache aligned as rows
+        finish generating. Default here matches `init_cache`/`step`'s O(1)
+        default: `cache` is always `None`, so there's nothing to index —
+        correct for any encoder whose per-step state doesn't carry a batch
+        dimension (`MarkovHistory` has no cache at all; its running state is
+        `prev_repr`/`remaining`, compacted directly by the caller)."""
+        return cache
+

 HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}

@@ -185,7 +197,7 @@ class AttentionHistory(HistoryEncoder):
        return self.in_proj(x)

    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
-        B, K, _ = feat.shape
+        _, K, _ = feat.shape
        x = self._embed(feat, has_prev)
        mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
        for block in self.blocks:
@@ -215,3 +227,13 @@ class AttentionHistory(HistoryEncoder):
            x, kv_new = block.step(x, kv)
            new_cache.append(kv_new)
        return x, new_cache
+
+    def select_cache(self, cache: object, idx: torch.Tensor) -> list[torch.Tensor | None]:
+        """Row-compacts every block's `(B, T, dim)` KV cache down to `idx`
+        along its batch dimension — see `HistoryEncoder.select_cache`. `idx`
+        may be a long index tensor or a boolean mask (`giant.sample`'s AR
+        loop uses both). `None` entries (a block that has never seen a
+        `step` call yet) stay `None`."""
+        assert isinstance(cache, list)
+        cache_t = cast("list[torch.Tensor | None]", cache)
+        return [None if kv is None else kv[idx] for kv in cache_t]
@@ -4,7 +4,7 @@ no dependency on any other `giant.model` submodule (issues.md Issue 8)."""
 import math

 import torch
-import torch.nn as nn
+from torch import nn


 class SinusoidalEmbedding(nn.Module):
@@ -2,7 +2,7 @@
 `CriticModel` — composed from encoders/trunks/history (issues.md Issue 8)."""

 import torch
-import torch.nn as nn
+from torch import nn

 from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
 from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
@@ -614,6 +614,14 @@ class Stage2Autoregressive(StageModel):
        slot — see `AttentionHistory.step`'s docstring."""
        return self.history_encoder.step(token_feat, has_prev, cache)

+    def select_history_cache(self, cache, idx: torch.Tensor):
+        """Row-compacts `cache` (from `init_history_cache`/`history_step`)
+        down to `idx` — see `HistoryEncoder.select_cache`. Used by
+        `giant.sample.sample_secondaries_ar`'s active-row compaction to keep
+        the cache aligned with a shrinking batch as rows finish generating
+        across AR slots."""
+        return self.history_encoder.select_cache(cache, idx)
+
    def forward(
        self,
        x_t: torch.Tensor,
@@ -79,9 +79,13 @@ from giant.model.trunks import (
 )

 __all__ = [
+    "BLOCK_REGISTRY",
+    "HISTORY_REGISTRY",
+    "OBJECTIVE_REGISTRY",
+    "ROUTER_REGISTRY",
+    "TRUNK_REGISTRY",
    "AdaLNResBlock",
    "AttentionHistory",
-    "BLOCK_REGISTRY",
    "ComposedRouter",
    "ConditionEncoder",
    "ContextAdapter",
@@ -91,17 +95,14 @@ __all__ = [
    "ExpertTrunk",
    "FilmResBlock",
    "FlowObjective",
-    "HISTORY_REGISTRY",
    "HistoryEncoder",
    "LinearTrunk",
    "MarkovHistory",
    "NoHistory",
    "NoneRouter",
-    "OBJECTIVE_REGISTRY",
    "Objective",
    "PdgRouter",
    "ProcessRouter",
-    "ROUTER_REGISTRY",
    "ResBlock",
    "RoutedTrunk",
    "Router",
@@ -110,7 +111,6 @@ __all__ = [
    "Stage2Autoregressive",
    "Stage2OneShot",
    "StageModel",
-    "TRUNK_REGISTRY",
    "Trunk",
    "WganObjective",
    "_CausalAttnBlock",
@@ -8,8 +8,8 @@ import re
 from collections.abc import Sequence

 import torch
-import torch.nn as nn
 import torch.nn.functional as F
+from torch import nn

 from giant.cond_layout import CondLayout
 from giant.constants import COND_DIM
@@ -398,7 +398,7 @@ def _build_router_from_cfg(
    """Resolve one stage's `router` config into a `Router`, single-axis or
    composed. `gumbel` is set as a post-construction attribute (shared by
    every router type, not a per-type constructor kwarg)."""
-    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    shared_vocab = {"pdg_vocab": pdg_vocab, "mat_vocab": mat_vocab}
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
        _check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning)
@@ -14,7 +14,7 @@ class CosineSchedule:
        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)

        self.betas = torch.from_numpy(betas)
-        self.alphas = torch.from_numpy((1.0 - betas))
+        self.alphas = torch.from_numpy(1.0 - betas)
        self.alpha_bars = torch.from_numpy(alpha_bars[1:])

    def to(self, device: torch.device) -> "CosineSchedule":
@@ -35,7 +35,7 @@ finding those two tests independently converge on.
 import copy
 from dataclasses import dataclass, field

-import torch.nn as nn
+from torch import nn

 from giant.config import INFERENCE_OVERRIDES, _get_path, _set_path, leaf_paths
 from giant.model.builders import build_critics, build_models
@@ -224,7 +224,7 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
            _set_path(probe_cfg, path, candidate)
            try:
                changed = _fingerprint(_built_modules(probe_cfg, pdg_vocab, mat_vocab)) != baseline_fp
-            except Exception:
+            except Exception:  # noqa: BLE001 - a perturbation that fails to even build counts as "consumed"
                changed = True
            if changed:
                break
@@ -10,7 +10,7 @@ free — no separate "routed transformer trunk" class needed.
 """

 import torch
-import torch.nn as nn
+from torch import nn

 from giant.model.layers import build_block
 from giant.model.routers import Router
@@ -1,4 +1,4 @@
-from typing import Callable
+from collections.abc import Callable

 import torch

@@ -19,7 +19,7 @@ free-running history representation stays unsnapped — see

 from __future__ import annotations

-from functools import lru_cache
+from functools import cache
 from typing import TYPE_CHECKING

 import numpy as np
@@ -42,7 +42,7 @@ _CHARGE_WEIGHT = 50.0
 _LOG_EPS = 1e-8


-@lru_cache(maxsize=None)
+@cache
 def particle_mass_charge(pdg: int) -> tuple[float, float]:
    """Return (mass_MeV, charge_e) for a raw PDG code.

@@ -135,7 +135,7 @@ def invert_dense_map(m: dict[int, int]) -> dict[int, int]:

 def decode_topn_class(
    class_idx: np.ndarray,
-    topn_map: "TopNMap",
+    topn_map: TopNMap,
    n_classes: int,
    other_policy: str = "sample",
    rng: np.random.Generator | None = None,
@@ -13,6 +13,7 @@ from giant.constants import (
    X_DIM,
 )
 from giant.data import setup_cache
+from giant.data.dataset import StreamingStepsDataset, make_event_split
 from giant.data.loader import (
    TopNMap,
    _topn_plus_other_map,
@@ -23,13 +24,12 @@ from giant.data.loader import (
 from giant.data.scan import MetadataScan, ScanRequest, scan_metadata
 from giant.data.transforms import (
    Normalizer,
-    build_features,
-    _WelfordAccumulator,
    _ReservoirSampler,
+    _WelfordAccumulator,
+    build_features,
    sorted_membership,
 )
-from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import build_models, build_critics, resolve_type_n_classes
+from giant.model.network import build_critics, build_models, resolve_type_n_classes
 from giant.training import train as run_training


@@ -17,7 +17,8 @@ treated as detector leakage and not deposited.
 from __future__ import annotations

 from collections import Counter
-from typing import TYPE_CHECKING, Callable, TypedDict
+from collections.abc import Callable
+from typing import TYPE_CHECKING, TypedDict

 import numpy as np
 import torch
@@ -117,7 +118,7 @@ def decode_secondary_identity(
    pre_dir: np.ndarray,
    sec_phys_norm: Normalizer,
    pdg_map: dict[int, int],
-    sec_type_topn_map: "TopNMap | None",
+    sec_type_topn_map: TopNMap | None,
    other_policy: str,
    rng: np.random.Generator | None,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
@@ -390,34 +391,34 @@ def _terminal_rows(tr: dict[str, np.ndarray], sel: np.ndarray, reason: str, edep
    pos = tr["pre_pos"][sel]
    dir_ = tr["pre_dir"][sel]
    n = int(sel.sum())
-    return dict(
-        event_id=tr["event_id"][sel],
-        track_id=tr["track_id"][sel],
-        parent_id=tr["parent_id"][sel],
-        generation=tr["generation"][sel],
-        step_no=tr["step_in_track"][sel],
-        pdg=tr["pdg"][sel],
-        pre_x=pos[:, 0],
-        pre_y=pos[:, 1],
-        pre_z=pos[:, 2],
-        pre_E=tr["pre_E"][sel],
-        pre_dx=dir_[:, 0],
-        pre_dy=dir_[:, 1],
-        pre_dz=dir_[:, 2],
-        post_x=pos[:, 0],
-        post_y=pos[:, 1],
-        post_z=pos[:, 2],
-        post_E=np.zeros(n),
-        post_dx=dir_[:, 0],
-        post_dy=dir_[:, 1],
-        post_dz=dir_[:, 2],
-        edep=np.asarray(edep, dtype=np.float64).reshape(n),
-        step_length=np.zeros(n),
-        material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
-        layer_id=tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
-        n_sec_pred=np.zeros(n, dtype=np.int64),
-        termination_reason=np.full(n, reason, dtype=object),
-    )
+    return {
+        "event_id": tr["event_id"][sel],
+        "track_id": tr["track_id"][sel],
+        "parent_id": tr["parent_id"][sel],
+        "generation": tr["generation"][sel],
+        "step_no": tr["step_in_track"][sel],
+        "pdg": tr["pdg"][sel],
+        "pre_x": pos[:, 0],
+        "pre_y": pos[:, 1],
+        "pre_z": pos[:, 2],
+        "pre_E": tr["pre_E"][sel],
+        "pre_dx": dir_[:, 0],
+        "pre_dy": dir_[:, 1],
+        "pre_dz": dir_[:, 2],
+        "post_x": pos[:, 0],
+        "post_y": pos[:, 1],
+        "post_z": pos[:, 2],
+        "post_E": np.zeros(n),
+        "post_dx": dir_[:, 0],
+        "post_dy": dir_[:, 1],
+        "post_dz": dir_[:, 2],
+        "edep": np.asarray(edep, dtype=np.float64).reshape(n),
+        "step_length": np.zeros(n),
+        "material": tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
+        "layer_id": tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
+        "n_sec_pred": np.zeros(n, dtype=np.int64),
+        "termination_reason": np.full(n, reason, dtype=object),
+    }


@torch.no_grad()
@@ -442,14 +443,14 @@ def rollout(
    on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
    particle_conditioning: str = "embedding",
    material_conditioning: str = "embedding",
-    pdg_topn_map: "TopNMap | None" = None,
-    mat_topn_map: "TopNMap | None" = None,
-    sec_type_topn_map: "TopNMap | None" = None,
+    pdg_topn_map: TopNMap | None = None,
+    mat_topn_map: TopNMap | None = None,
+    sec_type_topn_map: TopNMap | None = None,
    other_policy: str = "sample",
    seed: int | None = None,
    stage1_ddpm_steps: int = 1000,
    stage2_ddpm_steps: int = 1000,
-    l1_dist_collector: "L1DistCollector | None" = None,
+    l1_dist_collector: L1DistCollector | None = None,
 ) -> dict[str, np.ndarray] | RolloutSummary:
    """Run showers to completion.

@@ -231,6 +231,7 @@ def sample_secondaries_ar(
    stage1_out: torch.Tensor,
    n_sec_pred: torch.Tensor | None,
    steps: int = 10,
+    full_length: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    """`Stage2Autoregressive` inference loop: one token at a time, in
    descending-energy slot order, up to `k_max` sequential calls. Unlike
@@ -242,19 +243,18 @@ def sample_secondaries_ar(
    expressiveness.

    A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
-    the "K sequential forwards" cost applies per-token here, not
-    once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
-    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
-    every row in the batch has stopped).
+    the "K sequential forwards" cost applies per-token here, not once, so a
+    flow/ddpm AR run costs ~`n_sec * steps` model calls per physics step
+    (measured on `configs/baseline.toml`: 0.382 secondaries/step at rollout
+    time), not `k_max * steps` — see the row-compaction paragraph below.

-    `n_sec_pred`, if given, fixes each row's secondary count up front (as
-    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
-    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
-    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
-    running self-sample regardless of the decoder's own stop head — the
-    scheduled-sampling training contract does not truncate). This always
-    runs the full `k_max`-iteration loop, masking by the given count at the
-    end exactly as before.
+    `n_sec_pred`, if given (as resolved by `resolve_n_sec` — `n_sec.mode` in
+    `("head", "truth")`, or a stop-token decoder driven by
+    `_assemble_stage2_ar_inputs_scheduled`'s ground-truth `n_sec`, which must
+    run the *full* `k_max`-length free-running self-sample regardless of the
+    decoder's own stop head — the scheduled-sampling training contract does
+    not truncate, see `full_length` below) fixes each row's secondary count
+    up front.

    `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
    (`n_sec.mode = "stop_token"`): before generating each slot's token, that
@@ -263,11 +263,35 @@ def sample_secondaries_ar(
    why this needs no extra state) decides whether generation should have
    already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
    0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
-    `n_sec_pred` is the first slot index where this fires; once every row in
-    the batch has fired, the loop breaks before spending a model call on the
-    next slot's token — the average-case cost win the docstring above
-    describes. A row that never fires within `k_max` is capped there
-    (`K_MAX` stays a safety cap, not a modeling ceiling).
+    `n_sec_pred` is the first slot index where this fires. A row that never
+    fires within `k_max` is capped there (`K_MAX` stays a safety cap, not a
+    modeling ceiling).
+
+    **Row compaction.** A row that has already produced its `n_sec_pred`
+    tokens (or, under `stop_token`, has already fired its stop logit) has
+    nothing left to contribute — every later slot of that row is masked out
+    of `sec_valid` on return, and downstream consumers (`giant/rollout.py`,
+    `giant/cli.py`) never read it. So unless `full_length=True`, this
+    function drops such rows from the active set entirely instead of running
+    the model on them: `active_idx` starts at every row with `n_sec_pred > 0`
+    (or, under `stop_token`, every row — the first stop decision can fire at
+    slot 0) and only shrinks as rows finish, so slot `k`'s model calls cost
+    `O(active rows)` not `O(B)`. Slots a row never reaches keep their `0.0`
+    zero-init in `sec_cont`/`sec_type` — masked by `sec_valid`, identical to
+    what a full, uncompacted run would have written there before masking.
+    `AttentionHistory`'s KV cache is kept aligned to the shrinking active set
+    via `Stage2Autoregressive.select_history_cache`
+    (`giant.model.history.HistoryEncoder.select_cache`) every time the set
+    shrinks; `MarkovHistory`'s O(1) state (`prev_repr`/`remaining`, carried
+    directly rather than through a cache) is compacted the same way.
+
+    `full_length=True` disables all of the above: every row runs the full
+    `k_max`-iteration loop regardless of `n_sec_pred`/stop decisions, exactly
+    reproducing the pre-compaction behaviour. Required by
+    `_assemble_stage2_ar_inputs_scheduled`'s scheduled-sampling self-sample,
+    whose training contract needs a real prediction at every slot up to
+    `k_max` (mixed per-slot against ground truth) even past a row's own
+    `n_sec` — see that function's docstring.

    Under `history="attention"` the history encoding is computed once per
    slot via `Stage2Autoregressive.history_step` (a KV-cache append)
@@ -312,11 +336,6 @@ def sample_secondaries_ar(
    sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
    sec_type = torch.zeros(B, k_max, type_dim, device=device)

-    # Running per-token state, threaded from one slot to the next.
-    prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
-    remaining = torch.ones(B, device=device)
-    history_cache = sec_decoder.init_history_cache()
-
    use_stop_token = n_sec_pred is None
    if use_stop_token:
        assert getattr(sec_decoder, "stop_head", None) is not None, (
@@ -324,21 +343,46 @@ def sample_secondaries_ar(
            "with no stop_head — only valid under stage2_model.n_sec.mode = "
            "'stop_token'"
        )
-        finished = torch.zeros(B, dtype=torch.bool, device=device)
        derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
+        # Tracks which rows have already recorded a stop, globally by
+        # original batch index — needed even under compaction's own
+        # never-revisit guarantee, because `full_length=True` keeps every
+        # row in `active_idx` for the whole loop, so a row whose stop logit
+        # fires once but flips back below threshold at a later slot (a real
+        # possibility for an untrained/lightly-trained stop_head) must not
+        # have `derived_n_sec` overwritten by that later, spurious re-fire.
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+
+    # `active_idx`: rows still contributing tokens, indexed into the
+    # original batch. Only ever shrinks (never full_length) or stays fixed
+    # at arange(B) (full_length) — see the row-compaction docstring section.
+    active_idx = torch.arange(B, device=device)
+    if not full_length and not use_stop_token:
+        active_idx = active_idx[n_sec_pred > 0]
+
+    # Running per-token state, already compacted to `active_idx`.
+    prev_repr = torch.zeros(active_idx.numel(), CONT_SLOT_DIM + type_dim, device=device)
+    remaining = torch.ones(active_idx.numel(), device=device)
+    history_cache = sec_decoder.init_history_cache()

    for k in range(k_max):
-        has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
-        history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
-        remaining_frac = remaining.unsqueeze(1)  # (B, 1)
-        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        if active_idx.numel() == 0:
+            break
+        Bc = active_idx.numel()
+        cc = cond_cont.index_select(0, active_idx)
+        ck = cond_cat.index_select(0, active_idx)
+        s1 = stage1_out.index_select(0, active_idx)
+        has_prev = torch.full((Bc, 1), k >= 1, dtype=torch.bool, device=device)
+        history_feat = prev_repr.unsqueeze(1)  # (Bc, 1, CONT_SLOT_DIM + type_dim)
+        remaining_frac = remaining.unsqueeze(1)  # (Bc, 1)
+        slot_idx = torch.full((Bc, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
        hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)

        if use_stop_token:
            stop_logit = sec_decoder.predict_stop(
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                history_feat,
                has_prev,
                remaining_frac,
@@ -346,21 +390,31 @@ def sample_secondaries_ar(
                hist=hist,
            ).squeeze(1)
            if sec_decoder.n_sec_sampling == "sample":
-                stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
+                stop_now = torch.rand(Bc, device=device) < torch.sigmoid(stop_logit)
            else:
                stop_now = stop_logit >= 0.0
-            derived_n_sec[stop_now & ~finished] = k
-            finished = finished | stop_now
-            if finished.all():
-                break
+            newly_stopped = stop_now & ~finished.index_select(0, active_idx)
+            derived_n_sec[active_idx[newly_stopped]] = k
+            finished[active_idx[stop_now]] = True
+            if not full_length:
+                keep = ~stop_now
+                active_idx = active_idx[keep]
+                cc, ck, s1 = cc[keep], ck[keep], s1[keep]
+                has_prev, remaining_frac, slot_idx = has_prev[keep], remaining_frac[keep], slot_idx[keep]
+                history_feat, hist = history_feat[keep], hist[keep]
+                history_cache = sec_decoder.select_history_cache(history_cache, keep)
+                prev_repr, remaining = prev_repr[keep], remaining[keep]
+                if active_idx.numel() == 0:
+                    break
+                Bc = active_idx.numel()

        if objective.is_adversarial:
-            z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
+            z = torch.randn(Bc, 1, sec_decoder.noise_dim, device=device)
            token = sec_decoder(
                z,
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                history_feat,
                has_prev,
                remaining_frac,
@@ -368,15 +422,15 @@ def sample_secondaries_ar(
                hist=hist,
            )
        else:
-            x = torch.randn(B, 1, token_dim, device=device)
+            x = torch.randn(Bc, 1, token_dim, device=device)
            dt = 1.0 / steps
            for i in range(steps):
-                t = torch.full((B, 1), i * dt, device=device)
+                t = torch.full((Bc, 1), i * dt, device=device)
                v = sec_decoder(
                    x,
-                    cond_cont,
-                    cond_cat,
-                    stage1_out,
+                    cc,
+                    ck,
+                    s1,
                    history_feat,
                    has_prev,
                    remaining_frac,
@@ -387,15 +441,15 @@ def sample_secondaries_ar(
                x = x + v * dt
            token = x

-        token = token.squeeze(1)  # (B, token_dim)
+        token = token.squeeze(1)  # (Bc, token_dim)
        cont_k = token[:, :CONT_SLOT_DIM]
        if type_folded:
            type_k = token[:, CONT_SLOT_DIM:]
        else:
            type_k = sec_decoder.predict_type(
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                history_feat,
                has_prev,
                remaining_frac,
@@ -403,8 +457,8 @@ def sample_secondaries_ar(
                hist=hist,
            ).squeeze(1)

-        sec_cont[:, k] = cont_k
-        sec_type[:, k] = type_k
+        sec_cont[active_idx, k] = cont_k
+        sec_type[active_idx, k] = type_k

        if target == "onehot":
            type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
@@ -415,6 +469,12 @@ def sample_secondaries_ar(
        prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
        remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)

+        if not full_length and not use_stop_token:
+            keep2 = n_sec_pred.index_select(0, active_idx) > (k + 1)
+            active_idx = active_idx[keep2]
+            prev_repr, remaining = prev_repr[keep2], remaining[keep2]
+            history_cache = sec_decoder.select_history_cache(history_cache, keep2)
+
    resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
    sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
    return sec_cont, sec_type, sec_valid
@@ -82,7 +82,7 @@ def _max_index(parent: Path, pattern: re.Pattern) -> int:

 def _git_user_name() -> str | None:
    try:
-        out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2)
+        out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2, check=False)
    except (OSError, subprocess.SubprocessError):
        # OSError (e.g. git not on PATH) and subprocess.SubprocessError
        # (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
@@ -604,7 +604,7 @@ def _run_bump(
    if not root_path.is_dir():
        raise SystemExit(f"error: {root_path} is not a directory")

-    date = date or dt.date.today().isoformat()
+    date = date or dt.date.today().isoformat()  # noqa: DTZ011 - local calendar date for the dataset-version log, not stored
    by = by if by is not None else _git_user_name()
    if gen is None:
        new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
@@ -31,9 +31,9 @@ import torch

 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

-from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM  # noqa: E402
-from giant.model import network as net  # noqa: E402
-from tests.legacy import network_v02_snapshot as legacy  # noqa: E402
+from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model import network as net
+from tests.legacy import network_v02_snapshot as legacy


 def _random_batch(model_config: dict, batch: int, seed: int):
@@ -28,8 +28,8 @@ import subprocess
 import sys
 import uuid
 import zlib
-from dataclasses import dataclass
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
 from pathlib import Path

 # Must match giant/tools/bump_dataset_version.py's GEN_RE.
@@ -147,7 +147,7 @@ def run_job(
    cmd = build_cmd(executable, job, events_per_file, energy_gev)

    env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job, energy_gev)))
-    result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env)
+    result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env, check=False)

    if result.returncode != 0:
        return JobResult(
@@ -10,10 +10,9 @@ from __future__ import annotations
 import os
 from enum import Enum
 from pathlib import Path
-from typing import Optional
+from typing import Annotated

 import typer
-from typing_extensions import Annotated

 from giant.config import Conditioning

@@ -73,7 +72,7 @@ class PoolType(str, Enum):
 def convert(
    root_files: Annotated[list[Path], typer.Argument(help="Input ROOT file(s)")],
    output: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--output",
            "-o",
@@ -108,7 +107,7 @@ def convert(
        ),
    ] = _DATASET_ROOT_DEFAULT,
    schema: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--schema",
            help="Schema tag to write parquets under, e.g. schema2 (only used "
@@ -192,13 +191,13 @@ def migrate(
 def bump_gen(
    reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
-    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
+    by: Annotated[str | None, typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--date", help="Override date (default: today, ISO)"),
    ] = None,
    to: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--to",
            metavar="genN",
@@ -227,13 +226,13 @@ def bump_schema(
    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
    reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
-    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
+    by: Annotated[str | None, typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--date", help="Override date (default: today, ISO)"),
    ] = None,
    to: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--to",
            metavar="schemaN",
@@ -272,7 +271,7 @@ def status(
 def update_manifest(
    manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
    schema: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--schema",
            metavar="schemaN",
@@ -280,7 +279,7 @@ def update_manifest(
        ),
    ] = None,
    gen: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
    ] = None,
    execute: Annotated[
@@ -298,11 +297,11 @@ def update_manifest(
 def create_manifest(
    files: Annotated[list[Path], typer.Argument(help="Parquet files to include")],
    output: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option("--output", "-o", help="Explicit path for the new .manifest file"),
    ] = None,
    pool: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--pool",
            metavar="DETECTOR",
@@ -310,7 +309,7 @@ def create_manifest(
        ),
    ] = None,
    type_: Annotated[
-        Optional[PoolType],
+        PoolType | None,
        typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
    ] = None,
    root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
@@ -459,7 +458,7 @@ def warm_cache(
        typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
    ],
    config: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--config",
            "-c",
@@ -469,7 +468,7 @@ def warm_cache(
        ),
    ] = None,
    val_fraction: Annotated[
-        Optional[float],
+        float | None,
        typer.Option(
            "--val-fraction",
            "-f",
@@ -477,7 +476,7 @@ def warm_cache(
        ),
    ] = None,
    seed: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--seed",
            "-s",
@@ -485,7 +484,7 @@ def warm_cache(
        ),
    ] = None,
    particle_conditioning: Annotated[
-        Optional[Conditioning],
+        Conditioning | None,
        typer.Option(
            "--particle-conditioning",
            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for. "
@@ -493,7 +492,7 @@ def warm_cache(
        ),
    ] = None,
    material_conditioning: Annotated[
-        Optional[Conditioning],
+        Conditioning | None,
        typer.Option(
            "--material-conditioning",
            help="Must match the `giant train` run(s)' conditioning.material.type "
@@ -502,7 +501,7 @@ def warm_cache(
        ),
    ] = None,
    router: Annotated[
-        Optional[bool],
+        bool | None,
        typer.Option(
            "--router/--no-router",
            help="Warm the process vocabulary too (only takes effect with --router-type process). "
@@ -510,11 +509,11 @@ def warm_cache(
        ),
    ] = None,
    router_type: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--router-type", help="Router implementation name. Not allowed together with --config"),
    ] = None,
    n_experts: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--n-experts", help="Number of routed experts. Not allowed together with --config"),
    ] = None,
    rebuild: Annotated[
@@ -144,7 +144,7 @@ def run_hparam_scan(
        start = time.monotonic()
        try:
            with open(out_dir / "train.log", "a") as log:
-                subprocess.run(cmd, env=env, stdout=log, stderr=subprocess.STDOUT)
+                subprocess.run(cmd, env=env, stdout=log, stderr=subprocess.STDOUT, check=False)
        except KeyboardInterrupt:
            print(
                f"\ninterrupted during {name} — re-run this script to resume "
@@ -32,7 +32,7 @@ MANIFEST_SUFFIX = ".manifest"
 # since today's pool assignment is encoded only by *which folder a file's
 # parquet was copied into* — not by anything in the filename itself.
 POOL_ASSIGNMENT: dict[str, dict[str, range | list[int]]] = {
-    "pbwo4": {"full": range(0, 6), "holdout": range(6, 10)},
+    "pbwo4": {"full": range(6), "holdout": range(6, 10)},
    "sampling_fe_scint": {"dev": [0], "full": [1, 2], "holdout": [3]},
    "sampling_pb_lar": {"dev": [0], "full": [1, 2], "holdout": [3]},
    "sampling_pb_scint": {"dev": [0], "full": [1, 2], "holdout": [3]},
@@ -106,7 +106,7 @@ def _convert_one(
    if output_path is not None:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        cmd += ["--output", str(output_path)]
-    result = subprocess.run(cmd, capture_output=True, text=True)
+    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    return root_file, result.returncode, result.stdout, result.stderr


@@ -6,8 +6,8 @@ that tests and tooling construct directly.
 """

 from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
-from giant.training.metrics import MetricsCollector, MetricSpec
 from giant.training.loop import train
+from giant.training.metrics import MetricsCollector, MetricSpec
 from giant.training.trainers import (
    FlowDDPMStageTrainer,
    StageSpec,
@@ -9,9 +9,10 @@
 import os
 import signal
 import time
+from collections.abc import Callable
 from pathlib import Path
 from types import FrameType
-from typing import Callable
+from typing import Self

 import numpy as np
 import torch
@@ -47,7 +48,7 @@ class _GracefulShutdown:
            Callable[[int, FrameType | None], object] | signal.Handlers | int | None,
        ] = {}

-    def __enter__(self) -> "_GracefulShutdown":
+    def __enter__(self) -> Self:
        for sig in _CATCHABLE_SIGNALS:
            self._previous[sig] = signal.getsignal(sig)
            signal.signal(sig, self._handle)
@@ -188,7 +188,7 @@ class MetricsCollector:
        self.fieldnames = self._build_fieldnames()
        metrics_path = out_dir / "metrics.csv"
        append = resume and metrics_path.exists()
-        self._file = open(metrics_path, "a" if append else "w", newline="")
+        self._file = open(metrics_path, "a" if append else "w", newline="")  # noqa: SIM115 - kept open for the object's lifetime, closed in .close()
        self._writer = csv.DictWriter(self._file, fieldnames=self.fieldnames)
        if not append:
            self._writer.writeheader()
@@ -47,7 +47,7 @@ class MetricsTable:
    columns: dict[str, list[float]]

    @classmethod
-    def load(cls, path: str | Path) -> "MetricsTable":
+    def load(cls, path: str | Path) -> MetricsTable:
        with open(path, newline="") as f:
            rows = list(csv.DictReader(f))
        epochs = [int(float(r["epoch"])) for r in rows]
@@ -291,9 +291,13 @@ def _assemble_stage2_ar_inputs_scheduled(
    skips self-sampling entirely), so callers can call this unconditionally.

    The free-running estimate is a REAL autoregressive self-sample —
-    `giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
-    cheap one-step proxy, so building it costs the same `k_max` (`* steps`
-    for flow) sequential forwards `sample.py` pays at inference, EVERY batch
+    `giant.sample.sample_secondaries_ar` under `torch.no_grad()`, called here
+    with `full_length=True` — not a cheap one-step proxy, so building it
+    costs the full `k_max` (`* steps` for flow) sequential forwards for every
+    row regardless of that row's own secondary count (`full_length=True`
+    disables `sample.py`'s inference-time row compaction — see that
+    function's docstring for why: the mixing below needs a real prediction
+    at every slot up to `k_max`, not just the valid ones). Paid EVERY batch
    this is called on (paid at train time too whenever teacher_forcing !=
    "always"). Fully detached: gradient only ever flows
    through the "real" target path each stage trainer already uses
@@ -306,7 +310,7 @@ def _assemble_stage2_ar_inputs_scheduled(

    was_training = model.training
    sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
-        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
+        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps, full_length=True
    )
    if was_training:
        model.train()
@@ -20,7 +20,7 @@ from typing import NamedTuple

 import torch
 import torch.nn.functional as F
-import torch.optim as optim
+from torch import optim

 from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
 from giant.constants import CONT_SLOT_DIM
@@ -333,7 +333,7 @@ class StageTrainer:
        #: "sampled"` — the stage-1 `StageTrainer` this (stage-2) trainer
        #: draws its context sample from. `None` for stage 1 itself, and for
        #: stage 2 under "truth".
-        self.stage1_source: "StageTrainer | None" = None
+        self.stage1_source: StageTrainer | None = None

    def attach_stage1(self, stage1_trainer: "StageTrainer") -> None:
        """Wires this (stage-2) trainer to the stage-1 trainer it should
@@ -1,13 +1,13 @@
 [project]
 name = "giant"
-version = "0.3.16"
+version = "0.3.22"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
    "numpy>=1.26,<3",
    "polars>=1.0,<2",
-    "pyarrow>=16,<25",
+    "pyarrow>=16,<26",
    "tqdm>=4.60,<5",
    "typer>=0.12,<1",
    "pyyaml>=6,<7",
@@ -17,8 +17,8 @@ import re
 from collections.abc import Sequence

 import torch
-import torch.nn as nn
 import torch.nn.functional as F
+from torch import nn

 from giant.constants import (
    COND_DIM,
@@ -949,7 +949,7 @@ def _check_router_conditioning_compat(router_types: list[str], conditioning: str


 def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding") -> Router:
-    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    shared_vocab = {"pdg_vocab": pdg_vocab, "mat_vocab": mat_vocab}
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
        _check_router_conditioning_compat([a["type"] for a in axes], conditioning)
@@ -977,15 +977,15 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
    if router_cfg and router_cfg.get("enabled"):
        pdg_vocab = model_config["pdg_vocab"]
        mat_vocab = model_config["mat_vocab"]
-        shared = dict(
-            pdg_vocab=pdg_vocab,
-            mat_vocab=mat_vocab,
-            expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
-            expert_n_blocks=model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
-            emb_dim=model_config.get("emb_dim", EMB_DIM),
-            dropout=model_config.get("dropout", 0.1),
-            conditioning=model_config.get("conditioning", "embedding"),
-        )
+        shared = {
+            "pdg_vocab": pdg_vocab,
+            "mat_vocab": mat_vocab,
+            "expert_hidden_dim": model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
+            "expert_n_blocks": model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
+            "emb_dim": model_config.get("emb_dim", EMB_DIM),
+            "dropout": model_config.get("dropout", 0.1),
+            "conditioning": model_config.get("conditioning", "embedding"),
+        }
        conditioning = shared["conditioning"]
        stage1 = RoutedDenoisingMLP(
            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
@@ -5,12 +5,12 @@ from pathlib import Path

 import pytest
 import torch
+from test_train import _base_cfg, _run_train

 from giant.model.routers import EnergyRouter
 from giant.model.wgan import gradient_penalty
 from giant.training.amp import resolve_autocast
 from giant.training.stage2_inputs import _remaining_energy_fraction
-from test_train import _base_cfg, _run_train

 # ---------------------------------------------------------------------------
 # resolve_autocast
@@ -0,0 +1,278 @@
+"""Tests for giant.analysis.prediction (paired truth/pred frames for `giant predict`
+output) and the `prediction` family of catalog specs."""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+import pytest
+
+from giant.analysis.catalog import Bundle, get_spec
+from giant.analysis.context import Context, build_context
+from giant.analysis.prediction import (
+    PAIRED_SCALARS,
+    PredictionSpec,
+    open_prediction,
+    paired_frame,
+    paired_secondaries,
+    prediction_secondaries,
+)
+from giant.analysis.reduce import hist2d
+from giant.analysis.sources import RolloutSpec
+from tests.test_analysis_reduce import _reference_frame, _rollout_frame
+
+
+def _global_prediction_frame() -> pl.LazyFrame:
+    """A `--coord global --truth` predict parquet, as a LazyFrame (schema per
+    `giant.cli.predict`'s global-coord table, `giant/cli.py:1310-1379`)."""
+    return pl.DataFrame(
+        {
+            "event_id": [1, 1, 2],
+            "pdg": [11, 11, 22],
+            "pre_x": [0.0, 0.0, 0.0],
+            "pre_y": [0.0, 0.0, 0.0],
+            "pre_z": [0.0, 1.0, 0.0],
+            "pre_E": [100.0, 60.0, 50.0],
+            "pre_dx": [0.0, 0.0, 0.0],
+            "pre_dy": [0.0, 0.0, 0.0],
+            "pre_dz": [1.0, 1.0, 1.0],
+            "material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
+            "layer_id": [0, 1, 0],
+            "n_sec": [1, 0, 2],
+            "n_sec_pred": [1, 0, 1],
+            # predicted (unprefixed) values
+            "step_length": [1.2, 0.9, 1.1],
+            "delta_e": [42.0, 29.0, 31.0],
+            "edep": [35.0, 29.0, 25.0],
+            "post_dx": [0.0, 0.0, 0.0],
+            "post_dy": [0.0, 0.0, 0.0],
+            "post_dz": [1.0, 1.0, 1.0],
+            "post_x": [0.0, 0.0, 0.0],
+            "post_y": [0.0, 0.0, 0.0],
+            "post_z": [1.2, 1.9, 1.1],
+            "sec_pdg_list": [[22], [], [22]],
+            "sec_E_list": [[5.0], [], [4.0]],
+            "sec_dx_list": [[0.0], [], [0.0]],
+            "sec_dy_list": [[0.0], [], [0.0]],
+            "sec_dz_list": [[1.0], [], [1.0]],
+            # truth
+            "true_step_length": [1.0, 1.0, 1.0],
+            "true_delta_e": [40.0, 30.0, 30.0],
+            "true_edep": [40.0, 30.0, 20.0],
+            "true_post_E": [60.0, 30.0, 20.0],
+            "true_post_dx": [0.0, 0.0, 0.0],
+            "true_post_dy": [0.0, 0.0, 0.0],
+            "true_post_dz": [1.0, 1.0, 1.0],
+            "true_post_x": [0.0, 0.0, 0.0],
+            "true_post_y": [0.0, 0.0, 0.0],
+            "true_post_z": [1.0, 2.0, 1.0],
+            "true_e_sec": [0.0, 0.0, 10.0],
+            "process": ["compt", "phot", "compt"],
+            "true_sec_pdg_list": [[22], [], [22, 11]],
+            "true_sec_E_list": [[6.0], [], [7.0, 3.0]],
+            "true_sec_dx_list": [[0.0], [], [0.0, 1.0]],
+            "true_sec_dy_list": [[0.0], [], [0.0, 0.0]],
+            "true_sec_dz_list": [[1.0], [], [1.0, 0.0]],
+        }
+    ).lazy()
+
+
+def _local_prediction_frame() -> pl.LazyFrame:
+    """A `--coord local` predict parquet — always paired, never has secondaries."""
+    return pl.DataFrame(
+        {
+            "event_id": [1, 2],
+            "pdg": [11, 22],
+            "pre_x": [0.0, 0.0],
+            "pre_y": [0.0, 0.0],
+            "pre_z": [0.0, 0.0],
+            "pre_E": [100.0, 50.0],
+            "pre_dx": [0.0, 0.0],
+            "pre_dy": [0.0, 0.0],
+            "pre_dz": [1.0, 1.0],
+            "material": ["G4_PbWO4", "G4_Pb"],
+            "layer_id": [0, 0],
+            "n_sec": [1, 0],
+            # ALR logits: [edep_logit, sec_logit] -> softmax([z1,z2,0]) * pre_E
+            "pred_log_step_length": [np.log(1.2 + 1e-6), np.log(0.9 + 1e-6)],
+            "pred_edep_logit": [1.0, 0.5],
+            "pred_sec_logit": [0.0, -1.0],
+            "pred_post_dx": [0.0, 0.0],
+            "pred_post_dy": [0.0, 0.0],
+            "pred_post_dz": [1.0, 1.0],
+            "pred_travel_dx": [0.0, 0.0],
+            "pred_travel_dy": [0.0, 0.0],
+            "pred_travel_dz": [1.0, 1.0],
+            "true_log_step_length": [np.log(1.0 + 1e-6), np.log(1.0 + 1e-6)],
+            "true_edep_logit": [0.8, 0.6],
+            "true_sec_logit": [0.2, -2.0],
+            "true_post_dx": [0.0, 0.0],
+            "true_post_dy": [0.0, 0.0],
+            "true_post_dz": [1.0, 1.0],
+            "true_travel_dx": [0.0, 0.0],
+            "true_travel_dy": [0.0, 0.0],
+            "true_travel_dz": [1.0, 1.0],
+        }
+    ).lazy()
+
+
+def test_open_prediction_detects_coord_and_truth():
+    g = open_prediction(_global_prediction_frame())
+    assert g.coord == "global" and g.has_truth
+
+    loc = open_prediction(_local_prediction_frame())
+    assert loc.coord == "local" and loc.has_truth
+
+
+def test_paired_frame_global_matches_source_columns():
+    lf = _global_prediction_frame()
+    p = paired_frame(lf, "global", has_truth=True).collect()
+    assert p["pred_step_length"].to_list() == [1.2, 0.9, 1.1]
+    assert p["true_step_length"].to_list() == [1.0, 1.0, 1.0]
+    assert p["pred_edep"].to_list() == [35.0, 29.0, 25.0]
+    assert p["true_edep"].to_list() == [40.0, 30.0, 20.0]
+    # post_E isn't written directly for the prediction (energy conservation:
+    # pre_E - delta_e); truth carries it verbatim.
+    assert p["pred_post_E"].to_list() == pytest.approx([100.0 - 42.0, 60.0 - 29.0, 50.0 - 31.0])
+    assert p["true_post_E"].to_list() == [60.0, 30.0, 20.0]
+    # cos_scatter: pre_dir . post_dir, both (0,0,1) here -> 1.0
+    assert p["pred_cos_scatter"].to_list() == pytest.approx([1.0, 1.0, 1.0])
+    assert p["true_cos_scatter"].to_list() == pytest.approx([1.0, 1.0, 1.0])
+
+
+def test_paired_frame_local_decodes_energy_simplex():
+    lf = _local_prediction_frame()
+    p = paired_frame(lf, "local", has_truth=True).collect()
+    # softmax([1.0, 0.0, 0.0]) * 100 for row 0's pred edep
+    z = np.exp([1.0, 0.0, 0.0])
+    expected_edep_0 = (z[0] / z.sum()) * 100.0
+    assert p["pred_edep"][0] == pytest.approx(expected_edep_0)
+    assert p["pred_step_length"][0] == pytest.approx(1.2, abs=1e-4)
+    # local coord never has a meaningful cos_travel (no reconstructed post_pos)
+    assert "cos_travel" not in [c.rsplit("_", 1)[-1] for c in ["pred_cos_travel"] if c in p.columns] or True
+    assert "pred_cos_travel" not in p.columns
+
+
+def test_prediction_secondaries_and_pairing():
+    lf = _global_prediction_frame()
+    true_sec = prediction_secondaries(lf, "true").collect()
+    pred_sec = prediction_secondaries(lf, "pred").collect()
+    assert true_sec["pdg"].to_list() == [22, 22, 11]
+    assert pred_sec["pdg"].to_list() == [22, 22]
+
+    pairs = paired_secondaries(lf).collect()
+    # event 1: 1 true, 1 pred -> paired (22, 22); event 2: 2 true, 1 pred -> paired rank0 only (22, 22)
+    assert pairs["true_pdg"].to_list() == [22, 22]
+    assert pairs["pred_pdg"].to_list() == [22, 22]
+
+
+def test_hist2d_basic():
+    lf = pl.DataFrame({"x": [0.1, 0.5, 0.9, 0.5], "y": [0.1, 0.9, 0.9, 0.1]}).lazy()
+    edges = np.linspace(0.0, 1.0, 3)  # 2 bins: [0,0.5), [0.5,1]
+    mat = hist2d(lf, pl.col("x"), pl.col("y"), edges, edges)
+    assert mat.sum() == 4
+    assert mat.shape == (2, 2)
+
+
+def _ctx_with_predictions(n_marginal_bins: int = 10) -> Context:
+    return build_context(
+        [RolloutSpec("rollout", _rollout_frame())],
+        _reference_frame(),
+        predictions=[PredictionSpec("pred", _global_prediction_frame())],
+        n_energy_bins=2,
+        n_marginal_bins=n_marginal_bins,
+        top_k_pdg=3,
+        sample_rows=1000,
+    )
+
+
+def test_build_context_resolves_prediction_ranges():
+    ctx = _ctx_with_predictions()
+    assert "edep" in ctx.pred_var_ranges
+    assert "edep" in ctx.pred_residual_ranges
+    assert ctx.pred_top_sec_pdgs  # secondaries present in the fixture
+
+
+def test_prediction_specs_compute_valid_reduced():
+    ctx = _ctx_with_predictions()
+    bundle = Bundle.open(
+        [RolloutSpec("rollout", _rollout_frame())],
+        _reference_frame(),
+        ctx,
+        predictions=[PredictionSpec("pred", _global_prediction_frame())],
+    )
+    for spec_id in (
+        "pred_marginal_edep",
+        "pred_scatter_edep",
+        "pred_residual_edep",
+        "pred_relative_residual_edep",
+        "pred_residual_profile_edep",
+        "pred_ks_summary",
+        "pred_bias_summary",
+        "pred_rmse_summary",
+        "pred_n_sec_confusion",
+        "pred_sec_species_confusion",
+        "pred_dir_alignment_post",
+        "pred_dir_alignment_travel",
+        "pred_constraint_violations",
+        "pred_correlation_delta",
+    ):
+        spec = get_spec(spec_id)
+        r = spec.finalize([spec.compute_partial(bundle)], ctx)
+        assert r.id == spec_id
+        assert r.kind != "unavailable", f"{spec_id} unexpectedly unavailable"
+        assert "pred" in r.payload["series"]
+
+
+def test_prediction_specs_unavailable_without_predictions():
+    ctx = _ctx_with_predictions()
+    bundle = Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)
+    for spec_id in ("pred_marginal_edep", "pred_scatter_edep", "pred_n_sec_confusion", "pred_ks_summary"):
+        spec = get_spec(spec_id)
+        r = spec.finalize([spec.compute_partial(bundle)], ctx)
+        assert r.kind == "unavailable"
+        assert r.payload["note"]
+
+
+@pytest.mark.parametrize(
+    "spec_id",
+    ["pred_marginal_edep", "pred_scatter_edep", "pred_n_sec_confusion", "pred_ks_summary", "pred_correlation_delta"],
+)
+def test_prediction_chunked_matches_unchunked(spec_id: str):
+    ctx = _ctx_with_predictions()
+    specs = [RolloutSpec("rollout", _rollout_frame())]
+    preds = [PredictionSpec("pred", _global_prediction_frame())]
+    spec = get_spec(spec_id)
+
+    unchunked_bundle = Bundle.open(specs, _reference_frame(), ctx, predictions=preds)
+    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)
+
+    n_chunks = 2
+    parts = [
+        spec.compute_partial(Bundle.open(specs, _reference_frame(), ctx, chunk=(k, n_chunks), predictions=preds))
+        for k in range(n_chunks)
+    ]
+    chunked = spec.finalize(parts, ctx)
+
+    assert chunked.kind == unchunked.kind
+    _assert_close(unchunked.payload, chunked.payload)
+
+
+def _assert_close(a, b) -> None:
+    """Recursively compare two JSON-shaped payloads (float-tolerant)."""
+    if isinstance(a, dict):
+        assert set(a) == set(b)
+        for k in a:
+            _assert_close(a[k], b[k])
+    elif isinstance(a, list):
+        assert len(a) == len(b)
+        for x, y in zip(a, b):
+            _assert_close(x, y)
+    elif isinstance(a, float):
+        assert np.isclose(a, b, atol=1e-9) or (np.isnan(a) and np.isnan(b))
+    else:
+        assert a == b
+
+
+def test_paired_scalars_are_subset_of_all_vars():
+    assert set(PAIRED_SCALARS) <= {"step_length", "edep", "delta_e", "post_E"}
@@ -184,7 +184,7 @@ def test_weighted_profile_matches_manual_bincount():
    ea = R.entry_axis(lf)
    lf2 = R.attach_entry_axis(lf, ea)
    edges = np.linspace(0.0, 3.0, 4)  # depth bins along +z
-    mean, std = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
+    mean, _ = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
    assert mean.shape == (3,)
    # totals conserved: sum over bins == mean total edep per event
    assert np.isclose(mean.sum() * 1, (90.0 + 30.0) / 2)  # 2 events
@@ -197,7 +197,7 @@ def test_update_manifest_reports_missing_targets(tmp_path):
    # schema2 dir exists but the parquet file does not
    (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)

-    lines, missing = plan_update_manifest(manifest, "schema2")
+    _, missing = plan_update_manifest(manifest, "schema2")
    assert len(missing) == 1
    assert "schema2" in str(missing[0])

@@ -313,7 +313,7 @@ def test_create_manifest_writes_relative_paths(tmp_path):
 def test_create_manifest_reports_missing_files(tmp_path):
    ghost = tmp_path / "processed" / "gen1" / "schema2" / "shard-000.parquet"
    output = tmp_path / "pools" / "full.manifest"
-    lines, missing, _ = plan_create_manifest(output, [ghost])
+    _, missing, _ = plan_create_manifest(output, [ghost])
    assert len(missing) == 1
    assert missing[0] == ghost.resolve()

@@ -6,6 +6,7 @@ from typer.testing import CliRunner

 from giant.cli import (
    _CEPH_PREDICTIONS,
+    _build_predict_timing,
    _resolve_prediction_output,
    _write_prediction_ref,
    app,
@@ -103,6 +104,7 @@ def test_ref_yaml_contains_expected_fields(tmp_path):
    ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset)
    data = yaml.safe_load(ref_path.read_text())

+    assert data["kind"] == "prediction"
    assert data["prediction_id"] == pred_uuid
    assert data["output"] == str(out)
    assert data["dataset"] == str(dataset)
@@ -144,6 +146,46 @@ def test_ref_timestamp_is_iso_format(tmp_path):
    assert ts.tzinfo is not None


+def test_ref_yaml_merges_extra_after_base_fields(tmp_path):
+    ckpt_dir = tmp_path / "checkpoints"
+    ckpt_dir.mkdir()
+    checkpoint = ckpt_dir / "best.pt"
+    checkpoint.touch()
+
+    out = tmp_path / "pred.parquet"
+    dataset = tmp_path / "full.manifest"
+    pred_uuid = str(uuid.uuid4())
+
+    ref_path = _write_prediction_ref(
+        checkpoint,
+        pred_uuid,
+        out,
+        dataset,
+        extra={"coord": "global", "n_rows": 42, "timing": {"setup_s": 1.0}},
+    )
+    data = yaml.safe_load(ref_path.read_text())
+
+    # Base fields untouched, extras layered on top.
+    assert data["kind"] == "prediction"
+    assert data["prediction_id"] == pred_uuid
+    assert data["coord"] == "global"
+    assert data["n_rows"] == 42
+    assert data["timing"] == {"setup_s": 1.0}
+
+
+def test_ref_yaml_without_extra_matches_today(tmp_path):
+    ckpt_dir = tmp_path / "checkpoints"
+    ckpt_dir.mkdir()
+    checkpoint = ckpt_dir / "best.pt"
+    checkpoint.touch()
+
+    pred_uuid = str(uuid.uuid4())
+    ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
+    data = yaml.safe_load(ref_path.read_text())
+
+    assert set(data) == {"kind", "prediction_id", "output", "dataset", "checkpoint", "timestamp"}
+
+
 def test_ref_checkpoint_path_is_absolute(tmp_path):
    ckpt_dir = tmp_path / "checkpoints"
    ckpt_dir.mkdir()
@@ -212,3 +254,72 @@ def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):

    assert result.exit_code == 1
    assert "not an inference-safe override" in result.output
+
+
+# ---------------------------------------------------------------------------
+# schema v3 constants (truth-column tagging)
+# ---------------------------------------------------------------------------
+
+
+def test_predict_schema_version_is_v3():
+    from giant.constants import PREDICT_SCHEMA_VERSION
+
+    assert PREDICT_SCHEMA_VERSION == "3"
+
+
+def test_predict_truth_metadata_key_exists():
+    from giant.constants import PREDICT_TRUTH_METADATA_KEY
+
+    assert PREDICT_TRUTH_METADATA_KEY == "giant.predict.has_truth"
+
+
+# ---------------------------------------------------------------------------
+# _build_predict_timing
+# ---------------------------------------------------------------------------
+
+
+def test_build_predict_timing_computes_per_step_cost():
+    timing = _build_predict_timing(
+        setup_s=1.0,
+        predict_s=10.0,
+        write_s=2.0,
+        n_rows=100,
+        device="cpu",
+        torch_threads=4,
+    )
+    assert timing["n_rows"] == 100
+    assert timing["sample_s"] == 8.0  # predict_s - write_s
+    assert timing["us_per_step"] == 8.0 / 100 * 1e6
+    assert timing["write_us_per_step"] == 2.0 / 100 * 1e6
+    assert timing["rows_per_s"] == 10.0
+    assert timing["device"] == "cpu" and timing["torch_threads"] == 4
+
+
+def test_build_predict_timing_handles_zero_rows():
+    timing = _build_predict_timing(
+        setup_s=1.0,
+        predict_s=0.0,
+        write_s=0.0,
+        n_rows=0,
+        device="cpu",
+        torch_threads=1,
+    )
+    assert timing["us_per_step"] is None
+    assert timing["write_us_per_step"] is None
+    assert timing["rows_per_s"] is None
+
+
+def test_predict_has_truth_flag_default_on():
+    # Inspecting rendered --help text is brittle across terminal
+    # widths/color settings (wraps or re-colors mid-flag); go straight to
+    # the underlying click command's registered option instead.
+    from typing import cast
+
+    import typer
+    from click import Group
+
+    predict_cmd = cast(Group, typer.main.get_command(app)).commands["predict"]
+    truth_param = next(p for p in predict_cmd.params if p.name == "truth")
+    assert truth_param.opts == ["--truth"]
+    assert truth_param.secondary_opts == ["--no-truth"]
+    assert truth_param.default is True
@@ -9,7 +9,7 @@ from pathlib import Path

 from typer.testing import CliRunner

-import giant.cli as cli
+from giant import cli

 runner = CliRunner()

@@ -1,4 +1,5 @@
 import pytest
+
 from giant.cond_layout import AXIS_TYPES, CondLayout
 from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM

@@ -16,6 +16,8 @@ from giant.analysis import (
    compute_one,
    compute_reduced,
    derive_run_dir,
+    load_prediction_yaml,
+    load_prediction_yamls,
    load_rollout_yaml,
    load_rollout_yamls,
    merge_one,
@@ -25,7 +27,8 @@ from giant.analysis import (
 from giant.analysis.catalog import get_spec
 from giant.analysis.condor import Context
 from giant.analysis.reduced import Partial, Reduced
-from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
+from giant.constants import PREDICT_COORD_METADATA_KEY, PREDICT_TRUTH_METADATA_KEY, ROLLOUT_COORD_VALUE
+from tests.test_analysis_prediction import _global_prediction_frame
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame


@@ -86,6 +89,30 @@ def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
    return paths[0], paths[1]


+def _write_prediction(path: Path, coord: str = "global") -> None:
+    tbl = _global_prediction_frame().collect().to_arrow()
+    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: coord, PREDICT_TRUTH_METADATA_KEY: "1"})
+    pq.write_table(tbl, path)
+
+
+def _write_prediction_yaml(tmp_path: Path, reference: Path, tag: str = "p", coord: str = "global") -> Path:
+    pred = tmp_path / f"pred_{tag}.parquet"
+    _write_prediction(pred, coord=coord)
+    yaml_path = tmp_path / f"pred_{tag}.yaml"
+    yaml_path.write_text(
+        yaml.safe_dump(
+            {
+                "kind": "prediction",
+                "prediction_id": f"{tag}pred1234",
+                "output": str(pred),
+                "dataset": str(reference),
+                "checkpoint": f"/ckpt/{tag}.pt",
+            }
+        )
+    )
+    return yaml_path
+
+
 def _fake_venv(repo_dir: Path) -> None:
    """Stand in for a `uv sync`'d venv: write_submit checks `.venv/bin/giant` exists."""
    giant = repo_dir / ".venv" / "bin" / "giant"
@@ -94,13 +121,14 @@ def _fake_venv(repo_dir: Path) -> None:
    giant.chmod(0o755)


-def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
+def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None, prediction_yamls=()) -> Path:
    """``prep`` with small test-sized context bins/sampling."""
    return prep(
        rollout_yamls,
        run_dir,
        n_chunks=chunks,
        labels=labels,
+        prediction_yamls=prediction_yamls,
        n_energy_bins=2,
        n_marginal_bins=8,
        top_k_pdg=3,
@@ -161,6 +189,113 @@ def test_load_rollout_yamls_rejects_mismatched_reference(tmp_path: Path):
        load_rollout_yamls([a, c])


+def test_load_prediction_yaml_requires_paths(tmp_path: Path):
+    bad = tmp_path / "bad.yaml"
+    bad.write_text(yaml.safe_dump({"output": "x.parquet"}))  # no dataset
+    with pytest.raises(ValueError):
+        load_prediction_yaml(bad)
+
+
+def test_load_prediction_yaml_rejects_rollout_kind(tmp_path: Path):
+    y = tmp_path / "r.yaml"
+    y.write_text(yaml.safe_dump({"output": "x.parquet", "dataset": "d.parquet", "kind": "rollout"}))
+    with pytest.raises(ValueError, match="kind"):
+        load_prediction_yaml(y)
+
+
+def test_load_prediction_yamls_single_defaults_to_prediction_name(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    y = _write_prediction_yaml(tmp_path, reference)
+    loaded = load_prediction_yamls([y], str(reference))
+    assert [lp.name for lp in loaded] == ["prediction"]
+    assert loaded[0].coord == "global"
+
+
+def test_load_prediction_yamls_multi_defaults_to_stem_and_labels(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    a = _write_prediction_yaml(tmp_path, reference, tag="a")
+    b = _write_prediction_yaml(tmp_path, reference, tag="b")
+    loaded = load_prediction_yamls([a, b], str(reference))
+    assert [lp.name for lp in loaded] == ["pred_a", "pred_b"]
+    loaded = load_prediction_yamls([a, b], str(reference), labels=["ep20", "ep50"])
+    assert [lp.name for lp in loaded] == ["ep20", "ep50"]
+
+
+def test_load_prediction_yamls_rejects_mismatched_reference(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    other_ref = tmp_path / "other_reference.parquet"
+    _reference_frame().collect().write_parquet(other_ref)
+    y = _write_prediction_yaml(tmp_path, other_ref)
+    with pytest.raises(ValueError, match="same reference"):
+        load_prediction_yamls([y], str(reference))
+
+
+def test_load_prediction_yamls_rejects_mixed_coord(tmp_path: Path):
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+    a = _write_prediction_yaml(tmp_path, reference, tag="a", coord="global")
+    b = _write_prediction_yaml(tmp_path, reference, tag="b", coord="local")
+    with pytest.raises(ValueError, match="coord"):
+        load_prediction_yamls([a, b], str(reference))
+
+
+def test_prep_with_prediction_writes_run_meta(tmp_path: Path):
+    rollout_yaml = _write_inputs(tmp_path)
+    reference = load_rollout_yaml(rollout_yaml)["dataset"]
+    pred_yaml = _write_prediction_yaml(tmp_path, Path(reference))
+    run_dir = _prep([rollout_yaml], prediction_yamls=[pred_yaml])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert [p["name"] for p in meta.predictions] == ["prediction"]
+    assert meta.predictions[0]["plot_meta"]["checkpoint"] == "/ckpt/p.pt"
+
+    computed = compute_one("pred_marginal_edep", run_dir, chunk_index=0)
+    partial = Partial.load(computed)
+    assert partial.data["available"]
+
+
+def test_prep_forwards_predict_only_metadata_keys(tmp_path: Path):
+    """A rich `giant predict` sidecar's provenance/timing keys reach
+    run_meta.json's plot_meta, same as a rollout's do — a thin legacy
+    sidecar (no such keys) still loads fine (see _write_prediction_yaml)."""
+    rollout_yaml = _write_inputs(tmp_path)
+    reference = load_rollout_yaml(rollout_yaml)["dataset"]
+    pred = tmp_path / "pred_rich.parquet"
+    _write_prediction(pred, coord="global")
+    yaml_path = tmp_path / "pred_rich.yaml"
+    yaml_path.write_text(
+        yaml.safe_dump(
+            {
+                "kind": "prediction",
+                "prediction_id": "richpred12",
+                "output": str(pred),
+                "dataset": str(reference),
+                "checkpoint": "/ckpt/rich.pt",
+                "coord": "global",
+                "has_truth": True,
+                "schema_version": "3",
+                "n_input_rows": 1000,
+                "n_files": 1,
+                "n_skipped_rows": 3,
+                "unknown_pdg_counts": {"999999": 3},
+                "batch_size_auto": False,
+                "timing": {"us_per_step": 12.5},
+            }
+        )
+    )
+    run_dir = _prep([rollout_yaml], prediction_yamls=[yaml_path])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    plot_meta = meta.predictions[0]["plot_meta"]
+    assert plot_meta["coord"] == "global"
+    assert plot_meta["has_truth"] is True
+    assert plot_meta["n_input_rows"] == 1000
+    assert plot_meta["n_skipped_rows"] == 3
+    assert plot_meta["unknown_pdg_counts"] == {"999999": 3}
+    assert plot_meta["timing"] == {"us_per_step": 12.5}
+
+
 def test_derive_run_dir_next_to_rollout():
    y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
    assert derive_run_dir([y]) == Path("/data/analysis_abcd1234")
@@ -576,7 +576,7 @@ def test_save_config_round_trips_three_level_nesting(tmp_path):
 # default_out_dir_name
 # ---------------------------------------------------------------------------

-_NOW = datetime(2026, 7, 29, 14, 30)
+_NOW = datetime(2026, 7, 29, 14, 30)  # noqa: DTZ001 - naive, matching default_out_dir_name's naive datetime.now()


 def _cfg_with(**dotted_overrides):
@@ -90,7 +90,7 @@ def _collect_names(source: str, filename: str) -> set[str]:
            names.add(node.attr)
        elif isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstring_ids:
            names.add(node.value)
-        elif isinstance(node, ast.arg):
+        elif isinstance(node, ast.arg):  # noqa: SIM114 - kept separate so ty narrows node.arg to str, not str | None
            names.add(node.arg)
        elif isinstance(node, ast.keyword) and node.arg is not None:
            names.add(node.arg)
@@ -1,3 +1,4 @@
+from test_pipeline import _make_synthetic_steps
 from typer.testing import CliRunner

 from giant import cli as giant_cli
@@ -5,7 +6,6 @@ from giant.config import Conditioning
 from giant.data import setup_cache
 from giant.tools import dwarf
 from giant.tools.dwarf import app
-from test_pipeline import _make_synthetic_steps

 runner = CliRunner()

@@ -1,9 +1,10 @@
 import torch
+
 from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM
 from giant.model.network import Stage1Model
 from giant.model.schedule import CosineSchedule, flow_matching_loss
-from giant.sample import sample_flow, sample_ddim
+from giant.sample import sample_ddim, sample_flow

 PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
 MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
@@ -2,8 +2,9 @@ import copy

 import pytest
 import torch
+
 from giant import config as gconfig
-from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
+from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
    HISTORY_REGISTRY,
    AttentionHistory,
@@ -1217,18 +1218,18 @@ def test_stage_classes_are_stagemodel_subclasses(cls):
@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive])
@pytest.mark.parametrize("generator", ["flow", "ddpm", "wgan"])
 def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator):
-    kwargs = dict(
-        pdg_vocab=5,
-        mat_vocab=3,
-        particle_cfg=PARTICLE_CFG,
-        material_cfg=MATERIAL_CFG,
-        hidden_dim=_STAGE_HIDDEN_DIM,
-        n_res_blocks=_STAGE_N_BLOCKS,
-        cond_out_dim=_STAGE_COND_OUT_DIM,
-        generator=generator,
-        time_dim=8,
-        noise_dim=8,
-    )
+    kwargs = {
+        "pdg_vocab": 5,
+        "mat_vocab": 3,
+        "particle_cfg": PARTICLE_CFG,
+        "material_cfg": MATERIAL_CFG,
+        "hidden_dim": _STAGE_HIDDEN_DIM,
+        "n_res_blocks": _STAGE_N_BLOCKS,
+        "cond_out_dim": _STAGE_COND_OUT_DIM,
+        "generator": generator,
+        "time_dim": 8,
+        "noise_dim": 8,
+    }
    if cls is Stage1Model:
        kwargs["n_sec_head_k_max"] = 15
    else:
@@ -20,7 +20,6 @@ from giant.model.schedule import (
 )
 from giant.sample import sample_secondaries

-
 # ── helpers ──────────────────────────────────────────────────────────────────


@@ -341,7 +340,7 @@ def test_encode_secondaries_energy_conservation():
 def test_encode_secondaries_stick_logits_match_naive_reference():
    """Cumsum-based remaining-budget computation must match a naive
    per-row, per-slot Python reference (no cumsum) within float tolerance."""
-    from giant.data.transforms import encode_secondaries, _EPS, _STICK_LOGIT_CLIP
+    from giant.data.transforms import _EPS, _STICK_LOGIT_CLIP, encode_secondaries

    rng = np.random.default_rng(11)
    N = 25
@@ -569,7 +568,7 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
    e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)

-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+    sec_E, _sec_dir, _mass, _charge, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)

    for i, k in enumerate(n_sec):
        if k == 0:
@@ -593,7 +592,7 @@ def test_decode_secondaries_rescale_preserves_relative_shares():
    n_sec = np.array([4])
    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

-    sec_E_small, _, _, _, sec_valid = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
+    sec_E_small, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
    sec_E_large, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir)

    ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
@@ -9,8 +9,8 @@ import pytest

 pytest.importorskip("plotstyle")

-from giant.analysis import render as render_mod  # noqa: E402
-from giant.analysis.reduced import Reduced  # noqa: E402
+from giant.analysis import render as render_mod
+from giant.analysis.reduced import Reduced


 def _try_render(reduced: list[Reduced], out: Path):
@@ -329,6 +329,58 @@ def test_render_one_of_each_kind(tmp_path: Path):
                "log_color": True,
            },
        ),
+        Reduced(
+            "ph1",
+            "prediction",
+            "paired_hist",
+            "Paired hist (single prediction)",
+            "x",
+            {"edges": [0, 1, 2, 3], "series": {"pred": {"pred": [1, 2, 3], "true": [2, 2, 2]}}, "log_y": False},
+        ),
+        Reduced(
+            "ph2",
+            "prediction",
+            "paired_hist",
+            "Paired hist (two predictions)",
+            "x",
+            {
+                "edges": [0, 1, 2, 3],
+                "series": {"a": {"pred": [1, 2, 3], "true": [2, 2, 2]}, "b": {"pred": [3, 2, 1]}},
+                "log_y": False,
+            },
+        ),
+        Reduced(
+            "hm2d",
+            "prediction",
+            "heatmap2d",
+            "Scatter (truth vs pred)",
+            "true x",
+            {
+                "x_edges": [0, 1, 2],
+                "y_edges": [0, 1, 2],
+                "series": {"pred": [[2, 0], [1, 3]]},
+                "ylabel": "predicted x",
+                "cbar_label": "count",
+                "log_color": True,
+                "diagonal": True,
+            },
+        ),
+        Reduced(
+            "profile_noref",
+            "prediction",
+            "profile",
+            "Residual profile (no reference)",
+            "true x",
+            {"edges": [0, 1, 2], "series": {"pred": {"mean": [0.1, -0.1], "std": [0.2, 0.2]}}},
+        ),
+        Reduced(
+            "bar_noref",
+            "prediction",
+            "bar",
+            "Constraint violations (no reference)",
+            "check",
+            {"labels": ["a", "b"], "series": {"pred": [0.01, 0.0]}, "ylabel": "rate"},
+        ),
    ]
    try:
        pdfs = _try_render(reduced, tmp_path)
@@ -9,7 +9,7 @@ import pytest
 import torch

 from giant.config import ConditioningAxisConfig, ParticleTypeConfig
-from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX
+from giant.constants import K_MAX, TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG
 from giant.data.loader import TopNMap
 from giant.data.transforms import Normalizer
 from giant.model.network import (
@@ -21,7 +21,7 @@ from giant.model.network import (
 from giant.rollout import L1DistCollector, make_seed_frontier, rollout

 pytest.importorskip("sklearn")
-from giant import geometry as g  # noqa: E402
+from giant import geometry as g

 PDG_MAP = {22: 0, 11: 1, -11: 2}
 MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1}
@@ -1,5 +1,7 @@
 """Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""

+import itertools
+
 import pytest
 import torch

@@ -7,6 +9,7 @@ from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 from giant.model.network import (
    BLOCK_REGISTRY,
+    ROUTER_REGISTRY,
    TRUNK_REGISTRY,
    AdaLNResBlock,
    ComposedRouter,
@@ -17,7 +20,6 @@ from giant.model.network import (
    NoneRouter,
    PdgRouter,
    ProcessRouter,
-    ROUTER_REGISTRY,
    ResBlock,
    RoutedTrunk,
    Stage1Model,
@@ -397,7 +399,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
            router.raw_width[0] = raw
            shares.append(router.gate(cond_cont, cond_cat)[0, 0].item())

-    assert all(a <= b + 1e-6 for a, b in zip(shares, shares[1:]))
+    assert all(a <= b + 1e-6 for a, b in itertools.pairwise(shares))


 def test_build_router_threads_learn_width_kwargs_through():
@@ -1010,9 +1012,7 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
    assert n_sec_pred.shape == (B,)

-    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
-        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
-    )
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, K_MAX, 4)
    assert sec_valid.shape == (B, K_MAX)

@@ -1261,9 +1261,7 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow():
    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
    assert n_sec_pred.shape == (B,)

-    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
-        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
-    )
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, K_MAX, 4)
    assert sec_valid.shape == (B, K_MAX)

@@ -262,7 +262,7 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 1, 1])
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
+    sec_cont, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
    assert sec_valid.tolist() == [[False], [True], [True]]

@@ -281,7 +281,7 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_s
    _force_stop_head_logit(decoder, 50.0)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
    assert sec_valid.shape == (B, k_max)
    assert not sec_valid.any()

@@ -296,7 +296,7 @@ def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_
    _force_stop_head_logit(decoder, -50.0)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
    assert sec_valid.all()


@@ -341,6 +341,87 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
        sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)


+# ── Row compaction: full_length=False (default, inference) must agree with
+# full_length=True (the pre-compaction behaviour, still exercised by
+# _assemble_stage2_ar_inputs_scheduled's training-time self-sample) ────────
+
+
+def _zero_randn(*size, **kwargs):
+    """Drop-in replacement for `torch.randn` that returns zeros of the same
+    shape — makes the ODE/WGAN noise deterministic so a compacted run and a
+    full_length run can be compared row-for-row regardless of how many
+    `torch.randn` calls each makes (compaction changes the batch size, and
+    therefore the RNG stream position, at every slot)."""
+    device = kwargs.get("device")
+    dtype = kwargs.get("dtype")
+    return torch.zeros(*size, device=device, dtype=dtype)
+
+
+@pytest.mark.parametrize("history", ["markov", "attention"])
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_compaction_matches_full_length_head_mode(generator, history, monkeypatch):
+    B, k_max, emb_dim = 4, 5, 6
+    decoder = _stage2_ar("physical", generator, emb_dim=emb_dim, k_max=k_max, history=history)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 1, 3, k_max])
+
+    monkeypatch.setattr(torch, "randn", _zero_randn)
+    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=False
+    )
+    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
+    )
+
+    assert torch.equal(sec_valid_c, sec_valid_f)
+    assert torch.equal(sec_valid_c, torch.arange(k_max).unsqueeze(0) < n_sec_pred.unsqueeze(1))
+    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_compaction_matches_full_length_stop_token(generator, monkeypatch):
+    """`n_sec_sampling="greedy"` keeps the stop decision itself deterministic
+    (no `torch.rand` draw), so only `torch.randn` needs zeroing."""
+    B, k_max = 6, 5
+    decoder = _stage2_ar_stop_token("physical", generator, n_sec_sampling="greedy", k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+
+    monkeypatch.setattr(torch, "randn", _zero_randn)
+    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=False
+    )
+    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=True
+    )
+
+    assert torch.equal(sec_valid_c, sec_valid_f)
+    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+
+
+def test_sample_secondaries_ar_full_length_ignores_n_sec_pred_zero_rows():
+    """A row with n_sec_pred == 0 would be dropped from the active set at
+    slot 0 under compaction (full_length=False) — full_length=True must
+    still run the model for it at every slot (only masked by sec_valid at
+    the end), matching _assemble_stage2_ar_inputs_scheduled's contract."""
+    B, k_max, emb_dim = 3, 4, 6
+    decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 0, 0])
+    sec_cont, _, sec_valid = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
+    )
+    assert not sec_valid.any()
+    # every slot still ran the model (not left at the zero-init default) —
+    # a real flow ODE output from randn-initialized noise is essentially
+    # never exactly zero.
+    assert not torch.allclose(sec_cont, torch.zeros_like(sec_cont))
+
+
 # ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────


@@ -12,6 +12,7 @@ import pytest
 import torch
 import torch.nn.functional as F

+from giant.checkpoint_io import load_for_inference
 from giant.config import ParticleTypeConfig
 from giant.constants import (
    COND_DIM,
@@ -21,7 +22,6 @@ from giant.constants import (
    SEC_SLOT_DIM,
    X_DIM,
 )
-from giant.checkpoint_io import load_for_inference
 from giant.data.dataset import StepBatch
 from giant.data.transforms import Normalizer
 from giant.model.network import Stage2Autoregressive, build_critics, build_models
@@ -36,7 +36,6 @@ from giant.training import (
    train,
 )
 from giant.training.metrics import _wandb_run_config
-from giant.training.trainers import _type_class_weight_vector
 from giant.training.stage2_inputs import (
    _ar_has_prev,
    _assemble_stage2_ar_inputs,
@@ -51,6 +50,7 @@ from giant.training.stage2_inputs import (
    _stop_target_and_mask,
    _type_repr,
 )
+from giant.training.trainers import _type_class_weight_vector

 PDG_VOCAB = 6
 MAT_VOCAB = 3
@@ -543,19 +543,18 @@ def test_train_raises_when_no_active_stage():
    model_config = _model_config(cfg)
    models = build_models(model_config)
    critics = build_critics(model_config)
-    with tempfile.TemporaryDirectory() as tmp:
-        with pytest.raises(ValueError, match="no active stage"):
-            train(
-                cfg=cfg,
-                models=models,
-                critics=critics,
-                train_loader=_fake_batches(1, 8),
-                val_loader=_fake_batches(1, 8),
-                device=torch.device("cpu"),
-                out_dir=Path(tmp) / "run",
-                model_config=model_config,
-                total_train_batches=1,
-            )
+    with tempfile.TemporaryDirectory() as tmp, pytest.raises(ValueError, match="no active stage"):
+        train(
+            cfg=cfg,
+            models=models,
+            critics=critics,
+            train_loader=_fake_batches(1, 8),
+            val_loader=_fake_batches(1, 8),
+            device=torch.device("cpu"),
+            out_dir=Path(tmp) / "run",
+            model_config=model_config,
+            total_train_batches=1,
+        )


 def test_metrics_csv_columns_are_stage_prefixed():
@@ -11,8 +11,8 @@ import pytest

 pytest.importorskip("plotstyle")

-from giant.training import plots as plots_mod  # noqa: E402
-from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics  # noqa: E402
+from giant.training import plots as plots_mod
+from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics

 # --- fixtures ----------------------------------------------------------

@@ -2,9 +2,13 @@ import warnings

 import numpy as np
 import pytest
+
 from giant.cond_layout import AXIS_TYPES, CondLayout
 from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX
 from giant.data.transforms import (
+    Normalizer,
+    _vectorized_map_lookup,
+    _WelfordAccumulator,
    build_cond_features,
    build_features,
    encode_secondaries,
@@ -14,12 +18,9 @@ from giant.data.transforms import (
    inv_log_transform,
    local_frame_rotation,
    log_transform,
-    Normalizer,
    reconstruct_post_pos,
    sorted_membership,
    travel_direction,
-    _vectorized_map_lookup,
-    _WelfordAccumulator,
 )