chore: release v0.3.23

Reviewed-on: #99

.bumpversion.toml

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

.gitea/actions/setup/action.yml

@@ -0,0 +1,20 @@
+name: Setup uv + sync deps
+description: >-
+  Install uv, point its cache at the runner-local mount, and sync the
+  project with the cpu + dev extras. Every CI job does this identically;
+  the caller must still mount /srv/act-runner-cache/uv:/uv-cache on its
+  own container (a composite action can't set that) and check out the
+  repo before this runs.
+
+runs:
+  using: composite
+  steps:
+    - uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: false
+    - shell: bash
+      run: |
+        echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+        echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+    - shell: bash
+      run: uv sync --extra cpu --extra dev

.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/scripts/release-commit.sh

@@ -0,0 +1,63 @@
+#!/usr/bin/env bash
+# Build one complete release commit: version bump (or sync to a given
+# version, or no bump at all), changelog entry, and refreshed README
+# badges — all in a single commit, tagged at the end. Used by both the
+# merge-to-master release job and the hand-pushed-tag sync job, so "the
+# tag == a complete release commit" holds either way a release gets made.
+#
+# Usage: release-commit.sh [--no-bump|VERSION]
+#   (no argument)  bump the current patch version (normal merge path).
+#   --no-bump      don't touch the version — the branch already bumped it
+#                   (e.g. a manual minor/major bump); just build the
+#                   changelog/badge commit for whatever version is current.
+#   VERSION        sync the project to this exact version (tag-sync path).
+#
+# Preconditions: repo is checked out with full history (fetch-depth: 0),
+# `uv` is available and synced, and git user.name/user.email are configured.
+# Idempotent: if the resulting tag already exists, this is a no-op.
+set -euo pipefail
+
+VERSION_ARG="${1:-}"
+
+if [ "$VERSION_ARG" = "--no-bump" ]; then
+  echo "Version already bumped by this branch; using current version as-is"
+elif [ -n "$VERSION_ARG" ]; then
+  echo "Syncing project version to $VERSION_ARG"
+  uv version "$VERSION_ARG" --no-sync
+  uv lock
+else
+  CURRENT_VERSION=$(uv version --short)
+  echo "Bumping patch version from $CURRENT_VERSION"
+  uv run bump-my-version bump patch --current-version "$CURRENT_VERSION" --no-commit --no-tag
+  # bump-my-version's pre_commit_hooks (uv lock + git add uv.lock) only run
+  # on its own commit path, which we skipped with --no-commit — so do it here.
+  uv lock
+fi
+
+VERSION=$(uv version --short)
+TAG="v$VERSION"
+
+if git rev-parse "$TAG" >/dev/null 2>&1; then
+  echo "Tag $TAG already exists; nothing to do"
+  exit 0
+fi
+
+uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
+
+TEST_COUNT=$(uv run pytest --collect-only -q 2>/dev/null | grep -oE '^[0-9]+ tests? collected' | grep -oE '^[0-9]+')
+sed -i -E "s|badge/version-[^-]+-informational|badge/version-${VERSION}-informational|" README.md
+sed -i -E "s|badge/tests-[0-9]+%20passing-brightgreen|badge/tests-${TEST_COUNT}%20passing-brightgreen|" README.md
+
+# .bumpversion.toml stores its own current_version, which bump-my-version
+# rewrites even with --no-commit — must be staged or the next run sees a
+# dirty tree and bump-my-version refuses (allow_dirty = false).
+git add pyproject.toml .bumpversion.toml uv.lock CHANGELOG.md README.md
+if git diff --cached --quiet; then
+  echo "Nothing changed; skipping release commit"
+  exit 0
+fi
+
+git commit -m "chore: release $TAG"
+git tag -a "$TAG" -m "$TAG"
+
+echo "Created release commit and tag $TAG"

.gitea/workflows/ci.yml

@@ -12,6 +12,7 @@ env:
 jobs:
   ruff-check:
     name: Lint (ruff check)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
     runs-on: ubuntu-latest
     container:
       image: docker.gitea.com/runner-images:ubuntu-latest
@@ -19,17 +20,12 @@ jobs:
         - /srv/act-runner-cache/uv:/uv-cache
     steps:
       - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
       - run: uv run ruff check .
 
   ruff-format:
     name: Format (ruff format)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
     runs-on: ubuntu-latest
     container:
       image: docker.gitea.com/runner-images:ubuntu-latest
@@ -37,17 +33,12 @@ jobs:
         - /srv/act-runner-cache/uv:/uv-cache
     steps:
       - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
       - run: uv run ruff format --check .
 
   type-check:
     name: Type check (ty)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
     runs-on: ubuntu-latest
     container:
       image: docker.gitea.com/runner-images:ubuntu-latest
@@ -55,18 +46,13 @@ jobs:
         - /srv/act-runner-cache/uv:/uv-cache
     steps:
       - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
       - run: uv run ty check .
 
   test:
     name: Tests
-    needs: [ruff-check, type-check]
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
+    needs: [ruff-check, ruff-format, type-check]
     runs-on: ubuntu-latest
     container:
       image: docker.gitea.com/runner-images:ubuntu-latest
@@ -74,21 +60,15 @@ jobs:
         - /srv/act-runner-cache/uv:/uv-cache
     steps:
       - uses: actions/checkout@v4
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
       - run: uv run pytest --cov --cov-report=term-missing --cov-report=xml
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
         with:
           name: coverage-report
           path: coverage.xml
 
-  bump-version:
-    name: Bump version, tag, and update changelog on merge to master
+  release:
+    name: Release (bump, changelog, badges, tag) on merge to master
     needs: [ruff-check, ruff-format, type-check, test]
     if: github.ref == 'refs/heads/master' && github.event_name == 'push'
     runs-on: ubuntu-latest
@@ -98,7 +78,7 @@ jobs:
         - /srv/act-runner-cache/uv:/uv-cache
     steps:
       # CI_TOKEN needs write:repository scope (not just read) — this job
-      # pushes commits and tags to master, unlike ruff-check/ruff-format/
+      # pushes a commit and a tag to master, unlike ruff-check/ruff-format/
       # type-check/test above, which only need to check out the repo.
       - uses: actions/checkout@v4
         with:
@@ -114,22 +94,14 @@ jobs:
           else
             echo "is_merge=false" >> "$GITHUB_OUTPUT"
           fi
-      - uses: astral-sh/setup-uv@v5
-        if: steps.merge_check.outputs.is_merge == 'true'
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-        if: steps.merge_check.outputs.is_merge == 'true'
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
         if: steps.merge_check.outputs.is_merge == 'true'
       - name: Configure git identity
         if: steps.merge_check.outputs.is_merge == 'true'
         run: |
           git config user.name "gitea-actions"
           git config user.email "actions@git.larsbogner.de"
-      - name: Bump patch version if this merge didn't already bump it
+      - name: Build the release commit
         if: steps.merge_check.outputs.is_merge == 'true'
         run: |
           OLD_VERSION=$(git show "${{ github.event.before }}:pyproject.toml" 2>/dev/null | grep -m1 '^version = ' | sed -E 's/version = "(.*)"/\1/')
@@ -140,70 +112,64 @@ jobs:
           fi
           if [ "$OLD_VERSION" = "$CURRENT_VERSION" ]; then
             echo "Version unchanged by this merge ($CURRENT_VERSION); bumping patch"
-            uv run bump-my-version bump patch --current-version "$CURRENT_VERSION"
+            .gitea/scripts/release-commit.sh
           else
-            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); skipping auto-bump"
+            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); building release commit without bumping"
+            .gitea/scripts/release-commit.sh --no-bump
           fi
-      - name: Update changelog for the current version if not already tagged
+      - name: Push the release commit and its tag together
         if: steps.merge_check.outputs.is_merge == 'true'
         run: |
           VERSION=$(uv version --short)
           TAG="v$VERSION"
-          if git rev-parse "$TAG" >/dev/null 2>&1; then
-            echo "Tag $TAG already exists; skipping changelog update"
+          if git rev-parse "$TAG" >/dev/null 2>&1 && [ "$(git rev-parse "$TAG^{commit}")" = "$(git rev-parse HEAD)" ]; then
+            git push --atomic origin HEAD:master "refs/tags/$TAG"
           else
-            uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
-            git add CHANGELOG.md
-            if ! git diff --cached --quiet -- CHANGELOG.md; then
-              git commit -m "chore: update changelog for $TAG"
-            else
-              git restore --staged CHANGELOG.md
-            fi
-          fi
-      - name: Push commits and tag the current version
-        if: steps.merge_check.outputs.is_merge == 'true'
-        run: |
-          git push origin HEAD:master
-          VERSION=$(uv version --short)
-          TAG="v$VERSION"
-          if git rev-parse "$TAG" >/dev/null 2>&1; then
-            echo "Tag $TAG already exists"
-          else
-            git tag -a "$TAG" -m "$TAG"
-            git push origin "refs/tags/$TAG"
+            echo "No new release commit/tag to push (already released, or nothing changed)"
+            git push origin HEAD:master
           fi
 
   sync-version-on-tag:
-    name: Sync project version with tag
+    name: Sync project version with tag (hand-pushed tags only)
     if: startsWith(github.ref, 'refs/tags/')
     runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
     steps:
       - uses: actions/checkout@v4
         with:
           token: ${{ secrets.CI_TOKEN }}
-      - uses: astral-sh/setup-uv@v5
-      - name: Check tag against project version, update if they differ
+          fetch-depth: 0
+      - name: Require the tagged commit to already be on master
+        run: |
+          git fetch origin master
+          if ! git merge-base --is-ancestor "${{ github.sha }}" origin/master; then
+            echo "::error::Tag ${GITHUB_REF_NAME} points at a commit not on master; refusing to publish an unreviewed tree. Push the commit to master first, or delete and re-push the tag once it is."
+            exit 1
+          fi
+      - uses: ./.gitea/actions/setup
+      - name: Check tag against project version, build a release commit if they differ
         run: |
           TAG_VERSION="${GITHUB_REF_NAME#v}"
           CURRENT_VERSION=$(uv version --short)
-          if [ "$TAG_VERSION" != "$CURRENT_VERSION" ]; then
-            echo "Tag version ($TAG_VERSION) != project version ($CURRENT_VERSION); updating pyproject.toml"
-            uv version "$TAG_VERSION" --no-sync
+          if [ "$TAG_VERSION" = "$CURRENT_VERSION" ]; then
+            echo "Tag version matches project version ($CURRENT_VERSION)"
+          else
+            echo "Tag version ($TAG_VERSION) != project version ($CURRENT_VERSION); building a release commit"
             git config user.name "gitea-actions"
             git config user.email "actions@git.larsbogner.de"
-            git add pyproject.toml uv.lock
-            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME}"
+            .gitea/scripts/release-commit.sh "$TAG_VERSION"
             git push origin HEAD:master
             git push origin ":refs/tags/${GITHUB_REF_NAME}"
             git tag -f "${GITHUB_REF_NAME}" HEAD
             git push origin "refs/tags/${GITHUB_REF_NAME}"
-          else
-            echo "Tag version matches project version ($CURRENT_VERSION)"
           fi
 
   publish-package:
     name: Publish package to Gitea package registry
-    needs: [ruff-check, ruff-format, type-check, test, sync-version-on-tag]
+    needs: [sync-version-on-tag]
     if: startsWith(github.ref, 'refs/tags/')
     runs-on: ubuntu-latest
     container:
@@ -216,12 +182,7 @@ jobs:
       - uses: actions/checkout@v4
         with:
           ref: ${{ github.ref_name }}
-      - uses: astral-sh/setup-uv@v5
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+      - uses: ./.gitea/actions/setup
       - run: uv build
       # CI_TOKEN needs write:package scope (in addition to write:repository,
       # used elsewhere) for this upload to authenticate.

.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,61 @@
 # 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
+
+- Chore: update README badges (version 0.3.18, 1138 tests)
+
+- Ci: restructure release pipeline into a single atomic release commit
+
+## [0.3.18] - 2026-09-03
+
+### Changed
+
+- Docs: rewrite README, keep version/test badges live via Gitea Actions
+
+- Perf: compact Stage-2 AR inference loop to active rows only
+
+## [0.3.17] - 2026-09-02
+
+### Changed
+
+- Feat: add WGAN + AR stop-token config variant
+
+- Perf: replace pandas with polars in the setup-stage scan
+
 ## [0.3.16] - 2026-08-31
 
 ### Changed

CLAUDE.md

@@ -9,7 +9,6 @@ uv sync --extra cpu                # install dependencies with CPU-only torch (s
 uv sync --extra cuda               # install dependencies with CUDA 11.8 torch
 uv sync --extra cpu --extra dev    # add dev extras (pytest, ruff, ty, bump-my-version, git-cliff, + all runtime extras)
 uv sync --extra cpu --extra geometry  # add scikit-learn for the geometry oracle (giant rollout)
-uv sync --extra cpu --extra workflow  # add b2luigi for `giant workflow` pipeline orchestration
 pytest               # run tests
 giant new-run --hidden-dim 512 --lr 3e-4        # scaffold a config.toml + run dir ahead of training
 giant train path/to/steps.parquet               # train (defaults: stage 1 flow, stage 2 wgan + autoregressive)
@@ -19,9 +18,8 @@ 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 workflow run spec.toml --batch --workers 20          # whole pipeline (cache-warm -> train -> rollout -> analysis)
-giant analyze prep rollout.yaml --chunks 32                # lay out an analysis run dir (compute jobs come from the workflow)
-giant analyze render <run_dir> --gallery                   # render PDFs + HTML gallery (run_dir from prep)
+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,
                                                   # bump-schema, status, update-manifest, create-manifest,
@@ -47,13 +45,19 @@ 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:
 
 - **Local dev machines** (laptop + desktop, identical): repo at `~/Programming/giant`, no access to `/ceph` — datasets, training results, and models aren't reachable here.
 - **Portal machines** (`portal1`, `deepthought`, `deepthought2`, `bms1`, `bms2`, `bms3`): repo lives under `/work`, and `/ceph` holds ROOT/parquet files and trained models. **These are shared with other users** — stay strictly within `/work/lbogner` and `/ceph/lbogner`, and keep resource usage to roughly a quarter of CPU/RAM and a single GPU so as not to disturb other users' jobs.
-- **HTCondor worker nodes**: never run or SSH onto these directly — the only sanctioned path is `giant workflow run <spec.toml> --batch` (b2luigi, see the Workflow section), which submits and polls every job. `/ceph` is available there; `/work` is only sometimes mounted, depending on the node.
+- **HTCondor worker nodes**: never run or SSH onto these directly — the only sanctioned path is submitting jobs through condor (`giant analyze submit`, and the in-progress remote-GPU train/rollout submission on `condor-gpu-train-rollout`). `/ceph` is available there; `/work` is only sometimes mounted, depending on the node.
 
 ## Architecture
 
@@ -101,9 +105,9 @@ Secondary energies are a **stick-breaking partition of the `e_sec` budget** from
 
 **Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive `giant rollout` runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + `RolloutSpec`/`Side` — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `variables.py` (the per-step value expressions shared by range sizing and the plot registry), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json` over the union of the reference and every rollout, so every compute job is one pass with no range scan), `reduced.py` (`Partial`/`Reduced` — the compact self-describing JSON a compute job emits), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles/containment, species/leakage, secondaries, distance summaries, router and type-embedding diagnostics; `giant analyze list` prints every id), `runtime_estimate.py` (per-(plot, chunk) walltime estimates for the submit description), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`; each rollout gets a stable `ps.get_color(i)` slot by its position in `series`, the reference always draws in one fixed dashed-ink style). `Bundle.rollouts` is a name-keyed dict of `Side`, and every `compute_partial`/`finalize` builds a `Reduced.payload["series"]` dict keyed the same way, with `payload["reference"]` as the one distinguished non-rollout entry. The heatmap-shaped specs (`marginal_distance_summary`, `sec_count_per_step_by_species` — the latter also drawing the reference as its own panel) and the checkpoint-bound diagnostics (`router_gating.py`, `type_embedding_distance.py`) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout.
 
