chore: release v0.3.23

Reviewed-on: #99

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.19"
+current_version = "0.3.23"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

.gitea/scripts/check_dep_bounds.py

@@ -0,0 +1,184 @@
+#!/usr/bin/env python3
+"""Check (and optionally raise) pyproject.toml's dependency upper bounds
+against the latest release on PyPI.
+
+Used by the monthly deps-bounds workflow (see deps-bounds-pr.sh): a lockfile
+refresh (deps-lock.yml, weekly) only ever moves within the existing
+`>=..,<..` constraints, so a stale upper bound never gets touched by that
+job. This script closes that gap by proposing (or applying) a raised
+ceiling, always as its own reviewable PR — a deliberately excluded package
+(torch, pinned to a driver-compatible range; plotstyle, not on PyPI at all)
+never gets edited.
+
+Standalone stdlib + `packaging` script (no project deps needed to run it):
+    uv run --no-project --with packaging python check_dep_bounds.py --report -
+    uv run --no-project --with packaging python check_dep_bounds.py --apply --report bounds.md
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import re
+import sys
+import tomllib
+import urllib.request
+from dataclasses import dataclass
+from pathlib import Path
+
+from packaging.requirements import Requirement
+from packaging.version import Version
+
+ROOT = Path(__file__).resolve().parents[2]
+PYPROJECT = ROOT / "pyproject.toml"
+
+# Deliberately never touched by this script, with the reason surfaced in the
+# report footer.
+SKIP_REASONS = {
+    "giant": "self-reference (extras-of-extras), not a real upper bound to check",
+    "torch": "pinned <2.4 deliberately — newer torch needs newer NVIDIA drivers"
+    " than the shared portal machines have (see CLAUDE.md)",
+    "plotstyle": "served from the private `larsbogner` index, not PyPI",
+}
+
+
+@dataclass
+class Finding:
+    name: str
+    old_specifier: str
+    new_specifier: str
+    latest: str
+
+
+def iter_requirement_strings(pyproject: dict) -> list[str]:
+    project = pyproject["project"]
+    reqs = list(project.get("dependencies", []))
+    for group_reqs in project.get("optional-dependencies", {}).values():
+        reqs.extend(group_reqs)
+    return reqs
+
+
+def canonical_requirements(pyproject: dict) -> dict[str, Requirement]:
+    """One Requirement per distinct package name (specifiers are expected to
+    agree across groups — that's true today; a future mismatch would just
+    mean the last-seen group's specifier gets checked, which is fine for a
+    monthly advisory script)."""
+    out: dict[str, Requirement] = {}
+    for req_str in iter_requirement_strings(pyproject):
+        req = Requirement(req_str)
+        out[req.name] = req
+    return out
+
+
+def fetch_latest_version(name: str) -> str | None:
+    url = f"https://pypi.org/pypi/{name}/json"
+    try:
+        with urllib.request.urlopen(url, timeout=10) as resp:
+            data = json.load(resp)
+        return data["info"]["version"]
+    except Exception as exc:  # noqa: BLE001 - network hiccup: degrade, don't fail the job
+        print(f"warning: could not fetch latest version for {name}: {exc}", file=sys.stderr)
+        return None
+
+
+def next_ceiling(latest: Version) -> str:
+    """The upper-bound scheme already used by this repo's own pins:
+    `<{major+1}` once a package is past 0.x (e.g. numpy>=1.26,<3), but
+    `<0.{minor+1}` while still pre-1.0 (e.g. ruff>=0.15,<1 -> next would be
+    <0.17 style if ruff were still 0.x; ty>=0.0.50,<0.1 is the same idea one
+    level deeper). Only the major (or, pre-1.0, the minor) component of the
+    latest release matters here — the point is "next breaking-change
+    boundary", not "exactly latest.patch + epsilon".
+    """
+    if latest.major == 0:
+        return f"<0.{latest.minor + 1}"
+    return f"<{latest.major + 1}"
+
+
+def find_findings(requirements: dict[str, Requirement], skip: dict[str, str]) -> list[Finding]:
+    findings = []
+    for name, req in sorted(requirements.items()):
+        if name in skip:
+            continue
+        latest_str = fetch_latest_version(name)
+        if latest_str is None:
+            continue
+        latest = Version(latest_str)
+        if latest in req.specifier:
+            continue
+        new_upper = next_ceiling(latest)
+        lower_clauses = [str(s) for s in req.specifier if s.operator != "<"]
+        new_specifier = ",".join([*lower_clauses, new_upper])
+        findings.append(
+            Finding(name=name, old_specifier=str(req.specifier), new_specifier=new_specifier, latest=latest_str)
+        )
+    return findings
+
+
+def apply_findings(text: str, findings: list[Finding]) -> str:
+    """Textual, in-place replacement of each finding's specifier substring —
+    deliberately not a TOML round-trip, so comments/formatting/key order in
+    pyproject.toml survive untouched. Every occurrence is replaced (a name
+    like polars appears in three separate dependency groups with identical
+    specifiers, and all of them must move together)."""
+    for finding in findings:
+        # Matched by package name rather than the old specifier string
+        # verbatim: packaging.requirements.Requirement's str(specifier)
+        # doesn't preserve clause order (e.g. "numpy>=1.26,<3" round-trips
+        # as "<3,>=1.26"), so an exact-string match on the old requirement
+        # would rarely hit. subn with no count replaces every occurrence in
+        # one pass, which is what a multi-group dependency (e.g. polars)
+        # needs.
+        new = f'"{finding.name}{finding.new_specifier}"'
+        # Negative lookahead guards against matching a longer package name
+        # sharing this one as a prefix (e.g. "numpy" must not match
+        # "numpydoc>=...").
+        pattern = re.compile(rf'"{re.escape(finding.name)}(?![\w.-])[^"]*"')
+        text, n = pattern.subn(new, text)
+        if n == 0:
+            print(f"warning: could not locate {finding.name!r} requirement string to rewrite", file=sys.stderr)
+    return text
+
+
+def render_report(findings: list[Finding], skip: dict[str, str]) -> str:
+    lines = []
+    if findings:
+        lines.append("| package | old constraint | new constraint | latest on PyPI |")
+        lines.append("| --- | --- | --- | --- |")
+        for f in findings:
+            lines.append(f"| {f.name} | `{f.old_specifier}` | `{f.new_specifier}` | {f.latest} |")
+    else:
+        lines.append("All checked dependency upper bounds already cover the latest PyPI release.")
+    lines.append("")
+    lines.append("Skipped (never auto-raised):")
+    for name, reason in sorted(skip.items()):
+        lines.append(f"- `{name}` — {reason}")
+    return "\n".join(lines) + "\n"
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--apply", action="store_true", help="rewrite pyproject.toml in place")
+    parser.add_argument("--report", default="-", help="write the markdown report here ('-' for stdout, default: -)")
+    args = parser.parse_args()
+
+    pyproject_text = PYPROJECT.read_text()
+    pyproject = tomllib.loads(pyproject_text)
+    requirements = canonical_requirements(pyproject)
+    findings = find_findings(requirements, SKIP_REASONS)
+
+    report = render_report(findings, SKIP_REASONS)
+    if args.report == "-":
+        print(report, end="")
+    else:
+        Path(args.report).write_text(report)
+
+    if args.apply and findings:
+        new_text = apply_findings(pyproject_text, findings)
+        PYPROJECT.write_text(new_text)
+
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

.gitea/scripts/deps-bounds-pr.sh

@@ -0,0 +1,44 @@
+#!/usr/bin/env bash
+# Monthly job: raise pyproject.toml upper bounds that have fallen behind the
+# latest PyPI release (see check_dep_bounds.py — torch/plotstyle/giant are
+# deliberately excluded there), re-lock, and open/update a PR with the
+# result. Separate from the weekly uv.lock-only refresh (deps-lock.yml)
+# because this one can legitimately break CI (a new major version), which
+# should never block the routine weekly lockfile bump.
+#
+# Preconditions: repo checked out on master with fetch-depth: 0, uv synced
+# (./.gitea/actions/setup), CI_TOKEN/GITHUB_* env set by Gitea Actions.
+set -euo pipefail
+
+BRANCH="chore/dep-bounds"
+
+git config user.name "gitea-actions"
+git config user.email "actions@git.larsbogner.de"
+git checkout -B "$BRANCH" origin/master
+
+BODY_FILE=$(mktemp)
+uv run --no-project --with packaging python .gitea/scripts/check_dep_bounds.py \
+  --apply --report "$BODY_FILE"
+
+if git diff --quiet -- pyproject.toml; then
+  echo "No upper bounds out of date; nothing to propose"
+  .gitea/scripts/deps-pr.sh close "$BRANCH"
+  exit 0
+fi
+
+# Re-lock so the PR carries a pyproject + uv.lock that agree.
+uv lock
+
+N_RAISED=$(grep -c '^| ' "$BODY_FILE" || true)
+# Subtract the header + separator row from the markdown table, if present.
+if [ "$N_RAISED" -ge 2 ]; then
+  N_RAISED=$((N_RAISED - 2))
+else
+  N_RAISED=0
+fi
+
+git add pyproject.toml uv.lock
+.gitea/scripts/deps-pr.sh open "$BRANCH" \
+  "chore(deps): raise dependency upper bounds" \
+  "chore(deps): raise dependency upper bounds ($N_RAISED packages)" \
+  "$BODY_FILE"

.gitea/scripts/deps-lock-pr.sh

@@ -0,0 +1,48 @@
+#!/usr/bin/env bash
+# Weekly job: refresh uv.lock against the pinned constraints in
+# pyproject.toml (no constraint edits — see check_dep_bounds.py for the
+# separate monthly job that raises upper bounds) and open/update a PR with
+# the result. See deps-pr.sh for the commit/push/PR-upsert mechanics.
+#
+# Preconditions: repo checked out on master with fetch-depth: 0, uv synced
+# (./.gitea/actions/setup), CI_TOKEN/GITHUB_* env set by Gitea Actions.
+set -euo pipefail
+
+BRANCH="chore/uv-lock-upgrade"
+
+git config user.name "gitea-actions"
+git config user.email "actions@git.larsbogner.de"
+git checkout -B "$BRANCH" origin/master
+
+UPDATES_FILE=$(mktemp)
+uv lock --upgrade 2>&1 | tee "$UPDATES_FILE"
+
+if git diff --quiet -- uv.lock; then
+  echo "uv.lock already up to date; nothing to propose"
+  .gitea/scripts/deps-pr.sh close "$BRANCH"
+  exit 0
+fi
+
+N_PACKAGES=$(grep -c '^Update ' "$UPDATES_FILE" || true)
+
+BODY_FILE=$(mktemp)
+{
+  if [ "$N_PACKAGES" -gt 0 ]; then
+    echo "Weekly automated \`uv lock --upgrade\` — updates within the existing"
+    echo "\`pyproject.toml\` constraints:"
+    echo
+    grep '^Update ' "$UPDATES_FILE" | sed 's/^/- /'
+  else
+    echo "Weekly automated \`uv lock --upgrade\` refreshed the lockfile (e.g. hash"
+    echo "or metadata changes) without a visible version bump."
+  fi
+  echo
+  echo "CI on this PR (lint/format/type-check/tests) is the gate; merge normally"
+  echo "once green, which triggers the usual patch release."
+} > "$BODY_FILE"
+
+git add uv.lock
+.gitea/scripts/deps-pr.sh open "$BRANCH" \
+  "chore(deps): weekly uv.lock refresh" \
+  "chore(deps): weekly uv.lock refresh ($N_PACKAGES packages)" \
+  "$BODY_FILE"

.gitea/scripts/deps-pr.sh

@@ -0,0 +1,84 @@
+#!/usr/bin/env bash
+# Shared commit/push/PR-upsert mechanics for the scheduled dependency-bump
+# workflows (deps-lock.yml, deps-bounds.yml). Both jobs stage some changes,
+# then delegate here to turn them into a standing pull request against
+# master — one stable branch per job, force-pushed every run, so the PR
+# stays a single commit and a single open proposal across weeks/months
+# instead of accumulating history or duplicate PRs.
+#
+# Usage:
+#   deps-pr.sh open <branch> <commit-subject> <pr-title> <body-file>
+#     Commit the currently staged changes, force-push <branch>, and
+#     create-or-update an open PR from <branch> onto master.
+#   deps-pr.sh close <branch>
+#     Close any open PR from <branch> onto master (if one exists) and
+#     delete the remote branch. Used when a run finds nothing to change.
+#
+# Preconditions: repo checked out with fetch-depth: 0, git user.name/email
+# already configured, GITHUB_SERVER_URL/GITHUB_REPOSITORY/CI_TOKEN set (all
+# provided by Gitea Actions), and — for "open" — the changes to publish are
+# already `git add`-ed.
+set -euo pipefail
+
+MODE="${1:?usage: deps-pr.sh open|close ...}"
+BRANCH="${2:?branch name required}"
+
+API="${GITHUB_SERVER_URL}/api/v1/repos/${GITHUB_REPOSITORY}"
+AUTH_HEADER="Authorization: token ${CI_TOKEN}"
+
+# Look up the currently open PR (if any) from $BRANCH onto master.
+find_open_pr() {
+  curl -sf -H "$AUTH_HEADER" "${API}/pulls?state=open&base=master" \
+    | jq -r --arg ref "$BRANCH" '.[] | select(.head.ref == $ref) | .number' \
+    | head -n1
+}
+
+case "$MODE" in
+  open)
+    SUBJECT="${3:?commit subject required}"
+    TITLE="${4:?PR title required}"
+    BODY_FILE="${5:?PR body file required}"
+
+    git commit -F - <<EOF
+$SUBJECT
+
+$(cat "$BODY_FILE")
+EOF
+    git push --force origin "HEAD:refs/heads/$BRANCH"
+
+    BODY_JSON=$(jq -Rs '.' < "$BODY_FILE")
+    PR_NUMBER=$(find_open_pr || true)
+    if [ -n "$PR_NUMBER" ]; then
+      echo "Updating existing PR #$PR_NUMBER from $BRANCH"
+      curl -sf -X PATCH -H "$AUTH_HEADER" -H "Content-Type: application/json" \
+        "${API}/pulls/${PR_NUMBER}" \
+        -d "{\"title\": $(jq -Rs '.' <<<"$TITLE"), \"body\": ${BODY_JSON}}" \
+        > /dev/null
+    else
+      echo "Opening new PR from $BRANCH"
+      curl -sf -X POST -H "$AUTH_HEADER" -H "Content-Type: application/json" \
+        "${API}/pulls" \
+        -d "{\"head\": \"${BRANCH}\", \"base\": \"master\", \"title\": $(jq -Rs '.' <<<"$TITLE"), \"body\": ${BODY_JSON}}" \
+        > /dev/null
+    fi
+    ;;
+
+  close)
+    PR_NUMBER=$(find_open_pr || true)
+    if [ -n "$PR_NUMBER" ]; then
+      echo "Closing PR #$PR_NUMBER from $BRANCH (nothing to update)"
+      curl -sf -X PATCH -H "$AUTH_HEADER" -H "Content-Type: application/json" \
+        "${API}/pulls/${PR_NUMBER}" -d '{"state": "closed"}' > /dev/null
+    else
+      echo "No open PR from $BRANCH to close"
+    fi
+    if git ls-remote --exit-code --heads origin "$BRANCH" > /dev/null 2>&1; then
+      git push origin --delete "$BRANCH" || echo "Could not delete remote branch $BRANCH (already gone?)"
+    fi
+    ;;
+
+  *)
+    echo "Unknown mode: $MODE (expected 'open' or 'close')" >&2
+    exit 1
+    ;;
+esac

.gitea/workflows/ci.yml

@@ -122,7 +122,7 @@ jobs:
         run: |
           VERSION=$(uv version --short)
           TAG="v$VERSION"
-          if git rev-parse "$TAG" >/dev/null 2>&1 && [ "$(git rev-parse "$TAG")" = "$(git rev-parse HEAD)" ]; then
+          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
             echo "No new release commit/tag to push (already released, or nothing changed)"

.gitea/workflows/deps-bounds.yml

@@ -0,0 +1,29 @@
+name: Monthly dependency upper-bound raise
+
+"on":
+  schedule:
+    - cron: "0 5 1 * *"
+  workflow_dispatch: {}
+
+env:
+  UV_CACHE_DIR: /uv-cache
+
+jobs:
+  bounds-upgrade:
+    name: raise stale pyproject ceilings -> PR
+    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
+    steps:
+      # CI_TOKEN needs write:repository (same scope the release job uses)
+      # to push the refresh branch and open/update its PR.
+      - uses: actions/checkout@v4
+        with:
+          token: ${{ secrets.CI_TOKEN }}
+          fetch-depth: 0
+      - uses: ./.gitea/actions/setup
+      - run: .gitea/scripts/deps-bounds-pr.sh
+        env:
+          CI_TOKEN: ${{ secrets.CI_TOKEN }}

.gitea/workflows/deps-lock.yml

@@ -0,0 +1,29 @@
+name: Weekly dependency lock refresh
+
+"on":
+  schedule:
+    - cron: "0 4 * * 1"
+  workflow_dispatch: {}
+
+env:
+  UV_CACHE_DIR: /uv-cache
+
+jobs:
+  lock-upgrade:
+    name: uv lock --upgrade -> PR
+    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
+    steps:
+      # CI_TOKEN needs write:repository (same scope the release job uses)
+      # to push the refresh branch and open/update its PR.
+      - uses: actions/checkout@v4
+        with:
+          token: ${{ secrets.CI_TOKEN }}
+          fetch-depth: 0
+      - uses: ./.gitea/actions/setup
+      - run: .gitea/scripts/deps-lock-pr.sh
+        env:
+          CI_TOKEN: ${{ secrets.CI_TOKEN }}

CHANGELOG.md

@@ -1,5 +1,37 @@
 # Changelog
 
+## [0.3.23] - 2026-09-07
+
+### Changed
+
+- Fix(deps): silence polars explode() empty_as_null deprecation warnings
+
+- Feat(ci): add scheduled dependency-bump workflows (Renovate-lite)
+
+## [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

CLAUDE.md

@@ -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,
@@ -45,6 +45,12 @@ Part of the `dev` extra. Run these periodically (not just at commit time) to cat
 
 Merges to `master` auto-bump the patch version, tag, and update `CHANGELOG.md` via the Gitea workflow in `.gitea/workflows/ci.yml` (bump-my-version + git-cliff). Don't hand-edit the version in `pyproject.toml` or write changelog entries by hand.
 
+Two scheduled Gitea workflows keep dependencies current, each opening/updating one standing pull request (never editing `master` directly, never opening a Gitea issue) so the normal `pull_request` CI (lint/format/type-check/tests) gates every change before a human merges it:
+- **`deps-lock.yml`** (weekly, Mondays) — `uv lock --upgrade` within the existing `pyproject.toml` constraints, PR'd on branch `chore/uv-lock-upgrade` (`.gitea/scripts/deps-lock-pr.sh`).
+- **`deps-bounds.yml`** (monthly) — raises `pyproject.toml` upper bounds that have fallen behind the latest PyPI release, re-locks, and PRs on branch `chore/dep-bounds` (`.gitea/scripts/deps-bounds-pr.sh`, driving `.gitea/scripts/check_dep_bounds.py`). `torch` (pinned `<2.4` for portal-machine driver support), `plotstyle` (private index, not on PyPI), and the `giant[...]` self-references are never auto-raised — see `check_dep_bounds.py`'s `SKIP_REASONS`.
+
+Both scripts share PR-upsert mechanics in `.gitea/scripts/deps-pr.sh`; both branches are force-pushed fresh from `master` each run rather than accumulated, so at most one open PR exists per job at a time. Both can be triggered by hand via `workflow_dispatch` in the Gitea Actions UI.
+
 ## Compute environment
 
 Work on this repo happens across three kinds of machine:
@@ -101,6 +107,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

README.md

@@ -10,8 +10,8 @@ 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.19-informational)](CHANGELOG.md)
+[![version](https://img.shields.io/badge/version-0.3.23-informational)](CHANGELOG.md)
-[![tests](https://img.shields.io/badge/tests-1138%20passing-brightgreen)](tests/)
+[![tests](https://img.shields.io/badge/tests-1176%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)
 

giant/analysis/__init__.py

@@ -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",
+    "RUNTIME_SAFETY_MARGIN",
-    "catalog_ids",
+    "Context",
-    "get_spec",
+    "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",
 ]

giant/analysis/condor.py

@@ -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]:
+def _rows_per_chunk(
-    """Combined rollout+reference row count of each ``event_id % n_chunks`` 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
+    resolve the actual directory, ``load_rollout_yamls`` for how
-    ``labels``/YAML stems resolve each rollout's series name.
+    ``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
+    rollout/reference/prediction files changed) would otherwise let
-    merge stale partials against the new ``shared.json``.
+    ``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,
     )
 
 