-**Input is one or more `giant rollout` YAML sidecars** (`run.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 prep a.yaml [b.yaml ...] --chunks N` records `N` in `run_meta.json`, and the workflow's `AnalysisComputeTask` runs 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`.
+**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`.
 
-**Workflow orchestration** (`giant/workflow/`, `giant workflow run` CLI): b2luigi is the **only sanctioned way to run a multi-step pipeline**; `giant`/`dwarf` are single-step primitives the tasks invoke. One workflow TOML (`configs/workflow_example.toml`) parameterises a whole experiment — `[workflow]`/`[condor]`/`[dataset]`/`[geometry]` plus repeated `[[train]]`/`[[rollout]]`/`[[analysis]]` tables, each cross-referenced by name — and `spec.py` parses it into frozen dataclasses, rejecting unknown keys and dangling references. Every task's output directory is `<result_dir>/<kind>/name=<name>/spec_hash=<hash>/…`, where the 8-hex `spec_hash` covers that task's resolved sub-spec **and its transitive parents**, so an edited spec re-runs exactly the affected subtree instead of silently reusing stale outputs. The DAG (`tasks.py`): `DatasetTask` (external, fails fast if `/ceph` isn't mounted) → `WarmCacheTask` / `GeometryOracleTask` → `TrainEpochTask(name, milestone)` → `TrainTask` → `RolloutTask` → `AnalysisPrepTask` → `AnalysisComputeTask(name, plot_id, chunk)` → `AnalysisRenderTask` → `WorkflowTask`. Training is fanned out into **one short GPU job per epoch** (`epochs_per_job` trades queue waits back), chained by `--resume` on the previous job's `last.pt` — the loop already handles that unchanged — and `TrainTask` republishes `best.pt`/`last.pt`/a concatenated `metrics.csv` so nothing downstream sees the fan-out. `StreamingStepsDataset.set_epoch` and `config.epoch_seed` (both applied per epoch by `training/loop.py`) derive the batch order and the global RNG state from `(seed, epoch)`, so epoch *k* is bit-identical either way — verified by diffing a chained run's `metrics.csv` against a single 3-epoch `giant train`. `AnalysisRenderTask` is always local (the only step importing plotstyle/LaTeX); `htcondor.py` holds the CPU/GPU submit settings, with the GPU requirement strings (`TARGET.ProvidesEtpCeph` + device/memory pins) ported from the `condor-gpu-train-rollout` branch. `run.py` is the script b2luigi re-executes on workers (`--spec` forwarded via `task_cmd_additional_args`, so a worker resolves the identical graph); `giant workflow run` is a thin exec of it. Needs `uv sync --extra cpu --extra workflow`.
+**`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`.
 
@@ -127,4 +131,4 @@ v0.2 configs and checkpoints are auto-migrated (`config.migrate_config`, `model.
 
 A sampling-calorimeter (multi-material) dataset track is still open and unblocked, not yet started. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
 
-**Condor-submitted GPU training/rollout (`condor-gpu-train-rollout` branch, superseded):** its goal — moving `giant train`/`giant rollout` off the shared portal GPU dev machines onto remote-GPU HTCondor submission — is now met by the b2luigi workflow above. Its `train-submit`/`rollout-submit` commands are deliberately **not** ported and must not be revived when that branch is eventually merged; the only part that survived is `_gpu_requirements`, which moved into `giant/workflow/htcondor.py`.
+**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2. Partway between "needs major features" and feature-complete — not ready to merge yet.

README.md

@@ -1,221 +1,295 @@
-# giant
+<div align="center">
 
-**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
+# GIANT
 
-A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.
+### **G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate
+
+A conditional generative model that replaces the Geant4 step function — sample a
+post-step outcome instead of simulating one, then roll that out into full
+calorimeter showers.
+
+[![python](https://img.shields.io/badge/python-3.12%2B-3776AB?logo=python&logoColor=white)](pyproject.toml)
+[![torch](https://img.shields.io/badge/torch-2.3.x-EE4C2C?logo=pytorch&logoColor=white)](pyproject.toml)
+[![version](https://img.shields.io/badge/version-0.3.23-informational)](CHANGELOG.md)
+[![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)
+
+</div>
+
+---
+
+## The idea
+
+Geant4's step function is the innermost loop of detector simulation — for every
+particle, at every step, it stochastically samples where the particle goes next,
+how much energy it deposits, and what secondaries it spawns. GIANT learns that
+function instead of running it: given a pre-step particle state (position,
+energy, direction, particle species, material), a two-stage model samples a
+post-step outcome — including the variable-length list of secondaries — and
+autoregressively rolls that out into whole showers. It trains entirely from
+parquet dumps of a Geant4 steps tree; nothing downstream needs a Geant4 runtime.
+
+Two guarantees are architectural, not learned:
+
+- **Energy is conserved by construction.** Stage 1 decodes deposit / secondary /
+  post-step energy through a softmax simplex that sums to the pre-step energy
+  exactly; Stage 2's secondaries stick-break that same energy budget.
+- **No shower leaks across the train/val split.** Steps are split by `event_id`,
+  never by row, so correlated steps from the same shower can't appear on both
+  sides.
+
+## How a step becomes a shower
+
+```mermaid
+flowchart LR
+    A["pre-step state\nposition · energy · direction\nspecies · material"] --> B["ConditionEncoder\nphysical / embedding / onehot"]
+    B --> C["Stage 1\n9D post-step outcome"]
+    C --> D["Stage 2 (autoregressive)\nsecondaries, descending energy"]
+    D --> E["rollout step"]
+    E -->|"primary continues"| F["GeometryOracle\nposition to material, layer"]
+    E -->|"secondaries pushed"| G["track queue"]
+    F --> A
+    G --> A
+    E -->|"terminated"| H["deposited shower"]
+```
 
 ## Quick start
 
 ```bash
-uv sync --extra cpu               # install deps (CPU torch; use --extra cuda for GPU)
-
-giant new-run --hidden-dim 512 --lr 3e-4          # scaffold config.toml + run dir
-giant model summary --config config.toml          # parameter counts + which config keys actually bite
-giant train path/to/steps.parquet                 # train (flow + wgan by default)
-giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
+uv sync --extra cpu                                  # install deps (CPU torch; --extra cuda for GPU)
 
+giant new-run --hidden-dim 512 --lr 3e-4              # scaffold config.toml + run dir
+giant train path/to/steps.parquet                     # train (flow stage 1 + wgan stage 2, by default)
 dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # needed for rollout
 giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
+
+giant analyze prep rollout.yaml && giant analyze render <run_dir> --gallery   # rollout-vs-Geant4 diagnostics
 ```
 
-Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.
+Every command takes `--help` for its full flag list, and `--config config.toml`
+for anything not exposed as a flag.
 
-## Architecture
+---
 
-A **two-stage model**, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (`--stage1-generator`/`--stage2-generator`, or `--mode` to set both at once): `flow` (conditional flow matching, ODE-sampled in ~10 steps), `ddpm` (denoising diffusion), or `wgan` (single-pass WGAN-GP generator/critic).
+<details>
+<summary><h2 style="display:inline">Architecture</h2></summary>
 
-**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
+**Stage 1 — primary step.** Predicts the 9D post-step outcome
+(`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
 
 | Index | Variable | Encoding |
 |-------|----------|----------|
 | 0 | `step_length` [mm] | log |
-| 1–2 | `edep_logit`, `sec_logit` | ALR coords of the deposit/secondary/post-energy simplex |
-| 3–5 | `post_dir` in local frame | unit vector |
-| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
+| 1–2 | `edep_logit`, `sec_logit` | ALR coordinates of the deposit / secondary / post-energy simplex |
+| 3–5 | `post_dir` | unit vector, local frame (`pre_dir = ẑ`) |
+| 6–8 | `travel_dir` (`post_pos − pre_pos`) | unit vector, local frame |
 
-- Energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — conservation is architectural, not learned.
-- `post_dir`/`travel_dir` live in the frame where `pre_dir = ẑ`. `post_pos` isn't a target — it's reconstructed as `pre_pos + step_length * world_frame(travel_dir)`.
+Energy logits decode via `softmax([edep_logit, sec_logit, 0]) × pre_E`, so
+`edep + e_sec + post_E == pre_E` holds exactly. `post_pos` is not itself a
+target — it's reconstructed as `pre_pos + step_length · world_frame(travel_dir)`,
+since duplicating that magnitude in a second target would let the two drift out
+of sync.
 
-**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (`--stage2-decoder`):
+**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's
+outcome, it generates the variable-length secondary list, one token at a time in
+descending-energy order (`autoregressive`, default) or all `K_MAX` slots in one
+masked pass (`one_shot`). Autoregressive tokens condition on a running history —
+`markov` (previous token only) or `attention` (causal self-attention, KV-cached
+at inference). Either way, secondary energies stick-break the `e_sec` budget
+handed down from Stage 1, so the whole chain conserves energy. A secondary's
+species is represented `onehot` (categorical, top-N PDG codes + "other"),
+`physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour
+lookup).
 
-- `autoregressive` — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (`markov`: previous token only, or `attention`: causal self-attention, KV-cached at inference)
-- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
+**Conditioning (15D).** Pre-step position / energy / direction / layer, plus
+particle mass/charge and material Z_eff/A_eff/density/X0/λ_int — encoded the
+same three ways as secondary species above, configured *independently* per axis
+(`conditioning.particle.type` / `conditioning.material.type`). `physical`
+computes rather than looks up, so it generalizes to species and materials
+outside the training menu; that's the default. `n_sec`/`e_sec` are always model
+outputs, never conditioning inputs.
 
-Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).
+**Composable by design** — every stage assembles from small registries, so
+swapping one axis doesn't touch the others:
 
-**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above. The particle and material axes are configured independently (`conditioning.particle.type` / `conditioning.material.type`; `--conditioning` sets both at once) and may mix — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
+| Registry | Choices |
+|---|---|
+| Objective | `flow` (matching, ~10-step ODE sample) · `ddpm` (denoising diffusion) · `wgan` (single-pass GAN) |
+| Trunk | `resmlp` · `none`, optionally MoE-routed (`RoutedTrunk`) |
+| Router | `energy` · `pdg` · `process` · `composed` · `none` — soft-mixed at train time, **top-1 dispatched at eval time**, which is the actual inference-speed win |
+| History (stage 2 AR) | `markov` · `attention` · `none` |
 
-**MoE routing** (`--router`): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk. The CLI flags configure Stage 1's router; Stage 2 has its own `stage2_model.router` block, config-file only.
+</details>
 
-## Data
+<details>
+<summary><h2 style="display:inline">Configuration</h2></summary>
 
-- Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from ROOT via `dwarf convert`. One row = one Geant4 step.
-- **Conditioning (pre-step) columns:** `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz` (direction), `material`, `layer_id`.
-- **Primary outcome (post-step) columns:** `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep` (energy deposited in this step), `e_sec` (total energy carried off by secondaries), `child_track_ids` (its length gives `n_sec`).
-- **Secondary columns**, one variable-length list per step: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX` (15) slots on load, ordered by descending energy.
-- **Optional:** `process` — the physics process that produced the step (e.g. `compt`, `phot`, `eBrem`); a post-step label used only as classifier supervision (`ProcessRouter`), never as conditioning.
-- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
-- Loading a directory or `.manifest` of multiple parquet files (each one Geant4 job, `event_id` restarting from 0) offsets each file's `event_id`s by a fixed per-file stride so ids stay globally unique across files.
+Every default lives in one place: frozen dataclasses in `giant/config.py`,
+composed into `GiantConfig` (`conditioning` / `stage1_model` / `stage2_model` /
+`train`). `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather
+than hand-maintained, so the dataclasses can't drift from what actually gets
+merged. TOML config keys are validated against that shape — an unknown key is
+rejected with a did-you-mean suggestion. Precedence: CLI flag > `--config` file
+> default.
 
-## Project structure
+```toml
+# config.toml — resolved shape of the four blocks
+[conditioning]
+particle.type = "physical"
+material.type = "physical"
 
-```
-giant/
-├── giant/
-│   ├── data/
-│   │   ├── loader.py       # parquet → numpy arrays (incl. streaming/chunked reads)
-│   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
-│   │   ├── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
-│   │   └── setup_cache.py  # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
-│   ├── model/
-│   │   ├── models.py       # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
-│   │   ├── builders.py     # build_models / build_critics — config dict → assembled stage models
-│   │   ├── encoders.py     # ConditionEncoder (physical / embedding / onehot, per axis)
-│   │   ├── layers.py       # ResBlock/AdaLNResBlock registry, SinusoidalEmbedding, MLP heads
-│   │   ├── trunks.py       # trunk registry (resmlp, none) + RoutedTrunk (MoE expert bodies)
-│   │   ├── routers.py      # Router registry: energy / pdg / process / composed / none
-│   │   ├── history.py      # stage-2 AR history encoders: markov / attention (KV-cached) / none
-│   │   ├── objectives.py   # flow / ddpm / wgan objective registry
-│   │   ├── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
-│   │   ├── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
-│   │   ├── summary.py      # build-only introspection behind `giant model summary`
-│   │   ├── _legacy.py      # v0.2 checkpoint model_config/state-dict migration
-│   │   └── network.py      # re-export shim over all of the above
-│   ├── constants.py        # output/conditioning dims, K_MAX, secondary slot layout, schema keys
-│   ├── cond_layout.py      # single source of truth for the cond_cont/cond_cat column layout
-│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
-│   ├── materials.py        # material name → (Z_eff, A_eff, density, X0, λ_int)
-│   ├── config.py           # default hyperparameters, TOML config merging, device autodetect
-│   ├── pipeline.py         # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
-│   ├── training/           # two-stage training: loop, per-stage trainers, metrics, checkpointing
-│   │   ├── loop.py         #   epoch loop, graceful shutdown, best-checkpoint selection
-│   │   ├── trainers.py     #   StageSpec + flow/ddpm and WGAN-GP per-stage trainers
-│   │   ├── stage2_inputs.py#   ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
-│   │   ├── metrics.py      #   MetricsCollector: metrics.csv columns, W&B logging, progress/summary
-│   │   ├── amp.py          #   bf16 autocast (`train.precision`)
-│   │   ├── plots.py        #   training-progress plots (`giant analyze metrics`)
-│   │   └── checkpoint.py   #   checkpoint assembly/restore (format unchanged since v0.2)
-│   ├── sample.py           # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
-│   ├── checkpoint_io.py    # checkpoint → ready-to-run models/normalizers (predict + rollout)
-│   ├── geometry.py         # GeometryOracle: position → (material, layer_id, escaped) for rollout
-│   ├── rollout.py          # autoregressive shower rollout driver
-│   ├── validate.py         # step-level marginal + KL-divergence validation
-│   ├── _migration.py       # shared v0.2 → v0.3 facts used by both migration surfaces
-│   ├── analysis/           # rollout-vs-reference analysis pipeline (see `giant analyze` below)
-│   │   ├── sources.py      #   canonical LazyFrames + secondary view
-│   │   ├── variables.py    #   per-step value expressions shared by range sizing and the catalog
-│   │   ├── reduce.py       #   streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
-│   │   ├── grouping.py     #   fixed bin edges + energy/pdg/material group sets
-│   │   ├── context.py      #   resolves grouping into `shared.json` once per run
-│   │   ├── reduced.py      #   Partial/Reduced — the compact JSON a compute job emits
-│   │   ├── catalog.py      #   declarative PlotSpec registry (`giant analyze list`)
-│   │   ├── router_gating.py / type_embedding_distance.py   # checkpoint-bound diagnostics
-│   │   ├── runtime_estimate.py  #   per-(plot, chunk) walltime estimates for the job requests
-│   │   ├── run.py          #   prep / compute-one / merge plumbing
-│   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
-│   ├── workflow/           # b2luigi pipeline orchestration (`giant workflow run spec.toml`)
-│   │   ├── spec.py         #   workflow TOML -> frozen dataclasses, validation, per-task spec hashes
-│   │   ├── htcondor.py     #   CPU/GPU submit settings (docker image, +RemoteJob, GPU requirements)
-│   │   ├── tasks.py        #   the task graph: cache-warm -> train (one job/epoch) -> rollout -> analysis
-│   │   └── run.py          #   the script b2luigi re-executes on every worker
-│   └── cli.py              # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze` / `workflow`
-├── giant/tools/                        # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
-│   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
-│   │                                    #   update-manifest, create-manifest, make-root,
-│   │                                    #   build-geometry-oracle, warm-cache, hparam-scan
-│   ├── steps_to_parquet.py             # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
-│   ├── steps_to_parquet_parallel.py    # fan out conversion over several ROOT files — `dwarf convert --jobs N`
-│   ├── migrate_geant_steps.py          # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
-│   ├── bump_dataset_version.py         # cut a new raw gen or parquet schema, with a logged reason —
-│   │                                    #   `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
-│   ├── create_root_files.py            # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
-│   ├── geometry_oracle.py              # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
-│   ├── warm_setup_cache.py             # precompute `giant train`'s setup-stage sidecar — `dwarf warm-cache`
-│   ├── hparam_scan.py                  # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
-│   └── profile_analysis_costs.py       # profiling helper for the `giant analyze` reduction pipeline
-└── tests/
+[stage1_model]
+generator = "flow"
+
+[stage2_model]
+generator = "wgan"
+decoder   = "autoregressive"
+
+[train]
+epochs     = 100
+batch_size = 4096
+lr         = 3e-4
 ```
 
-## Setup
+Some knobs only exist in the config file, with no CLI flag:
+`stage2_model.autoregressive.teacher_forcing`/`.history`,
+`stage2_model.particle_type.target`/`.class_weighting`,
+`stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`,
+`stage2_model.router.*`, and the finer `router` knobs (`lambda_balance`,
+`gumbel`, `learn_width`, …).
+
+`configs/` holds kept reference configs — `baseline.toml` is the fixed
+comparison point every experimental variant (routed trunk, WGAN, attention
+history, embedding conditioning) is a single edit away from. v0.2 flat-schema
+configs and checkpoints load and auto-migrate.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Data</h2></summary>
+
+Input is parquet — one row per Geant4 step — from
+[miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from
+ROOT via `dwarf convert`.
+
+| Group | Columns |
+|---|---|
+| **Conditioning (pre-step)** | `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz`, `material`, `layer_id` |
+| **Primary outcome (post-step)** | `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep`, `e_sec`, `child_track_ids` (length → `n_sec`) |
+| **Secondaries** (variable-length lists) | `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX = 15` slots, descending energy |
+| **Optional** | `process` — physics-process label, classifier supervision only (`ProcessRouter`), never conditioning |
+
+Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so a
+shower's correlated steps never straddle the split. Loading a directory or
+`.manifest` of several parquet files offsets each file's `event_id`s by a
+per-file stride so ids stay globally unique. The pre-epoch setup scan (vocab
+maps, event split, normalizer stats) persists to a sidecar cache
+(`--cache-setup`/`--rebuild-setup-cache`), precomputable ahead of time via
+`dwarf warm-cache`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">CLI reference</h2></summary>
+
+**`giant`** — train, run, and analyze the surrogate:
+
+| Command | Does |
+|---|---|
+| `new-run` | scaffold a `config.toml` + run directory from flags |
+| `train DATA` | train the two-stage model |
+| `model summary` | build-only parameter counts, without training |
+| `predict DATA --checkpoint …` | per-step predictions from a checkpoint |
+| `rollout DATA --checkpoint … --geometry …` | full autoregressive shower rollout |
+| `analyze prep/submit` | build a run dir; `submit` also queues HTCondor compute jobs |
+| `analyze compute-one` / `merge-one` | one plot × chunk reduction / merge (what a condor job runs) |
+| `analyze render <run_dir> --gallery` | merge chunks → styled PDFs + HTML gallery (local, needs LaTeX) |
+| `analyze metrics <train_run_dir>` | training-progress plots from `metrics.csv` |
+| `analyze list` | every catalog plot id |
+
+**`dwarf`** — dataset/tooling CLI:
+
+| Command | Does |
+|---|---|
+| `convert` | ROOT Steps tree → parquet (`--jobs N` fans out) |
+| `migrate` | one-time move into the raw/processed/pools/derived layout |
+| `bump-gen` / `bump-schema` / `status` | dataset versioning |
+| `update-manifest` / `create-manifest` | point/build a manifest of parquet files |
+| `make-root` | generate new ROOT shards via a minicalosim executable |
+| `build-geometry-oracle` | fit position → (material, layer_id) for rollout |
+| `warm-cache` | precompute `giant train`'s setup-stage sidecar |
+| `hparam-scan` | grid-scan dropout × n_blocks × hidden_dim |
+
+Worth knowing on `giant train` (full surface behind `--help`):
+`--mode {flow,ddpm,wgan}` / `--stage1-generator` / `--stage2-generator`,
+`--stage2-decoder {autoregressive,one_shot}`, `--conditioning
+{physical,embedding,onehot}`, `--router` / `--router-type` / `--n-experts` /
+`--router-axis`, `--stage{1,2}-init-from` + `--stage{1,2}-freeze` (retrain one
+stage against a fixed other one), `--precision {fp32,bf16}`, `--wandb`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Rollout &amp; analysis</h2></summary>
+
+`giant rollout` seeds showers from each event's highest-energy entry step, then
+autoregressively steps the model to completion — advancing all active tracks
+breadth-first, batched — pushing secondaries as new tracks and looking up
+`material`/`layer_id` from the geometry oracle each step. Tracks terminate on
+one of six reasons (energy cutoff, max steps, detector escape, natural end,
+unknown pdg, max tracks); every reason but escape deposits the remaining energy
+locally, so showers conserve energy by construction — only `escaped` counts as
+leakage.
+
+`giant analyze` compares one or more rollouts against a single held-out
+reference: `prep` resolves shared bin edges/groups once, `submit`/`compute-one`
+run each (plot, `event_id`-disjoint chunk) pair as a polars/numpy-only HTCondor
+job, `render` merges the chunks and produces the styled PDFs + HTML gallery
+locally (the only step that needs LaTeX). Each rollout gets its own colored
+series against one shared reference line. `giant analyze metrics` is a separate
+entry point — training-progress plots straight from a run's `metrics.csv`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Install</h2></summary>
+
+| Extra | Adds | For |
+|---|---|---|
+| `cpu` **or** `cuda` | torch 2.3.x | required — mutually exclusive, pick one |
+| `geometry` | scikit-learn | `dwarf build-geometry-oracle`, rollout |
+| `analysis` | matplotlib, plotstyle | `giant analyze render` |
+| `convert` | uproot, awkward | `dwarf convert` |
+| `wandb` | wandb | `giant train --wandb` |
+| `dev` | pytest, ruff, ty, + all of the above | development |
 
 ```bash
-uv sync --extra cpu                   # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
-uv sync --extra cpu --extra dev       # add dev tools (pytest, ruff, ty)
-uv sync --extra cpu --extra geometry  # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
-uv sync --extra cpu --extra analysis  # matplotlib/polars/plotstyle, for `giant analyze render`
-uv sync --extra cpu --extra convert   # uproot/awkward/polars, for `dwarf convert`
-uv sync --extra cpu --extra wandb     # W&B logging (`giant train --wandb`)
-uv sync --extra cpu --extra workflow  # b2luigi, for `giant workflow run`
+uv sync --extra cpu --extra dev     # everything needed to develop
 ```
 
-The `dev` extra pulls in `convert`, `analysis`, `geometry`, `wandb` and `workflow` as well.
+Plain `uv sync` with no extra installs **no torch at all** — always include
+`--extra cpu` or `--extra cuda`.
 
-`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
+</details>
 
-## Training, prediction, rollout
+<details>
+<summary><h2 style="display:inline">Development</h2></summary>
 
 ```bash
-giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir
-giant train path/to/steps.parquet                          # train (flow stage 1 + wgan stage 2, default)
-giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
-
-dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # position → material/layer_id
-giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
-```
-
-Useful flags on `giant train`:
-
-- `--mode {flow,ddpm,wgan}` sets both stages' objective at once; `--stage1-generator`/`--stage2-generator` override per stage
-- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
-- `--conditioning {physical,embedding,onehot}` — conditioning representation
-- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
-- `--stage2-stage1-context {truth,sampled}` — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via `stage2_model.ctx_p_start`/`ctx_p_end`)
-- `--precision {fp32,bf16}` — bf16 autocast in the training loop
-- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
-- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
-- `--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
-
-Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`/`.class_weighting`, `stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`, `stage*_model.trunk.*` and the finer `router` knobs (`lambda_balance`, `gumbel`, `learn_width`, …). `configs/` holds kept reference configs. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
-
-`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.
-
-## Validation and analysis
-
-- `giant.validate.validate_marginals` — step-level marginal + KL-divergence checks during training (`--validate-every`)
-- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
-
-```bash
-giant analyze prep rollout.yaml --chunks 8                 # lay out the run directory
-giant analyze prep a.yaml b.yaml --label flow --label wgan  # N rollouts vs one shared reference
-giant analyze render <run_dir> --gallery                   # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
-
-giant analyze list                                         # every catalog plot id
-giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0   # what a condor job runs
-giant analyze merge-one --id marginal_edep --run-dir <run_dir>               # merge one plot's chunks (debugging)
-```
-
-The per-(plot, chunk) compute jobs themselves are submitted by the workflow (below), not by `giant analyze` — these commands are the single-step primitives it calls. `<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep` prints it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
-
-## Workflow orchestration
-
-Multi-step pipelines run through [b2luigi](https://github.com/belle2/b2luigi) — one spec file describes a whole experiment, and every step's outputs are files on `/ceph` that are only recomputed when their spec (or an upstream one) changes:
-
-```bash
-uv sync --extra cpu --extra workflow
-giant workflow run configs/workflow_example.toml --mode dry-run   # what would run
-giant workflow run configs/workflow_example.toml --mode show-output  # where every output goes
-giant workflow run configs/workflow_example.toml --batch --workers 20  # submit to HTCondor and wait
-```
-
-The spec holds `[workflow]`/`[condor]`/`[dataset]`/`[geometry]` plus repeated `[[train]]`, `[[rollout]]` and `[[analysis]]` tables cross-referenced by name (see `configs/workflow_example.toml`). The task graph is `DatasetTask → WarmCacheTask/GeometryOracleTask → TrainEpochTask… → TrainTask → RolloutTask → AnalysisPrepTask → AnalysisComputeTask(plot, chunk) → AnalysisRenderTask`. Training is split into one short GPU job per epoch (chained by `--resume`), which schedules better on a busy farm and survives preemption; `TrainTask` then publishes one `best.pt`/`last.pt`/`metrics.csv` for everything downstream. Rendering always runs locally, since it is the only step that needs LaTeX.
-
-Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.
-
-## Development
-
-```bash
-uv run pytest            # run tests
+uv run pytest            # 964 tests
 uv run ruff check .      # lint
 uv run ruff format .     # format
 uv run ty check .        # type check
 ```
+
+Gitea Actions (`.gitea/workflows/ci.yml`) runs lint + format-check + type-check
++ tests on every push and PR; merges to `master` auto-bump the patch version
+and regenerate `CHANGELOG.md` — don't hand-edit either.
+
+</details>
+
+## License
+
+Not yet decided — treat this repository as all-rights-reserved until a
+`LICENSE` file is added.

cliff.toml

@@ -37,7 +37,7 @@ commit_preprocessors = [
 protect_breaking_commits = false
 commit_parsers = [
     { message = "^Merge ", skip = true },
-    { message = "^chore: (bump version|update changelog|sync project version)", skip = true },
+    { message = "^chore: (release|bump version|update changelog|sync project version)", skip = true },
     { message = "^Add", group = "<!-- 0 -->Added" },
     { message = "^(Fix|Clamp|Clip)", group = "<!-- 1 -->Fixed" },
     { message = "^(Remove|Drop|Deprecate)", group = "<!-- 2 -->Removed" },

configs/baseline.toml

@@ -84,15 +84,30 @@ dropout      = 0.0
 # secondary-species failure. Flow (not the schema default wgan) so the
 # baseline varies only the decoder relative to the best v0.2 result.
 #
-# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
-# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
-# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
-# k_max * steps = 150 stage-2 calls per physics step. That makes this block
-# the dominant cost on both sides:
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated — but see
+# the row-compaction note below, which changes the INFERENCE side of this:
 #   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
 #   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
 # Accepted deliberately: one-shot is the configuration whose secondary
 # species distribution failed, and that failure is what v0.3 exists to fix.
+#
+# Row compaction (landed after the above measurement): at inference,
+# sample.sample_secondaries_ar used to loop `for k in range(k_max)`
+# unconditionally — all 15 slots regardless of predicted n_sec — so a flow
+# AR token cost k_max * steps = 150 stage-2 calls per physics step. It now
+# drops a row from the batch the moment its own secondary count is
+# exhausted, so the real inference cost is ~n_sec * steps stage-2 calls
+# (this checkpoint's own rollout measured 0.382 secondaries/step — see
+# giant-baseline-flow-ar-rollout-validation.md), not k_max * steps. A CPU
+# micro-benchmark at that multiplicity (giant/model/history.py's
+# hidden_dim=512/6-block shape, k_max=15, batch 512) measured 17.6-22.9x
+# fewer wall-clock seconds for the AR loop alone (markov/attention history
+# respectively) — directional only (CPU, synthetic n_sec distribution, not
+# an end-to-end rollout); the 8.1x inference ratio above is now stale and
+# should be re-measured on GPU via a real rollout + `eval_cost_per_step`
+# once one is run against this checkpoint. Training cost (the 6.5x/29.5k
+# figures) is untouched by this: teacher_forcing = "always" here never
+# calls the AR sampler at train time (see [stage2_model.autoregressive]).
 decoder      = "autoregressive"
 generator    = "flow"
 hidden_dim   = 512

configs/wgan_ar_stoptoken.toml

@@ -0,0 +1,122 @@
+# GIANT WGAN-GP + AR stop-token variant of configs/baseline.toml.
+#
+# Two roadmap axes, combined into one run: WGAN-GP generators for both
+# stages (unbenchmarked since the 2026-08-03 pre-v0.3.0 failure, which was
+# secondary-species mode collapse — the failure v0.3.0's AR/categorical
+# pivot exists to fix) and the AR stop-token multiplicity mode
+# (stage2_model.n_sec.mode = "stop_token", never benchmarked at all).
+# Everything else is byte-identical to baseline.toml so a rollout compared
+# against baseline's analysis_341dfb14 is attributable to these two axes
+# alone: conditioning (physical/physical), hidden_dim 512 / n_res_blocks 6 /
+# dropout 0.0 per stage, k_max 15, history "markov", teacher_forcing
+# "always", particle_type.target "onehot" (n_classes 32, other_policy
+# "sample"), lr 3e-4, warmup_epochs 3, weight_decay 0.01, ema_decay 0.9999,
+# val_fraction 0.1, num_workers 4, seed 0, validate_steps 10, W&B on.
+#
+# No [stage1_model.wgan] / [stage2_model.wgan] block: the dataclass defaults
+# (noise_dim 64, n_critic 5, gp_weight 10.0, critic_lr 0.0 = inherit
+# train.lr, critic_hidden_dim/critic_n_res_blocks 0 = inherit the stage's
+# 512/6, stage 2's gumbel_tau_start/_end 1.0/0.1) are what the earlier WGAN
+# runs used — writing them out would add keys that don't vary.
+#
+# particle_type.class_weighting stays "none" (the default): config.py's
+# validate_config rejects any other value under stage2_model.generator =
+# "wgan", since that path feeds the type slice to the critic via a
+# straight-through Gumbel relaxation instead of a weighted cross-entropy.
+#
+# Prior WGAN writeup (pre-v0.3.0, describes the failure this run re-tests):
+# /home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md
+
+[meta]
+# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+# rewrites it from V02_FIXED_FACTS — silently forcing decoder = "one_shot",
+# particle_type.target = "physical" and the v0.2 default sizes, while still
+# passing validate_config.
+config_version = 3
+
+[conditioning]
+# Physical-property MLPs rather than learned vocab embeddings: computable for
+# any PDG code / material, which is what the held-out-species and
+# held-out-material generalization comparisons need.
+out_dim      = 128
+share_stages = false
+
+# n_layers = 2 rather than the v0.3 default of 1: v0.2's conditioning MLP was
+# always 2 deep (see _migration.V02_FIXED_FACTS), so this keeps the encoder
+# identical to baseline.toml.
+[conditioning.particle]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[conditioning.material]
+type     = "physical"
+emb_dim  = 16
+n_layers = 2
+
+[stage1_model]
+generator    = "wgan"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+
+[stage2_model]
+# Autoregressive in descending-energy order, as baseline.toml — this variant
+# only swaps the generator (flow -> wgan) and the multiplicity mode
+# (head -> stop_token), not the decoder shape.
+decoder      = "autoregressive"
+generator    = "wgan"
+hidden_dim   = 512
+n_res_blocks = 6
+dropout      = 0.0
+k_max        = 15
+
+[stage2_model.autoregressive]
+history         = "markov"
+teacher_forcing = "always"
+
+[stage2_model.n_sec]
+# EOS-style per-slot stop head on the AR secondary decoder, replacing the
+# n_sec classifier entirely (mutually exclusive — see NSecConfig's
+# docstring in giant/config.py). Requires decoder = "autoregressive" and
+# owner = "stage2" (both already true above/by default); validate_config
+# enforces this.
+mode = "stop_token"
+
+[stage2_model.particle_type]
+target = "onehot"
+# Decoupled from conditioning.particle.emb_dim (gitea #29). 32 classes + the
+# "other" bucket keeps essentially all real secondary species out of "other"
+# without making the head expensive.
+n_classes    = 32
+other_policy = "sample"
+
+[train]
+epochs = 30
+# Halved from baseline's 36864. That figure came from a measured linear fit
+# of the *flow-AR* training step (peak reserved MiB = 0.9736 * batch_size +
+# 115); WGAN invalidates it twice over — each stage gains a critic that by
+# default inherits the stage's own 512/6 body, and gradient_penalty
+# (giant/model/wgan.py, forced fp32 internally) runs a double-backward every
+# batch. 18432 is a conservative choice pending a real memory measurement on
+# this exact config, not a re-derived fit. Throughput is already flat above
+# bs~4096 on the 4070, so this costs occupancy on the L40S, not step
+# efficiency.
+batch_size    = 18432
+lr            = 3e-4
+warmup_epochs = 3
+weight_decay  = 0.01
+ema_decay     = 0.9999
+val_fraction  = 0.1
+num_workers   = 4
+seed          = 0
+# Tightened from baseline's 10: WGANStageTrainer.supports_val_loss = False,
+# and with both stages adversarial there is no per-epoch val loss at all, so
+# validate_every's marginal-KL pass (giant/training/trainers.py's
+# val_objective) is the only comparable-across-epochs best-checkpoint
+# selection signal available. 5 gives 6 evaluations over 30 epochs instead
+# of baseline's 3, at ~6x5000s of extra walltime.
+validate_every = 5
+validate_steps = 10
+wandb          = true
+wandb_project  = "giant"

configs/workflow_example.toml

@@ -1,68 +0,0 @@
-# Example GIANT workflow spec — `giant workflow run configs/workflow_example.toml`.
-#
-# One file parameterises a whole experiment: the datasets, the geometry oracle,
-# N trainings, N rollouts, and the analyses comparing them. Every task's output
-# directory carries a hash of its resolved sub-spec (plus its parents), so
-# editing anything here re-runs exactly the affected subtree and nothing else.
-#
-# result_dir/log_dir must be visible from both the submit host and the workers
-# (i.e. on /ceph) — there is deliberately no HTCondor file transfer.
-
-[workflow]
-name       = "baseline-vs-router"
-result_dir = "/ceph/lbogner/workflows/baseline-vs-router"
-log_dir    = "/ceph/lbogner/workflows/baseline-vs-router/logs"
-
-[condor]
-accounting_group = "cms"
-repo_dir         = "/work/lbogner/giant"   # also b2luigi's working_dir
-env_script       = "/work/lbogner/giant/condor_env.sh"
-docker_image_cpu = "cverstege/alma9-gridjob"
-docker_image_gpu = "mschnepf/slc7-condocker"
-remote           = true
-
-[dataset]
-steps     = "/ceph/lbogner/geant_steps/train/"     # training data
-reference = "/ceph/lbogner/geant_steps/holdout/"   # rollout seeds + analysis truth
-
-[geometry]
-method    = "slab"
-subsample = 500_000
-
-[[train]]
-name           = "baseline"
-config         = "configs/baseline.toml"
-epochs         = 200
-epochs_per_job = 1            # one short GPU job per epoch, chained
-request_gpus   = 1
-gpu_memory_mb  = 20000
-overrides      = { lr = 3e-4 }   # `giant train` flag names
-
-[[train]]
-name           = "router-balanced"
-config         = "configs/router.toml"
-epochs         = 200
-epochs_per_job = 1
-request_gpus   = 1
-gpu_memory_mb  = 20000
-
-[[rollout]]
-name          = "baseline"
-train         = "baseline"     # -> [[train]].name
-n_events      = 2000
-energy_cutoff = 0.1
-
-[[rollout]]
-name          = "router-balanced"
-train         = "router-balanced"
-n_events      = 2000
-energy_cutoff = 0.1
-
-[[analysis]]
-name        = "baseline-vs-router"
-rollouts    = ["baseline", "router-balanced"]
-chunks      = 32
-energy_bins = 4
-bins        = 50
-top_pdg     = 6
-gallery     = true

giant/analysis/__init__.py

@@ -4,7 +4,7 @@ Compares one or more autoregressive ``giant rollout`` runs against a single
 held-out miniCaloSim reference file shared by all of them, producing
 publication-styled comparison plots (one colored series per rollout, one
 reference line) generated in parallel on HTCondor (one job per plot x data
-chunk, compute/merge/render split) — orchestrated by ``giant/workflow``.
+chunk, compute/merge/render split).
 
 Only ``render`` (and the ``render`` CLI path) imports plotstyle/LaTeX; everything
 re-exported here is plotstyle-free so it runs on a compute worker. Import
@@ -12,43 +12,55 @@ 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.run import (
+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,
     merge_one,
     prep,
+    write_submit,
 )
 from giant.analysis.context import Context, build_context
+from giant.analysis.prediction import PredictionSpec
 from giant.analysis.reduced import Partial, Reduced
 from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
 from giant.analysis.sources import RolloutSpec, Side
 
 __all__ = [
-    "build_catalog",
-    "catalog_ids",
-    "get_spec",
+    "RUNTIME_SAFETY_MARGIN",
+    "Context",
+    "LoadedPrediction",
     "LoadedRollout",
+    "Partial",
+    "PredictionSpec",
+    "Reduced",
+    "RolloutSpec",
     "RunMeta",
+    "Side",
+    "SubmitConfig",
+    "build_catalog",
+    "build_context",
+    "catalog_ids",
     "compute_one",
     "compute_reduced",
     "derive_run_dir",
+    "estimate_runtime_s",
+    "get_spec",
+    "load_prediction_yaml",
+    "load_prediction_yamls",
     "load_rollout_yaml",
     "load_rollout_yamls",
     "merge_all",
     "merge_one",
     "prep",
-    "Context",
-    "build_context",
-    "Partial",
-    "Reduced",
-    "RolloutSpec",
-    "Side",
-    "RUNTIME_SAFETY_MARGIN",
-    "estimate_runtime_s",
+    "write_submit",
 ]

giant/analysis/condor.py

@@ -1,6 +1,4 @@
-"""Analysis run directories: prep, per-(plot, chunk) compute, and merge.
-
-Driven by one or more ``giant rollout`` YAML sidecars.
+"""HTCondor orchestration driven by one or more ``giant rollout`` YAML sidecars.
 
 A rollout writes a YAML sidecar (``giant/cli.py:_write_prediction_ref`` +
 rollout extras) that already names both files we need and carries the run's
@@ -24,16 +22,13 @@ everything out under it:
     <run_dir>/reduced/<id>.json   merged, per plot
     <run_dir>/plots/<family>/<id>.pdf   rendered locally
 
-Job model (one job per (plot, chunk), compute/merge/render split). Job
-submission itself is b2luigi's (``giant/workflow/tasks.py`` — ``AnalysisPrepTask``
-/ ``AnalysisComputeTask`` / ``AnalysisRenderTask``); this module only provides
-the three steps they call:
+Job model (one condor job per (plot, chunk), compute/merge/render split):
 
-1. ``prep`` runs once locally — reads the YAML, resolves the shared
+1. ``prep`` runs once on the submit node — reads the YAML, resolves the shared
    context from a subsample, writes ``shared.json`` + ``run_meta.json``
    (including the run's configured ``n_chunks``).
 2. one job per catalog id x chunk index runs ``giant analyze compute-one
-   --run-dir`` (or ``compute_one`` in-process) on a worker — a single streaming pass over that
+   --run-dir`` on a worker — a single streaming pass over that
    ``event_id``-disjoint chunk, writing ``reduced_partial/<id>__<chunk>.json``
    (polars/numpy only, no LaTeX). Specs marked ``chunkable=False``
    (``PlotSpec``, ``catalog.py``) always run as a single chunk.
@@ -43,15 +38,15 @@ the three steps they call:
    ``reduced/<id>.json``, then renders those into the styled PDF + gallery tree
    (that step imports plotstyle/LaTeX).
 
-Files on ``/ceph`` or ``/work`` are reached directly (see
-``giant/workflow/htcondor.py``); no HTCondor file transfer of the multi-GB
-inputs.
+Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no
+HTCondor file transfer of the multi-GB inputs.
 """
 
 from __future__ import annotations
 
 import json
 import shutil
+import sys
 from collections.abc import Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
@@ -61,10 +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",
@@ -88,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,
@@ -167,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]
@@ -227,17 +302,26 @@ class RunMeta:
     # Empty/0 on run directories written before this field existed.
     rows_per_chunk: list[int] = field(default_factory=list)
     total_rows: int = 0
+    # `giant predict` inputs (the paired-truth "prediction" family) — same
+    # shape as `rollouts`. Empty on a run with no --prediction, so old
+    # run_meta.json files still load.
+    predictions: list[dict] = field(default_factory=list)
 
     def save(self, path: str | Path) -> None:
         Path(path).write_text(json.dumps(self.__dict__, indent=2))
 
     @classmethod
-    def load(cls, path: str | Path) -> "RunMeta":
+    def load(cls, path: str | Path) -> RunMeta:
         return cls(**json.loads(Path(path).read_text()))
 
 
-def _rows_per_chunk(rollouts: list[str | Path], reference: str | Path, n_chunks: int) -> list[int]:
-    """Combined rollout+reference row count of each ``event_id % n_chunks`` chunk.
+def _rows_per_chunk(
+    rollouts: list[str | Path],
+    reference: str | Path,
+    n_chunks: int,
+    predictions: Sequence[str | Path] = (),
+) -> list[int]:
+    """Combined rollout+reference+prediction row count of each ``event_id % n_chunks`` chunk.
 
     One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
     the sizing input every job's estimated walltime
@@ -253,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()):
@@ -267,26 +355,33 @@ def prep(
     n_chunks: int = 1,
     default_base: str | Path | None = None,
     labels: Sequence[str] | None = None,
+    prediction_yamls: Sequence[str | Path] = (),
+    prediction_labels: Sequence[str] | None = None,
     **ctx_kwargs,
 ) -> Path:
-    """Read the rollout YAML(s), build the shared context, and lay out the run dir.
+    """Read the rollout (+ optional prediction) YAML(s), build the shared
+    context, and lay out the run dir.
 
     Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
     ``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one``
     job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is
     resolved once here rather than re-passed (and risking disagreement) at every
     later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
-    resolve the actual directory, and ``load_rollout_yamls`` for how
-    ``labels``/YAML stems resolve each rollout's series name.
+    resolve the actual directory, ``load_rollout_yamls`` for how
+    ``labels``/YAML stems resolve each rollout's series name, and
+    ``load_prediction_yamls`` for the same on ``prediction_yamls`` (which,
+    unlike rollouts, is optional — the ``prediction`` plot family degrades to
+    ``kind="unavailable"`` when it's empty).
 
     Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
     this same ``run_dir``: partial files carry no record of what context
     (``n_chunks``, bin edges, group sets) they were computed under, so
     re-prepping with a different ``n_chunks``/``**ctx_kwargs`` (or after the
-    rollout/reference files changed) would otherwise let ``merge_one`` silently
-    merge stale partials against the new ``shared.json``.
+    rollout/reference/prediction files changed) would otherwise let
+    ``merge_one`` silently merge stale partials against the new ``shared.json``.
     """
     loaded, reference = load_rollout_yamls(list(rollout_yamls), labels)
+    loaded_preds = load_prediction_yamls(list(prediction_yamls), reference, prediction_labels)
     run_path = derive_run_dir([lr.yaml for lr in loaded], run_dir, default_base=default_base)
     run_path.mkdir(parents=True, exist_ok=True)
 
@@ -296,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(
@@ -314,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
 
@@ -331,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``
@@ -361,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,
@@ -387,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,
@@ -395,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,
     )
 
 
@@ -436,3 +544,136 @@ def merge_one(spec_id: str, run_dir: str | Path) -> Path:
 def merge_all(run_dir: str | Path) -> list[Path]:
     """Merge every catalog plot's chunk partials into ``reduced/<id>.json``."""
     return [merge_one(spec_id, run_dir) for spec_id in catalog_ids()]
+
+
+# ---------------------------------------------------------------------------
+# submit description
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class SubmitConfig:
+    run_dir: Path
+    accounting_group: str
+    repo_dir: Path
+    docker_image: str = "cverstege/alma9-gridjob"
+    request_memory_mb: int = 8192
+    request_cpus: int = 1
+    remote: bool = False  # +RemoteJob (grid I/O) vs ProvidesETPResources (local files)
+    n_chunks: int = 1  # per-plot data chunks; ignored for chunkable=False specs
+
+
+_WRAPPER = """#!/bin/bash
+set -euo pipefail
+cd {repo_dir}
+exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
+"""
+
+
+def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
+    reqs_attrs = "+RemoteJob = True\n" if cfg.remote else "requirements = TARGET.ProvidesETPResources\n"
+    return (
+        "universe = docker\n"
+        f"docker_image = {cfg.docker_image}\n"
+        f"executable = {wrapper}\n"
+        "arguments = $(plotid) $(chunk)\n"
+        "should_transfer_files = YES\n"
+        "when_to_transfer_output = ON_EXIT\n"
+        f"request_memory = {cfg.request_memory_mb}\n"
+        f"request_cpus = {cfg.request_cpus}\n"
+        "+RequestWalltime = $(walltime)\n"
+        f"accounting_group = {cfg.accounting_group}\n"
+        f"{reqs_attrs}"
+        f"output = {cfg.run_dir}/logs/$(plotid)__$(chunk).out\n"
+        f"error = {cfg.run_dir}/logs/$(plotid)__$(chunk).err\n"
+        f"log = {cfg.run_dir}/logs/condor.log\n"
+        f"queue plotid,chunk,walltime from {jobs_file}\n"
+    )
+
+
+def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]:
+    """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.
+
+    Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
+    ``chunkable=False`` specs (router diagnostics) always use the dataset
+    total since they run as a single job regardless of ``n_chunks``.
+    """
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    jobs: list[tuple[str, int, int]] = []
+    for spec_id in ids:
+        chunkable = get_spec(spec_id).chunkable
+        chunks = range(n_chunks) if chunkable else [0]
+        for chunk in chunks:
+            n_rows = meta.rows_per_chunk[chunk] if chunkable else meta.total_rows
+            jobs.append((spec_id, chunk, estimate_runtime_s(spec_id, n_rows)))
+    return jobs
+
+
+def _resolve_giant_executable(repo_dir: Path) -> Path:
+    """Path to the ``giant`` entry point to bake into the condor wrapper script.
+
+    Prefers the venv currently running this process (``sys.executable``'s
+    sibling ``giant``) so a submit from a non-default venv (e.g. ``--extra
+    cuda`` on a dev box) doesn't silently pick up a different one; falls back
+    to ``repo_dir/.venv/bin/giant`` for the case this is invoked from outside
+    any venv (e.g. a system Python).
+    """
+    active = Path(sys.executable).parent / "giant"
+    if active.exists():
+        return active
+    venv_giant = repo_dir / ".venv" / "bin" / "giant"
+    if not venv_giant.exists():
+        raise FileNotFoundError(
+            f"no `giant` executable found next to {sys.executable} or at "
+            f"{venv_giant} — condor jobs run it directly (no `uv` on the "
+            f"worker image), so run `uv sync --extra cpu` in {repo_dir} "
+            "before submitting."
+        )
+    return venv_giant
+
+
+def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
+    """Write the wrapper script, (plot, chunk) job list, and HTCondor submit
+    description.
+
+    Each catalog id gets ``cfg.n_chunks`` jobs, except ``chunkable=False``
+    specs (the router diagnostics), which always get exactly one regardless of
+    ``cfg.n_chunks``. Every job's ``+RequestWalltime`` is estimated from its
+    chunk's row count (``runtime_estimate.estimate_runtime_s``, requires
+    ``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``).
+    Returns the submit description path (``<run_dir>/analyze.sub``). Does not
+    submit — call ``condor_submit`` on the returned file.
+
+    ``cfg.n_chunks`` and the run directory's own ``RunMeta.n_chunks`` (fixed by
+    ``prep``, and what ``RunMeta.rows_per_chunk`` was sized against) are two
+    independent values — checked equal up front so a mismatch is a clear error
+    here rather than an ``IndexError`` out of ``_job_walltimes``.
+    """
+    giant_exe = _resolve_giant_executable(cfg.repo_dir)
+
+    ids = ids or catalog_ids()
+    run_dir = cfg.run_dir
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    if cfg.n_chunks != meta.n_chunks:
+        raise ValueError(
+            f"SubmitConfig.n_chunks={cfg.n_chunks} does not match the "
+            f"n_chunks this run directory was prepped with "
+            f"(RunMeta.n_chunks={meta.n_chunks} in {run_dir}/run_meta.json) — "
+            "re-run `prep` with the desired n_chunks, or fix cfg.n_chunks to "
+            "match it."
+        )
+    (run_dir / "logs").mkdir(parents=True, exist_ok=True)
+    (run_dir / "reduced").mkdir(parents=True, exist_ok=True)
+    (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)
+
+    wrapper = run_dir / "run_compute.sh"
+    wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir))
+    wrapper.chmod(0o755)
+
+    jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
+    jobs_file = run_dir / "jobs.txt"
+    jobs_file.write_text("\n".join(f"{i},{k},{w}" for i, k, w in jobs) + "\n")
+
+    sub = run_dir / "analyze.sub"
+    sub.write_text(_submit_description(cfg, wrapper, jobs_file))
+    return sub

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)
-#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint)
+#                    distance scorecard) or per prediction (a confusion matrix)
+#   "paired_hist"    per-prediction true/pred density histogram over shared
+#                    edges (giant predict's paired truth, not a rollout)
+#   "heatmap2d"      numeric x/y-binned true-vs-pred count matrix + colorbar,
+#                    one panel per prediction, with a y=x diagonal guide
+#   "unavailable"    plot not applicable to this run (e.g. no MoE checkpoint,
+#                    or no --prediction given)
 
 
 @dataclass
@@ -46,7 +51,7 @@ class Reduced:
         Path(path).write_text(json.dumps(asdict(self)))
 
     @classmethod
-    def load(cls, path: str | Path) -> "Reduced":
+    def load(cls, path: str | Path) -> Reduced:
         return cls(**json.loads(Path(path).read_text()))
 
 
@@ -70,5 +75,5 @@ class Partial:
         Path(path).write_text(json.dumps(asdict(self)))
 
     @classmethod
-    def load(cls, path: str | Path) -> "Partial":
+    def load(cls, path: str | Path) -> Partial:
         return cls(**json.loads(Path(path).read_text()))

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,
 }
 
@@ -545,7 +619,7 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
     (checkpoint, paths, cutoffs) from ``run_meta.json`` into every plot's
     gallery metadata and renders.
     """
-    from giant.analysis.run import RunMeta, merge_all
+    from giant.analysis.condor import RunMeta, merge_all
 
     run_dir = Path(run_dir)
     merge_all(run_dir)

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/runtime_estimate.py

@@ -1,4 +1,4 @@
-"""Per-(plot, chunk) HTCondor walltime estimates for the analysis compute jobs.
+"""Per-(plot, chunk) HTCondor walltime estimates for `giant analyze submit`.
 
 Each catalog spec's compute cost is close to linear in the number of input
 rows a `compute-one` job streams over — every spec is one (or a couple of)

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:
@@ -1888,19 +1888,6 @@ def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path
     return out_dir
 
 
-def epoch_seed(seed: int, epoch: int) -> int:
-    """Per-epoch derivative of the run seed.
-
-    Reseeding the global RNGs from this at the top of every epoch makes epoch
-    *k* draw the same noise whether it runs inside one long `giant train` or
-    as its own resumed job in a per-epoch workflow chain
-    (`giant/workflow/tasks.py:TrainEpochTask`) — without it, a fresh process
-    would restart the stream at epoch 1's state. Mirrors what
-    `StreamingStepsDataset.set_epoch` does for the batch order.
-    """
-    return (int(seed) * 1_000_003 + int(epoch)) % (2**32)
-
-
 def seed_everything(seed: int) -> None:
     import numpy as np
     import torch
@@ -1964,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/dataset.py

@@ -97,7 +97,6 @@ class StreamingStepsDataset(IterableDataset):
         mat_topn_map: dict[str, int] | None = None,
         sec_type_class_map: dict | None = None,
         k_max: int = K_MAX,
-        seed: int = 0,
     ) -> None:
         self.files = list(files)
         self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
@@ -118,35 +117,16 @@ class StreamingStepsDataset(IterableDataset):
         self.mat_topn_map = mat_topn_map
         self.sec_type_class_map = sec_type_class_map
         self.k_max = k_max
-        self.seed = seed
-        self.epoch = 0
-        self._rng = np.random.default_rng()
-
-    def set_epoch(self, epoch: int) -> None:
-        """Select the shuffle stream for `epoch` (the DistributedSampler convention).
-
-        The training loop calls this at the top of every epoch. Shuffling is
-        seeded from `(seed, epoch, worker_id)` rather than the global numpy
-        state so epoch *k*'s batch order is the same whether it runs as epoch
-        *k* of one long `giant train`, or as its own resumed job in a
-        per-epoch workflow chain (`giant/workflow/tasks.py:TrainEpochTask`).
-        Workers are re-forked from this object each epoch (no
-        `persistent_workers`), so setting it here reaches them.
-        """
-        self.epoch = int(epoch)
 
     def __iter__(self):
         worker_info = torch.utils.data.get_worker_info()
         files = self.files
-        worker_id = worker_info.id if worker_info is not None else 0
         if worker_info is not None:
-            files = files[worker_id :: worker_info.num_workers]
-
-        self._rng = np.random.default_rng([self.seed, self.epoch, worker_id])
+            files = files[worker_info.id :: worker_info.num_workers]
 
         if self.shuffle:
             files = list(files)
-            self._rng.shuffle(files)
+            np.random.default_rng().shuffle(files)
 
         buf_cont: list[np.ndarray] = []
         buf_cat: list[np.ndarray] = []
@@ -242,7 +222,7 @@ class StreamingStepsDataset(IterableDataset):
         styp = np.concatenate(buf_type)
 
         if self.shuffle:
-            idx = self._rng.permutation(len(cont))
+            idx = np.random.permutation(len(cont))
             cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
             nsec, sec, proc, styp = nsec[idx], sec[idx], proc[idx], styp[idx]
 

giant/data/loader.py

@@ -1,13 +1,17 @@
+from collections.abc import Iterator, Mapping
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Iterator
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
-import pandas as pd
+import polars as pl
 import pyarrow.parquet as pq
 
 from giant.constants import K_MAX
 
+if TYPE_CHECKING:
+    from giant.data.scan import ValueStat
+
 # A manifest is a plain text file listing one parquet path per line, used to
 # name a curated subset of files (e.g. a train/holdout pool) without copying
 # or symlinking the underlying parquet files. Lines are resolved relative to
@@ -77,88 +81,75 @@ def find_parquet_files(path: str | Path) -> list[Path]:
     return [p]
 
 
-def _pad_list_col(series: pd.Series, K: int, fill: float = 0.0) -> np.ndarray:
-    """Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
-    out = np.full((len(series), K), fill, dtype=np.float32)
-    for i, lst in enumerate(series):
-        if lst is not None and len(lst) > 0:
-            n = min(len(lst), K)
-            out[i, :n] = lst[:n]
-    return out
+def _pad_list_column(df: pl.DataFrame, col: str, k: int, fill, dtype: type[pl.DataType] | pl.DataType) -> np.ndarray:
+    """Pad / truncate a list-valued column to fixed width `k` → (N, k) numpy array.
 
-
-def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
-    """Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
-    out = np.full((len(series), K), fill, dtype=np.int64)
-    for i, lst in enumerate(series):
-        if lst is not None and len(lst) > 0:
-            n = min(len(lst), K)
-            out[i, :n] = lst[:n]
-    return out
-
-
-def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
-    """Pad three list-valued direction columns → (N, K, 3) float32.
-
-    Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
+    Concatenating `k` fill values before truncating to `k` guarantees every
+    row ends up with exactly `k` non-null elements regardless of how short
+    (including empty) or long the original list was, so `list.to_array(k)`
+    (a fixed-size-array dtype) converts to a plain 2D numpy array with a
+    single vectorized expression — no per-row Python loop.
     """
-    N = len(dx)
-    out = np.zeros((N, K, 3), dtype=np.float32)
-    out[:, :, 2] = 1.0
-    for i in range(N):
-        lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
-        if lx is not None and len(lx) > 0:
-            n = min(len(lx), K)
-            out[i, :n, 0] = lx[:n]
-            out[i, :n, 1] = ly[:n]
-            out[i, :n, 2] = lz[:n]
-    return out
+    fill_tail = pl.lit([fill] * k, dtype=pl.List(dtype))
+    out = df.select(pl.col(col).cast(pl.List(dtype)).list.concat(fill_tail).list.head(k).list.to_array(k).alias("_p"))
+    return out["_p"].to_numpy()
 
 
-def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
+def _pad_dir_col(df: pl.DataFrame, dx: str, dy: str, dz: str, k: int) -> np.ndarray:
+    """Pad three list-valued direction columns → (N, k, 3) float32.
+
+    Padding direction defaults to (0, 0, 1) (forward) so it is a valid unit vector.
+    """
+    px = _pad_list_column(df, dx, k, 0.0, pl.Float64)
+    py = _pad_list_column(df, dy, k, 0.0, pl.Float64)
+    pz = _pad_list_column(df, dz, k, 1.0, pl.Float64)
+    return np.stack([px, py, pz], axis=-1).astype(np.float32)
+
+
+def _df_to_dict(df: pl.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
     has_sec_lists = "sec_E_list" in df.columns
 
     d: dict[str, np.ndarray] = {
         "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
-        "pdg": df["pdg"].to_numpy(dtype=np.int32),
-        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
-        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
-        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
-        "material": df["material"].to_numpy(dtype=object),
-        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
-        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
-        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
+        "pdg": df["pdg"].to_numpy().astype(np.int32),
+        "pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
+        "pre_E": df["pre_E"].to_numpy().astype(np.float32),
+        "pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
+        "material": df["material"].to_numpy().astype(object),
+        "layer_id": df["layer_id"].to_numpy().astype(np.int32),
+        "n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
+        "e_sec": df["e_sec"].to_numpy().astype(np.float32),
         # The physics process that ended the step (e.g. "compt", "phot",
         # "eBrem") — a post-step outcome, so it's a router/classifier
         # supervision label only, never conditioning (see build_process_map*
         # / ProcessRouter). Guarded like has_sec_lists: older parquet
         # conversions predating this column still load fine.
         "process": (
-            df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
+            df["process"].to_numpy().astype(object) if "process" in df.columns else np.full(len(df), "", dtype=object)
         ),
-        "step_length": df["step_length"].to_numpy(dtype=np.float32),
-        "post_E": df["post_E"].to_numpy(dtype=np.float32),
-        "delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
-        "edep": df["edep"].to_numpy(dtype=np.float32),
-        "post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
-        "post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
+        "step_length": df["step_length"].to_numpy().astype(np.float32),
+        "post_E": df["post_E"].to_numpy().astype(np.float32),
+        "delta_e": (df["pre_E"] - df["post_E"]).to_numpy().astype(np.float32),
+        "edep": df["edep"].to_numpy().astype(np.float32),
+        "post_dir": df.select(["post_dx", "post_dy", "post_dz"]).to_numpy().astype(np.float32),
+        "post_pos": df.select(["post_x", "post_y", "post_z"]).to_numpy().astype(np.float32),
     }
 
     if has_sec_lists:
-        d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
-        d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
-        d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)
+        d["sec_E_list"] = _pad_list_column(df, "sec_E_list", k_max, 0.0, pl.Float64).astype(np.float32)
+        d["sec_pdg_list"] = _pad_list_column(df, "sec_pdg_list", k_max, 0, pl.Int64).astype(np.int64)
+        d["sec_dir_list"] = _pad_dir_col(df, "sec_dx_list", "sec_dy_list", "sec_dz_list", k_max)
 
     return d
 
 
 def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
-    return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
+    return _df_to_dict(pl.read_parquet(path), offset=offset, k_max=k_max)
 
 
 def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
     """Read only the event_id column — cheap scan for split assignment."""
-    ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
+    ids = pl.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
     return _offset_event_id(ids, offset)
 
 
@@ -170,7 +161,7 @@ def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> I
     module constant for callers that don't care (e.g. Stage-1-only reads)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max)
+        yield _df_to_dict(pl.DataFrame(pf.read_row_group(i)), offset=offset, k_max=k_max)
 
 
 _COND_COLS = [
@@ -190,17 +181,17 @@ _COND_COLS = [
 ]
 
 
-def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
+def _cond_df_to_dict(df: pl.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
     return {
         "event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
-        "pdg": df["pdg"].to_numpy(dtype=np.int32),
-        "pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
-        "pre_E": df["pre_E"].to_numpy(dtype=np.float32),
-        "pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
-        "material": df["material"].to_numpy(dtype=object),
-        "layer_id": df["layer_id"].to_numpy(dtype=np.int32),
-        "n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
-        "e_sec": df["e_sec"].to_numpy(dtype=np.float32),
+        "pdg": df["pdg"].to_numpy().astype(np.int32),
+        "pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
+        "pre_E": df["pre_E"].to_numpy().astype(np.float32),
+        "pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
+        "material": df["material"].to_numpy().astype(object),
+        "layer_id": df["layer_id"].to_numpy().astype(np.int32),
+        "n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
+        "e_sec": df["e_sec"].to_numpy().astype(np.float32),
     }
 
 
@@ -208,7 +199,7 @@ def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np
     """Yield conditioning-only row-groups (no post-step columns read from disk)."""
     pf = pq.ParquetFile(path)
     for i in range(pf.num_row_groups):
-        yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)
+        yield _cond_df_to_dict(pl.DataFrame(pf.read_row_group(i, columns=_COND_COLS)), offset=offset)
 
 
 def build_index_maps(
@@ -225,42 +216,28 @@ def build_index_maps(
 def build_index_maps_from_files(
     files: list[Path],
 ) -> tuple[dict[int, int], dict[str, int]]:
-    """Scan only pdg and material columns across all files (2-column read)."""
-    pdg_vals: set[int] = set()
-    mat_vals: set[str] = set()
-    for path in files:
-        df = pd.read_parquet(path, columns=["pdg", "material"])
-        pdg_vals.update(int(v) for v in df["pdg"].unique())
-        mat_vals.update(str(v) for v in df["material"].unique())
+    """Scan only pdg and material columns across all files (fused single-pass scan)."""
+    from giant.data.scan import ScanRequest, scan_metadata
+
+    result = scan_metadata(files, ScanRequest(pdg=True, material=True))
+    assert result.pdg is not None and result.material is not None
     return (
-        {v: i for i, v in enumerate(sorted(pdg_vals))},
-        {v: i for i, v in enumerate(sorted(mat_vals))},
+        {v: i for i, v in enumerate(sorted(result.pdg))},
+        {v: i for i, v in enumerate(sorted(result.material))},
     )
 
 
-def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None:
-    for name, count in series.value_counts().items():
-        name = cast(name)
-        counts[name] = counts.get(name, 0) + int(count)
-
-
-def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
-    """Scan `column` across `files` and return `{cast(value): total_count}`,
-    accumulated in file order (see `fingerprint_files`'s docstring on why
-    scan order — not a normalized/sorted order — is preserved: it drives
-    tie-breaking in the frequency ranking below)."""
-    counts: dict = {}
-    for path in files:
-        df = pd.read_parquet(path, columns=[column])
-        _accumulate_value_counts(counts, df[column], cast)
-    return counts
-
-
-def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]:
+def _topn_plus_other_map(counts: "Mapping[Any, ValueStat]", n_classes: int) -> tuple[dict, dict, dict]:
     """Frequency-capped value->index map: the `n_classes - 1` most frequent
     keys get their own index; every rarer key is bucketed into a shared
     "other" index (`n_classes - 1`).
 
+    `counts` maps each key to something with `.count` and `.first_seen`
+    attributes (`giant.data.scan.ValueStat`) — ties in `.count` are broken by
+    `.first_seen` (whichever value was scanned first: file order, then row
+    order within a file — see `giant.data.scan`'s module docstring). This is
+    an explicit, documented contract, not an accident of iteration order.
+
     Returns `(class_map, other_members, class_counts)` — `other_members` is
     `{key: count}` for every key bucketed into "other" (the empirical
     within-bucket distribution, for `other_policy = "sample"` at rollout);
@@ -270,15 +247,15 @@ def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict
     (gitea #44) needs and that would otherwise be dropped once `counts` is
     collapsed into `class_map`.
     """
-    ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
+    ranked = sorted(counts, key=lambda k: (-counts[k].count, counts[k].first_seen))
     keep = ranked[: max(n_classes - 1, 0)]
     class_map = {k: i for i, k in enumerate(keep)}
-    class_counts = {i: counts[k] for i, k in enumerate(keep)}
+    class_counts = {i: counts[k].count for i, k in enumerate(keep)}
     other_idx = n_classes - 1
     other_members: dict = {}
     for k in ranked[len(keep) :]:
         class_map[k] = other_idx
-        other_members[k] = counts[k]
+        other_members[k] = counts[k].count
     if other_members:
         class_counts[other_idx] = sum(other_members.values())
     return class_map, other_members, class_counts
@@ -294,8 +271,11 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
     mirrors how `build_features` clamps the n_sec label to K_MAX for the
     fixed-width n_sec_head classifier.
     """
-    counts = _rank_by_frequency_from_files(files, "process", str)
-    class_map, _, _ = _topn_plus_other_map(counts, n_experts)
+    from giant.data.scan import ScanRequest, scan_metadata
+
+    result = scan_metadata(files, ScanRequest(process=True))
+    assert result.process is not None
+    class_map, _, _ = _topn_plus_other_map(result.process, n_experts)
     return class_map
 
 
@@ -327,8 +307,13 @@ def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, ca
     later for `other_policy = "sample"` at rollout — computed now since it's
     free during this same scan.
     """
-    counts = _rank_by_frequency_from_files(files, column, cast)
-    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    from giant.data.scan import ScanRequest, scan_metadata
+
+    if column != "material":
+        raise ValueError(f"build_topn_map_from_files only supports column='material', got {column!r}")
+    result = scan_metadata(files, ScanRequest(material=True))
+    assert result.material is not None
+    class_map, other_members, class_counts = _topn_plus_other_map(result.material, n_classes)
     return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
 
 
@@ -349,16 +334,9 @@ def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
     join (see `_df_to_dict`'s `has_sec_lists` guard) — silently skipped for
     those, same convention as elsewhere in this module.
     """
-    counts: dict = {}
-    for path in files:
-        columns = ["pdg"]
-        has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names
-        if has_sec:
-            columns.append("sec_pdg_list")
-        df = pd.read_parquet(path, columns=columns)
-        _accumulate_value_counts(counts, df["pdg"], int)
-        if has_sec:
-            exploded = df["sec_pdg_list"].explode().dropna()
-            _accumulate_value_counts(counts, exploded, int)
-    class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
+    from giant.data.scan import ScanRequest, scan_metadata
+
+    result = scan_metadata(files, ScanRequest(pooled_pdg=True))
+    assert result.pooled_pdg is not None
+    class_map, other_members, class_counts = _topn_plus_other_map(result.pooled_pdg, n_classes)
     return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)

giant/data/scan.py

@@ -0,0 +1,175 @@
+"""Fused metadata scan over one or more parquet files.
+
+`giant.pipeline.run_setup_stage` needs several distinct frequency summaries
+before training can start — the event-id → row-count index (for the train/val
+split), the pdg/material vocabularies, an optional physics-process count, and
+a pooled pdg count (primary + secondary species, for onehot conditioning).
+Each of those used to be its own full `pd.read_parquet(path, columns=[...])`
+per file (`giant.data.loader`'s old `_rank_by_frequency_from_files` /
+`build_index_maps_from_files` / `build_pdg_topn_map_from_files`) — up to five
+separate reads of the same file. `scan_metadata` answers all of them in one
+`pl.collect_all` per file instead, sharing the file open/decompress cost.
+
+Every requested count comes back keyed by value, as a `ValueStat(count,
+first_seen)`. `first_seen` is the value's row ordinal — file order (as given
+in `files`), then row order within a file — via `row_index_name` on the
+per-file lazy scan plus a running row offset across files. This is what
+`giant.data.loader._topn_plus_other_map`'s frequency-ranking tie-break keys
+on: among equally-frequent values, whichever was scanned first wins its own
+class slot. That is an explicit, documented contract (this module is where
+it's implemented), not an accident of iteration order.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+import polars as pl
+
+from giant.data.loader import _offset_event_id, event_id_offset
+
+
+@dataclass(frozen=True)
+class ScanRequest:
+    """Which aggregations to compute. Every field defaults off so a caller
+    only pays for what it actually needs."""
+
+    event_index: bool = False
+    pdg: bool = False
+    material: bool = False
+    process: bool = False
+    pooled_pdg: bool = False
+    """pdg ∪ exploded sec_pdg_list — both roles a PDG code plays (primary
+    species and secondary species), pooled into one count per code. See
+    `giant.data.loader.build_pdg_topn_map_from_files`'s docstring for why."""
+
+
+@dataclass(frozen=True)
+class ValueStat:
+    count: int
+    first_seen: int
+
+
+@dataclass
+class MetadataScan:
+    event_index: tuple[np.ndarray, np.ndarray] | None = None
+    """(unique_ids, counts), ids ascending — matches
+    `setup_cache.compute_event_index_from_files`'s return shape."""
+    pdg: dict[int, ValueStat] | None = None
+    material: dict[str, ValueStat] | None = None
+    process: dict[str, ValueStat] | None = None
+    pooled_pdg: dict[int, ValueStat] | None = None
+
+
+def _group_lazy(path: Path, column: str) -> pl.LazyFrame:
+    return (
+        pl.scan_parquet(path, row_index_name="__row")
+        .select(column, "__row")
+        .group_by(column)
+        .agg(pl.len().alias("__count"), pl.col("__row").min().alias("__first_row"))
+    )
+
+
+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", 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"))
+
+
+def _merge_counts(acc: dict, df: pl.DataFrame, column: str, row_offset: int, cast) -> None:
+    for key, count, first_row in zip(
+        df[column].to_list(), df["__count"].to_list(), df["__first_row"].to_list(), strict=True
+    ):
+        key = cast(key)
+        first_seen = row_offset + int(first_row)
+        if key in acc:
+            prev_count, prev_first = acc[key]
+            acc[key] = (prev_count + int(count), min(prev_first, first_seen))
+        else:
+            acc[key] = (int(count), first_seen)
+
+
+def scan_metadata(files: list[Path], request: ScanRequest) -> MetadataScan:
+    """Scan `files` once (one `pl.collect_all` per file) and return every
+    aggregation `request` asks for. Files with zero rows contribute nothing
+    but still advance nothing (no row_offset change, nothing to merge)."""
+    event_id_parts: list[tuple[np.ndarray, np.ndarray]] = []
+    pdg_acc: dict[int, tuple[int, int]] = {}
+    material_acc: dict[str, tuple[int, int]] = {}
+    process_acc: dict[str, tuple[int, int]] = {}
+    pooled_pdg_acc: dict[int, tuple[int, int]] = {}
+
+    row_offset = 0
+    for file_idx, path in enumerate(files):
+        keys: list[str] = []
+        lazies: list[pl.LazyFrame] = []
+
+        if request.event_index:
+            keys.append("event_id")
+            lazies.append(_group_lazy(path, "event_id"))
+        if request.pdg:
+            keys.append("pdg")
+            lazies.append(_group_lazy(path, "pdg"))
+        if request.material:
+            keys.append("material")
+            lazies.append(_group_lazy(path, "material"))
+        if request.process:
+            keys.append("process")
+            lazies.append(_group_lazy(path, "process"))
+        if request.pooled_pdg:
+            has_sec = "sec_pdg_list" in pl.scan_parquet(path).collect_schema().names()
+            keys.append("pooled_pdg")
+            lazies.append(_pooled_pdg_lazy(path, has_sec))
+
+        keys.append("__n")
+        lazies.append(pl.scan_parquet(path).select(pl.len().alias("__n")))
+
+        results = dict(zip(keys, pl.collect_all(lazies, engine="streaming"), strict=True))
+        n_rows = int(results["__n"].item()) if len(results["__n"]) else 0
+
+        if request.event_index:
+            df = results["event_id"]
+            ids = _offset_event_id(df["event_id"].to_numpy(), event_id_offset(file_idx))
+            counts = df["__count"].to_numpy().astype(np.int64)
+            if ids.size:
+                event_id_parts.append((ids, counts))
+        if request.pdg:
+            _merge_counts(pdg_acc, results["pdg"], "pdg", row_offset, int)
+        if request.material:
+            _merge_counts(material_acc, results["material"], "material", row_offset, str)
+        if request.process:
+            _merge_counts(process_acc, results["process"], "process", row_offset, str)
+        if request.pooled_pdg:
+            _merge_counts(pooled_pdg_acc, results["pooled_pdg"], "__val", row_offset, int)
+
+        row_offset += n_rows
+
+    event_index = None
+    if request.event_index:
+        if event_id_parts:
+            all_ids = np.concatenate([p[0] for p in event_id_parts])
+            all_counts = np.concatenate([p[1] for p in event_id_parts])
+            order = np.argsort(all_ids, kind="stable")
+            event_index = (all_ids[order], all_counts[order])
+        else:
+            event_index = (np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64))
+
+    def _to_stats(acc: dict) -> dict:
+        return {k: ValueStat(*v) for k, v in acc.items()}
+
+    return MetadataScan(
+        event_index=event_index,
+        pdg=_to_stats(pdg_acc) if request.pdg else None,
+        material=_to_stats(material_acc) if request.material else None,
+        process=_to_stats(process_acc) if request.process else None,
+        pooled_pdg=_to_stats(pooled_pdg_acc) if request.pooled_pdg else None,
+    )

giant/data/setup_cache.py

@@ -23,7 +23,7 @@ import numpy as np
 
 from giant import config
 from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
-from giant.data.loader import TopNMap, event_id_offset, load_event_ids
+from giant.data.loader import TopNMap
 from giant.data.transforms import Normalizer, sorted_membership
 
 # Bump manually on a change to the data-encoding semantics (e.g. a future
@@ -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
@@ -349,12 +349,21 @@ def save(
 
 
 def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.ndarray]:
-    """Unique event ids + per-event row (step) counts, across all `files`."""
+    """Unique event ids + per-event row (step) counts, across all `files`.
+
+    Computed via a streaming per-file `group_by("event_id")` (see
+    `giant.data.scan.scan_metadata`) rather than concatenating every row's
+    raw event_id across every file before `np.unique` — the latter's peak
+    memory is 8 bytes x total row count; this is bounded by the (much
+    smaller) unique event count instead.
+    """
+    from giant.data.scan import ScanRequest, scan_metadata
+
     if not files:
         return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
-    all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
-    unique_ids, counts = np.unique(all_ids, return_counts=True)
-    return unique_ids, counts
+    result = scan_metadata(files, ScanRequest(event_index=True))
+    assert result.event_index is not None
+    return result.event_index
 
 
 def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:

giant/geometry.py

@@ -23,12 +23,13 @@ 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 pandas as pd
+import polars as pl
 import pyarrow.parquet as pq
 
 _INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"
@@ -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
 
@@ -301,14 +302,18 @@ def _fit_slab_lookup(
     edges = np.linspace(z_min, z_max, n_bins + 1)
     bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
 
+    # pandas' groupby(...).size() sorts group keys ascending by default, so
+    # a tie in `n` for the same bin (equal counts split between two
+    # material/layer_id combos) resolves to the lexicographically-first
+    # combo — matched here by sorting on the keys first, then a
+    # maintain_order-stable sort on `n` so ties keep that key order.
     counts = (
-        pd.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
-        .groupby(["bin", "material", "layer_id"])
-        .size()
-        .to_frame("n")
-        .reset_index()
-        .sort_values("n", ascending=False)
-        .drop_duplicates("bin")
+        pl.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
+        .group_by(["bin", "material", "layer_id"])
+        .agg(pl.len().alias("n"))
+        .sort(["bin", "material", "layer_id"])
+        .sort("n", descending=True, maintain_order=True)
+        .unique(subset="bin", keep="first", maintain_order=True)
     )
 
     bin_material = np.full(n_bins, "", dtype=object)
@@ -349,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,
@@ -408,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,
@@ -454,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

@@ -4,9 +4,10 @@ for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
 `giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""
 
 import inspect
+from typing import cast
 
 import torch
-import torch.nn as nn
+from torch import nn
 
 
 class HistoryEncoder(nn.Module):
@@ -34,6 +35,17 @@ class HistoryEncoder(nn.Module):
     def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
         return self.forward(feat, has_prev), cache
 
+    def select_cache(self, cache: object, idx: torch.Tensor) -> object:
+        """Row-compacts an inference cache (`init_cache`/`step`'s state) down
+        to `idx` — used by `giant.sample.sample_secondaries_ar`'s row
+        compaction to keep a shrinking active-row set's cache aligned as rows
+        finish generating. Default here matches `init_cache`/`step`'s O(1)
+        default: `cache` is always `None`, so there's nothing to index —
+        correct for any encoder whose per-step state doesn't carry a batch
+        dimension (`MarkovHistory` has no cache at all; its running state is
+        `prev_repr`/`remaining`, compacted directly by the caller)."""
+        return cache
+
 
 HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}
 
@@ -185,7 +197,7 @@ class AttentionHistory(HistoryEncoder):
         return self.in_proj(x)
 
     def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
-        B, K, _ = feat.shape
+        _, K, _ = feat.shape
         x = self._embed(feat, has_prev)
         mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
         for block in self.blocks:
@@ -215,3 +227,13 @@ class AttentionHistory(HistoryEncoder):
             x, kv_new = block.step(x, kv)
             new_cache.append(kv_new)
         return x, new_cache
+
+    def select_cache(self, cache: object, idx: torch.Tensor) -> list[torch.Tensor | None]:
+        """Row-compacts every block's `(B, T, dim)` KV cache down to `idx`
+        along its batch dimension — see `HistoryEncoder.select_cache`. `idx`
+        may be a long index tensor or a boolean mask (`giant.sample`'s AR
+        loop uses both). `None` entries (a block that has never seen a
+        `step` call yet) stay `None`."""
+        assert isinstance(cache, list)
+        cache_t = cast("list[torch.Tensor | None]", cache)
+        return [None if kv is None else kv[idx] for kv in cache_t]

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
@@ -614,6 +614,14 @@ class Stage2Autoregressive(StageModel):
         slot — see `AttentionHistory.step`'s docstring."""
         return self.history_encoder.step(token_feat, has_prev, cache)
 
+    def select_history_cache(self, cache, idx: torch.Tensor):
+        """Row-compacts `cache` (from `init_history_cache`/`history_step`)
+        down to `idx` — see `HistoryEncoder.select_cache`. Used by
+        `giant.sample.sample_secondaries_ar`'s active-row compaction to keep
+        the cache aligned with a shrinking batch as rows finish generating
+        across AR slots."""
+        return self.history_encoder.select_cache(cache, idx)
+
     def forward(
         self,
         x_t: torch.Tensor,

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,25 +13,23 @@ 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,
     event_id_offset,
     find_parquet_files,
     iter_file_chunks,
-    build_index_maps_from_files,
-    build_pdg_topn_map_from_files,
-    build_process_map_from_files,
-    build_topn_map_from_files,
 )
+from giant.data.scan import MetadataScan, ScanRequest, scan_metadata
 from giant.data.transforms import (
     Normalizer,
-    build_features,
-    _WelfordAccumulator,
     _ReservoirSampler,
+    _WelfordAccumulator,
+    build_features,
     sorted_membership,
 )
-from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import build_models, build_critics, resolve_type_n_classes
+from giant.model.network import build_critics, build_models, resolve_type_n_classes
 from giant.training import train as run_training
 
 
@@ -127,12 +125,71 @@ def run_setup_stage(
             loaded = setup_cache.load(data, files, echo=echo)
         cache = loaded if loaded is not None else setup_cache.SetupCache.empty(files)
 
-    if cache is not None and cache.event_index is not None:
+    # Every section below first asks the cache; whatever's missing is
+    # collected into one ScanRequest and answered by a single fused scan
+    # (giant.data.scan.scan_metadata), instead of a separate full pass per
+    # section (event index, vocab, process counts, pdg/material top-N counts
+    # used to each re-open and re-read every file on their own).
+    particle_cfg = cfg["conditioning"]["particle"]
+    material_cfg = cfg["conditioning"]["material"]
+    particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
+    particle_type_target = particle_type_cfg.target
+
+    # A process map is needed if either stage's router reads the physics
+    # process label (type="process"). Only one map is built even if both
+    # stages want one — see the module-level note in giant/cli.py's
+    # _router_total_experts for why composed-router n_experts isn't a plain
+    # int; process routers are never composed in practice, so this doesn't
+    # need that generality.
+    process_router_cfg = next(
+        (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
+        None,
+    )
+    process_n_experts = process_router_cfg["n_experts"] if process_router_cfg is not None else None
+
+    need_pdg_onehot = particle_cfg["type"] == "onehot"
+    need_sec_type_onehot = particle_type_target == "onehot"
+    sec_type_n_classes = (
+        resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"]) if need_sec_type_onehot else None
+    )
+    need_material_onehot = material_cfg["type"] == "onehot"
+    material_n_classes = material_cfg["emb_dim"] if need_material_onehot else None
+
+    def _topn_cached(axis: str, n_classes: int) -> TopNMap | None:
+        return cache.topn_maps.get(setup_cache.topn_key(axis, n_classes)) if cache is not None else None
+
+    need_event_index = cache is None or cache.event_index is None
+    need_vocab = cache is None or cache.vocab is None
+    need_process = process_n_experts is not None and (cache is None or cache.proc_maps.get(process_n_experts) is None)
+    # The PDG axis is used independently by conditioning.particle.type="onehot"
+    # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
+    # (secondary-species decode) — their class counts can now differ (gitea
+    # #29: stage2_model.particle_type.n_classes, 0 = inherit
+    # conditioning.particle.emb_dim), but both are built from the same
+    # pooled pdg-count scan, so a cache miss on either one asks for it.
+    need_pdg_pooled = (need_pdg_onehot and _topn_cached("pdg", particle_cfg["emb_dim"]) is None) or (
+        need_sec_type_onehot and sec_type_n_classes is not None and _topn_cached("pdg", sec_type_n_classes) is None
+    )
+    need_material_topn = (
+        need_material_onehot and material_n_classes is not None and _topn_cached("material", material_n_classes) is None
+    )
+
+    request = ScanRequest(
+        event_index=need_event_index,
+        pdg=need_vocab,
+        material=need_vocab or need_material_topn,
+        process=need_process,
+        pooled_pdg=need_pdg_pooled,
+    )
+    scan = scan_metadata(files, request) if request != ScanRequest() else MetadataScan()
+
+    if not need_event_index:
         unique_ids, counts = cache.event_index
         echo(f"event index: cache hit ({len(unique_ids):,} unique events)")
     else:
         echo("scanning event IDs …")
-        unique_ids, counts = setup_cache.compute_event_index_from_files(files)
+        assert scan.event_index is not None
+        unique_ids, counts = scan.event_index
         if cache is not None:
             cache.event_index = (unique_ids, counts)
 
@@ -141,55 +198,34 @@ def run_setup_stage(
     n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
     echo(f"  {int(counts.sum()):,} steps  |  {len(train_events)} train events  |  {len(val_events)} val events")
 
-    if cache is not None and cache.vocab is not None:
+    if not need_vocab:
         pdg_map, mat_map = cache.vocab
         echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
     else:
         echo("building vocabulary maps …")
-        pdg_map, mat_map = build_index_maps_from_files(files)
+        assert scan.pdg is not None and scan.material is not None
+        pdg_map = {v: i for i, v in enumerate(sorted(scan.pdg))}
+        mat_map = {v: i for i, v in enumerate(sorted(scan.material))}
         echo(f"  {len(pdg_map)} PDG codes  |  {len(mat_map)} materials")
         if cache is not None:
             cache.vocab = (pdg_map, mat_map)
 
-    # A process map is needed if either stage's router reads the physics
-    # process label (type="process"). Only one map is built even if both
-    # stages want one — see the module-level note in giant/cli.py's
-    # _router_total_experts for why composed-router n_experts isn't a plain
-    # int; process routers are never composed in practice, so this doesn't
-    # need that generality.
     proc_map: dict[str, int] | None = None
-    process_router_cfg = next(
-        (r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
-        None,
-    )
     if process_router_cfg is not None:
-        n_experts = process_router_cfg["n_experts"]
-        cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
-        if cached_proc_map is not None:
+        assert process_n_experts is not None
+        if not need_process:
+            assert cache is not None
+            cached_proc_map = cache.proc_maps.get(process_n_experts)
+            assert cached_proc_map is not None
             proc_map = cached_proc_map
-            echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
+            echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {process_n_experts} experts)")
         else:
             echo("building process vocabulary …")
-            proc_map = build_process_map_from_files(files, n_experts=n_experts)
-            echo(f"  {len(proc_map)} process labels mapped to {n_experts} experts")
+            assert scan.process is not None
+            proc_map, _, _ = _topn_plus_other_map(scan.process, process_n_experts)
+            echo(f"  {len(proc_map)} process labels mapped to {process_n_experts} experts")
             if cache is not None:
-                cache.proc_maps[n_experts] = proc_map
-
-    # Top-N-plus-other maps for onehot conditioning/type axes.
-    # The PDG axis is used independently by conditioning.particle.type="onehot"
-    # (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
-    # (secondary-species decode) — their class counts can now differ (gitea
-    # #29: stage2_model.particle_type.n_classes, 0 = inherit
-    # conditioning.particle.emb_dim), so each is resolved and built
-    # independently via _pdg_topn below. cache.topn_maps is keyed by
-    # (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
-    # the same N the second call is a cache hit against the first — no extra
-    # scan in the common case where they still match. The material axis is
-    # independent of both.
-    particle_cfg = cfg["conditioning"]["particle"]
-    material_cfg = cfg["conditioning"]["material"]
-    particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
-    particle_type_target = particle_type_cfg.target
+                cache.proc_maps[process_n_experts] = proc_map
 
     def _pdg_topn(n_classes: int) -> TopNMap:
         cache_key = setup_cache.topn_key("pdg", n_classes)
@@ -198,33 +234,36 @@ def run_setup_stage(
             echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
             return cached
         echo("building pdg top-N map …")
-        topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
+        assert scan.pooled_pdg is not None
+        class_map, other_members, class_counts = _topn_plus_other_map(scan.pooled_pdg, n_classes)
+        topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
         echo(f"  {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
         if cache is not None:
             cache.topn_maps[cache_key] = topn_map
         return topn_map
 
-    pdg_topn_map: TopNMap | None = None
-    if particle_cfg["type"] == "onehot":
-        pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
-
+    pdg_topn_map: TopNMap | None = _pdg_topn(particle_cfg["emb_dim"]) if need_pdg_onehot else None
     sec_type_topn_map: TopNMap | None = None
-    if particle_type_target == "onehot":
-        sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
+    if need_sec_type_onehot:
+        assert sec_type_n_classes is not None
         sec_type_topn_map = _pdg_topn(sec_type_n_classes)
 
     mat_topn_map: TopNMap | None = None
-    if material_cfg["type"] == "onehot":
-        n_classes = material_cfg["emb_dim"]
-        cache_key = setup_cache.topn_key("material", n_classes)
+    if need_material_onehot:
+        assert material_n_classes is not None
+        cache_key = setup_cache.topn_key("material", material_n_classes)
         cached = cache.topn_maps.get(cache_key) if cache is not None else None
         if cached is not None:
             mat_topn_map = cached
-            echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
+            echo(
+                f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {material_n_classes} classes)"
+            )
         else:
             echo("building material top-N map …")
-            mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
-            echo(f"  {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
+            assert scan.material is not None
+            class_map, other_members, class_counts = _topn_plus_other_map(scan.material, material_n_classes)
+            mat_topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
+            echo(f"  {len(mat_topn_map.class_map)} materials mapped to {material_n_classes} classes")
             if cache is not None:
                 cache.topn_maps[cache_key] = mat_topn_map
 
@@ -435,7 +474,6 @@ def run_train_job(
         mat_topn_map=cond_mat_topn,
         sec_type_class_map=sec_type_class_map,
         k_max=k_max,
-        seed=t["seed"],
     )
     val_ds = StreamingStepsDataset(
         files=files,
@@ -457,17 +495,30 @@ def run_train_job(
     )
 
     pin = device.type == "cuda"
+    # DataLoader worker subprocesses default to fork() on Linux, but by the
+    # time they're created this process has already run polars queries
+    # (run_setup_stage's fused metadata scan, above) — polars' native
+    # (rayon) thread pool doesn't survive a fork: a worker that inherits it
+    # mid-fork deadlocks the instant it touches polars itself, which
+    # StreamingStepsDataset's iter_file_chunks now does on every row group.
+    # "spawn" starts each worker as a fresh interpreter with no inherited
+    # thread-pool state, avoiding that hazard entirely. Only matters when
+    # workers actually exist — num_workers=0 runs the dataset in-process and
+    # never forks.
+    mp_context = "spawn" if num_workers > 0 else None
     train_loader = DataLoader(
         train_ds,
         batch_size=None,
         num_workers=num_workers,
         pin_memory=pin,
+        multiprocessing_context=mp_context,
     )
     val_loader = DataLoader(
         val_ds,
         batch_size=None,
         num_workers=num_workers,
         pin_memory=pin,
+        multiprocessing_context=mp_context,
     )
 
     model_config = {

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(
-        event_id=tr["event_id"][sel],
-        track_id=tr["track_id"][sel],
-        parent_id=tr["parent_id"][sel],
-        generation=tr["generation"][sel],
-        step_no=tr["step_in_track"][sel],
-        pdg=tr["pdg"][sel],
-        pre_x=pos[:, 0],
-        pre_y=pos[:, 1],
-        pre_z=pos[:, 2],
-        pre_E=tr["pre_E"][sel],
-        pre_dx=dir_[:, 0],
-        pre_dy=dir_[:, 1],
-        pre_dz=dir_[:, 2],
-        post_x=pos[:, 0],
-        post_y=pos[:, 1],
-        post_z=pos[:, 2],
-        post_E=np.zeros(n),
-        post_dx=dir_[:, 0],
-        post_dy=dir_[:, 1],
-        post_dz=dir_[:, 2],
-        edep=np.asarray(edep, dtype=np.float64).reshape(n),
-        step_length=np.zeros(n),
-        material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
-        layer_id=tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
-        n_sec_pred=np.zeros(n, dtype=np.int64),
-        termination_reason=np.full(n, reason, dtype=object),
-    )
+    return {
+        "event_id": tr["event_id"][sel],
+        "track_id": tr["track_id"][sel],
+        "parent_id": tr["parent_id"][sel],
+        "generation": tr["generation"][sel],
+        "step_no": tr["step_in_track"][sel],
+        "pdg": tr["pdg"][sel],
+        "pre_x": pos[:, 0],
+        "pre_y": pos[:, 1],
+        "pre_z": pos[:, 2],
+        "pre_E": tr["pre_E"][sel],
+        "pre_dx": dir_[:, 0],
+        "pre_dy": dir_[:, 1],
+        "pre_dz": dir_[:, 2],
+        "post_x": pos[:, 0],
+        "post_y": pos[:, 1],
+        "post_z": pos[:, 2],
+        "post_E": np.zeros(n),
+        "post_dx": dir_[:, 0],
+        "post_dy": dir_[:, 1],
+        "post_dz": dir_[:, 2],
+        "edep": np.asarray(edep, dtype=np.float64).reshape(n),
+        "step_length": np.zeros(n),
+        "material": tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
+        "layer_id": tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
+        "n_sec_pred": np.zeros(n, dtype=np.int64),
+        "termination_reason": np.full(n, reason, dtype=object),
+    }
 
 
 @torch.no_grad()
@@ -442,14 +443,14 @@ def rollout(
     on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
     particle_conditioning: str = "embedding",
     material_conditioning: str = "embedding",
-    pdg_topn_map: "TopNMap | None" = None,
-    mat_topn_map: "TopNMap | None" = None,
-    sec_type_topn_map: "TopNMap | None" = None,
+    pdg_topn_map: TopNMap | None = None,
+    mat_topn_map: TopNMap | None = None,
+    sec_type_topn_map: TopNMap | None = None,
     other_policy: str = "sample",
     seed: int | None = None,
     stage1_ddpm_steps: int = 1000,
     stage2_ddpm_steps: int = 1000,
-    l1_dist_collector: "L1DistCollector | None" = None,
+    l1_dist_collector: L1DistCollector | None = None,
 ) -> dict[str, np.ndarray] | RolloutSummary:
     """Run showers to completion.
 

giant/sample.py

@@ -231,6 +231,7 @@ def sample_secondaries_ar(
     stage1_out: torch.Tensor,
     n_sec_pred: torch.Tensor | None,
     steps: int = 10,
+    full_length: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """`Stage2Autoregressive` inference loop: one token at a time, in
     descending-energy slot order, up to `k_max` sequential calls. Unlike
@@ -242,19 +243,18 @@ def sample_secondaries_ar(
     expressiveness.
 
     A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
-    the "K sequential forwards" cost applies per-token here, not
-    once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
-    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
-    every row in the batch has stopped).
+    the "K sequential forwards" cost applies per-token here, not once, so a
+    flow/ddpm AR run costs ~`n_sec * steps` model calls per physics step
+    (measured on `configs/baseline.toml`: 0.382 secondaries/step at rollout
+    time), not `k_max * steps` — see the row-compaction paragraph below.
 
-    `n_sec_pred`, if given, fixes each row's secondary count up front (as
-    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
-    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
-    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
-    running self-sample regardless of the decoder's own stop head — the
-    scheduled-sampling training contract does not truncate). This always
-    runs the full `k_max`-iteration loop, masking by the given count at the
-    end exactly as before.
+    `n_sec_pred`, if given (as resolved by `resolve_n_sec` — `n_sec.mode` in
+    `("head", "truth")`, or a stop-token decoder driven by
+    `_assemble_stage2_ar_inputs_scheduled`'s ground-truth `n_sec`, which must
+    run the *full* `k_max`-length free-running self-sample regardless of the
+    decoder's own stop head — the scheduled-sampling training contract does
+    not truncate, see `full_length` below) fixes each row's secondary count
+    up front.
 
     `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
     (`n_sec.mode = "stop_token"`): before generating each slot's token, that
@@ -263,11 +263,35 @@ def sample_secondaries_ar(
     why this needs no extra state) decides whether generation should have
     already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
     0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
-    `n_sec_pred` is the first slot index where this fires; once every row in
-    the batch has fired, the loop breaks before spending a model call on the
-    next slot's token — the average-case cost win the docstring above
-    describes. A row that never fires within `k_max` is capped there
-    (`K_MAX` stays a safety cap, not a modeling ceiling).
+    `n_sec_pred` is the first slot index where this fires. A row that never
+    fires within `k_max` is capped there (`K_MAX` stays a safety cap, not a
+    modeling ceiling).
+
+    **Row compaction.** A row that has already produced its `n_sec_pred`
+    tokens (or, under `stop_token`, has already fired its stop logit) has
+    nothing left to contribute — every later slot of that row is masked out
+    of `sec_valid` on return, and downstream consumers (`giant/rollout.py`,
+    `giant/cli.py`) never read it. So unless `full_length=True`, this
+    function drops such rows from the active set entirely instead of running
+    the model on them: `active_idx` starts at every row with `n_sec_pred > 0`
+    (or, under `stop_token`, every row — the first stop decision can fire at
+    slot 0) and only shrinks as rows finish, so slot `k`'s model calls cost
+    `O(active rows)` not `O(B)`. Slots a row never reaches keep their `0.0`
+    zero-init in `sec_cont`/`sec_type` — masked by `sec_valid`, identical to
+    what a full, uncompacted run would have written there before masking.
+    `AttentionHistory`'s KV cache is kept aligned to the shrinking active set
+    via `Stage2Autoregressive.select_history_cache`
+    (`giant.model.history.HistoryEncoder.select_cache`) every time the set
+    shrinks; `MarkovHistory`'s O(1) state (`prev_repr`/`remaining`, carried
+    directly rather than through a cache) is compacted the same way.
+
+    `full_length=True` disables all of the above: every row runs the full
+    `k_max`-iteration loop regardless of `n_sec_pred`/stop decisions, exactly
+    reproducing the pre-compaction behaviour. Required by
+    `_assemble_stage2_ar_inputs_scheduled`'s scheduled-sampling self-sample,
+    whose training contract needs a real prediction at every slot up to
+    `k_max` (mixed per-slot against ground truth) even past a row's own
+    `n_sec` — see that function's docstring.
 
     Under `history="attention"` the history encoding is computed once per
     slot via `Stage2Autoregressive.history_step` (a KV-cache append)
@@ -312,11 +336,6 @@ def sample_secondaries_ar(
     sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
     sec_type = torch.zeros(B, k_max, type_dim, device=device)
 
-    # Running per-token state, threaded from one slot to the next.
-    prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
-    remaining = torch.ones(B, device=device)
-    history_cache = sec_decoder.init_history_cache()
-
     use_stop_token = n_sec_pred is None
     if use_stop_token:
         assert getattr(sec_decoder, "stop_head", None) is not None, (
@@ -324,21 +343,46 @@ def sample_secondaries_ar(
             "with no stop_head — only valid under stage2_model.n_sec.mode = "
             "'stop_token'"
         )
-        finished = torch.zeros(B, dtype=torch.bool, device=device)
         derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
+        # Tracks which rows have already recorded a stop, globally by
+        # original batch index — needed even under compaction's own
+        # never-revisit guarantee, because `full_length=True` keeps every
+        # row in `active_idx` for the whole loop, so a row whose stop logit
+        # fires once but flips back below threshold at a later slot (a real
+        # possibility for an untrained/lightly-trained stop_head) must not
+        # have `derived_n_sec` overwritten by that later, spurious re-fire.
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+
+    # `active_idx`: rows still contributing tokens, indexed into the
+    # original batch. Only ever shrinks (never full_length) or stays fixed
+    # at arange(B) (full_length) — see the row-compaction docstring section.
+    active_idx = torch.arange(B, device=device)
+    if not full_length and not use_stop_token:
+        active_idx = active_idx[n_sec_pred > 0]
+
+    # Running per-token state, already compacted to `active_idx`.
+    prev_repr = torch.zeros(active_idx.numel(), CONT_SLOT_DIM + type_dim, device=device)
+    remaining = torch.ones(active_idx.numel(), device=device)
+    history_cache = sec_decoder.init_history_cache()
 
     for k in range(k_max):
-        has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
-        history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
-        remaining_frac = remaining.unsqueeze(1)  # (B, 1)
-        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        if active_idx.numel() == 0:
+            break
+        Bc = active_idx.numel()
+        cc = cond_cont.index_select(0, active_idx)
+        ck = cond_cat.index_select(0, active_idx)
+        s1 = stage1_out.index_select(0, active_idx)
+        has_prev = torch.full((Bc, 1), k >= 1, dtype=torch.bool, device=device)
+        history_feat = prev_repr.unsqueeze(1)  # (Bc, 1, CONT_SLOT_DIM + type_dim)
+        remaining_frac = remaining.unsqueeze(1)  # (Bc, 1)
+        slot_idx = torch.full((Bc, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
         hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
         if use_stop_token:
             stop_logit = sec_decoder.predict_stop(
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -346,21 +390,31 @@ def sample_secondaries_ar(
                 hist=hist,
             ).squeeze(1)
             if sec_decoder.n_sec_sampling == "sample":
-                stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
+                stop_now = torch.rand(Bc, device=device) < torch.sigmoid(stop_logit)
             else:
                 stop_now = stop_logit >= 0.0
-            derived_n_sec[stop_now & ~finished] = k
-            finished = finished | stop_now
-            if finished.all():
-                break
+            newly_stopped = stop_now & ~finished.index_select(0, active_idx)
+            derived_n_sec[active_idx[newly_stopped]] = k
+            finished[active_idx[stop_now]] = True
+            if not full_length:
+                keep = ~stop_now
+                active_idx = active_idx[keep]
+                cc, ck, s1 = cc[keep], ck[keep], s1[keep]
+                has_prev, remaining_frac, slot_idx = has_prev[keep], remaining_frac[keep], slot_idx[keep]
+                history_feat, hist = history_feat[keep], hist[keep]
+                history_cache = sec_decoder.select_history_cache(history_cache, keep)
+                prev_repr, remaining = prev_repr[keep], remaining[keep]
+                if active_idx.numel() == 0:
+                    break
+                Bc = active_idx.numel()
 
         if objective.is_adversarial:
-            z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
+            z = torch.randn(Bc, 1, sec_decoder.noise_dim, device=device)
             token = sec_decoder(
                 z,
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -368,15 +422,15 @@ def sample_secondaries_ar(
                 hist=hist,
             )
         else:
-            x = torch.randn(B, 1, token_dim, device=device)
+            x = torch.randn(Bc, 1, token_dim, device=device)
             dt = 1.0 / steps
             for i in range(steps):
-                t = torch.full((B, 1), i * dt, device=device)
+                t = torch.full((Bc, 1), i * dt, device=device)
                 v = sec_decoder(
                     x,
-                    cond_cont,
-                    cond_cat,
-                    stage1_out,
+                    cc,
+                    ck,
+                    s1,
                     history_feat,
                     has_prev,
                     remaining_frac,
@@ -387,15 +441,15 @@ def sample_secondaries_ar(
                 x = x + v * dt
             token = x
 
-        token = token.squeeze(1)  # (B, token_dim)
+        token = token.squeeze(1)  # (Bc, token_dim)
         cont_k = token[:, :CONT_SLOT_DIM]
         if type_folded:
             type_k = token[:, CONT_SLOT_DIM:]
         else:
             type_k = sec_decoder.predict_type(
-                cond_cont,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -403,8 +457,8 @@ def sample_secondaries_ar(
                 hist=hist,
             ).squeeze(1)
 
-        sec_cont[:, k] = cont_k
-        sec_type[:, k] = type_k
+        sec_cont[active_idx, k] = cont_k
+        sec_type[active_idx, k] = type_k
 
         if target == "onehot":
             type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
@@ -415,6 +469,12 @@ def sample_secondaries_ar(
         prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
         remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
 
+        if not full_length and not use_stop_token:
+            keep2 = n_sec_pred.index_select(0, active_idx) > (k + 1)
+            active_idx = active_idx[keep2]
+            prev_repr, remaining = prev_repr[keep2], remaining[keep2]
+            history_cache = sec_decoder.select_history_cache(history_cache, keep2)
+
     resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
     sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
     return sec_cont, sec_type, sec_valid

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

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/profile_analysis_costs.py

@@ -26,7 +26,7 @@ import numpy as np
 import polars as pl
 
 from giant.analysis.catalog import catalog_ids, get_spec
-from giant.analysis.run import compute_reduced
+from giant.analysis.condor import compute_reduced
 from giant.analysis.context import build_context
 from giant.analysis.sources import RolloutSpec
 

giant/tools/profile_setup_scan.py

@@ -0,0 +1,162 @@
+"""Benchmark `giant.pipeline.run_setup_stage`'s cold-cache scan against synthetic data.
+
+Generates a schema-complete synthetic steps parquet (matching
+`tests/test_pipeline.py`'s `_make_synthetic_steps`, but built with vectorized
+numpy instead of a per-row Python loop so it scales to millions of rows) at a
+few row counts, times `run_setup_stage` with `cache_setup=False` (so every
+call is a genuine cold scan, never served from the sidecar), and prints a
+before/after-style table. Run this on `master` before a change and again
+after to see what a step actually bought — see the "speed up dwarf
+warm-cache" plan for the pass-by-pass breakdown this benchmark is meant to
+attribute (giant/data/loader.py, giant/data/scan.py, giant/pipeline.py).
+
+Usage: ``uv run python giant/tools/profile_setup_scan.py``
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from tempfile import TemporaryDirectory
+
+import numpy as np
+import polars as pl
+
+from giant import config as gconfig
+from giant.pipeline import run_setup_stage
+
+ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000]
+
+_MATERIALS = ["G4_AIR", "G4_Fe"]
+_PDGS = [11, 22]
+_PROCESSES = ["eIoni", "phot", "compt"]
+
+
+def _unit_vectors(n: int, rng: np.random.Generator) -> np.ndarray:
+    v = rng.normal(size=(n, 3))
+    return v / np.linalg.norm(v, axis=1, keepdims=True)
+
+
+def _ragged_lists(k: np.ndarray, rng: np.random.Generator, lo: float, hi: float) -> list[list[float]]:
+    total = int(k.sum())
+    flat = rng.uniform(lo, hi, size=total)
+    idx = np.cumsum(k)[:-1]
+    return [arr.tolist() for arr in np.split(flat, idx)]
+
+
+def _make_synthetic_steps(n: int, seed: int = 0) -> pl.DataFrame:
+    """Vectorized equivalent of tests/test_pipeline.py's `_make_synthetic_steps`.
+
+    event_id is assigned so each event gets 2-3 steps (matching that
+    fixture's structure), and pdg/material/process cycle deterministically
+    by row index rather than being drawn at random, same as the original.
+    """
+    rng = np.random.default_rng(seed)
+    n_events = max(n // 3, 1)
+
+    pre_E = rng.uniform(50.0, 500.0, size=n)
+    n_sec = rng.integers(0, 3, size=n)
+    frac_dep = rng.uniform(0.05, 0.3, size=n)
+    frac_sec = np.where(n_sec > 0, rng.uniform(0.05, 0.2, size=n), 0.0)
+    frac_post = 1.0 - frac_dep - frac_sec
+    edep = pre_E * frac_dep
+    e_sec = pre_E * frac_sec
+    post_E = pre_E * frac_post
+    pre_pos = rng.uniform(-10, 10, size=(n, 3))
+    step_length = rng.uniform(0.1, 5.0, size=n)
+    pre_dir = np.zeros((n, 3))
+    pre_dir[:, 2] = 1.0
+    post_dir = _unit_vectors(n, rng)
+    post_pos = pre_pos + step_length[:, None] * pre_dir
+
+    row_idx = np.arange(n)
+    event_id = row_idx % n_events
+
+    sec_E = _ragged_lists(n_sec, rng, 0.1, 1.0)  # placeholder magnitude, rescaled below
+    sec_dx = _ragged_lists(n_sec, rng, -1.0, 1.0)
+    sec_dy = _ragged_lists(n_sec, rng, -1.0, 1.0)
+    sec_dz = _ragged_lists(n_sec, rng, -1.0, 1.0)
+    total_sec = int(n_sec.sum())
+    flat_pdg = [_PDGS[(row_idx[i] + j) % 2] for i in range(n) for j in range(n_sec[i])]
+    idx = np.cumsum(n_sec)[:-1]
+    sec_pdg = (
+        [list(x) for x in np.split(np.array(flat_pdg, dtype=np.int64), idx)] if total_sec else [[] for _ in range(n)]
+    )
+    # Rescale each row's secondary energies to sum to that row's e_sec (a
+    # Dirichlet split, like the original fixture) rather than the raw
+    # uniform placeholder.
+    sec_E_scaled = []
+    for i in range(n):
+        vals = np.array(sec_E[i])
+        if vals.size:
+            sec_E_scaled.append((vals / vals.sum() * e_sec[i]).tolist())
+        else:
+            sec_E_scaled.append([])
+
+    return pl.DataFrame(
+        {
+            "event_id": event_id,
+            "pdg": np.array(_PDGS)[row_idx % 2],
+            "pre_x": pre_pos[:, 0],
+            "pre_y": pre_pos[:, 1],
+            "pre_z": pre_pos[:, 2],
+            "pre_E": pre_E,
+            "pre_dx": pre_dir[:, 0],
+            "pre_dy": pre_dir[:, 1],
+            "pre_dz": pre_dir[:, 2],
+            "material": np.array(_MATERIALS)[row_idx % 2],
+            "layer_id": row_idx % 5,
+            "child_track_ids": [list(range(int(k))) for k in n_sec],
+            "e_sec": e_sec,
+            "process": np.array(_PROCESSES)[row_idx % 3],
+            "step_length": step_length,
+            "post_E": post_E,
+            "edep": edep,
+            "post_dx": post_dir[:, 0],
+            "post_dy": post_dir[:, 1],
+            "post_dz": post_dir[:, 2],
+            "post_x": post_pos[:, 0],
+            "post_y": post_pos[:, 1],
+            "post_z": post_pos[:, 2],
+            "sec_E_list": sec_E_scaled,
+            "sec_pdg_list": sec_pdg,
+            "sec_dx_list": sec_dx,
+            "sec_dy_list": sec_dy,
+            "sec_dz_list": sec_dz,
+        }
+    )
+
+
+def _time_setup_stage(data: Path) -> float:
+    cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {})
+    gconfig.validate_config(cfg)
+    t0 = time.perf_counter()
+    run_setup_stage(
+        data,
+        val_fraction=cfg["train"]["val_fraction"],
+        seed=cfg["train"]["seed"],
+        cfg=cfg,
+        cache_setup=False,
+        echo=lambda *a, **k: None,
+    )
+    return time.perf_counter() - t0
+
+
+def main() -> None:
+    with TemporaryDirectory(prefix="giant-setup-scan-profile-") as tmp:
+        tmp_path = Path(tmp)
+        print(f"{'n_rows':>10s}  {'time (s)':>10s}  {'rows/s':>12s}")
+        for n in ROW_COUNTS:
+            path = tmp_path / f"steps_{n}.parquet"
+            _make_synthetic_steps(n).write_parquet(path)
+            # warm the OS page cache so the timed pass measures compute, not
+            # the one-time cold read of a freshly-written file.
+            pl.scan_parquet(path).select(pl.len()).collect()
+
+            dt = _time_setup_stage(path)
+            print(f"{n:>10,d}  {dt:>10.3f}  {n / dt:>12,.0f}")
+            path.unlink()
+
+
+if __name__ == "__main__":
+    main()

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,16 +9,16 @@
 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
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from giant import config
 from giant.data.loader import TopNMap
 from giant.data.setup_cache import topnmap_to_json
 from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
@@ -48,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)
@@ -185,17 +185,6 @@ def train(
             if device.type == "cuda":
                 torch.cuda.reset_peak_memory_stats(device)
             collector.start_epoch(epoch)
-            # Epoch-aware RNG: same noise (and, below, same batch order) for
-            # epoch k whether the run is one process or a chain of per-epoch
-            # jobs. See giant.config.epoch_seed.
-            config.seed_everything(config.epoch_seed(t["seed"], epoch))
-            # Epoch-aware shuffle stream (see StreamingStepsDataset.set_epoch):
-            # keeps epoch k's batch order identical whether it runs here or as
-            # its own resumed per-epoch job in a b2luigi workflow.
-            # (tests hand `train` a plain list of batches, which has neither)
-            set_epoch = getattr(getattr(train_loader, "dataset", None), "set_epoch", None)
-            if callable(set_epoch):
-                set_epoch(epoch)
             for trainer in trainers.values():
                 trainer.train_mode()
 

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]
@@ -115,7 +115,7 @@ def derive_metrics_dir(
 
     Precedence: an explicit `out_dir` always wins. Otherwise
     `default_base / f"metrics_{run_dir.name}"` (the CLI passes the repo's
-    gitignored `analysis_runs/`, matching `giant.analysis.run.derive_run_dir`'s
+    gitignored `analysis_runs/`, matching `giant.analysis.condor.derive_run_dir`'s
     convention) — training-progress plots live alongside rollout-vs-reference
     analysis runs, not inside the training run directory itself.
     """

giant/training/stage2_inputs.py

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

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

giant/workflow/__init__.py

@@ -1,40 +0,0 @@
-"""b2luigi orchestration of the full GIANT pipeline.
-
-One workflow TOML (``spec.py``) parameterises an entire experiment — dataset,
-geometry oracle, N trainings, N rollouts, N analyses — and ``giant workflow
-run <spec.toml>`` turns it into a b2luigi task graph (``tasks.py``) whose
-targets are files on ``/ceph``: nothing is recomputed that already exists,
-every step waits for its inputs, and HTCondor submission/polling is b2luigi's
-job rather than a hand-rolled submit-file generator.
-
-This is the only sanctioned way to run a multi-step pipeline; ``giant`` and
-``dwarf`` stay single-step primitives that these tasks invoke.
-
-``tasks``/``run`` import b2luigi, so they are *not* imported here — a plain
-``import giant.workflow`` (or ``giant.workflow.spec``) works without the
-``workflow`` extra installed.
-"""
-
-from giant.workflow.spec import (
-    AnalysisSpec,
-    CondorSpec,
-    DatasetSpec,
-    GeometrySpec,
-    RolloutSpec,
-    TrainSpec,
-    WorkflowSpec,
-    load_spec,
-    spec_hash,
-)
-
-__all__ = [
-    "AnalysisSpec",
-    "CondorSpec",
-    "DatasetSpec",
-    "GeometrySpec",
-    "RolloutSpec",
-    "TrainSpec",
-    "WorkflowSpec",
-    "load_spec",
-    "spec_hash",
-]

giant/workflow/htcondor.py

@@ -1,82 +0,0 @@
-"""HTCondor job descriptions for the workflow tasks.
-
-b2luigi writes every key of a task's ``htcondor_settings`` dict straight into
-that job's submit description, so these helpers are just the ETP-specific
-resource/requirement conventions in one place:
-
-* **CPU jobs** (setup cache, geometry oracle, analysis compute) keep what
-  the deleted ``giant analyze submit`` used: ``+RemoteJob`` for grid I/O, or
-  ``TARGET.ProvidesETPResources`` when the files are local to the cluster.
-* **GPU jobs** (training epochs, rollout) are remote-only, so they always
-  carry ``+RemoteJob`` and reach ``/ceph`` through
-  ``TARGET.ProvidesEtpCeph`` — the requirement strings are ported from the
-  ``condor-gpu-train-rollout`` branch's ``giant/condor.py`` rather than
-  rewritten, since they encode what the ETP HTCondor wiki documents for
-  TOpAS/NEMO2 GPU workers.
-"""
-
-from __future__ import annotations
-
-from giant.workflow.spec import CondorSpec
-
-__all__ = ["cpu_settings", "gpu_settings", "gpu_requirements"]
-
-
-def cpu_settings(
-    condor: CondorSpec,
-    *,
-    request_memory_mb: int | None = None,
-    request_cpus: int | None = None,
-    walltime_s: int | None = None,
-) -> dict:
-    settings: dict = {
-        "universe": "docker",
-        "docker_image": condor.docker_image_cpu,
-        "request_memory": request_memory_mb if request_memory_mb is not None else condor.request_memory_mb,
-        "request_cpus": request_cpus if request_cpus is not None else condor.request_cpus,
-        "accounting_group": condor.accounting_group,
-        "should_transfer_files": "YES",
-        "when_to_transfer_output": "ON_EXIT",
-    }
-    if condor.remote:
-        settings["+RemoteJob"] = "True"
-    else:
-        settings["requirements"] = "TARGET.ProvidesETPResources"
-    if walltime_s is not None:
-        settings["+RequestWalltime"] = int(walltime_s)
-    return settings
-
-
-def gpu_requirements(gpu_type: str | None = None, gpu_memory_mb: int | None = None) -> str:
-    """``TARGET.ProvidesEtpCeph`` (remote /ceph access) ANDed with any GPU pin."""
-    clauses = ["TARGET.ProvidesEtpCeph =?= True"]
-    if gpu_type is not None:
-        clauses.append(f'TARGET.GPUs_DeviceName =?= "{gpu_type}"')
-    if gpu_memory_mb is not None:
-        clauses.append(f"TARGET.GPUs_GlobalMemoryMb >= {gpu_memory_mb}")
-    return " && ".join(clauses)
-
-
-def gpu_settings(
-    condor: CondorSpec,
-    *,
-    request_gpus: int = 1,
-    gpu_type: str | None = None,
-    gpu_memory_mb: int | None = None,
-    request_memory_mb: int = 16384,
-    request_cpus: int = 4,
-    walltime_s: int = 86400,
-) -> dict:
-    return {
-        "universe": "docker",
-        "docker_image": condor.docker_image_gpu,
-        "request_memory": request_memory_mb,
-        "request_cpus": request_cpus,
-        "RequestGPUs": request_gpus,
-        "+RequestWalltime": int(walltime_s),
-        "accounting_group": condor.accounting_group,
-        "should_transfer_files": "YES",
-        "when_to_transfer_output": "ON_EXIT",
-        "+RemoteJob": "True",
-        "requirements": f"({gpu_requirements(gpu_type, gpu_memory_mb)})",
-    }

giant/workflow/run.py

@@ -1,103 +0,0 @@
-#!/usr/bin/env python
-"""Entry point b2luigi re-executes on every worker.
-
-Locally this is what ``giant workflow run <spec.toml>`` execs; on a batch
-worker it is what the generated wrapper script runs (after ``cd repo_dir`` and
-sourcing ``env_script``), with ``--spec`` forwarded via the
-``task_cmd_additional_args`` setting so the worker resolves exactly the same
-spec — and therefore the same task graph and output paths — as the submitter.
-
-b2luigi needs a real script path for that re-execution, which is why this is a
-script rather than a ``python -m`` module.
-"""
-
-from __future__ import annotations
-
-import argparse
-import sys
-from pathlib import Path
-
-# Allow `python giant/workflow/run.py` from a checkout that isn't installed.
-sys.path.insert(0, str(Path(__file__).resolve().parents[2]))
-
-import b2luigi  # noqa: E402
-
-from giant.workflow.spec import WorkflowSpec, load_spec  # noqa: E402
-from giant.workflow.tasks import WorkflowTask, set_spec  # noqa: E402
-
-
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="Run a GIANT workflow spec with b2luigi.")
-    parser.add_argument("--spec", required=True, help="Workflow TOML (see configs/workflow_example.toml)")
-    parser.add_argument("--workers", type=int, default=1, help="Concurrent luigi workers")
-    parser.add_argument(
-        "--batch",
-        action="store_true",
-        help="Submit batch-system tasks to HTCondor (otherwise everything runs locally)",
-    )
-    parser.add_argument(
-        "--mode",
-        choices=("run", "dry-run", "show-output", "remove"),
-        default="run",
-        help="run (default), dry-run (print pending tasks), show-output (print every target), remove (delete outputs)",
-    )
-    parser.add_argument("--scheduler-host", default=None, help="luigid host (default: local scheduler)")
-    parser.add_argument("--scheduler-port", type=int, default=None, help="luigid port")
-    return parser
-
-
-def configure(spec: WorkflowSpec, spec_path: Path, batch: bool) -> None:
-    """Wire b2luigi's settings from the spec.
-
-    ``/ceph`` is shared between submit host and workers, so there is
-    deliberately no ``transfer_files``: ``result_dir``/``log_dir`` must live
-    somewhere both sides can see.
-    """
-    set_spec(spec)
-
-    b2luigi.set_setting("result_dir", spec.result_dir)
-    b2luigi.set_setting("log_dir", spec.log_dir)
-    b2luigi.set_setting("task_file_dir", str(Path(spec.result_dir) / "task_files"))
-    b2luigi.set_setting("use_parameter_name_in_output", True)
-    b2luigi.set_setting("batch_system", "htcondor" if batch else "local")
-    b2luigi.set_setting("working_dir", spec.condor.repo_dir)
-    b2luigi.set_setting("job_name", spec.name)
-    if spec.condor.env_script:
-        b2luigi.set_setting("env_script", spec.condor.env_script)
-    # The worker command is `<executable> [<basename of this file>] --batch-runner
-    # --task-id ...`, run after `cd working_dir`. Only the *basename* would be
-    # used, so the filename is dropped and the repo-relative script path is
-    # made part of the executable instead.
-    b2luigi.set_setting("add_filename_to_cmd", False)
-    b2luigi.set_setting("executable", [".venv/bin/python", "giant/workflow/run.py"])
-    b2luigi.set_setting("task_cmd_additional_args", ["--spec", str(spec_path)])
-
-
-def main(argv: list[str] | None = None) -> None:
-    args, _ = build_parser().parse_known_args(argv)
-    spec_path = Path(args.spec).resolve()
-    spec = load_spec(spec_path)
-    configure(spec, spec_path, batch=args.batch)
-
-    kwargs: dict = {}
-    if args.scheduler_host:
-        kwargs["scheduler_host"] = args.scheduler_host
-    if args.scheduler_port:
-        kwargs["scheduler_port"] = args.scheduler_port
-
-    b2luigi.process(
-        WorkflowTask(workflow_name=spec.name),
-        workers=args.workers,
-        batch=args.batch,
-        dry_run=args.mode == "dry-run",
-        show_output=args.mode == "show-output",
-        remove=args.mode == "remove",
-        auto_confirm=args.mode == "remove",
-        # run.py owns --spec/--mode/...; b2luigi must not choke on them.
-        ignore_additional_command_line_args=True,
-        **kwargs,
-    )
-
-
-if __name__ == "__main__":
-    main()

giant/workflow/spec.py

@@ -1,383 +0,0 @@
-"""Workflow TOML -> frozen dataclasses, validation, and per-task spec hashes.
-
-One spec file is the only place a pipeline is parameterised (see
-``configs/workflow_example.toml``):
-
-    [workflow]  name / result_dir / log_dir
-    [condor]    accounting group, repo dir, env script, docker images
-    [dataset]   steps (training) + reference (rollout seeds & analysis truth)
-    [geometry]  geometry-oracle build options
-    [[train]]   one per training run      (name, config, epochs, overrides, ...)
-    [[rollout]] one per rollout           (name, train = <a [[train]].name>, ...)
-    [[analysis]] one per comparison       (name, rollouts = [<[[rollout]].name>, ...])
-
-Every task carries its ``name`` plus a short ``spec_hash`` — 8 hex of the
-canonical JSON of its own resolved sub-spec **including its transitive
-parents**. That is what makes an edited spec produce a fresh result directory
-instead of silently reusing outputs computed under different settings: change
-the dataset and every hash downstream of it changes too.
-
-Unknown keys are rejected (with the valid ones listed), in the same spirit as
-``giant.config.validate_config_keys`` — a typo in a workflow spec would
-otherwise be a silently ignored setting on a multi-day pipeline.
-"""
-
-from __future__ import annotations
-
-import hashlib
-import json
-import tomllib
-from dataclasses import MISSING, dataclass, field, fields, is_dataclass
-from pathlib import Path
-from typing import Any
-
-__all__ = [
-    "AnalysisSpec",
-    "CondorSpec",
-    "DatasetSpec",
-    "GeometrySpec",
-    "RolloutSpec",
-    "TrainSpec",
-    "WorkflowSpec",
-    "load_spec",
-    "spec_hash",
-]
-
-
-class WorkflowSpecError(ValueError):
-    """Raised for any malformed workflow spec (unknown key, bad reference, ...)."""
-
-
-# ---------------------------------------------------------------------------
-# sub-specs
-# ---------------------------------------------------------------------------
-
-
-@dataclass(frozen=True)
-class CondorSpec:
-    """Where and how jobs run — the batch-system half of the spec.
-
-    ``repo_dir`` doubles as b2luigi's ``working_dir`` (jobs ``cd`` there before
-    running ``giant/workflow/run.py``), and ``env_script`` is sourced first,
-    since submit and worker machines don't share an environment.
-    """
-
-    accounting_group: str
-    repo_dir: str
-    env_script: str = ""
-    docker_image_cpu: str = "cverstege/alma9-gridjob"
-    docker_image_gpu: str = "mschnepf/slc7-condocker"
-    remote: bool = True
-    request_cpus: int = 1
-    request_memory_mb: int = 8192
-
-
-@dataclass(frozen=True)
-class DatasetSpec:
-    """The two datasets every pipeline needs.
-
-    ``steps`` is what training reads; ``reference`` is the held-out file
-    rollouts are seeded from and the analysis compares against (the "one
-    ground truth" premise of ``giant.analysis``).
-    """
-
-    steps: str
-    reference: str
-
-
-@dataclass(frozen=True)
-class GeometrySpec:
-    """``dwarf build-geometry-oracle`` options (see giant/tools/geometry_oracle.py)."""
-
-    method: str = "slab"
-    k: int = 1
-    subsample: int = 500_000
-    escape_factor: float = 5.0
-    seed: int = 0
-    depth_axis: int = 2
-    n_bins: int = 2000
-
-
-@dataclass(frozen=True)
-class TrainSpec:
-    """One training run, fanned out into ``ceil(epochs / epochs_per_job)`` jobs.
-
-    ``overrides`` are ``[train]``/model config keys merged on top of ``config``
-    exactly as ``giant train``'s flags are (``giant.config.merge_cli_overrides``),
-    so anything expressible on the CLI is expressible here.
-    """
-
-    name: str
-    config: str | None = None
-    epochs: int = 1
-    epochs_per_job: int = 1
-    overrides: dict[str, Any] = field(default_factory=dict)
-    request_gpus: int = 1
-    gpu_type: str | None = None
-    gpu_memory_mb: int | None = None
-    request_memory_mb: int = 16384
-    request_cpus: int = 4
-    walltime_s: int = 86400
-    num_workers: int = 4
-    shuffle_buffer: int = 65536
-    device: str | None = None
-
-
-@dataclass(frozen=True)
-class RolloutSpec:
-    """One ``giant rollout`` run against the checkpoint of ``train``."""
-
-    name: str
-    train: str
-    n_events: int | None = None
-    energy_cutoff: float = 0.1
-    max_steps: int = 1000
-    steps: int = 10
-    batch_size: int = 4096
-    max_tracks_per_event: int | None = None
-    escape_threshold: float | None = None
-    weights: str = "raw"
-    seed: int | None = None
-    request_gpus: int = 1
-    gpu_type: str | None = None
-    gpu_memory_mb: int | None = None
-    request_memory_mb: int = 16384
-    request_cpus: int = 2
-    walltime_s: int = 86400
-    device: str | None = None
-
-
-@dataclass(frozen=True)
-class AnalysisSpec:
-    """One rollout-vs-reference comparison (N rollout series, one reference)."""
-
-    name: str
-    rollouts: tuple[str, ...]
-    chunks: int = 1
-    energy_bins: int = 4
-    bins: int = 50
-    top_pdg: int = 6
-    gallery: bool = False
-    request_memory_mb: int = 8192
-    request_cpus: int = 1
-
-
-@dataclass(frozen=True)
-class WorkflowSpec:
-    """A whole pipeline: the parsed spec file plus name-keyed lookups."""
-
-    name: str
-    result_dir: str
-    log_dir: str
-    condor: CondorSpec
-    dataset: DatasetSpec
-    geometry: GeometrySpec
-    trains: tuple[TrainSpec, ...]
-    rollouts: tuple[RolloutSpec, ...]
-    analyses: tuple[AnalysisSpec, ...]
-    path: str = ""
-
-    # -- lookups ----------------------------------------------------------
-    def train(self, name: str) -> TrainSpec:
-        return _lookup(self.trains, name, "train")
-
-    def rollout(self, name: str) -> RolloutSpec:
-        return _lookup(self.rollouts, name, "rollout")
-
-    def analysis(self, name: str) -> AnalysisSpec:
-        return _lookup(self.analyses, name, "analysis")
-
-    # -- hashes -----------------------------------------------------------
-    # Each one folds in everything upstream of it, so a change anywhere in a
-    # task's ancestry moves its result directory (and only the affected
-    # subtree's).
-    def dataset_hash(self) -> str:
-        return spec_hash(self.dataset)
-
-    def warm_cache_hash(self, train_name: str) -> str:
-        # The setup cache depends on the dataset and on what this training's
-        # config asks of it (val split, conditioning, router) — not on how
-        # many epochs it runs for, so epochs/resources are deliberately left
-        # out and two trainings sharing a config share one warm-cache job.
-        t = self.train(train_name)
-        return spec_hash(self.dataset, t.config, t.overrides)
-
-    def geometry_hash(self) -> str:
-        return spec_hash(self.dataset, self.geometry)
-
-    def train_hash(self, name: str) -> str:
-        return spec_hash(self.dataset, self.train(name))
-
-    def rollout_hash(self, name: str) -> str:
-        ro = self.rollout(name)
-        return spec_hash(self.dataset, self.geometry, self.train(ro.train), ro)
-
-    def analysis_hash(self, name: str) -> str:
-        an = self.analysis(name)
-        parents = [self.rollout(r) for r in an.rollouts]
-        train_parents = [self.train(r.train) for r in parents]
-        return spec_hash(self.dataset, self.geometry, train_parents, parents, an)
-
-
-def _lookup(items, name: str, kind: str):
-    for item in items:
-        if item.name == name:
-            return item
-    known = ", ".join(sorted(i.name for i in items)) or "(none defined)"
-    raise WorkflowSpecError(f"no [[{kind}]] named {name!r} in this workflow — defined: {known}")
-
-
-# ---------------------------------------------------------------------------
-# hashing
-# ---------------------------------------------------------------------------
-
-
-def spec_hash(*parts: Any, length: int = 8) -> str:
-    """Short stable hash of one or more (sub-)specs.
-
-    Canonical JSON (sorted keys, dataclasses expanded) so the value depends
-    only on the resolved settings — not on key order in the TOML, nor on
-    which defaults were written out explicitly.
-    """
-    payload = json.dumps([_canonical(p) for p in parts], sort_keys=True, separators=(",", ":"))
-    return hashlib.sha256(payload.encode()).hexdigest()[:length]
-
-
-def _canonical(value: Any) -> Any:
-    if is_dataclass(value) and not isinstance(value, type):
-        return {f.name: _canonical(getattr(value, f.name)) for f in fields(value)}
-    if isinstance(value, dict):
-        return {str(k): _canonical(v) for k, v in value.items()}
-    if isinstance(value, (list, tuple)):
-        return [_canonical(v) for v in value]
-    if isinstance(value, Path):
-        return str(value)
-    return value
-
-
-# ---------------------------------------------------------------------------
-# parsing
-# ---------------------------------------------------------------------------
-
-
-def _build(cls, data: dict, where: str):
-    """Instantiate a frozen sub-spec, rejecting unknown/missing keys loudly."""
-    valid = {f.name for f in fields(cls)}
-    unknown = sorted(set(data) - valid)
-    if unknown:
-        raise WorkflowSpecError(f"{where}: unknown key(s) {unknown} — valid keys: {sorted(valid)}")
-    required = {f.name for f in fields(cls) if f.default is MISSING and f.default_factory is MISSING}
-    missing = sorted(required - set(data))
-    if missing:
-        raise WorkflowSpecError(f"{where}: missing required key(s) {missing}")
-    return cls(**data)
-
-
-def load_spec(path: str | Path) -> WorkflowSpec:
-    """Parse and validate a workflow TOML file."""
-    path = Path(path)
-    try:
-        raw = tomllib.loads(path.read_text())
-    except tomllib.TOMLDecodeError as exc:
-        raise WorkflowSpecError(f"{path}: not valid TOML — {exc}") from exc
-    return parse_spec(raw, path=path)
-
-
-def parse_spec(raw: dict, path: str | Path = "") -> WorkflowSpec:
-    """Validate an already-parsed workflow spec mapping."""
-    top_valid = {"workflow", "condor", "dataset", "geometry", "train", "rollout", "analysis"}
-    unknown = sorted(set(raw) - top_valid)
-    if unknown:
-        raise WorkflowSpecError(
-            f"{path or '<spec>'}: unknown top-level table(s) {unknown} — valid: {sorted(top_valid)}"
-        )
-
-    for required in ("workflow", "condor", "dataset"):
-        if required not in raw:
-            raise WorkflowSpecError(f"{path or '<spec>'}: missing required [{required}] table")
-
-    wf = dict(raw["workflow"])
-    wf_valid = {"name", "result_dir", "log_dir"}
-    wf_unknown = sorted(set(wf) - wf_valid)
-    if wf_unknown:
-        raise WorkflowSpecError(f"[workflow]: unknown key(s) {wf_unknown} — valid keys: {sorted(wf_valid)}")
-    if "name" not in wf or "result_dir" not in wf:
-        raise WorkflowSpecError("[workflow]: 'name' and 'result_dir' are required")
-    result_dir = str(Path(wf["result_dir"]).expanduser())
-    log_dir = str(Path(wf.get("log_dir", Path(result_dir) / "logs")).expanduser())
-
-    condor = _build(CondorSpec, dict(raw["condor"]), "[condor]")
-    dataset = _build(DatasetSpec, dict(raw["dataset"]), "[dataset]")
-    geometry = _build(GeometrySpec, dict(raw.get("geometry", {})), "[geometry]")
-
-    trains = tuple(_build(TrainSpec, dict(t), f"[[train]] #{i}") for i, t in enumerate(raw.get("train", [])))
-    rollouts = tuple(_build(RolloutSpec, dict(r), f"[[rollout]] #{i}") for i, r in enumerate(raw.get("rollout", [])))
-    analyses = tuple(
-        _build(AnalysisSpec, {**a, "rollouts": tuple(a.get("rollouts", ()))}, f"[[analysis]] #{i}")
-        for i, a in enumerate(raw.get("analysis", []))
-    )
-
-    _check_unique(trains, "train")
-    _check_unique(rollouts, "rollout")
-    _check_unique(analyses, "analysis")
-
-    train_names = {t.name for t in trains}
-    for ro in rollouts:
-        if ro.train not in train_names:
-            raise WorkflowSpecError(
-                f"[[rollout]] {ro.name!r}: train={ro.train!r} names no [[train]] — defined: {sorted(train_names)}"
-            )
-    rollout_names = {r.name for r in rollouts}
-    for an in analyses:
-        if not an.rollouts:
-            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: 'rollouts' must name at least one [[rollout]]")
-        for r in an.rollouts:
-            if r not in rollout_names:
-                raise WorkflowSpecError(
-                    f"[[analysis]] {an.name!r}: rollout {r!r} is not defined — "
-                    f"defined: {sorted(rollout_names) or '(none)'}"
-                )
-        if len(set(an.rollouts)) != len(an.rollouts):
-            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: repeated rollout name(s) in 'rollouts'")
-        if an.chunks < 1:
-            raise WorkflowSpecError(f"[[analysis]] {an.name!r}: chunks must be >= 1, got {an.chunks}")
-
-    for t in trains:
-        if t.epochs < 1:
-            raise WorkflowSpecError(f"[[train]] {t.name!r}: epochs must be >= 1, got {t.epochs}")
-        if t.epochs_per_job < 1:
-            raise WorkflowSpecError(f"[[train]] {t.name!r}: epochs_per_job must be >= 1, got {t.epochs_per_job}")
-
-    return WorkflowSpec(
-        name=wf["name"],
-        result_dir=result_dir,
-        log_dir=log_dir,
-        condor=condor,
-        dataset=dataset,
-        geometry=geometry,
-        trains=trains,
-        rollouts=rollouts,
-        analyses=analyses,
-        path=str(path),
-    )
-
-
-def _check_unique(items, kind: str) -> None:
-    names = [i.name for i in items]
-    dupes = sorted({n for n in names if names.count(n) > 1})
-    if dupes:
-        raise WorkflowSpecError(f"[[{kind}]] names must be unique — repeated: {dupes}")
-
-
-def epoch_milestones(train: TrainSpec) -> list[int]:
-    """Cumulative epoch counts, one per chained ``TrainEpochTask``.
-
-    ``epochs_per_job`` trades queue waits against job length: with
-    ``epochs=10, epochs_per_job=3`` this is ``[3, 6, 9, 10]``, i.e. job *k*
-    resumes job *k-1*'s ``last.pt`` and trains up to its own milestone.
-    """
-    step = train.epochs_per_job
-    milestones = list(range(step, train.epochs + 1, step))
-    if not milestones or milestones[-1] != train.epochs:
-        milestones.append(train.epochs)
-    return milestones

giant/workflow/tasks.py

@@ -1,647 +0,0 @@
-"""The b2luigi task graph: cache-warm -> train -> rollout -> analysis.
-
-    DatasetTask (external)  ─┬─> WarmCacheTask(train) ─> TrainEpochTask(train, 1..N) ─> TrainTask(train) ─┐
-                             └─> GeometryOracleTask ──┐                                                   │
-                                                      └──> RolloutTask(rollout) <────────────────────────┘
-                                                                       │
-    AnalysisPrepTask(analysis) ─> AnalysisComputeTask(analysis, plot, chunk) ─> AnalysisRenderTask(analysis)
-                                                                                        ^
-    WorkflowTask (wrapper) ─────────────────────────────────────────────────────────────┘
-
-Every task's output directory is ``<result_dir>/<kind>/name=<name>/spec_hash=
-<hash>/…`` — the hash covers the task's resolved sub-spec *and its transitive
-parents* (``giant/workflow/spec.py``), so editing the spec produces a fresh
-directory for exactly the affected subtree instead of silently reusing stale
-outputs.
-
-Task bodies never reimplement anything: they call the same entry points the
-CLIs do (``run_warm_setup_cache``, ``run_build_geometry_oracle``,
-``run_train_job``, ``giant.analysis.prep``/``compute_one``/``merge_all``,
-``render_run``), or shell out to ``giant rollout``, which has no library-level
-entry point of its own.
-
-Training is fanned out into **one short GPU job per epoch** (or per
-``epochs_per_job`` epochs): job *k* runs ``run_train_job`` with ``epochs = k``
-and ``resume = <job k-1>/last.pt``, which the training loop already handles
-(``giant/training/loop.py`` sets ``start_epoch = ckpt["epoch"] + 1`` and
-returns early when the checkpoint already covers ``epochs``). A 200-epoch run
-then becomes 200 schedulable jobs that survive preemption and give luigi a
-real progress signal, at the cost of one (cache-warmed) setup scan and one
-queue wait per job.
-"""
-
-from __future__ import annotations
-
-import json
-import shutil
-import subprocess
-import sys
-from pathlib import Path
-
-import b2luigi
-
-from giant.workflow.htcondor import cpu_settings, gpu_settings
-from giant.workflow.spec import WorkflowSpec, epoch_milestones
-
-__all__ = [
-    "AnalysisComputeTask",
-    "AnalysisPrepTask",
-    "AnalysisRenderTask",
-    "DatasetTask",
-    "GeometryOracleTask",
-    "RolloutTask",
-    "TrainEpochTask",
-    "TrainTask",
-    "WarmCacheTask",
-    "WorkflowTask",
-    "analysis_dir",
-    "analysis_jobs",
-    "get_spec",
-    "set_spec",
-]
-
-
-# ---------------------------------------------------------------------------
-# the active spec
-# ---------------------------------------------------------------------------
-# luigi parameters must be simple scalars, so tasks carry only `name` +
-# `spec_hash` and read the rest out of the one spec this process was started
-# with. Batch workers re-execute `run.py --spec <same file>` (see
-# `task_cmd_additional_args` there), so they resolve the identical spec.
-
-_SPEC: WorkflowSpec | None = None
-
-
-def set_spec(spec: WorkflowSpec) -> None:
-    global _SPEC
-    _SPEC = spec
-
-
-def get_spec() -> WorkflowSpec:
-    if _SPEC is None:
-        raise RuntimeError("no workflow spec loaded — call giant.workflow.tasks.set_spec() first")
-    return _SPEC
-
-
-def _result_dir(*parts: str) -> Path:
-    return Path(get_spec().result_dir).joinpath(*parts)
-
-
-def _task_dir(kind: str, name: str, spec_hash: str) -> Path:
-    """``<result_dir>/<kind>/name=<name>/spec_hash=<hash>``."""
-    return _result_dir(kind, f"name={name}", f"spec_hash={spec_hash}")
-
-
-def analysis_dir(spec: WorkflowSpec, name: str) -> Path:
-    """The analysis run directory — what ``prep`` lays out and every later step reads."""
-    return Path(spec.result_dir) / "analysis" / f"name={name}" / f"spec_hash={spec.analysis_hash(name)}"
-
-
-def analysis_jobs(spec: WorkflowSpec, name: str) -> list[tuple[str, int]]:
-    """Every ``(plot_id, chunk)`` compute job of one analysis.
-
-    ``chunkable=False`` specs (the checkpoint-bound diagnostics, already
-    bounded/subsampled) always run as a single chunk — the same rule the
-    deleted ``_job_walltimes`` applied.
-    """
-    from giant.analysis.catalog import catalog_ids, get_spec as get_plot_spec
-
-    chunks = spec.analysis(name).chunks
-    jobs: list[tuple[str, int]] = []
-    for plot_id in catalog_ids():
-        n = chunks if get_plot_spec(plot_id).chunkable else 1
-        jobs.extend((plot_id, chunk) for chunk in range(n))
-    return jobs
-
-
-def _giant_cmd() -> list[str]:
-    """How to invoke the ``giant`` CLI from inside a task (worker or locally)."""
-    return [sys.executable, "-m", "giant.cli"]
-
-
-# ---------------------------------------------------------------------------
-# inputs
-# ---------------------------------------------------------------------------
-
-
-class DatasetTask(b2luigi.ExternalTask):
-    """A steps parquet file or directory that must already exist.
-
-    Nothing produces it, so a missing path is a hard, immediate error rather
-    than a job that fails hours later — the usual cause being ``/ceph`` not
-    mounted on the machine the workflow was started from.
-    """
-
-    path = b2luigi.Parameter()
-
-    def output(self):
-        return b2luigi.LocalTarget(str(self.path))
-
-    def complete(self):
-        if not Path(str(self.path)).exists():
-            raise FileNotFoundError(
-                f"dataset {self.path!r} does not exist — is /ceph mounted on this machine? "
-                "(see CLAUDE.md's Compute environment section)"
-            )
-        return True
-
-
-# ---------------------------------------------------------------------------
-# setup stage
-# ---------------------------------------------------------------------------
-
-
-class WarmCacheTask(b2luigi.Task):
-    """Precompute one training's setup-stage sidecar (vocab maps, event split,
-    normalizer stats) so every per-epoch job is a cache hit instead of a
-    full rescan.
-
-    The real product (``<data>.giant_train_cache.json``) lives next to the
-    dataset, not under ``result_dir``, so the target here is a small stamp
-    recording that sidecar's path/mtime/size.
-    """
-
-    name = b2luigi.Parameter()
-    spec_hash = b2luigi.Parameter()
-
-    @property
-    def htcondor_settings(self):
-        spec = get_spec()
-        return cpu_settings(spec.condor, request_memory_mb=32768, request_cpus=4, walltime_s=21600)
-
-    def requires(self):
-        yield DatasetTask(path=get_spec().dataset.steps)
-
-    def output(self):
-        return b2luigi.LocalTarget(str(_task_dir("warm_cache", str(self.name), str(self.spec_hash)) / "stamp.json"))
-
-    def run(self):
-        from giant.data.setup_cache import sidecar_path
-        from giant.tools.warm_setup_cache import run_warm_setup_cache
-
-        spec = get_spec()
-        train = spec.train(str(self.name))
-        run_warm_setup_cache(
-            data=spec.dataset.steps,
-            config_path=Path(train.config) if train.config else None,
-        )
-        sidecar = Path(sidecar_path(spec.dataset.steps))
-        stamp = {
-            "sidecar": str(sidecar),
-            "mtime": sidecar.stat().st_mtime if sidecar.exists() else None,
-            "size": sidecar.stat().st_size if sidecar.exists() else None,
-        }
-        out = Path(self.output().path)
-        out.parent.mkdir(parents=True, exist_ok=True)
-        out.write_text(json.dumps(stamp, indent=2))
-
-
-class GeometryOracleTask(b2luigi.Task):
-    """Build the position -> (material, layer_id) oracle every rollout needs."""
-
-    spec_hash = b2luigi.Parameter()
-
-    @property
-    def htcondor_settings(self):
-        spec = get_spec()
-        return cpu_settings(spec.condor, request_memory_mb=32768, request_cpus=4, walltime_s=21600)
-
-    def requires(self):
-        yield DatasetTask(path=get_spec().dataset.steps)
-
-    def output(self):
-        return b2luigi.LocalTarget(
-            str(_result_dir("geometry", f"spec_hash={self.spec_hash}") / "oracle.pkl"),
-        )
-
-    def run(self):
-        from giant.tools.geometry_oracle import run_build_geometry_oracle
-
-        spec = get_spec()
-        g = spec.geometry
-        out = Path(self.output().path)
-        out.parent.mkdir(parents=True, exist_ok=True)
-        run_build_geometry_oracle(
-            data=Path(spec.dataset.steps),
-            out=out,
-            method=g.method,
-            k=g.k,
-            subsample=g.subsample,
-            escape_factor=g.escape_factor,
-            seed=g.seed,
-            depth_axis=g.depth_axis,
-            n_bins=g.n_bins,
-        )
-
-
-# ---------------------------------------------------------------------------
-# training
-# ---------------------------------------------------------------------------
-
-
-def _train_cfg(spec: WorkflowSpec, name: str, epochs: int) -> dict:
-    """The merged config one training job runs, resolved exactly as `giant train` does."""
-    from giant import config as gconfig
-
-    train = spec.train(name)
-    flags = {**train.overrides, "epochs": epochs}
-    overrides = gconfig.overrides_from_flags(flags)
-    cfg = gconfig.merge_cli_overrides(
-        gconfig.DEFAULT_CONFIG,
-        Path(train.config) if train.config else None,
-        overrides,
-    )
-    gconfig.validate_config(cfg, resume=True)
-    return cfg
-
-
-class TrainEpochTask(b2luigi.Task):
-    """Epochs up to ``milestone`` of one training, resuming the previous job.
-
-    Target is ``last.pt``. ``best.pt`` is written by the loop *only when that
-    epoch improved*, and ``best_val_loss`` travels inside the checkpoint, so
-    the global best comparison stays correct across jobs: "``best.pt`` exists
-    in milestone dir *k*" means exactly "one of that job's epochs was the best
-    so far".
-    """
-
-    name = b2luigi.Parameter()
-    spec_hash = b2luigi.Parameter()
-    milestone = b2luigi.IntParameter()
-
-    @property
-    def htcondor_settings(self):
-        spec = get_spec()
-        train = spec.train(str(self.name))
-        return gpu_settings(
-            spec.condor,
-            request_gpus=train.request_gpus,
-            gpu_type=train.gpu_type,
-            gpu_memory_mb=train.gpu_memory_mb,
-            request_memory_mb=train.request_memory_mb,
-            request_cpus=train.request_cpus,
-            walltime_s=train.walltime_s,
-        )
-
-    @property
-    def _dir(self) -> Path:
-        return _task_dir("train_epoch", str(self.name), str(self.spec_hash)) / f"epochs={int(self.milestone)}"
-
-    def _previous_milestone(self) -> int | None:
-        spec = get_spec()
-        milestones = epoch_milestones(spec.train(str(self.name)))
-        index = milestones.index(int(self.milestone))
-        return milestones[index - 1] if index > 0 else None
-
-    def requires(self):
-        previous = self._previous_milestone()
-        if previous is None:
-            yield WarmCacheTask(name=self.name, spec_hash=get_spec().warm_cache_hash(str(self.name)))
-        else:
-            yield TrainEpochTask(name=self.name, spec_hash=self.spec_hash, milestone=previous)
-
-    def output(self):
-        return b2luigi.LocalTarget(str(self._dir / "last.pt"))
-
-    def run(self):
-        import torch
-
-        from giant import config as gconfig
-        from giant.pipeline import run_train_job
-
-        spec = get_spec()
-        train = spec.train(str(self.name))
-        cfg = _train_cfg(spec, str(self.name), int(self.milestone))
-
-        previous = self._previous_milestone()
-        resume = None
-        if previous is not None:
-            resume = _task_dir("train_epoch", str(self.name), str(self.spec_hash)) / f"epochs={previous}" / "last.pt"
-
-        device = torch.device(train.device) if train.device else gconfig.auto_device()
-        out_dir = self._dir
-        out_dir.mkdir(parents=True, exist_ok=True)
-        run_train_job(
-            data=Path(spec.dataset.steps),
-            cfg=cfg,
-            out_dir=out_dir,
-            device=device,
-            shuffle_buffer=train.shuffle_buffer,
-            num_workers=train.num_workers,
-            resume=resume,
-            cache_setup=True,
-        )
-
-
-class TrainTask(b2luigi.Task):
-    """Publish one training's canonical outputs, hiding the epoch fan-out.
-
-    Everything downstream (``RolloutTask``, humans, ``giant analyze metrics``)
-    points here and never has to know which milestone directory happened to
-    hold the best checkpoint.
-    """
-
-    name = b2luigi.Parameter()
-    spec_hash = b2luigi.Parameter()
-    batch_system = "local"
-
-    @property
-    def _milestones(self) -> list[int]:
-        return epoch_milestones(get_spec().train(str(self.name)))
-
-    def requires(self):
-        yield TrainEpochTask(name=self.name, spec_hash=self.spec_hash, milestone=self._milestones[-1])
-
-    @property
-    def _dir(self) -> Path:
-        return _task_dir("train", str(self.name), str(self.spec_hash))
-
-    def output(self):
-        d = self._dir
-        return {
-            "best.pt": b2luigi.LocalTarget(str(d / "best.pt")),
-            "last.pt": b2luigi.LocalTarget(str(d / "last.pt")),
-            "metrics.csv": b2luigi.LocalTarget(str(d / "metrics.csv")),
-        }
-
-    def run(self):
-        epoch_base = _task_dir("train_epoch", str(self.name), str(self.spec_hash))
-        milestone_dirs = [epoch_base / f"epochs={m}" for m in self._milestones]
-
-        best_dirs = [d for d in milestone_dirs if (d / "best.pt").exists()]
-        if not best_dirs:
-            raise FileNotFoundError(
-                f"no best.pt in any milestone directory under {epoch_base} — "
-                "did every epoch job run with a validation split?"
-            )
-        out = self._dir
-        out.mkdir(parents=True, exist_ok=True)
-        shutil.copy2(best_dirs[-1] / "best.pt", out / "best.pt")
-        shutil.copy2(milestone_dirs[-1] / "last.pt", out / "last.pt")
-        for extra in ("config.toml", "run_meta.json"):
-            src = milestone_dirs[-1] / extra
-            if src.exists():
-                shutil.copy2(src, out / extra)
-
-        # One metrics.csv for the whole run: the first job's header, then
-        # every job's rows in epoch order, so `giant analyze metrics` sees a
-        # single continuous training curve.
-        lines: list[str] = []
-        header: str | None = None
-        for d in milestone_dirs:
-            csv = d / "metrics.csv"
-            if not csv.exists():
-                continue
-            rows = csv.read_text().splitlines()
-            if not rows:
-                continue
-            if header is None:
-                header = rows[0]
-            lines.extend(rows[1:])
-        (out / "metrics.csv").write_text("\n".join([header or ""] + lines) + "\n")
-
-
-# ---------------------------------------------------------------------------
-# rollout
-# ---------------------------------------------------------------------------
-
-
-class RolloutTask(b2luigi.Task):
-    """Roll one trained checkpoint forward into full showers.
-
-    ``giant rollout`` has no library-level entry point, so this shells out to
-    the CLI — with an explicit ``--out``, which puts the YAML sidecar at the
-    deterministic ``rollout.yaml`` next to the parquet (see
-    ``giant/cli.py:_write_prediction_ref``).
-    """
-
-    name = b2luigi.Parameter()
-    spec_hash = b2luigi.Parameter()
-
-    @property
-    def htcondor_settings(self):
-        spec = get_spec()
-        ro = spec.rollout(str(self.name))
-        return gpu_settings(
-            spec.condor,
-            request_gpus=ro.request_gpus,
-            gpu_type=ro.gpu_type,
-            gpu_memory_mb=ro.gpu_memory_mb,
-            request_memory_mb=ro.request_memory_mb,
-            request_cpus=ro.request_cpus,
-            walltime_s=ro.walltime_s,
-        )
-
-    @property
-    def _dir(self) -> Path:
-        return _task_dir("rollout", str(self.name), str(self.spec_hash))
-
-    def requires(self):
-        spec = get_spec()
-        ro = spec.rollout(str(self.name))
-        yield TrainTask(name=ro.train, spec_hash=spec.train_hash(ro.train))
-        yield GeometryOracleTask(spec_hash=spec.geometry_hash())
-        yield DatasetTask(path=spec.dataset.reference)
-
-    def output(self):
-        d = self._dir
-        return {
-            "rollout.parquet": b2luigi.LocalTarget(str(d / "rollout.parquet")),
-            "rollout.yaml": b2luigi.LocalTarget(str(d / "rollout.yaml")),
-        }
-
-    def run(self):
-        spec = get_spec()
-        ro = spec.rollout(str(self.name))
-        out = self._dir / "rollout.parquet"
-        out.parent.mkdir(parents=True, exist_ok=True)
-
-        checkpoint = _task_dir("train", ro.train, spec.train_hash(ro.train)) / "best.pt"
-        oracle = _result_dir("geometry", f"spec_hash={spec.geometry_hash()}") / "oracle.pkl"
-
-        cmd = [
-            *_giant_cmd(),
-            "rollout",
-            spec.dataset.reference,
-            "--checkpoint",
-            str(checkpoint),
-            "--geometry",
-            str(oracle),
-            "--out",
-            str(out),
-            "--energy-cutoff",
-            str(ro.energy_cutoff),
-            "--max-steps",
-            str(ro.max_steps),
-            "--steps",
-            str(ro.steps),
-            "--batch-size",
-            str(ro.batch_size),
-            "--weights",
-            ro.weights,
-        ]
-        for flag, value in (
-            ("--n-events", ro.n_events),
-            ("--max-tracks-per-event", ro.max_tracks_per_event),
-            ("--escape-threshold", ro.escape_threshold),
-            ("--seed", ro.seed),
-            ("--device", ro.device),
-        ):
-            if value is not None:
-                cmd += [flag, str(value)]
-        subprocess.run(cmd, check=True)
-
-
-# ---------------------------------------------------------------------------
-# analysis
-# ---------------------------------------------------------------------------
-
-
-class AnalysisPrepTask(b2luigi.Task):
-    """Resolve the shared bin edges/group sets once, for every compute job.
-
-    Cheap and streaming, so it runs locally: everything after it needs
-    ``shared.json``/``run_meta.json`` to already exist.
-    """
-
-    name = b2luigi.Parameter()
-    spec_hash = b2luigi.Parameter()
-    batch_system = "local"
-
-    def requires(self):
-        spec = get_spec()
-        for rollout_name in spec.analysis(str(self.name)).rollouts:
-            yield RolloutTask(name=rollout_name, spec_hash=spec.rollout_hash(rollout_name))
-
-    @property
-    def _dir(self) -> Path:
-        return analysis_dir(get_spec(), str(self.name))
-
-    def output(self):
-        d = self._dir
-        return {
-            "shared.json": b2luigi.LocalTarget(str(d / "shared.json")),
-            "run_meta.json": b2luigi.LocalTarget(str(d / "run_meta.json")),
-        }
-
-    def run(self):
-        from giant.analysis import prep
-
-        spec = get_spec()
-        an = spec.analysis(str(self.name))
-        yamls = [_task_dir("rollout", r, spec.rollout_hash(r)) / "rollout.yaml" for r in an.rollouts]
-        prep(
-            yamls,
-            run_dir=self._dir,
-            n_chunks=an.chunks,
-            labels=list(an.rollouts),
-            n_energy_bins=an.energy_bins,
-            n_marginal_bins=an.bins,
-            top_k_pdg=an.top_pdg,
-        )
-
-
-class AnalysisComputeTask(b2luigi.Task):
-    """One (plot, chunk) streaming reduction — the replaced ``jobs.txt`` row.
-
-    The output path is the on-disk contract ``compute-one``/``merge_one``
-    already share (``reduced_partial/<id>__<chunk>.json``), declared
-    explicitly rather than through b2luigi's own output naming so that
-    contract is untouched.
-    """
-
-    name = b2luigi.Parameter()
-    spec_hash = b2luigi.Parameter()
-    plot_id = b2luigi.Parameter()
-    chunk = b2luigi.IntParameter()
-
-    @property
-    def htcondor_settings(self):
-        # A property, so it is evaluated at submit time — i.e. after prep has
-        # written run_meta.json, whose row counts size the walltime request.
-        from giant.analysis import RunMeta
-        from giant.analysis.runtime_estimate import estimate_runtime_s
-
-        spec = get_spec()
-        an = spec.analysis(str(self.name))
-        walltime = None
-        meta_path = analysis_dir(spec, str(self.name)) / "run_meta.json"
-        if meta_path.exists():
-            from giant.analysis.catalog import get_spec as get_plot_spec
-
-            meta = RunMeta.load(meta_path)
-            chunkable = get_plot_spec(str(self.plot_id)).chunkable
-            n_rows = meta.rows_per_chunk[int(self.chunk)] if chunkable and meta.rows_per_chunk else meta.total_rows
-            walltime = estimate_runtime_s(str(self.plot_id), n_rows)
-        return cpu_settings(
-            spec.condor,
-            request_memory_mb=an.request_memory_mb,
-            request_cpus=an.request_cpus,
-            walltime_s=walltime,
-        )
-
-    def requires(self):
-        yield AnalysisPrepTask(name=self.name, spec_hash=self.spec_hash)
-
-    def output(self):
-        run_dir = analysis_dir(get_spec(), str(self.name))
-        return b2luigi.LocalTarget(str(run_dir / "reduced_partial" / f"{self.plot_id}__{int(self.chunk)}.json"))
-
-    def run(self):
-        from giant.analysis import compute_one
-
-        compute_one(str(self.plot_id), analysis_dir(get_spec(), str(self.name)), chunk_index=int(self.chunk))
-
-
-class AnalysisRenderTask(b2luigi.Task):
-    """Merge every plot's chunk partials, then render the PDFs + gallery.
-
-    Always local — this is the only step that imports plotstyle/LaTeX, which
-    the compute worker images don't have.
-    """
-
-    name = b2luigi.Parameter()
-    spec_hash = b2luigi.Parameter()
-    batch_system = "local"
-
-    def requires(self):
-        spec = get_spec()
-        for plot_id, chunk in analysis_jobs(spec, str(self.name)):
-            yield AnalysisComputeTask(
-                name=self.name,
-                spec_hash=self.spec_hash,
-                plot_id=plot_id,
-                chunk=chunk,
-            )
-
-    def output(self):
-        run_dir = analysis_dir(get_spec(), str(self.name))
-        return b2luigi.LocalTarget(str(run_dir / "plots" / "metadata.yaml"))
-
-    def run(self):
-        # render_run joins every plot's chunk partials (merge_all) before
-        # rendering, so this one call is the whole merge+render step.
-        from giant.analysis.render import render_run
-
-        spec = get_spec()
-        render_run(analysis_dir(spec, str(self.name)), run_gallery=spec.analysis(str(self.name)).gallery)
-
-
-class WorkflowTask(b2luigi.WrapperTask):
-    """The whole pipeline: every analysis in the spec, rendered."""
-
-    workflow_name = b2luigi.Parameter()
-
-    def requires(self):
-        spec = get_spec()
-        if not spec.analyses:
-            # A spec with no [[analysis]] still has work to do — fall back to
-            # the deepest tasks it does define.
-            for ro in spec.rollouts:
-                yield RolloutTask(name=ro.name, spec_hash=spec.rollout_hash(ro.name))
-            if not spec.rollouts:
-                for tr in spec.trains:
-                    yield TrainTask(name=tr.name, spec_hash=spec.train_hash(tr.name))
-            return
-        for an in spec.analyses:
-            yield AnalysisRenderTask(name=an.name, spec_hash=spec.analysis_hash(an.name))

pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "giant"
-version = "0.4.0"
+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",
-    "pandas>=2.2,<4",
-    "pyarrow>=16,<25",
+    "polars>=1.43,<2",
+    "pyarrow>=16,<26",
     "tqdm>=4.60,<5",
     "typer>=0.12,<1",
     "pyyaml>=6,<7",
@@ -28,7 +28,11 @@ dev = [
     "ty>=0.0.50,<0.1",
     "bump-my-version>=1.2,<2",
     "git-cliff>=2,<3",
-    "giant[convert,analysis,geometry,wandb,workflow]",
+    # Only used by test fixtures (writing small parquet files) — not a
+    # runtime dependency of giant itself since the pandas -> polars
+    # data-loading rewrite.
+    "pandas>=2.2,<4",
+    "giant[convert,analysis,geometry,wandb]",
 ]
 geometry = [
     "scikit-learn>=1.4,<2",
@@ -39,21 +43,16 @@ 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.
     "plotstyle>=1.0.0",
 ]
-# b2luigi pulls luigi + tenacity; the only sanctioned way to chain a
-# multi-step pipeline (see giant/workflow/).
-workflow = [
-    "b2luigi>=1.0,<2",
-]
 
 [project.scripts]
 giant = "giant.cli:app"
@@ -108,12 +107,3 @@ explicit = true
 name = "larsbogner"
 url = "https://git.larsbogner.de/api/packages/lars/pypi/simple/"
 explicit = true
-
-# luigi builds task constructors from class-level Parameter descriptors, so a
-# static checker sees no keyword parameters at all on `Task(name=..., ...)`.
-# The workflow code is written against that API; nothing else in the repo is.
-[[tool.ty.overrides]]
-include = ["giant/workflow/**", "tests/test_workflow_tasks.py"]
-
-[tool.ty.overrides.rules]
-unknown-argument = "ignore"

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

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

@@ -1,7 +1,8 @@
-"""Tests for the rollout-YAML(s) → run-directory flow, compute, and merge."""
+"""Tests for the rollout-YAML(s) → run-directory flow, compute, and submit."""
 
 from __future__ import annotations
 
+import sys
 from pathlib import Path
 
 import pyarrow.parquet as pq
@@ -10,17 +11,24 @@ import yaml
 
 from giant.analysis import (
     RunMeta,
+    SubmitConfig,
+    catalog_ids,
     compute_one,
     compute_reduced,
     derive_run_dir,
+    load_prediction_yaml,
+    load_prediction_yamls,
     load_rollout_yaml,
     load_rollout_yamls,
     merge_one,
     prep,
+    write_submit,
 )
-from giant.analysis.run import Context
+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
 
 
@@ -81,13 +89,46 @@ def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
     return paths[0], paths[1]
 
 
-def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
+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"
+    giant.parent.mkdir(parents=True, exist_ok=True)
+    giant.write_text("#!/bin/bash\n")
+    giant.chmod(0o755)
+
+
+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,
@@ -148,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")
@@ -316,6 +464,73 @@ def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
         compute_one("marginal_edep", run_dir, chunk_index=1)
 
 
+def test_write_submit_description(tmp_path: Path):
+    run_dir = _prep([_write_inputs(tmp_path)])
+    _fake_venv(tmp_path)
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
+    txt = write_submit(cfg).read_text()
+    assert "universe = docker" in txt
+    assert "docker_image = cverstege/alma9-gridjob" in txt
+    assert "requirements = TARGET.ProvidesETPResources" in txt
+    assert "accounting_group = cms" in txt
+    assert "+RequestWalltime = $(walltime)" in txt
+    assert "queue plotid,chunk,walltime from" in txt
+    jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
+    assert [i for i, _, _ in jobs] == catalog_ids()
+    assert all(k == "0" for _, k, _ in jobs)  # n_chunks=1 default
+    assert all(int(w) > 0 for _, _, w in jobs)
+    wrapper = run_dir / "run_compute.sh"
+    assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
+    body = wrapper.read_text()
+    assert "giant analyze compute-one --id" in body
+    assert "--chunk" in body and "--run-dir" in body
+
+
+def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
+    run_dir = _prep([_write_inputs(tmp_path)])
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
+    # No `giant` next to the (fake) active interpreter, so this falls through
+    # to repo_dir/.venv/bin/giant, which _write_inputs/_prep also didn't create.
+    monkeypatch.setattr(sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python"))
+    with pytest.raises(FileNotFoundError, match="uv sync"):
+        write_submit(cfg)
+
+
+def test_write_submit_remote_flag(tmp_path: Path):
+    run_dir = _prep([_write_inputs(tmp_path)])
+    _fake_venv(tmp_path)
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True)
+    txt = write_submit(cfg).read_text()
+    assert "+RemoteJob = True" in txt
+    assert "ProvidesETPResources" not in txt
+
+
+def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
+    assert get_spec("router_gating").chunkable is False
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=4)
+    _fake_venv(tmp_path)
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
+    write_submit(cfg)
+    jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
+    counts: dict[str, int] = {}
+    for spec_id, _, _ in jobs:
+        counts[spec_id] = counts.get(spec_id, 0) + 1
+    assert counts["marginal_edep"] == 4
+    assert counts["router_gating"] == 1  # chunkable=False, ignores n_chunks
+
+
+def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
+    """cfg.n_chunks must match the n_chunks the run_dir was actually prepped
+    with — RunMeta.rows_per_chunk is sized to the prepped value, so a
+    mismatch would otherwise surface as a confusing IndexError deep inside
+    _job_walltimes instead of a clear error here."""
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
+    _fake_venv(tmp_path)
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
+    with pytest.raises(ValueError, match="n_chunks"):
+        write_submit(cfg)
+
+
 def test_estimate_runtime_s_scales_with_rows_and_margin():
     from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
     from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S
@@ -325,3 +540,18 @@ def test_estimate_runtime_s_scales_with_rows_and_margin():
     large = estimate_runtime_s("marginal_edep", 100_000_000)
     assert small >= (1 + RUNTIME_SAFETY_MARGIN) * _FIXED_OVERHEAD_S
     assert large > small  # bigger chunk -> longer estimate
+
+
+def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
+    """A chunked run's later job walltimes track that chunk's row count."""
+    from giant.analysis.runtime_estimate import estimate_runtime_s
+
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    _fake_venv(tmp_path)
+    cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2)
+    write_submit(cfg)
+    jobs = {(i, int(k)): int(w) for i, k, w in (line.split(",") for line in (run_dir / "jobs.txt").read_text().split())}
+    for chunk in range(2):
+        expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
+        assert jobs[("marginal_edep", chunk)] == expected

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_loader.py

@@ -108,12 +108,12 @@ def test_build_process_map_from_files_spans_multiple_files(tmp_path):
 
 def test_build_process_map_from_files_tie_breaking_pins_first_seen_order(tmp_path):
     """When two processes end up with equal total counts, ranking falls back
-    to whichever was accumulated first (`sorted(..., reverse=True)` is stable,
-    and `counts` is built in file/row-scan order) — this is implementation-
-    defined, not a documented contract, so pin it explicitly: a future
-    rewrite (e.g. a polars-based single-scan) that ties differently would
-    silently reshuffle which processes get their own expert slot across a
-    retrain, and this test is what should catch that."""
+    to whichever was scanned first — file order, then row order within a
+    file (`giant.data.scan`'s `first_seen` ordinal, ranked by
+    `giant.data.loader._topn_plus_other_map`'s `(-count, first_seen)` key).
+    This is an explicit, documented contract (not an accident of iteration
+    order), pinned here so a future change to the ranking can't silently
+    reshuffle which processes get their own expert slot across a retrain."""
     path = tmp_path / "a.parquet"
     pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path)
 
@@ -233,6 +233,23 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path)
     assert m.class_counts == {0: 11, 1: 5}
 
 
+def test_build_pdg_topn_map_from_files_pooled_tie_breaks_by_row_position(tmp_path):
+    """Pooled pdg counting merges the primary `pdg` column and the exploded
+    `sec_pdg_list` column via one `group_by` over both (see
+    `giant.data.scan._pooled_pdg_lazy`), keyed by row position regardless of
+    which role (primary or secondary) a code was seen in — not "all
+    primaries before all secondaries" the way a two-pass accumulation would.
+    11 (primary, row 0), 33 (primary, row 1) and 22 (secondary, row 1) all
+    end up with count 1; 11's strictly earlier row wins the tie over both,
+    whatever order 33/22 (tied with each other, same row) land in."""
+    path = tmp_path / "a.parquet"
+    pd.DataFrame({"pdg": [11, 33], "sec_pdg_list": [[], [22]]}).to_parquet(path)
+
+    m = build_pdg_topn_map_from_files([path], n_classes=4)
+
+    assert m.class_map[11] == 0
+
+
 def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):
     """Files predating the parent->child join have no sec_pdg_list column —
     must not raise, just count the primary pdg column alone."""

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

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_pipeline.py

@@ -142,7 +142,7 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
     def _forbidden(*a, **k):
         raise AssertionError("should be served from cache, not recomputed")
 
-    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+    monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden)
     monkeypatch.setattr("giant.pipeline.iter_file_chunks", _forbidden)
 
     echo2 = _run(data, tmp_path / "out2")
@@ -283,7 +283,7 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
     def _forbidden(*a, **k):
         raise AssertionError("vocab should be served from cache")
 
-    monkeypatch.setattr("giant.pipeline.build_index_maps_from_files", _forbidden)
+    monkeypatch.setattr("giant.pipeline.scan_metadata", _forbidden)
 
     echo2 = _run(data, tmp_path / "out2", cfg=_tiny_cfg(val_fraction=0.3))
     joined = "\n".join(echo2)

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.reduced import Reduced  # noqa: E402
+from giant.analysis import render as render_mod
+from giant.analysis.reduced import Reduced
 
 
 def _try_render(reduced: list[Reduced], out: Path):
@@ -206,21 +206,21 @@ def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
     assert kwargs == {"check": True}
 
 
-def test_render_run_glues_run_meta_into_render_all(tmp_path: Path, monkeypatch):
-    from giant.analysis import run as run_mod
+def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkeypatch):
+    from giant.analysis import condor as condor_mod
 
     run_dir = tmp_path / "run"
     (run_dir / "reduced").mkdir(parents=True)
 
     merge_calls = []
-    monkeypatch.setattr(run_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
-    meta = run_mod.RunMeta(
+    monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
+    meta = condor_mod.RunMeta(
         rollouts=[{"name": "rollout", "path": "rollout.parquet", "plot_meta": {"checkpoint": "ckpt/best.pt"}}],
         reference="reference.parquet",
         run_dir=str(run_dir),
         title="my-run",
     )
-    monkeypatch.setattr(run_mod.RunMeta, "load", classmethod(lambda cls, p: meta))
+    monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))
 
     Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "series": {"rollout": [1]}}).save(
         run_dir / "reduced" / "s.json"
@@ -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(
-        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
-    )
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_valid.shape == (B, K_MAX)
 
@@ -1261,9 +1261,7 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow():
     n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
     assert n_sec_pred.shape == (B,)
 
-    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
-        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
-    )
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
     assert sec_cont.shape == (B, K_MAX, 4)
     assert sec_valid.shape == (B, K_MAX)
 

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

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

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.plots import MetricsTable, derive_metrics_dir, render_metrics  # noqa: E402
+from giant.training import plots as plots_mod
+from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics
 
 # --- fixtures ----------------------------------------------------------
 

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