giant/analysis/context.py

@@ -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={

giant/analysis/grouping.py

@@ -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:

giant/analysis/prediction.py

@@ -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, empty_as_null=False)
+        .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", empty_as_null=False)
+        .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"),
+        )
+    )

giant/analysis/reduce.py

@@ -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).
 

giant/analysis/reduced.py

@@ -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)
+#                    distance scorecard) or per prediction (a confusion matrix)
-#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint)
+#   "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()))

giant/analysis/render.py

@@ -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,
 }
 

giant/analysis/router_gating.py

@@ -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.

giant/analysis/sources.py

@@ -230,7 +230,7 @@ def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
     lists = ["sec_E_list", "sec_pdg_list", "sec_dx_list", "sec_dy_list", "sec_dz_list"]
     return (
         lf.select("event_id", *lists)
-        .explode(lists)
+        .explode(lists, empty_as_null=False)
         .drop_nulls("sec_E_list")
         .select(
             "event_id",
@@ -268,7 +268,7 @@ def secondaries_by_step(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
     return (
         lf.select("sec_pdg_list")
         .with_row_index("_row")
-        .explode("sec_pdg_list")
+        .explode("sec_pdg_list", empty_as_null=False)
         .drop_nulls("sec_pdg_list")
         .select(pl.struct("_row").alias("step_key"), pl.col("sec_pdg_list").cast(pl.Int64).alias("pdg"))
     )

giant/analysis/type_embedding_distance.py

@@ -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.
 

giant/config.py

@@ -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),

giant/constants.py

@@ -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".

giant/data/loader.py

@@ -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

giant/data/scan.py

@@ -76,7 +76,11 @@ def _pooled_pdg_lazy(path: Path, has_sec_pdg_list: bool) -> pl.LazyFrame:
     lf = pl.scan_parquet(path, row_index_name="__row")
     parts = [lf.select(pl.col("pdg").alias("__val"), "__row")]
     if has_sec_pdg_list:
-        parts.append(lf.select(pl.col("sec_pdg_list").alias("__val"), "__row").explode("__val").drop_nulls("__val"))
+        parts.append(
+            lf.select(pl.col("sec_pdg_list").alias("__val"), "__row")
+            .explode("__val", empty_as_null=False)
+            .drop_nulls("__val")
+        )
     combined = pl.concat(parts)
     return combined.group_by("__val").agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row"))
 

giant/data/setup_cache.py

@@ -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

giant/geometry.py

@@ -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),

giant/model/builders.py

@@ -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

giant/model/encoders.py

@@ -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

giant/model/history.py

@@ -7,7 +7,7 @@ import inspect
 from typing import cast
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 
 class HistoryEncoder(nn.Module):
@@ -197,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:

giant/model/layers.py

@@ -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):

giant/model/models.py

@@ -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

giant/model/network.py

@@ -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",

giant/model/routers.py

@@ -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)

giant/model/schedule.py

@@ -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":

giant/model/summary.py

@@ -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

giant/model/trunks.py

@@ -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

giant/model/wgan.py

@@ -1,4 +1,4 @@
-from typing import Callable
+from collections.abc import Callable
 
 import torch
 

giant/particles.py

@@ -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,

giant/pipeline.py

@@ -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_critics, build_models, resolve_type_n_classes
-from giant.model.network import build_models, build_critics, resolve_type_n_classes
 from giant.training import train as run_training
 
 

giant/rollout.py

@@ -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(
+    return {
-        event_id=tr["event_id"][sel],
+        "event_id": tr["event_id"][sel],
-        track_id=tr["track_id"][sel],
+        "track_id": tr["track_id"][sel],
-        parent_id=tr["parent_id"][sel],
+        "parent_id": tr["parent_id"][sel],
-        generation=tr["generation"][sel],
+        "generation": tr["generation"][sel],
-        step_no=tr["step_in_track"][sel],
+        "step_no": tr["step_in_track"][sel],
-        pdg=tr["pdg"][sel],
+        "pdg": tr["pdg"][sel],
-        pre_x=pos[:, 0],
+        "pre_x": pos[:, 0],
-        pre_y=pos[:, 1],
+        "pre_y": pos[:, 1],
-        pre_z=pos[:, 2],
+        "pre_z": pos[:, 2],
-        pre_E=tr["pre_E"][sel],
+        "pre_E": tr["pre_E"][sel],
-        pre_dx=dir_[:, 0],
+        "pre_dx": dir_[:, 0],
-        pre_dy=dir_[:, 1],
+        "pre_dy": dir_[:, 1],
-        pre_dz=dir_[:, 2],
+        "pre_dz": dir_[:, 2],
-        post_x=pos[:, 0],
+        "post_x": pos[:, 0],
-        post_y=pos[:, 1],
+        "post_y": pos[:, 1],
-        post_z=pos[:, 2],
+        "post_z": pos[:, 2],
-        post_E=np.zeros(n),
+        "post_E": np.zeros(n),
-        post_dx=dir_[:, 0],
+        "post_dx": dir_[:, 0],
-        post_dy=dir_[:, 1],
+        "post_dy": dir_[:, 1],
-        post_dz=dir_[:, 2],
+        "post_dz": dir_[:, 2],
-        edep=np.asarray(edep, dtype=np.float64).reshape(n),
+        "edep": np.asarray(edep, dtype=np.float64).reshape(n),
-        step_length=np.zeros(n),
+        "step_length": np.zeros(n),
-        material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
+        "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],
+        "layer_id": tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
-        n_sec_pred=np.zeros(n, dtype=np.int64),
+        "n_sec_pred": np.zeros(n, dtype=np.int64),
-        termination_reason=np.full(n, reason, dtype=object),
+        "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,
+    pdg_topn_map: TopNMap | None = None,
-    mat_topn_map: "TopNMap | None" = None,
+    mat_topn_map: TopNMap | None = None,
-    sec_type_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.
 

giant/tools/bump_dataset_version.py

@@ -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)

giant/tools/check_migration_v02_v03.py

@@ -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.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
-from giant.model import network as net  # noqa: E402
+from giant.model import network as net
-from tests.legacy import network_v02_snapshot as legacy  # noqa: E402
+from tests.legacy import network_v02_snapshot as legacy
 
 
 def _random_batch(model_config: dict, batch: int, seed: int):

giant/tools/create_root_files.py

@@ -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(

giant/tools/dwarf.py

@@ -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[

giant/tools/hparam_scan.py

@@ -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 "

giant/tools/migrate_geant_steps.py

@@ -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]},

giant/tools/steps_to_parquet.py

@@ -58,7 +58,7 @@ def _add_secondary_attributes(df: pl.DataFrame) -> tuple[pl.DataFrame, int]:
     exploded = (
         df.select(["event_id", "child_track_ids"])
         .with_row_index("_step_row")
-        .explode("child_track_ids")
+        .explode("child_track_ids", empty_as_null=False)
         .rename({"child_track_ids": "child_track_id"})
         .drop_nulls("child_track_id")
     )

giant/tools/steps_to_parquet_parallel.py

@@ -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
 
 

giant/training/__init__.py

@@ -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,

giant/training/loop.py

@@ -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)

giant/training/metrics.py

@@ -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()

giant/training/plots.py

@@ -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]

giant/training/trainers.py

@@ -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

pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "giant"
-version = "0.3.19"
+version = "0.3.23"
 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",
+    "polars>=1.43,<2",
-    "pyarrow>=16,<25",
+    "pyarrow>=16,<26",
     "tqdm>=4.60,<5",
     "typer>=0.12,<1",
     "pyyaml>=6,<7",
@@ -43,11 +43,11 @@ wandb = [
 convert = [
     "uproot>=5.3,<6",
     "awkward>=2.6,<3",
-    "polars>=1.0,<2",
+    "polars>=1.43,<2",
 ]
 analysis = [
     "matplotlib>=3.8,<4",
-    "polars>=1.0,<2",
+    "polars>=1.43,<2",
     "ipykernel>=7.3.0",
     # KIT matplotlib theme, published from git.larsbogner.de. Only the local
     # `giant analyze render` step imports it; compute workers never do.

tests/legacy/network_v02_snapshot.py

@@ -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(
+        shared = {
-            pdg_vocab=pdg_vocab,
+            "pdg_vocab": pdg_vocab,
-            mat_vocab=mat_vocab,
+            "mat_vocab": mat_vocab,
-            expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
+            "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),
+            "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),
+            "emb_dim": model_config.get("emb_dim", EMB_DIM),
-            dropout=model_config.get("dropout", 0.1),
+            "dropout": model_config.get("dropout", 0.1),
-            conditioning=model_config.get("conditioning", "embedding"),
+            "conditioning": model_config.get("conditioning", "embedding"),
-        )
+        }
         conditioning = shared["conditioning"]
         stage1 = RoutedDenoisingMLP(
             router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),

tests/test_amp.py

@@ -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

tests/test_analysis_prediction.py

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

tests/test_analysis_reduce.py

@@ -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

tests/test_bump_dataset_version.py

@@ -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()
 

tests/test_cli_predict.py

@@ -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

tests/test_cli_train_overrides.py

@@ -9,7 +9,7 @@ from pathlib import Path
 
 from typer.testing import CliRunner
 
-import giant.cli as cli
+from giant import cli
 
 runner = CliRunner()
 

tests/test_cond_layout.py

@@ -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
 

tests/test_condor.py

@@ -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")

tests/test_config.py

@@ -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):

tests/test_config_consumed_keys.py

@@ -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)

tests/test_dependency_automation.py

@@ -0,0 +1,147 @@
+"""Tests for the scheduled dependency-bump automation (weekly uv.lock
+refresh, monthly pyproject upper-bound raise). The workflow YAMLs themselves
+can only be exercised by a real scheduled/dispatched run (same reasoning as
+tests/test_release_tooling.py for the release workflow), so this checks the
+script logic they drive plus the one piece of cross-file coupling that would
+silently misbehave if it drifted: each workflow's script must target the
+same standing branch that deps-pr.sh (invoked by that script) manages.
+"""
+
+import importlib.util
+import re
+import sys
+import tomllib
+from pathlib import Path
+
+import pytest
+import yaml
+from packaging.requirements import Requirement
+from packaging.version import Version
+
+_ROOT = Path(__file__).resolve().parents[1]
+_SCRIPTS = _ROOT / ".gitea" / "scripts"
+_WORKFLOWS = _ROOT / ".gitea" / "workflows"
+
+
+def _load_check_dep_bounds():
+    spec = importlib.util.spec_from_file_location("check_dep_bounds", _SCRIPTS / "check_dep_bounds.py")
+    assert spec is not None and spec.loader is not None
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+@pytest.fixture(scope="module")
+def cdb():
+    return _load_check_dep_bounds()
+
+
+def test_canonical_requirements_finds_real_dependencies_and_skips_excluded(cdb):
+    pyproject = tomllib.loads((_ROOT / "pyproject.toml").read_text())
+    requirements = cdb.canonical_requirements(pyproject)
+
+    assert "numpy" in requirements
+    assert "polars" in requirements
+    # polars appears in [project.dependencies], [convert], and [analysis]
+    # with the same specifier — canonicalization must dedupe by name.
+    assert isinstance(requirements["polars"], Requirement)
+
+    for skipped in ("torch", "plotstyle", "giant"):
+        assert skipped in cdb.SKIP_REASONS, f"{skipped} must stay on the never-auto-raise skip list"
+
+
+def test_next_ceiling_pre_and_post_1_0():
+    from importlib import import_module
+
+    cdb = import_module("check_dep_bounds")
+    # Post-1.0: next ceiling is the major above latest.
+    assert cdb.next_ceiling(Version("2.5.3")) == "<3"
+    assert cdb.next_ceiling(Version("25.0.1")) == "<26"
+    # Pre-1.0: next ceiling is the minor above latest (matches this repo's
+    # own pins, e.g. ruff>=0.15,<1 and ty>=0.0.50,<0.1).
+    assert cdb.next_ceiling(Version("0.16.2")) == "<0.17"
+
+
+def test_find_findings_flags_out_of_range_and_skips_excluded(cdb, monkeypatch):
+    requirements = {
+        "numpy": Requirement("numpy>=1.26,<3"),
+        "widget": Requirement("widget>=1,<2"),
+        "torch": Requirement("torch>=2.3,<2.4"),
+    }
+    fake_latest = {"numpy": "2.5.3", "widget": "3.1.0", "torch": "2.9.0"}
+    monkeypatch.setattr(cdb, "fetch_latest_version", lambda name: fake_latest[name])
+
+    findings = cdb.find_findings(requirements, {"torch": "pinned deliberately"})
+
+    assert len(findings) == 1
+    finding = findings[0]
+    assert finding.name == "widget"
+    assert set(finding.old_specifier.split(",")) == {">=1", "<2"}
+    # SpecifierSet doesn't guarantee clause order, so compare as a set.
+    assert set(finding.new_specifier.split(",")) == {">=1", "<4"}
+    assert finding.latest == "3.1.0"
+
+
+def test_apply_findings_rewrites_every_occurrence_and_nothing_else(cdb):
+    text = (
+        "[project]\n"
+        "dependencies = [\n"
+        '    "polars>=1.43,<2",\n'
+        "]\n\n"
+        "[project.optional-dependencies]\n"
+        "convert = [\n"
+        '    "polars>=1.43,<2",\n'
+        "]\n"
+        "other = [\n"
+        '    "numpy>=1.26,<3",\n'
+        "]\n"
+    )
+    finding = cdb.Finding(name="polars", old_specifier=">=1.43,<2", new_specifier=">=1.43,<3", latest="2.0.0")
+
+    new_text = cdb.apply_findings(text, [finding])
+
+    assert new_text.count('"polars>=1.43,<3"') == 2
+    assert '"polars>=1.43,<2"' not in new_text
+    # Untouched dependency (numpy) survives byte-for-byte.
+    assert '"numpy>=1.26,<3"' in new_text
+
+
+def test_apply_findings_on_real_pyproject_is_a_noop_when_no_findings(cdb):
+    text = (_ROOT / "pyproject.toml").read_text()
+    assert cdb.apply_findings(text, []) == text
+
+
+@pytest.mark.parametrize(
+    ("workflow_file", "script_file"),
+    [
+        ("deps-lock.yml", "deps-lock-pr.sh"),
+        ("deps-bounds.yml", "deps-bounds-pr.sh"),
+    ],
+)
+def test_workflow_yaml_parses_and_calls_matching_script(workflow_file, script_file):
+    workflow = yaml.safe_load((_WORKFLOWS / workflow_file).read_text())
+    assert "schedule" in workflow["on"]
+    assert "workflow_dispatch" in workflow["on"]
+
+    workflow_text = (_WORKFLOWS / workflow_file).read_text()
+    assert script_file in workflow_text, f"{workflow_file} must invoke .gitea/scripts/{script_file}"
+
+
+@pytest.mark.parametrize(
+    ("script_file", "branch"),
+    [
+        ("deps-lock-pr.sh", "chore/uv-lock-upgrade"),
+        ("deps-bounds-pr.sh", "chore/dep-bounds"),
+    ],
+)
+def test_driver_script_branch_matches_deps_pr_invocations(script_file, branch):
+    """Each driver script must define BRANCH as the expected literal, then
+    route every git/deps-pr.sh call through that one $BRANCH variable — a
+    hardcoded mismatch would silently leave a stray branch/PR or open a
+    second one each run."""
+    text = (_SCRIPTS / script_file).read_text()
+    assert f'BRANCH="{branch}"' in text, f"{script_file} should set BRANCH={branch!r}"
+    assert 'git checkout -B "$BRANCH"' in text
+    assert re.search(r'deps-pr\.sh open "\$BRANCH"', text)
+    assert re.search(r'deps-pr\.sh close "\$BRANCH"', text)

tests/test_dwarf.py

@@ -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()
 

tests/test_flow.py

@@ -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)

tests/test_network.py

@@ -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(
+    kwargs = {
-        pdg_vocab=5,
+        "pdg_vocab": 5,
-        mat_vocab=3,
+        "mat_vocab": 3,
-        particle_cfg=PARTICLE_CFG,
+        "particle_cfg": PARTICLE_CFG,
-        material_cfg=MATERIAL_CFG,
+        "material_cfg": MATERIAL_CFG,
-        hidden_dim=_STAGE_HIDDEN_DIM,
+        "hidden_dim": _STAGE_HIDDEN_DIM,
-        n_res_blocks=_STAGE_N_BLOCKS,
+        "n_res_blocks": _STAGE_N_BLOCKS,
-        cond_out_dim=_STAGE_COND_OUT_DIM,
+        "cond_out_dim": _STAGE_COND_OUT_DIM,
-        generator=generator,
+        "generator": generator,
-        time_dim=8,
+        "time_dim": 8,
-        noise_dim=8,
+        "noise_dim": 8,
-    )
+    }
     if cls is Stage1Model:
         kwargs["n_sec_head_k_max"] = 15
     else:

tests/test_phase2.py

@@ -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]

tests/test_render.py

@@ -9,8 +9,8 @@ import pytest
 
 pytest.importorskip("plotstyle")
 
-from giant.analysis import render as render_mod  # noqa: E402
+from giant.analysis import render as render_mod
-from giant.analysis.reduced import Reduced  # noqa: E402
+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)

tests/test_rollout.py

@@ -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}

tests/test_router.py

@@ -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(
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
-        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(
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
-        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)
 

tests/test_sample.py

@@ -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()
 
 
@@ -412,7 +412,7 @@ def test_sample_secondaries_ar_full_length_ignores_n_sec_pred_zero_rows():
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
     n_sec_pred = torch.tensor([0, 0, 0])
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
+    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()

tests/test_train.py

@@ -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 tempfile.TemporaryDirectory() as tmp, pytest.raises(ValueError, match="no active stage"):
-        with pytest.raises(ValueError, match="no active stage"):
+        train(
-            train(
+            cfg=cfg,
-                cfg=cfg,
+            models=models,
-                models=models,
+            critics=critics,
-                critics=critics,
+            train_loader=_fake_batches(1, 8),
-                train_loader=_fake_batches(1, 8),
+            val_loader=_fake_batches(1, 8),
-                val_loader=_fake_batches(1, 8),
+            device=torch.device("cpu"),
-                device=torch.device("cpu"),
+            out_dir=Path(tmp) / "run",
-                out_dir=Path(tmp) / "run",
+            model_config=model_config,
-                model_config=model_config,
+            total_train_batches=1,
-                total_train_batches=1,
+        )
-            )
 
 
 def test_metrics_csv_columns_are_stage_prefixed():

tests/test_training_plots.py

@@ -11,8 +11,8 @@ import pytest
 
 pytest.importorskip("plotstyle")
 
-from giant.training import plots as plots_mod  # noqa: E402
+from giant.training import plots as plots_mod
-from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics  # noqa: E402
+from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics
 
 # --- fixtures ----------------------------------------------------------
 

tests/test_transforms.py

@@ -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,
 )