docs: record analysis_341dfb14 baseline rollout benchmark results

Replaces the extrapolated pre-v0.3 weak-spot claims in baseline.toml's
header (which had the wrong sign on step-count error) with measured
numbers from the first full rollout validation of this exact config,
and adds a matching Roadmap entry in CLAUDE.md.

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.8"
+current_version = "0.3.15"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"
@@ -8,7 +8,7 @@ regex = false
 allow_dirty = false
 commit = true
 tag = false
-message = "chore: bump version {current_version} -> {new_version} [skip ci]"
+message = "chore: bump version {current_version} -> {new_version}"
 pre_commit_hooks = ["uv lock", "git add uv.lock"]
 
 [[tool.bumpversion.files]]

.gitea/workflows/ci.yml

@@ -2,10 +2,9 @@ name: CI
 
 "on":
   push:
-    branches: ["**"]
+    branches: ["master"]
     tags: ["**"]
-  pull_request:
-    branches: [master]
+  pull_request: {}
 
 env:
   UV_CACHE_DIR: /uv-cache
@@ -156,7 +155,7 @@ jobs:
             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 [skip ci]"
+              git commit -m "chore: update changelog for $TAG"
             else
               git restore --staged CHANGELOG.md
             fi
@@ -193,7 +192,7 @@ jobs:
             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} [skip ci]"
+            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME}"
             git push origin HEAD:master
             git push origin ":refs/tags/${GITHUB_REF_NAME}"
             git tag -f "${GITHUB_REF_NAME}" HEAD
@@ -201,3 +200,33 @@ jobs:
           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]
+    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:
+      # Check out by tag name (not the triggering SHA) since sync-version-on-tag
+      # may have force-moved the tag to a version-corrected commit.
+      - 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"
+      - run: uv build
+      # CI_TOKEN needs write:package scope (in addition to write:repository,
+      # used elsewhere) for this upload to authenticate.
+      - run: |
+          uv publish \
+            --publish-url "https://git.larsbogner.de/api/packages/lars/pypi" \
+            --username gitea-actions \
+            --password "${{ secrets.CI_TOKEN }}"

CHANGELOG.md

@@ -1,5 +1,56 @@
 # Changelog
 
+## [0.3.15] - 2026-08-28
+
+### Changed
+
+- Perf: defer heavy imports in giant/dwarf CLIs until commands run
+
+## [0.3.14] - 2026-08-28
+
+### Changed
+
+- Ci: give automated commits visible checks, scope CI triggers, publish releases
+
+- Ci: fix pull_request trigger not registering
+
+## [0.3.13] - 2026-08-28
+
+### Added
+
+- Add inference-time model_config overrides with a sampling-key allowlist [gitea #87](https://git.larsbogner.de/lars/giant/issues/87)
+
+## [0.3.12] - 2026-08-28
+
+### Added
+
+- Add sampled n_sec under n_sec.mode = 'head' [gitea #86](https://git.larsbogner.de/lars/giant/issues/86)
+
+## [0.3.11] - 2026-08-26
+
+### Changed
+
+- Feat(analysis): per-step secondary multiplicity plots
+
+## [0.3.10] - 2026-08-26
+
+### Changed
+
+- Backfill CHANGELOG.md for v0.2.0-v0.3.2
+
+- Docs: bring README and CLAUDE.md in line with v0.3.9
+
+## [0.3.9] - 2026-08-24
+
+### Added
+
+- Add multi-rollout support to giant analyze [gitea #77](https://git.larsbogner.de/lars/giant/issues/77)
+
+
+### Changed
+
+- Escape LaTeX-special characters in plot titles/xlabels [gitea #81](https://git.larsbogner.de/lars/giant/issues/81)
+
 ## [0.3.8] - 2026-08-24
 
 ### Added
@@ -54,4 +105,479 @@
 
 - Document CI_TOKEN's write:repository scope requirement [gitea #50](https://git.larsbogner.de/lars/giant/issues/50)
 
-# Changelog
+## [0.3.2] - 2026-08-17
+
+### Added
+
+- Add configs/baseline.toml as the kept reference model
+
+
+### Fixed
+
+- Clamp analysis histogram bins before the i32 cast, not after [gitea #61](https://git.larsbogner.de/lars/giant/issues/61)
+
+- Clip raw predicted log_mass in decode_secondaries [gitea #54](https://git.larsbogner.de/lars/giant/issues/54)
+
+
+### Changed
+
+- Let dwarf warm-cache take --config so it can't under-warm a config's cache keys [gitea #59](https://git.larsbogner.de/lars/giant/issues/59)
+
+- Implement n_sec.mode = "stop_token" for the AR secondary decoder [gitea #40](https://git.larsbogner.de/lars/giant/issues/40)
+
+- Bump patch version to 0.3.2
+
+## [0.3.1] - 2026-08-14
+
+### Added
+
+- Add an Objective registry for the flow/ddpm/wgan generator choice [gitea #32](https://git.larsbogner.de/lars/giant/issues/32)
+
+
+### Changed
+
+- Make trunk architecture selectable via a registry [gitea #33](https://git.larsbogner.de/lars/giant/issues/33)
+
+- Make ResBlock's conditioning-injection mechanism selectable [gitea #34](https://git.larsbogner.de/lars/giant/issues/34)
+
+- Make HistoryEncoder a pluggable registry, like Router/Objective [gitea #35](https://git.larsbogner.de/lars/giant/issues/35)
+
+- Deduplicate n_sec_head/type_head MLPs into build_mlp_head [gitea #36](https://git.larsbogner.de/lars/giant/issues/36)
+
+- Give the cond_cat/cond_cont column layout one owner [gitea #37](https://git.larsbogner.de/lars/giant/issues/37)
+
+- Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base [gitea #39](https://git.larsbogner.de/lars/giant/issues/39)
+
+- Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts [gitea #38](https://git.larsbogner.de/lars/giant/issues/38)
+
+- Bump patch version to 0.3.1
+
+## [0.3.0] - 2026-08-13
+
+### Added
+
+- Add v0.3.0 design doc: Stage-2 autoregressive redesign
+
+- Add pytest-cov to dev deps and run coverage in CI
+
+- Add coverage for router-center seeding, geometry batch reader, material topN cache, and setup-cache corruption paths
+
+- Add render.py coverage: figure params, router diagnostics plots, gallery/condor glue
+
+- Add unknown-key validation to config.toml merge (issues.md Issue 2)
+
+- Add consumed-keys audit test (issues.md Issue 5)
+
+
+### Fixed
+
+- Fix test_render_all_run_gallery_invokes_subprocess clobbering LaTeX's own subprocess.run
+
+
+### Removed
+
+- Remove issues.md
+
+
+### Changed
+
+- Refine v0.3.0 design: defaults, deferred scope, open questions
+
+- Document the differentiability position and its validation obligation
+
+- V0.3.0 step 1: new nested config schema, v0.2 migration shim
+
+- V0.3.0 step 2: network.py refactor to composable stage models
+
+- V0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite
+
+- V0.3.0 step 4: type map + particle_type.target = "onehot"/"embedding"
+
+- V0.3.0 step 5: Stage2Autoregressive (history=markov) + §11.4 grad instrumentation
+
+- V0.3.0 step 6: sample.py/rollout.py AR generation + class->PDG decode
+
+- V0.3.0 step 7: AttentionHistory (KV-cached) + scheduled/never teacher forcing
+
+- V0.3.0 post-implementation audit: resolve all 9 tracked discrepancies
+
+- Refactor train.py into giant/training/ around a metrics collector
+
+- Silence the fork-safety warning from num_workers>0 pipeline tests
+
+- Deduplicate giant/training/trainers.py shared per-stage logic
+
+- Rewrite README for v0.3.0 architecture, quick start, and data columns
+
+- Bump version to 0.3.0
+
+- Delete docs/v0.3.0-design.md and strip all references to it
+
+- Apply ruff format
+
+- Downgrade coverage-report upload to actions/upload-artifact@v3
+
+- Bump ruff line-length to 120 and reformat
+
+- Make config dataclasses the single source of truth for DEFAULT_CONFIG
+
+- Extract giant train/new-run's CLI override mapping into a table-driven function (issues.md Issues 3 & 4)
+
+- Mark issues.md Issues 3 & 4 as fixed
+
+- Extract predict/rollout's duplicated inference bootstrap into giant.checkpoint_io (issues.md Issue 5)
+
+- Mark issues.md Issue 5 as fixed
+
+- Unify the two v0.2->v0.3 migration surfaces (issues.md Issue 6)
+
+- Type the data/model/training batch contracts with NamedTuples (issues.md Issue 7)
+
+- Split giant/model/network.py into giant/model/ (issues.md Issue 8)
+
+- Move scripts/ to giant/tools/ (issues.md Issue 9)
+
+- Reject stage2_model.stage1_context = 'sampled' as unimplemented (issues.md Issue 1)
+
+- Honour wgan.critic_hidden_dim/critic_n_res_blocks in build_critics [gitea #28](https://git.larsbogner.de/lars/giant/issues/28)
+
+- Validate stage2_model.autoregressive.order in validate_config [gitea #30](https://git.larsbogner.de/lars/giant/issues/30)
+
+- Decouple secondary-species vocabulary from conditioning.particle.emb_dim [gitea #29](https://git.larsbogner.de/lars/giant/issues/29)
+
+- Skip router auxiliary loss compute when their lambda is 0 [gitea #31](https://git.larsbogner.de/lars/giant/issues/31)
+
+## [0.2.0] - 2026-08-04
+
+### Added
+
+- Add CLAUDE.md with architecture overview and dev commands
+
+- Add streaming data pipeline and giant CLI entry point
+
+- Add giant predict command
+
+- Add ROOT-to-parquet conversion script with convert dependency group
+
+- Add post_pos as a model target via travel_dir decomposition
+
+- Add --coord local mode to predict for raw-space prediction debugging
+
+- Add KL divergence to marginal validation and hook it into the training loop
+
+- Add graceful shutdown on SIGINT/SIGTERM
+
+- Add configurable dropout to ResBlocks
+
+- Add giant.analysis module for notebook-based model quality diagnostics
+
+- Add lazy polars I/O and duplicate KL/constraint checks for giant.analysis
+
+- Add ruff and ty as dev dependencies, fix lint/type findings
+
+- Add linear warmup before cosine LR decay
+
+- Add --batch-size auto to estimate batch size from free GPU memory
+
+- Add hyperparameter scan
+
+- Add --batch-size auto to predict, matching train
+
+- Add tqdm progress bar to predict
+
+- Add KL bar plots and sample_frac to load_predicted_local; ignore root parquet scratch files
+
+- Add event-level shower observables to giant.analysis
+
+- Add total length traveled per event to event observables
+
+- Add pdg energy/length contribution pie plots
+
+- Add export script for Tier 4 event-level/pdg-share plots
+
+- Add mean/median deposited energy and step length plots per event
+
+- Add export script for ETP group-update presentation plots
+
+- Add photon edep export scripts and per-step presentation plots
+
+- Add tooling for a versioned geant_steps dataset layout
+
+- Add --copy mode to migrate_geant_steps.py
+
+- Add update-manifest and create-manifest subcommands to bump_dataset_version
+
+- Add --to flag for bump-gen/bump-schema and --gen flag for update-manifest
+
+- Add disk usage summary to dwarf status
+
+- Add file counts and reference tracking to dwarf status
+
+- Add --comment option to predict, recorded in YAML sidecar
+
+- Add energy-conservation PoC ODE-step comparison scripts
+
+- Add autoregressive shower rollout driver
+
+- Add fast slab lookup for the GeometryOracle, replacing knn as the default
+
+- Add load_rollout_vs_truth to compare rollouts against held-out truth data
+
+- Add mixture-of-experts routing prototype for Stage 1 and Stage 2
+
+- Add ProcessRouter for physics-process-based expert gating
+
+- Add PdgRouter for particle-type-based expert gating
+
+- Add ComposedRouter for multi-axis MoE gating
+
+- Add EMA weights, weight decay, step-based LR schedule, and grad-norm logging to training
+
+- Add WGAN-GP mode as a throwaway fast-eval experiment
+
+- Add router gating diagnostic for MoE checkpoints
+
+- Add Gitea Actions CI pipeline
+
+- Add configs for router energy (embedding/physical) and WGAN baseline runs
+
+- Add opt-in Weights & Biases logging for the training loop
+
+- Add test coverage for resolve_expert_dims
+
+- Add regression coverage for vocab/process index-map builders
+
+- Add dwarf warm-cache to precompute the setup-stage sidecar
+
+- Add giant new-run to scaffold a config.toml + run dir ahead of training
+
+- Add learnable per-expert width and shared temperature to EnergyRouter
+
+- Add opt-in straight-through Gumbel-softmax combine weights to MoE router
+
+- Add gumbel router configs sweeping learnable-knob combinations
+
+- Add gumbel/learn_centers/learn_width/learn_temperature to out-dir naming
+
+- Add bigger WGAN config (hidden_dim=512, n_blocks=6)
+
+- Add data-integrity guards against silent NaN/Inf propagation and races
+
+
+### Fixed
+
+- Fix column names to match actual parquet schema
+
+- Fix installed torch version to be compatible with cuda drivers
+
+- Fix miniCaloSim link in README
+
+- Fix giant.analysis import after Phase 2 dataset API changes
+
+- Fix silent failure modes surfaced by extensive code review
+
+- Fix ruff, ty, and pytest failures; apply ruff format
+
+- Clamp n_sec classification label to K_MAX
+
+- Fix rollout edep mismatch and add truth overlay to Tier 4 observables
+
+- Fix crashes in physical-property conditioning edge cases
+
+- Fix router experts silently ignoring --hidden-dim/--n-blocks
+
+- Fix conditioning="physical" so it can actually generalize past training vocab
+
+- Fix training-loop checkpoint/resume and WGAN bugs
+
+- Fix stale-partial reuse and n_chunks mismatch in analysis condor pipeline
+
+- Fix CLI/tooling robustness gaps and dedupe the Conditioning enum
+
+- Fix test_write_submit_requires_synced_venv for active-venv resolution
+
+
+### Removed
+
+- Remove scripts/train.py in favor of the giant train CLI
+
+- Drop orphaned child tracks instead of nulling secondary targets
+
+
+### Changed
+
+- Initial commit: giant surrogate model with two-phase roadmap in README
+
+- Implement Phase 1: full data pipeline, model, training, and config support
+
+- Handle material column as string type
+
+- Rename pre_energy/post_energy columns to pre_E/post_E
+
+- Rename direction columns from pre_dir_x/y/z to pre_dx/dy/dz
+
+- Batch StreamingStepsDataset internally instead of per-row collate
+
+- Dedup training pipeline, add seeding/resume and per-epoch metrics logging
+
+- Split torch into cpu/cuda extras and pin dependency version bounds
+
+- Apply ruff format and document lint/type tooling in CLAUDE.md
+
+- Update README to match current architecture and tooling
+
+- Make sampler step count configurable for validation
+
+- Calibrate auto batch size separately for inference vs training
+
+- Skip rows with unknown PDG codes during predict
+
+- Buffer predict rows across row-group boundaries before inference
+
+- Export plots for knowledge base
+
+- Rework validation notebook with markdown sections and Tier 4 plots
+
+- Allow steps_to_parquet.py to accept multiple ROOT input files
+
+- Encode edep/secondary/post energy as a conservation-constrained simplex
+
+- Expose dataset/conversion scripts as uv entry points
+
+- Restrict holdout overlap check to holdout vs dev/full only
+
+- Route predict output to UUID-named parquet with YAML reference sidecar
+
+- Implement Phase 2: secondary particle prediction
+
+- Unify dataset/tooling scripts into a single `dwarf` Typer CLI
+
+- Fold --to/--gen dataset-versioning flags into the dwarf CLI
+
+- Prefix default train output dir with current date
+
+- Color-code dwarf status output by tree level
+
+- Show VERSIONS.md reason extracts in dwarf status
+
+- Wire up predict CLI to load and run the Stage-2 sec_decoder
+
+- Wire up n_sec/species/energy-fraction validation for Stage 2
+
+- Detach Stage-2 type-embedding target to stop self-referential collapse
+
+- Weight Stage-2 secondary loss equally between direction and type-embedding dims
+
+- Recalibrate batch-size estimate for the post-Phase-2 model size
+
+- Update CLAUDE.md and README for the implemented Phase 2 model
+
+- Error on missing secondary lists instead of silently zeroing Stage-2 targets
+
+- Derive a unique per-job seed for minicalosim shard generation
+
+- Rescale secondary energies to exactly consume the e_sec budget
+
+- Support --energy-gev in dwarf make-root for the new minicalosim energy arg
+
+- Stream giant rollout output instead of buffering the whole run
+
+- Scale auto batch-size estimate by MoE expert count during training
+
+- Rewrite analysis module as a lean, fully-streaming pipeline
+
+- Reimplement rollout-vs-truth comparison on the streaming analysis module
+
+- Condition on material/particle physical properties instead of learned embeddings
+
+- Ignore the scratchpad working directory
+
+- Quote the on: key in the CI workflow
+
+- Split CI lint stage into parallel jobs
+
+- Rewrite analysis as streaming rollout-vs-reference plotting pipeline
+
+- Analyze: drive prep/submit from the rollout YAML sidecar
+
+- Analyze: show model/training params on rendered figures
+
+- Deps: install plotstyle from git.larsbogner.de package index
+
+- Analyze: drop stale ty:ignore on plotstyle import
+
+- Test: replace prep(**_CTX) splat with a typed _prep helper
+
+- Analyze: add MoE router gating/share diagnostic plots
+
+- Chore: remove stray CUDA sanity script and stale Phase 2 planning doc
+
+- Docs: document compute environment, WGAN/MoE status, and condor-gpu-train-rollout
+
+- Analyze: normalize pdg dtype in open_side to fix rollout/reference concat
+
+- Analyze: chunk per-plot aggregation across HTCondor jobs
+
+- Analyze: expose bin/pdg options on `analyze submit`
+
+- Analyze: estimate per-job HTCondor walltime from chunk row count
+
+- Analyze: run condor compute jobs via .venv/bin/giant, not uv run
+
+- Analyze: default condor docker image to alma9-gridjob
+
+- Analyze: raise default condor job memory request to 8192 MB
+
+- Analyze: recalibrate condor walltime model from real cluster timings
+
+- Transforms: pad legacy cond normalizers for pre-physical-conditioning checkpoints
+
+- Analyze: default run directory to <repo>/analysis_runs, gitignored
+
+- Docs: record first MoE router rollout benchmark result in the roadmap
+
+- Router: seed EnergyRouter centers from data quantiles instead of a fixed linspace
+
+- Docs: note the EnergyRouter centers_init fix in the roadmap
+
+- Analyze: thread full model/training/rollout/dataset params to plots
+
+- Ci: share one uv sync across jobs, gate tests on lint+type-check, sync tag/version on release tags
+
+- Ci: replace unsupported artifact sharing with a bind-mounted uv cache
+
+- Ci: stop setup-uv from overriding UV_CACHE_DIR
+
+- Ci: re-pin UV_CACHE_DIR after setup-uv, which exports its own value regardless of enable-cache
+
+- Ci: set UV_LINK_MODE=copy to silence the cross-filesystem hardlink warning
+
+- Log batch-level metrics to W&B, not just per-epoch summaries
+
+- Log router health, WGAN grad-norm split, n_sec accuracy, GPU/throughput to W&B
+
+- Persist global_step across --resume so W&B step stays monotonic
+
+- Timestamp default checkpoint dir to avoid W&B run-id collisions
+
+- Skip empty-slice mean/std in sec phys validation print
+
+- Speed up giant train's setup stage
+
+- Speed up _WelfordAccumulator's per-chunk update
+
+- Make default checkpoint out_dir name reflect only non-default hyperparams
+
+- Cache giant train's setup stage in a sidecar file
+
+- Pass --seed through to the train/val event split
+
+- Offset event_id per file to avoid cross-file collisions
+
+- Store a quantile grid instead of a raw reservoir sample in the setup cache
+
+- Scope wandb run config to only-active hyperparameters
+
+- Resolve giant condor wrapper from the active venv, not a hardcoded path
+
+- Bump version to 0.2.0

CLAUDE.md

@@ -7,18 +7,20 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
 ```bash
 uv sync --extra cpu                # install dependencies with CPU-only torch (standard/default)
 uv sync --extra cuda               # install dependencies with CUDA 11.8 torch
-uv sync --extra cpu --extra dev    # add dev extras (pytest, etc.)
+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)
 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 --mode flow   # train (flow matching)
-giant train path/to/steps.parquet --mode ddpm   # train (DDPM baseline)
-giant train path/to/steps.parquet --mode wgan   # train (WGAN-GP, single-pass eval; implemented, not yet tested)
-giant train path/to/steps.parquet --router --router-type energy  # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap)
+giant train path/to/steps.parquet               # train (defaults: stage 1 flow, stage 2 wgan + autoregressive)
+giant train path/to/steps.parquet --mode flow   # set both stages' generative objective at once
+giant train path/to/steps.parquet --stage1-generator flow --stage2-generator wgan   # per-stage override
+giant train path/to/steps.parquet --router --router-type energy  # MoE routing trunk (see Roadmap for status)
+giant model summary --config config.toml         # build-only: parameter counts + which config keys actually bite
 giant predict path/to/steps.parquet --checkpoint ckpt/best.pt   # per-step predictions
 giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl  # full showers
 giant analyze submit rollout.yaml --accounting-group cms   # parallel rollout-vs-reference analysis on HTCondor
 giant analyze 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,
                                                   # make-root, build-geometry-oracle, warm-cache, hparam-scan
@@ -27,6 +29,8 @@ dwarf --help                                     # dataset/tooling CLI: convert,
 
 `cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.
 
+`configs/` holds kept reference configs (`baseline.toml`, `default.toml`, the router/WGAN scan configs) — pass them with `--config`.
+
 ### Lint and type checking
 
 ```bash
@@ -37,6 +41,10 @@ uv run ty check .       # type check
 
 Part of the `dev` extra. Run these periodically (not just at commit time) to catch drift early.
 
+### Release tooling
+
+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.
+
 ## Compute environment
 
 Work on this repo happens across three kinds of machine:
@@ -47,50 +55,72 @@ Work on this repo happens across three kinds of machine:
 
 ## Architecture
 
-GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — now including the variable-length list of secondary particles the step produces (Phase 2, see Roadmap).
+GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — including the variable-length list of secondary particles the step produces.
 
-**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`). Train/val split is by `event_id` to avoid leaking correlated steps from the same shower.
+**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`, streaming variant included). Train/val split is by `event_id` (`--seed`-controlled) to avoid leaking correlated steps from the same shower. 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. `setup_cache.py` persists the pre-epoch setup scan (vocab maps, event split, process maps, normalizer stats) as a sidecar so repeated runs over the same `data` path don't rescan (`--cache-setup`/`--rebuild-setup-cache`, precomputable with `dwarf warm-cache --config ...`).
 
 **Stage-1 output space (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):** `log_step_length`, two additive-log-ratio (ALR) coordinates `edep_logit`/`sec_logit` of a **deposit / secondary / post-energy simplex**, `post_dir` (post-scattering momentum direction, unit vector in the local frame), and `travel_dir` (direction of `post_pos - pre_pos`, unit vector in the local frame). The energy simplex decodes via softmax over `[edep_logit, sec_logit, 0]` × `pre_E` so `edep + e_sec + post_E == pre_E` holds by construction — energy conservation is architectural, not learned (see `energy_simplex_decode`). `post_pos` is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, since `step_length` already encodes that displacement's magnitude and duplicating it would let the two become inconsistent.
 
-**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus, since particle/material physical-property conditioning (`model.conditioning`, see below), 7 more columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model predicts them.
+**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus 7 physical-property columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** — the model predicts them. `giant/cond_layout.py` is the single source of truth for the `cond_cont`/`cond_cat` column layout shared by `giant.data.transforms`, `giant.model.encoders`, and `giant.model.routers`.
 
-`ConditionEncoder`/`SecondaryConditionEncoder` (`giant/model/network.py`) support two mutually exclusive `conditioning` modes, selected per-checkpoint (`model_config["conditioning"]`, defaulting to `"embedding"` for old checkpoints without the key, `"physical"` for new `giant train` runs — see `--conditioning`):
-- **`"embedding"`** (original Phase 2 design): a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab (`pdg_map`/`mat_map`). Memorizes the training menu.
-- **`"physical"`** (default): the 7 physical-property columns above are each routed through a small MLP (`particle_mlp`/`material_mlp`) to the same `emb_dim` width the embedding tables would have produced — a drop-in replacement computable for any PDG code / material name, not just ones seen in training, which is what lets the surrogate generalize to a held-out material or species. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships real Geant4-11.4.1-derived `z_eff`/`a_eff`/`density`/`x0`/`lambda_int` values for every material the detector geometry actually produces; the sole exception is `G4_LYSO` (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays `MaterialProperties(None, ...)` and raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting if it's ever requested.
+`ConditionEncoder` (`giant/model/encoders.py`) configures the particle and material identity axes **independently** (`conditioning.particle` / `conditioning.material`, each a `ConditioningAxisConfig` with `type`/`emb_dim`/`n_layers`), so they may mix freely. Three per-axis modes:
+- **`"physical"`** (default): the axis's raw physical properties routed through a small MLP — computable for any PDG code / material name, which is what lets the surrogate generalize beyond the training menu. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships real Geant4-11.4.1-derived values for every material the detector geometry actually produces; the sole exception is `G4_LYSO` (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays `MaterialProperties(None, ...)` and raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting.
+- **`"embedding"`**: a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab. Memorizes the training menu; the generalization-comparison baseline, and the only mode compatible with `stage2_model.particle_type.target = "embedding"`.
+- **`"onehot"`**: a fixed, unlearned vector over the top `emb_dim - 1` codes by training-set count plus one "other" bin. Not a reparameterization of `"embedding"` — the vocabulary cap is the real difference.
 
-**Model** (`giant/model/network.py`): a two-stage model, both checkpointed together.
-- **Stage 1 — `DenoisingMLP`:** `ResBlock` stack with a `SinusoidalEmbedding` for the flow/diffusion time variable and a `ConditionEncoder` fusing the conditioning. Predicts the 9D primary vector field, plus an `n_sec_head` classifier over `{0..K_MAX}` (`K_MAX=15`) that runs on the condition encoding alone (no diffusion noise), callable via `predict_n_sec`.
-- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, log-mass, charge)` = `SEC_SLOT_DIM=6`, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. A secondary's mass/charge are regressed directly against a fixed physics-derived target (its ground-truth PDG code's `giant.particles.particle_mass_charge`) — not a learned/moving embedding target, so nothing needs detaching. **No snapping at inference**: the predicted (mass, charge) are used as-is as the secondary's physical identity, including for its own future conditioning if it goes on to take further steps in a rollout. A separate, reporting-only nearest-known-PDG lookup (`giant.particles.nearest_known_pdg`) is used purely to populate a nominal `pdg` label for output rows / `"embedding"`-mode fallback conditioning — it never feeds back into the model.
+`conditioning.share_stages` decides whether the two stages get one shared encoder instance or two identically-configured independent ones.
 
-`schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
+**Model** (`giant/model/`, both stages checkpointed together). `network.py` is only a re-export shim now; the real code is split by concern:
+- `layers.py` — `ResBlock`/`AdaLNResBlock` + `BLOCK_REGISTRY` (conditioning-injection mechanism is selectable), `SinusoidalEmbedding`, `ContextAdapter`, `build_mlp_head`.
+- `encoders.py` — `ConditionEncoder` (above).
+- `trunks.py` — `TRUNK_REGISTRY`/`build_trunk`: everything downstream of the fused conditioning vector, as a registrable expert *body* (`resmlp` default, plus a `none` variant). `RoutedTrunk` builds `router.n_experts` instances of whichever body is named, so mixing is orthogonal to which body is mixed.
+- `routers.py` — `Router` base + `ROUTER_REGISTRY`: `energy`/`pdg`/`process`/`composed`/`none`. Soft-mixed at train time, **top-1 dispatched at eval time** (each row runs exactly one small expert), which is the actual inference-speed win. `EnergyRouter`/`PdgRouter` gate on a quantity known at inference; `ProcessRouter` runs its own small classifier (process isn't known upfront); `ComposedRouter` gates jointly over outer-product expert cells via repeated `--router-axis "type:key=val,..."`. The `--router*`/`--n-experts` CLI flags target `stage1_model.router` only; stage 2's router is config-file-only (`stage2_model.router`). `EnergyRouter` accepts `centers_init`, which `giant/pipeline.py` auto-populates from real data quantiles via a reservoir sample collected during the normalizer-fitting pass.
+- `history.py` — `HISTORY_REGISTRY`/`build_history`: `markov` (previous token only), `attention` (causal self-attention, KV-cached at inference via `init_cache`/`step`), `none`. Stage-2 autoregressive only.
+- `objectives.py` — `Objective` base + registry for `flow`/`ddpm`/`wgan`: answers in one place whether a stage needs a time embedding, is adversarial, folds the secondary type slice into its trunk output, what its trunk input is, and which loss it trains against.
+- `models.py` — the composed stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel`, all on a shared `StageModel` base.
+- `builders.py` — `build_models`/`build_critics`, assembling the above from a config dict.
+- `schedule.py` (`CosineSchedule` for DDPM + conditional-flow-matching losses), `wgan.py` (gradient penalty / critic / generator losses, Gulrajani et al. 2017), `summary.py` (`giant model summary`), `_legacy.py` (v0.2 checkpoint migration).
 
-**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
+**Stage 1 — primary step.** Trunk (routed or not) over the fused conditioning, plus a `SinusoidalEmbedding` of the flow/diffusion time for non-adversarial objectives, predicting the 9D vector field. An `n_sec` classifier head over `{0..k_max}` runs on the condition encoding alone; `stage2_model.n_sec.owner` decides whether it lives on stage 1 (v0.2 checkpoints) or stage 2 (default).
 
-**WGAN-GP mode (`--mode wgan`, implemented, not yet tested):** a throwaway fast-eval alternative to the flow/DDPM samplers above — single forward pass instead of ~10 ODE steps. Dedicated noise-conditioned generators (`WGANGenerator`/`WGANSecondaryGenerator`, `giant/model/network.py`) stand in for `DenoisingMLP`/`SecondaryDecoder`, trained against `Critic`/`SecondaryCritic` discriminators with the gradient-penalty loss in `giant/model/wgan.py` (Gulrajani et al. 2017); `sample_wgan` (`giant/sample.py`) does the single-pass draw at inference. Not yet validated against the flow-matching baseline.
+**Stage 2 — secondaries.** Conditioned on the pre-step state plus a projected stage-1 outcome (`stage2_model.context_dim`; `stage1_context` selects ground-truth vs sampled context, annealable via `ctx_p_start`/`ctx_p_end`). Two decoders (`stage2_model.decoder`):
+- **`autoregressive`** (default): one secondary at a time in descending-energy order, each token conditioned on a `HistoryEncoder` summary of prior tokens, with teacher forcing (`always`/`scheduled`/`never`, `tf_p_start`/`tf_p_end`). `n_sec.mode = "stop_token"` lets the length be emitted by the sequence itself instead of the classifier head.
+- **`one_shot`**: all `k_max` slots in one pass, masked past the predicted `n_sec` (the v0.2 behaviour).
 
-**MoE routing trunk (`--router`, implemented; first rollout benchmark shows the experts don't specialize — see Roadmap):** an alternative to `DenoisingMLP`'s monolithic `ResBlock` trunk — a `Router` (`giant/model/network.py`, `ROUTER_REGISTRY`/`build_router`) gates between small per-expert `ResBlock` stacks (`Expert`), soft-mixed over all experts at train time but **top-1 dispatched at eval time** (each row runs exactly one small expert), which is the actual inference-speed win. Router types gate on different conditioning axes: `EnergyRouter`/`PdgRouter` read a quantity already known at inference time, `ProcessRouter` runs its own small classifier over pre-step conditioning (since process isn't known upfront); `ComposedRouter` gates jointly over multiple axes (outer-product expert cells) via repeated `--router-axis "type:key=val,..."` flags. Config lives under `model.router` (`giant/config.py`), deep-merged one level so `router.enabled` alone doesn't drop the rest of the defaults.
+Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. Particle identity is set by `stage2_model.particle_type.target`: `"onehot"` (default — categorical over the top `n_classes - 1` PDG codes by training count plus "other", with configurable `other_policy` and `class_weighting`), `"physical"` (continuous `(log-mass, charge)` regressed against `giant.particles.particle_mass_charge`), or `"embedding"` (nearest-row snap into the conditioning embedding table; requires `conditioning.particle.type = "embedding"`).
 
-**Validation** (`giant/validate.py`): step-level marginal comparisons.
+**Samplers** (`giant/sample.py`): DDPM, DDIM, flow matching (ODE integration, ~10 steps), and single-pass WGAN, plus the stage-2 secondary sampling loop (one-shot and autoregressive).
 
-**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. 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 + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `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`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_<id>/`) holding `shared.json`, `run_meta.json`, `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit rollout.yaml --chunks N` runs `prep` (recording the run's chunk count `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`) and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` (`catalog.py`) splits into a `compute_partial`/`finalize` pair so a plot's chunks can be summed/concatenated back together correctly (`chunkable=False` specs — the router diagnostics, already bounded/subsampled — always run as a single chunk regardless of `N`). 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`), then turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
+**Training** (`giant/training/`): `loop.py` (epoch loop, graceful shutdown, best-checkpoint selection), `trainers.py` (`StageSpec` + per-stage flow/ddpm and WGAN-GP trainers, and the `MetricSpec` declarations that define `metrics.csv`'s columns), `stage2_inputs.py` (ground-truth stage-2 targets + teacher-forcing inputs), `metrics.py` (`MetricsCollector`: `metrics.csv`, W&B logging, progress/summary), `checkpoint.py`, `amp.py` (`train.precision = fp32|bf16` autocast), `plots.py` (`giant analyze metrics`). Per-stage `init_from`/`freeze` lets one stage be retrained against a fixed, known-good other stage while still producing a complete rollout-capable checkpoint.
 
-**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — 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 energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
+**Config** (`giant/config.py`): frozen dataclasses are the single source of truth; `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather than hand-maintained. Blocks: `[conditioning]`, `[stage1_model]`, `[stage2_model]`, `[train]`, `[meta]`. Unknown keys are rejected on merge (with a did-you-mean suggestion), and `tests/test_config_consumed_keys.py` audits that every key is actually read somewhere.
+
+**Validation** (`giant/validate.py`): step-level marginal + KL-divergence comparisons during training (`--validate-every`).
+
+**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** (`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`.
+
+**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower, advancing tracks breadth-first (every sweep steps all active tracks once, in `batch_size` chunks, so many tracks share each forward pass). Each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on one of the `TERM_*` reasons in `constants.py` (energy cutoff, max steps, escape, natural end, unknown pdg, max tracks); energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. `giant/checkpoint_io.py` is the shared checkpoint → ready-to-run-models path used by both `predict` and `rollout`.
 
 ## Roadmap
 
 **Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
 
-**Phase 2 (implemented — baseline):** the two-stage model above jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected). This is the "get a baseline out" track agreed with Jan & Tobias (2026-07-07).
+**Phase 2 (done):** the two-stage model jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected).
 
-**Physical-property conditioning (implemented):** `model.conditioning = "physical" | "embedding"` (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. `"embedding"` stays available as the generalization-comparison baseline. `giant/materials.py`'s table is already filled with real values for every material the geometry produces. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.
+**Physical-property conditioning (implemented, default):** `conditioning.particle.type` / `conditioning.material.type` = `physical | embedding | onehot`. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.
 
-**Faster-eval architectures (implemented, validation in progress):** both tracks below target a ~10× native-Geant4 eval budget and are now wired into `giant train`/`giant/model/network.py`, but neither has a validated result yet — treat both as unproven until the corresponding analysis run says otherwise:
-- **WGAN-GP** (`--mode wgan`, see Architecture above): implemented, **not yet tested** — no rollout-vs-reference analysis run against it yet.
-- **MoE routing trunk** (`--router`, see Architecture above): implemented, **first rollout benchmark done (2026-07-22), result: needs retraining with a different router config, not abandoned.** A 10-expert `EnergyRouter` run (`n_experts=10`, `temperature=0.5`, `learn_centers=true`, **`lambda_balance=0.0`**, only 20 fine-tuning epochs resumed from a non-routed checkpoint) diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. The `router_gating` diagnostic plot points at the likely cause: the ten experts overlap heavily across ~5 decades of pre-step energy instead of partitioning it — even the top-energy expert only reaches ~60–65% gate weight at the highest energies plotted — so eval-time top-1 (Voronoi) dispatch is choosing among near-ties rather than real specialists. Two contributors were identified: the missing load-balancing loss (`lambda_balance=0.0`), and `EnergyRouter`'s center init (`torch.linspace(-2, 2, n_experts)`) assuming a roughly uniform z-normalized energy distribution, which real energy spectra don't match. **Fixed (2026-07-27):** `EnergyRouter` now accepts an optional `centers_init` (backward compatible — omitting it keeps the old linspace), and `giant train` auto-populates it from real data quantiles via a reservoir sample collected during the existing normalizer-fitting pass in `giant/pipeline.py` (no extra file scan), for `--router-type energy` only. The routing *strategy* itself may still be sound, but the specific benchmarked config wasn't. **Next step before further evaluation: retrain with `lambda_balance > 0` and the new quantile-seeded centers (and consider more epochs / a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Full writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
+**v0.3.0 — Stage-2 autoregressive redesign (implemented, released; on `master` since 2026-08-13):** motivated by the 2026-08-03 WGAN rollout benchmark, which failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). Stage 2 became autoregressive in descending-energy order with teacher forcing, and the particle-type representation went back to **categorical** (`particle_type.target = "onehot"`), reversing the 2026-07-17 continuous `(log-mass, charge)` target. The config break (`[conditioning]`/`[stage1_model]`/`[stage2_model]`/`[train]` replacing the flat `train.mode` + `[model]`) makes per-stage generators, stage-2-only training, and one-shot-vs-autoregressive comparison all expressible, and the `network.py` refactor into composable parts (encoder × trunk × objective) also makes routed WGAN work for the first time.
 
-A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
+**Baseline benchmark (done, 2026-08-26):** `configs/baseline.toml`'s first full rollout-vs-Geant4 validation (`analysis_341dfb14`, checkpoint `20260814_1743_s2-flow_h512_s2h512_bs36864_ep50/best.pt`, epoch 50/50). Confirms the v0.3.0 pivot fixed the species collapse — zero photon secondaries / hallucinated `-14` muon antineutrinos are both gone (γ at 95% of truth, no `-14` in the top species) — and rules out `conditioning.*.type = "physical"` as the cause, since this checkpoint pairs it with `flow`/no-router and still doesn't collapse. Bulk shower observables are close to Geant4 (total deposited energy +1.9%, containment depth-90%/95% both 0.986×), but steps/event now *over*-shoots by 1.32× (the opposite sign from every pre-v0.3.0 checkpoint), no hadronic/nuclear secondaries are produced at all, and event-to-event energy variance is ~16× too narrow. Writeup: `/home/lars/knowledge-base/experiments/giant-baseline-flow-ar-rollout-validation.md`.
 
-**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`) — v0.2 checkpoint-loading support has **no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
+v0.2 configs and checkpoints are auto-migrated (`config.migrate_config`, `model._legacy._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`). **v0.2 checkpoint-loading support has no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
 
-**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 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.
+**Faster-eval architectures — both implemented, neither validated.** Target is a ~10× native-Geant4 eval budget; no eval-latency number exists for any configuration yet, so that budget is unverified across the board.
+- **WGAN-GP** (`--stage2-generator wgan`, now the stage-2 default): first rollout benchmark 2026-08-03 failed with secondary-species mode collapse — the failure v0.3.0 was designed to address. **No post-v0.3.0 benchmark has been run.** Writeup: `/home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md`.
+- **MoE routing trunk** (`--router`): first rollout benchmark 2026-07-22 diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. Cause identified as a bad config, not a bad idea: `lambda_balance=0.0` (no load-balancing loss) plus `EnergyRouter`'s `torch.linspace(-2, 2, n_experts)` center init assuming a roughly uniform z-normalized energy distribution — so the ten experts overlapped across ~5 decades of energy instead of partitioning it, and eval-time top-1 dispatch chose among near-ties rather than real specialists. Both prerequisites are fixed in code (quantile-seeded `centers_init` from `pipeline.py`, `lambda_balance` exposed). **Next step: retrain with `lambda_balance > 0` and quantile-seeded centers (consider a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
+
+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 (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

@@ -10,6 +10,7 @@ A conditional generative model that replaces the Geant4 step function: given a p
 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
 
@@ -42,9 +43,9 @@ A **two-stage model**, checkpointed together. Either stage's outcome can be prod
 
 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).
 
-**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 (`--conditioning`) — 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.
+**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.
 
-**MoE routing** (`--router`, either stage): 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.
+**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.
 
 ## Data
 
@@ -64,12 +65,24 @@ 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)
+│   │   ├── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
+│   │   └── setup_cache.py  # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
 │   ├── model/
-│   │   ├── network.py      # ConditionEncoder, Stage1Model, Stage2OneShot/Stage2Autoregressive, Router/MoE, CriticModel
+│   │   ├── 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
+│   │   ├── 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
@@ -79,20 +92,28 @@ giant/
 │   │   ├── 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
-│   │   ├── catalog.py      #   declarative PlotSpec registry
-│   │   ├── condor.py       #   prep / compute-one / submit-description plumbing
+│   │   ├── reduced.py      #   Partial/Reduced — the compact JSON a compute job emits
+│   │   ├── catalog.py      #   declarative PlotSpec registry (`giant analyze list`)
+│   │   ├── router_gating.py / type_embedding_distance.py   # checkpoint-bound diagnostics
+│   │   ├── runtime_estimate.py  #   per-(plot, chunk) walltime estimates for submit
+│   │   ├── condor.py       #   prep / compute-one / merge / submit-description plumbing
 │   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
-│   └── cli.py              # `giant train` / `new-run` / `predict` / `rollout` / `analyze` Typer app
+│   └── cli.py              # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze`
 ├── giant/tools/                        # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
 │   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
 │   │                                    #   update-manifest, create-manifest, make-root,
@@ -116,8 +137,13 @@ giant/
 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`)
 ```
 
+The `dev` extra pulls in `convert`, `analysis`, `geometry` and `wandb` as well.
+
 `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.
 
 ## Training, prediction, rollout
@@ -137,11 +163,13 @@ Useful flags on `giant train`:
 - `--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`. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
+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.
 
@@ -151,11 +179,19 @@ Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_mod
 - `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
 
 ```bash
-giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot (compute only)
-giant analyze render <run_dir> --gallery                   # local: styled PDFs + HTML gallery (needs LaTeX)
+giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot × chunk (compute only)
+giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan  # N rollouts vs one shared reference
+giant analyze render <run_dir> --gallery                   # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
+
+giant analyze list                                         # every catalog plot id
+giant analyze prep rollout.yaml --chunks 8                 # just the run directory, no submission
+giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0   # what a condor job runs
+giant analyze merge-one --id marginal_edep --run-dir <run_dir>               # merge one plot's chunks (debugging)
 ```
 
-`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
+`<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
+
+Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.
 
 ## Development
 

cliff.toml

@@ -37,7 +37,7 @@ commit_preprocessors = [
 protect_breaking_commits = false
 commit_parsers = [
     { message = "^Merge ", skip = true },
-    { message = "\\[skip ci\\]", skip = true },
+    { message = "^chore: (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

@@ -29,11 +29,21 @@
 #     capacity overfitting is not the binding constraint, and every recent
 #     run used 0.0.
 #
-# Known weak spots this baseline is expected to *exhibit* (they are the
-# reason for the comparisons, not a reason to retune this file): every model
-# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
-# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
-# head accuracy sits at 0.863-0.867 regardless of size or objective.
+# Known weak spots, now measured against this exact config rather than
+# extrapolated from the pre-v0.3 field (analysis_341dfb14, best.pt @ epoch
+# 50/50, full writeup: knowledge-base/experiments/
+# giant-baseline-flow-ar-rollout-validation.md). Unlike every pre-v0.3
+# checkpoint (which under-produced steps/event by 1.6-5x), this baseline
+# OVER-produces steps/event by 1.32x (1.86e5 vs Geant4 1.41e5) and
+# under-produces secondaries/event by 0.84x (5.97e4 vs 7.14e4) — the sign on
+# steps flipped with the v0.3 autoregressive pivot, so don't assume it still
+# undershoots. Secondary-species hallucination (zero photons, hallucinated
+# `-14` muon antineutrinos) that broke every prior checkpoint is gone; the
+# remaining species gap is a total absence of hadronic/nuclear secondaries
+# (protons, neutrons, ion recoils), not miscalibration of the ones produced.
+# Total deposited energy/event is +1.9% high but its event-to-event spread is
+# ~16x too narrow (31 MeV vs Geant4's 491 MeV). Per-step deposited energy is
+# the worst per-step marginal (KS 0.179 vs 0.004-0.071 for the others).
 
 [meta]
 # REQUIRED. Without it config.migrate_config reads this file as v0.2 and

giant/analysis/__init__.py

@@ -1,9 +1,10 @@
 """Rollout-vs-reference analysis: streaming compute + plotstyle rendering.
 
-Compares one autoregressive ``giant rollout`` against a held-out miniCaloSim
-reference file, producing publication-styled comparison plots generated in
-parallel on HTCondor (one job per plot x data chunk, compute/merge/render
-split).
+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).
 
 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,12 +13,14 @@ re-exported here is plotstyle-free so it runs on a compute worker. Import
 
 from giant.analysis.catalog import build_catalog, catalog_ids, get_spec
 from giant.analysis.condor import (
+    LoadedRollout,
     RunMeta,
     SubmitConfig,
     compute_one,
     compute_reduced,
     derive_run_dir,
     load_rollout_yaml,
+    load_rollout_yamls,
     merge_all,
     merge_one,
     prep,
@@ -26,18 +29,20 @@ from giant.analysis.condor import (
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial, Reduced
 from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
-from giant.analysis.sources import Side
+from giant.analysis.sources import RolloutSpec, Side
 
 __all__ = [
     "build_catalog",
     "catalog_ids",
     "get_spec",
+    "LoadedRollout",
     "RunMeta",
     "SubmitConfig",
     "compute_one",
     "compute_reduced",
     "derive_run_dir",
     "load_rollout_yaml",
+    "load_rollout_yamls",
     "merge_all",
     "merge_one",
     "prep",
@@ -46,6 +51,7 @@ __all__ = [
     "build_context",
     "Partial",
     "Reduced",
+    "RolloutSpec",
     "Side",
     "RUNTIME_SAFETY_MARGIN",
     "estimate_runtime_s",

giant/analysis/catalog.py

@@ -4,7 +4,7 @@ Each spec knows its stable ``id`` (used for the reduced-data filename, the PDF
 stem and the condor queue item), its gallery ``family`` (subdirectory), and a
 ``compute_partial(bundle) -> dict`` / ``finalize(parts, ctx) -> Reduced`` pair
 that together run the streaming reduction. ``compute_partial`` runs once per
-``(plot, chunk)`` condor job against a ``Bundle`` whose four LazyFrames are
+``(plot, chunk)`` condor job against a ``Bundle`` whose LazyFrames are
 already filtered to that chunk (see ``Bundle.open``'s ``chunk`` argument); it
 returns a small JSON-safe partial artifact — either a raw sum-mergeable count
 dict (histograms/species sums against fixed edges) or a raw per-event/
@@ -16,6 +16,19 @@ exactly what a single unchunked pass would produce. Specs marked
 ``chunkable=False`` (the router ones) always run as a single chunk regardless
 of the configured chunk count.
 
+Every ``compute_partial`` here returns ``{"r": {rollout_name: <shape>}, "t":
+<shape>}`` — one entry per rollout in ``Bundle.rollouts`` (insertion order,
+which is the order rollouts were given on the CLI) plus the single reference.
+``finalize`` merges each rollout's chunks independently and assembles a
+``Reduced.payload`` keyed the same way: ``"series": {name: ...}`` for the
+rollouts, ``"reference": ...`` as one distinguished entry (omitted on
+rollout-only plots like ``leakage_fraction``). The heatmap-shaped specs
+(``marginal_distance_summary``, ``sec_count_per_step_by_species``) and the
+router diagnostics are inherently one-matrix/one-checkpoint per rollout, so their
+``"series"`` entries are whole per-rollout artifacts (a matrix, a gating
+dict) rather than a single number/array — ``render.py`` draws those as one
+panel per rollout instead of one line/bar per rollout.
+
 Rendering lives in ``render.py`` and dispatches on ``Reduced.kind`` — the
 catalog itself never imports plotstyle, so ``compute-one`` jobs stay LaTeX-free.
 
@@ -47,7 +60,6 @@ from giant.analysis.reduce import (
     leakage_fraction,
     profile_finalize,
     profile_partial,
-    sec_count_by_event,
     species_share,
     sum_merge,
     transverse_expr,
@@ -59,7 +71,15 @@ from giant.analysis.router_gating import (
     compute_router_share_by_process,
     compute_router_specialization,
 )
-from giant.analysis.sources import Side, open_side, physical_steps, secondaries
+from giant.analysis.sources import (
+    RolloutSide,
+    RolloutSpec,
+    Side,
+    open_side,
+    physical_steps,
+    secondaries,
+    secondaries_by_step,
+)
 from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
 from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
 
@@ -69,51 +89,44 @@ class Bundle:
     """Everything a compute runs against — built once per ``compute-one`` job."""
 
     ctx: Context
-    r_all: pl.LazyFrame  # rollout, all rows (incl. synthetic termination rows)
+    rollouts: dict[str, RolloutSide]  # name -> frames, insertion order = CLI order
     t_all: pl.LazyFrame  # reference, all rows
-    r_phys: pl.LazyFrame  # rollout, physical steps only
     t_phys: pl.LazyFrame  # reference, physical steps only
-    checkpoint: str | None = None  # from the rollout YAML; router_gating only
-    # Diagnostic pre-aggregated at rollout time (giant.rollout.
-    # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
-    # only. Unlike checkpoint/router_gating, this needs no live model: it's
-    # already a finished histogram, just passed through.
-    type_embedding_l1_dist: dict | None = None
 
     @classmethod
     def open(
         cls,
-        rollout,
+        rollouts: list[RolloutSpec],
         reference,
         ctx: Context,
-        checkpoint=None,
         chunk: tuple[int, int] | None = None,
-        type_embedding_l1_dist: dict | None = None,
     ) -> "Bundle":
-        """Open both sides, optionally restricted to one event-disjoint chunk.
+        """Open the reference + every rollout, optionally restricted to one event-disjoint chunk.
 
-        ``chunk = (chunk_index, n_chunks)`` filters both sides to
+        ``chunk = (chunk_index, n_chunks)`` filters every side to
         ``event_id % n_chunks == chunk_index`` *before* deriving the physical/
         secondary views, so every downstream reduction (which is either
         row-local or a ``group_by("event_id")``) sees a self-contained,
         event-disjoint slice — no cross-chunk lookups are ever needed.
         """
-        r_all = open_side(rollout, Side.rollout)
         t_all = open_side(reference, Side.reference)
+        pred = None
         if chunk is not None:
             idx, n = chunk
             pred = pl.col("event_id") % n == idx
-            r_all = r_all.filter(pred)
             t_all = t_all.filter(pred)
-        return cls(
-            ctx=ctx,
-            r_all=r_all,
-            t_all=t_all,
-            r_phys=physical_steps(r_all, Side.rollout),
-            t_phys=physical_steps(t_all, Side.reference),
-            checkpoint=checkpoint,
-            type_embedding_l1_dist=type_embedding_l1_dist,
-        )
+        sides: dict[str, RolloutSide] = {}
+        for rs in rollouts:
+            r_all = open_side(rs.source, Side.rollout)
+            if pred is not None:
+                r_all = r_all.filter(pred)
+            sides[rs.name] = RolloutSide(
+                all=r_all,
+                phys=physical_steps(r_all, Side.rollout),
+                checkpoint=rs.checkpoint,
+                type_embedding_l1_dist=rs.type_embedding_l1_dist,
+            )
+        return cls(ctx=ctx, rollouts=sides, t_all=t_all, t_phys=physical_steps(t_all, Side.reference))
 
 
 @dataclass
@@ -146,8 +159,10 @@ def _unchunkable(
 # small numpy/hist helpers
 # ---------------------------------------------------------------------------
 
-_ROLL = "rollout"
-_REF = "reference"
+
+def _per_rollout(b: Bundle, fn: Callable[[RolloutSide], object]) -> dict[str, object]:
+    """``{name: fn(rollout_side)}`` over every rollout, preserving CLI order."""
+    return {name: fn(rs) for name, rs in b.rollouts.items()}
 
 
 def _counts(h: dict, key, nbins: int) -> list[int]:
@@ -168,14 +183,20 @@ def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]:
     return list(merged.get(str(key), [0] * nbins))
 
 
-def _np_hist_pair(r: np.ndarray, t: np.ndarray, nbins: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
-    """Shared-edge histogram of two small per-event arrays (robust range)."""
-    both = np.concatenate([r, t]) if (len(r) or len(t)) else np.array([0.0, 1.0])
+def _np_hist_shared_edges(arrays: list[np.ndarray], nbins: int) -> tuple[np.ndarray, list[np.ndarray]]:
+    """Shared-edge histogram of several small per-event arrays (robust range).
+
+    The edges are sized from the union of every array (reference + all
+    rollouts), so every series in the resulting overlay is directly
+    comparable on one axis.
+    """
+    non_empty = [a for a in arrays if len(a)]
+    both = np.concatenate(non_empty) if non_empty else np.array([0.0, 1.0])
     lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
     if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
         lo, hi = lo - 0.5, hi + 0.5
     edges = np.linspace(lo, hi, nbins + 1)
-    return edges, np.histogram(r, edges)[0], np.histogram(t, edges)[0]
+    return edges, [np.histogram(a, edges)[0] for a in arrays]
 
 
 def _ks_statistic(r_counts, t_counts) -> float:
@@ -197,24 +218,6 @@ def _ks_statistic(r_counts, t_counts) -> float:
     return float(np.max(np.abs(r_cdf - t_cdf)))
 
 
-def _integer_confusion(t: np.ndarray, r: np.ndarray, max_bins: int = 21) -> tuple[list[str], np.ndarray]:
-    """Confusion matrix of two paired small-integer arrays (e.g. secondary counts).
-
-    Bins are consecutive integers ``0..cap``, with the last bin an overflow
-    ``"cap+"`` bucket, so an occasional pathological count doesn't blow up the
-    heatmap. Returns ``(labels, matrix)`` with ``matrix[i, j]`` counting pairs
-    with ``t == i`` and ``r == j`` (both clipped into ``[0, cap]``).
-    """
-    cap = min(max(int(t.max()) if len(t) else 0, int(r.max()) if len(r) else 0, 1), max_bins - 1)
-    t_c = np.clip(t.astype(np.int64), 0, cap)
-    r_c = np.clip(r.astype(np.int64), 0, cap)
-    n = cap + 1
-    mat = np.zeros((n, n), dtype=np.int64)
-    np.add.at(mat, (t_c, r_c), 1)
-    labels = [str(i) for i in range(cap)] + [f"{cap}+"]
-    return labels, mat
-
-
 def _containment_depths(mat: np.ndarray, edges: np.ndarray, quantile: float) -> np.ndarray:
     """Per-event depth containing ``quantile`` of that event's deposited energy.
 
@@ -274,7 +277,7 @@ def _marginal_overall_partial(b: Bundle, var: str) -> dict:
     _, expr = _var(var)
     edges = _marginal_edges(b.ctx, var)
     return {
-        "r": _partial_hist(b.r_phys, expr, edges),
+        "r": _per_rollout(b, lambda rs: _partial_hist(rs.phys, expr, edges)),
         "t": _partial_hist(b.t_phys, expr, edges),
     }
 
@@ -283,7 +286,8 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
     label, _ = _var(var)
     edges = _marginal_edges(ctx, var)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
+    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
     t = sum_merge([p["t"] for p in parts])
     return Reduced(
         id=f"marginal_{var}",
@@ -293,8 +297,8 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
         xlabel=label,
         payload={
             "edges": edges.tolist(),
-            _ROLL: _finalize_counts(r, 0, nb),
-            _REF: _finalize_counts(t, 0, nb),
+            "series": series,
+            "reference": _finalize_counts(t, 0, nb),
             "log_y": True,
         },
     )
@@ -305,23 +309,24 @@ def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
     return pl.col("event_id").replace_strict(ids, bins, default=-1, return_dtype=pl.Int64)
 
 
+def _grouped_hist_dict(lf: pl.LazyFrame, expr: pl.Expr, edges: np.ndarray, axis: str, ctx: Context, nb: int) -> dict:
+    if axis == "pdg":
+        h = hist1d(lf, expr, edges, group=pl.col("pdg"))
+    elif axis == "material":
+        h = hist1d(lf, expr, edges, group=pl.col("material"))
+    else:  # energy
+        e_edges = np.asarray(ctx.energy_edges)
+        h = hist1d(lf, expr, edges, group=_energy_group_expr(lf, e_edges))
+    return {str(k): _counts(h, k, nb) for k in h}
+
+
 def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
     _, expr = _var(var)
     edges = _marginal_edges(b.ctx, var)
-    if axis == "pdg":
-        r = hist1d(b.r_phys, expr, edges, group=pl.col("pdg"))
-        t = hist1d(b.t_phys, expr, edges, group=pl.col("pdg"))
-    elif axis == "material":
-        r = hist1d(b.r_phys, expr, edges, group=pl.col("material"))
-        t = hist1d(b.t_phys, expr, edges, group=pl.col("material"))
-    else:  # energy
-        e_edges = np.asarray(b.ctx.energy_edges)
-        r = hist1d(b.r_phys, expr, edges, group=_energy_group_expr(b.r_phys, e_edges))
-        t = hist1d(b.t_phys, expr, edges, group=_energy_group_expr(b.t_phys, e_edges))
     nb = len(edges) - 1
     return {
-        "r": {str(k): _counts(r, k, nb) for k in r},
-        "t": {str(k): _counts(t, k, nb) for k in t},
+        "r": _per_rollout(b, lambda rs: _grouped_hist_dict(rs.phys, expr, edges, axis, b.ctx, nb)),
+        "t": _grouped_hist_dict(b.t_phys, expr, edges, axis, b.ctx, nb),
     }
 
 
@@ -329,29 +334,24 @@ def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis:
     label, _ = _var(var)
     edges = _marginal_edges(ctx, var)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
-    t = sum_merge([p["t"] for p in parts])
-    groups: dict[str, dict] = {}
+    names = list(parts[0]["r"])
+    r_merged = {name: sum_merge([p["r"][name] for p in parts]) for name in names}
+    t_merged = sum_merge([p["t"] for p in parts])
 
     if axis == "pdg":
-        for k in ctx.top_pdgs:
-            groups[pdg_label(k)] = {
-                _ROLL: _finalize_counts(r, k, nb),
-                _REF: _finalize_counts(t, k, nb),
-            }
+        keys, labels = ctx.top_pdgs, [pdg_label(k) for k in ctx.top_pdgs]
     elif axis == "material":
-        for m in ctx.materials:
-            groups[material_label(m)] = {
-                _ROLL: _finalize_counts(r, m, nb),
-                _REF: _finalize_counts(t, m, nb),
-            }
+        keys, labels = ctx.materials, [material_label(m) for m in ctx.materials]
     else:  # energy
         e_edges = np.asarray(ctx.energy_edges)
-        for bi, lbl in enumerate(energy_bin_labels(e_edges)):
-            groups[lbl] = {
-                _ROLL: _finalize_counts(r, bi, nb),
-                _REF: _finalize_counts(t, bi, nb),
-            }
+        keys, labels = list(range(len(e_edges) - 1)), energy_bin_labels(e_edges)
+
+    groups: dict[str, dict] = {}
+    for k, lbl in zip(keys, labels):
+        groups[lbl] = {
+            "series": {name: _finalize_counts(r_merged[name], k, nb) for name in names},
+            "reference": _finalize_counts(t_merged, k, nb),
+        }
 
     return Reduced(
         id=f"marginal_{var}_by_{axis}",
@@ -364,7 +364,7 @@ def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis:
 
 
 # ---------------------------------------------------------------------------
-# distance summary: a var x group-axis scorecard, reusing the marginal hists
+# distance summary: a var x group-axis scorecard per rollout, reusing the marginal hists
 # ---------------------------------------------------------------------------
 
 
@@ -379,29 +379,37 @@ def _distance_summary_partial(b: Bundle) -> dict:
 
 def _distance_summary_finalize(parts: list[dict], ctx: Context) -> Reduced:
     col_labels = ["overall", *GROUPING_AXES]
-    matrix: list[list[float]] = []
+    names = list(parts[0][MARGINAL_VARS[0]]["overall"]["r"])
+    matrices: dict[str, list[list[float]]] = {name: [] for name in names}
+
     for var in MARGINAL_VARS:
         edges = _marginal_edges(ctx, var)
         nb = len(edges) - 1
-        row: list[float] = []
 
-        r = sum_merge([p[var]["overall"]["r"] for p in parts])
-        t = sum_merge([p[var]["overall"]["t"] for p in parts])
-        row.append(_ks_statistic(_finalize_counts(r, 0, nb), _finalize_counts(t, 0, nb)))
+        t_overall = sum_merge([p[var]["overall"]["t"] for p in parts])
+        r_overall = {name: sum_merge([p[var]["overall"]["r"][name] for p in parts]) for name in names}
+        row: dict[str, list[float]] = {name: [] for name in names}
+        for name in names:
+            row[name].append(
+                _ks_statistic(_finalize_counts(r_overall[name], 0, nb), _finalize_counts(t_overall, 0, nb))
+            )
 
         for axis in GROUPING_AXES:
-            r = sum_merge([p[var][axis]["r"] for p in parts])
-            t = sum_merge([p[var][axis]["t"] for p in parts])
-            dists, weights = [], []
-            for k in _group_keys(ctx, axis):
-                rc, tc = _finalize_counts(r, k, nb), _finalize_counts(t, k, nb)
-                w = sum(rc) + sum(tc)
-                if w == 0:
-                    continue
-                dists.append(_ks_statistic(rc, tc))
-                weights.append(w)
-            row.append(float(np.average(dists, weights=weights)) if dists else float("nan"))
-        matrix.append(row)
+            t_grp = sum_merge([p[var][axis]["t"] for p in parts])
+            r_grp = {name: sum_merge([p[var][axis]["r"][name] for p in parts]) for name in names}
+            for name in names:
+                dists, weights = [], []
+                for k in _group_keys(ctx, axis):
+                    rc, tc = _finalize_counts(r_grp[name], k, nb), _finalize_counts(t_grp, k, nb)
+                    w = sum(rc) + sum(tc)
+                    if w == 0:
+                        continue
+                    dists.append(_ks_statistic(rc, tc))
+                    weights.append(w)
+                row[name].append(float(np.average(dists, weights=weights)) if dists else float("nan"))
+
+        for name in names:
+            matrices[name].append(row[name])
 
     return Reduced(
         id="marginal_distance_summary",
@@ -410,7 +418,7 @@ def _distance_summary_finalize(parts: list[dict], ctx: Context) -> Reduced:
         title="Marginal distance summary (KS statistic, rollout vs reference)",
         xlabel="grouping axis",
         payload={
-            "matrix": matrix,
+            "series": matrices,
             "row_labels": [_TITLE_NAMES[v] for v in MARGINAL_VARS],
             "col_labels": col_labels,
             "ylabel": "marginal variable",
@@ -427,16 +435,24 @@ def _distance_summary_finalize(parts: list[dict], ctx: Context) -> Reduced:
 
 
 def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
-    r_lf, t_lf = (b.r_all, b.t_all) if use_all else (b.r_phys, b.t_phys)
-    r = event_scalars(r_lf)[col].to_numpy()
-    t = event_scalars(t_lf)[col].to_numpy()
-    return {"r": r.tolist(), "t": t.tolist()}
+    t_lf = b.t_all if use_all else b.t_phys
+
+    def _vals(rs: RolloutSide) -> list[float]:
+        lf = rs.all if use_all else rs.phys
+        return event_scalars(lf)[col].to_numpy().tolist()
+
+    return {
+        "r": _per_rollout(b, _vals),
+        "t": event_scalars(t_lf)[col].to_numpy().tolist(),
+    }
 
 
 def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str) -> Reduced:
-    r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
+    names = list(parts[0]["r"])
+    r_arrays = {name: np.concatenate([np.asarray(p["r"][name], dtype=float) for p in parts]) for name in names}
     t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
-    edges, rc, tc = _np_hist_pair(r, t, ctx.n_marginal_bins)
+    edges, counts = _np_hist_shared_edges([t, *(r_arrays[n] for n in names)], ctx.n_marginal_bins)
+    t_counts, *r_counts = counts
     return Reduced(
         id=spec_id,
         family="event",
@@ -445,40 +461,44 @@ def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title:
         xlabel=xlabel,
         payload={
             "edges": edges.tolist(),
-            _ROLL: rc.astype(np.int64).tolist(),
-            _REF: tc.astype(np.int64).tolist(),
+            "series": {name: c.astype(np.int64).tolist() for name, c in zip(names, r_counts)},
+            "reference": t_counts.astype(np.int64).tolist(),
             "log_y": False,
         },
     )
 
 
 def _event_total_edep_by_energy_partial(b: Bundle) -> dict:
-    r = event_scalars(b.r_all)
     t = event_scalars(b.t_all)
+
+    def _vals(rs: RolloutSide) -> dict:
+        r = event_scalars(rs.all)
+        return {"incident": r["incident_E"].to_list(), "edep": r["total_edep"].to_list()}
+
     return {
-        "r_incident": r["incident_E"].to_list(),
-        "r_edep": r["total_edep"].to_list(),
-        "t_incident": t["incident_E"].to_list(),
-        "t_edep": t["total_edep"].to_list(),
+        "r": _per_rollout(b, _vals),
+        "t": {"incident": t["incident_E"].to_list(), "edep": t["total_edep"].to_list()},
     }
 
 
 def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
     e_edges = np.asarray(ctx.energy_edges)
-    r_inc = np.concatenate([np.asarray(p["r_incident"], dtype=float) for p in parts])
-    r_val = np.concatenate([np.asarray(p["r_edep"], dtype=float) for p in parts])
-    t_inc = np.concatenate([np.asarray(p["t_incident"], dtype=float) for p in parts])
-    t_val = np.concatenate([np.asarray(p["t_edep"], dtype=float) for p in parts])
-    r_bin = np.clip(np.digitize(r_inc, e_edges[1:-1]), 0, len(e_edges) - 2)
+    names = list(parts[0]["r"])
+    t_inc = np.concatenate([np.asarray(p["t"]["incident"], dtype=float) for p in parts])
+    t_val = np.concatenate([np.asarray(p["t"]["edep"], dtype=float) for p in parts])
+    r_inc = {n: np.concatenate([np.asarray(p["r"][n]["incident"], dtype=float) for p in parts]) for n in names}
+    r_val = {n: np.concatenate([np.asarray(p["r"][n]["edep"], dtype=float) for p in parts]) for n in names}
+
+    edges, _ = _np_hist_shared_edges([t_val, *(r_val[n] for n in names)], ctx.n_marginal_bins)
     t_bin = np.clip(np.digitize(t_inc, e_edges[1:-1]), 0, len(e_edges) - 2)
-    edges, _, _ = _np_hist_pair(r_val, t_val, ctx.n_marginal_bins)
+    r_bin = {n: np.clip(np.digitize(r_inc[n], e_edges[1:-1]), 0, len(e_edges) - 2) for n in names}
+
     groups: dict[str, dict] = {}
     for bi, lbl in enumerate(energy_bin_labels(e_edges)):
-        rc = np.histogram(r_val[r_bin == bi], edges)[0]
         tc = np.histogram(t_val[t_bin == bi], edges)[0]
         groups[lbl] = {
-            _ROLL: rc.astype(np.int64).tolist(),
-            _REF: tc.astype(np.int64).tolist(),
+            "series": {n: np.histogram(r_val[n][r_bin[n] == bi], edges)[0].astype(np.int64).tolist() for n in names},
+            "reference": tc.astype(np.int64).tolist(),
         }
     return Reduced(
         id="event_total_edep_by_energy",
@@ -497,15 +517,15 @@ def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Red
 
 def _profile_partial(b: Bundle, coord_fn, edges_key: str) -> dict:
     edges = np.asarray(getattr(b.ctx, edges_key))
-    r_lf = attach_entry_axis(b.r_all, entry_axis(b.r_all))
-    t_lf = attach_entry_axis(b.t_all, entry_axis(b.t_all))
-    r_ids, r_mat = profile_partial(r_lf, coord_fn(), edges, pl.col("edep"))
-    t_ids, t_mat = profile_partial(t_lf, coord_fn(), edges, pl.col("edep"))
+
+    def _mat(lf: pl.LazyFrame) -> dict:
+        lf2 = attach_entry_axis(lf, entry_axis(lf))
+        ids, mat = profile_partial(lf2, coord_fn(), edges, pl.col("edep"))
+        return {"ids": ids.tolist(), "mat": mat.tolist()}
+
     return {
-        "r_ids": r_ids.tolist(),
-        "r_mat": r_mat.tolist(),
-        "t_ids": t_ids.tolist(),
-        "t_mat": t_mat.tolist(),
+        "r": _per_rollout(b, lambda rs: _mat(rs.all)),
+        "t": _mat(b.t_all),
     }
 
 
@@ -537,12 +557,19 @@ def _profile_finalize(
 ) -> Reduced:
     edges = np.asarray(getattr(ctx, edges_key))
     nb = len(edges) - 1
-    _assert_event_disjoint([p["r_ids"] for p in parts], spec_id, "rollout")
-    _assert_event_disjoint([p["t_ids"] for p in parts], spec_id, "reference")
-    r_mats = [np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts]
-    t_mats = [np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts]
-    r_mean, r_std = profile_finalize(r_mats)
+    names = list(parts[0]["r"])
+
+    _assert_event_disjoint([p["t"]["ids"] for p in parts], spec_id, "reference")
+    t_mats = [np.asarray(p["t"]["mat"], dtype=float).reshape(-1, nb) for p in parts]
     t_mean, t_std = profile_finalize(t_mats)
+
+    series: dict[str, dict] = {}
+    for name in names:
+        _assert_event_disjoint([p["r"][name]["ids"] for p in parts], spec_id, name)
+        mats = [np.asarray(p["r"][name]["mat"], dtype=float).reshape(-1, nb) for p in parts]
+        mean, std = profile_finalize(mats)
+        series[name] = {"mean": mean.tolist(), "std": std.tolist()}
+
     return Reduced(
         id=spec_id,
         family="shower",
@@ -551,10 +578,8 @@ def _profile_finalize(
         xlabel=xlabel,
         payload={
             "edges": edges.tolist(),
-            "rollout_mean": r_mean.tolist(),
-            "rollout_std": r_std.tolist(),
-            "reference_mean": t_mean.tolist(),
-            "reference_std": t_std.tolist(),
+            "series": series,
+            "reference": {"mean": t_mean.tolist(), "std": t_std.tolist()},
             "ylabel": "mean deposited energy per event [MeV]",
         },
     )
@@ -573,13 +598,20 @@ _CONTAINMENT_QUANTILES: list[tuple[float, str]] = [
 def _containment_finalize(parts: list[dict], ctx: Context, spec_id: str, quantile: float) -> Reduced:
     edges = np.asarray(ctx.depth_edges)
     nb = len(edges) - 1
-    _assert_event_disjoint([p["r_ids"] for p in parts], spec_id, "rollout")
-    _assert_event_disjoint([p["t_ids"] for p in parts], spec_id, "reference")
-    r_full = np.concatenate([np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
-    t_full = np.concatenate([np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
-    r_depth = _containment_depths(r_full, edges, quantile)
+    names = list(parts[0]["r"])
+
+    _assert_event_disjoint([p["t"]["ids"] for p in parts], spec_id, "reference")
+    t_full = np.concatenate([np.asarray(p["t"]["mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
     t_depth = _containment_depths(t_full, edges, quantile)
-    hedges, rc, tc = _np_hist_pair(r_depth, t_depth, ctx.n_marginal_bins)
+
+    r_depths: dict[str, np.ndarray] = {}
+    for name in names:
+        _assert_event_disjoint([p["r"][name]["ids"] for p in parts], spec_id, name)
+        full = np.concatenate([np.asarray(p["r"][name]["mat"], dtype=float).reshape(-1, nb) for p in parts], axis=0)
+        r_depths[name] = _containment_depths(full, edges, quantile)
+
+    hedges, counts = _np_hist_shared_edges([t_depth, *(r_depths[n] for n in names)], ctx.n_marginal_bins)
+    t_counts, *r_counts = counts
     return Reduced(
         id=spec_id,
         family="shower",
@@ -588,8 +620,8 @@ def _containment_finalize(parts: list[dict], ctx: Context, spec_id: str, quantil
         xlabel=f"depth containing {quantile:.0%} of deposited energy [mm]",
         payload={
             "edges": hedges.tolist(),
-            _ROLL: rc.astype(np.int64).tolist(),
-            _REF: tc.astype(np.int64).tolist(),
+            "series": {name: c.astype(np.int64).tolist() for name, c in zip(names, r_counts)},
+            "reference": t_counts.astype(np.int64).tolist(),
             "log_y": False,
         },
     )
@@ -601,20 +633,30 @@ def _containment_finalize(parts: list[dict], ctx: Context, spec_id: str, quantil
 
 
 def _species_share_partial(b: Bundle) -> dict:
-    r = species_share(b.r_all)
     t = species_share(b.t_all)
+
+    def _map(rs: RolloutSide) -> dict[str, float]:
+        r = species_share(rs.all)
+        return {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())}
+
     return {
-        "r": {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())},
+        "r": _per_rollout(b, _map),
         "t": {str(k): v for k, v in zip(t["pdg"].to_list(), t["total_edep"].to_list())},
     }
 
 
 def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    r_map = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
+    r_maps = {n: sum_merge([p["r"][n] for p in parts]) for n in names}
     t_map = sum_merge([p["t"] for p in parts])
-    r_tot = sum(r_map.values()) or 1.0
     t_tot = sum(t_map.values()) or 1.0
     labels = [pdg_label(k) for k in ctx.top_pdgs]
+
+    series: dict[str, list[float]] = {}
+    for n in names:
+        r_tot = sum(r_maps[n].values()) or 1.0
+        series[n] = [r_maps[n].get(str(k), 0.0) / r_tot for k in ctx.top_pdgs]
+
     return Reduced(
         id="species_edep_share",
         family="species",
@@ -623,22 +665,23 @@ def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="species",
         payload={
             "labels": labels,
-            _ROLL: [r_map.get(str(k), 0.0) / r_tot for k in ctx.top_pdgs],
-            _REF: [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs],
+            "series": series,
+            "reference": [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs],
             "ylabel": "fraction of total deposited energy",
         },
     )
 
 
 def _leakage_partial(b: Bundle) -> dict:
-    frac = leakage_fraction(b.r_all)
-    return {"frac": frac.tolist()}
+    return {"r": _per_rollout(b, lambda rs: leakage_fraction(rs.all).tolist())}
 
 
 def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts])
-    edges = np.linspace(0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1)
-    counts = np.histogram(frac, edges)[0]
+    names = list(parts[0]["r"])
+    arrays = {n: np.concatenate([np.asarray(p["r"][n], dtype=float) for p in parts]) for n in names}
+    max_val = max((float(a.max()) for a in arrays.values() if len(a)), default=1e-3)
+    edges = np.linspace(0.0, max(max_val, 1e-3), ctx.n_marginal_bins + 1)
+    series = {n: np.histogram(arrays[n], edges)[0].astype(np.int64).tolist() for n in names}
     return Reduced(
         id="leakage_fraction",
         family="species",
@@ -647,7 +690,7 @@ def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="escaped energy fraction",
         payload={
             "edges": edges.tolist(),
-            _ROLL: counts.astype(np.int64).tolist(),
+            "series": series,
             "log_y": True,
             "note": "rollout only; the reference has no detector-escape concept",
         },
@@ -659,24 +702,36 @@ def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
 # ---------------------------------------------------------------------------
 
 
-def _sec_frames(b: Bundle):
-    return (
-        secondaries(b.r_phys, Side.rollout),
-        secondaries(b.t_all, Side.reference),
-    )
+def _t_sec(b: Bundle) -> pl.LazyFrame:
+    return secondaries(b.t_all, Side.reference)
+
+
+def _r_sec(rs: RolloutSide) -> pl.LazyFrame:
+    return secondaries(rs.phys, Side.rollout)
 
 
 def _sec_count_per_event_partial(b: Bundle) -> dict:
-    r_sec, t_sec = _sec_frames(b)
-    r = r_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
-    t = t_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
-    return {"r": r.tolist(), "t": t.tolist()}
+    t = _t_sec(b).group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
+
+    def _r(rs: RolloutSide) -> list[float]:
+        return (
+            _r_sec(rs)
+            .group_by("event_id")
+            .agg(pl.len().alias("n"))
+            .collect(engine="streaming")["n"]
+            .to_numpy()
+            .tolist()
+        )
+
+    return {"r": _per_rollout(b, _r), "t": t.tolist()}
 
 
 def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
+    names = list(parts[0]["r"])
     t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
-    edges, rc, tc = _np_hist_pair(r, t, min(ctx.n_marginal_bins, 40))
+    r = {n: np.concatenate([np.asarray(p["r"][n], dtype=float) for p in parts]) for n in names}
+    edges, counts = _np_hist_shared_edges([t, *(r[n] for n in names)], min(ctx.n_marginal_bins, 40))
+    t_c, *r_cs = counts
     return Reduced(
         id="sec_count_per_event",
         family="secondaries",
@@ -685,8 +740,8 @@ def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="secondaries per event",
         payload={
             "edges": edges.tolist(),
-            _ROLL: rc.astype(np.int64).tolist(),
-            _REF: tc.astype(np.int64).tolist(),
+            "series": {n: c.astype(np.int64).tolist() for n, c in zip(names, r_cs)},
+            "reference": t_c.astype(np.int64).tolist(),
             "log_y": False,
         },
     )
@@ -698,14 +753,22 @@ def _counts_by_pdg(sec_lf: pl.LazyFrame) -> dict[str, int]:
 
 
 def _sec_count_per_species_partial(b: Bundle) -> dict:
-    r_sec, t_sec = _sec_frames(b)
-    return {"r": _counts_by_pdg(r_sec), "t": _counts_by_pdg(t_sec)}
+    return {"r": _per_rollout(b, lambda rs: _counts_by_pdg(_r_sec(rs))), "t": _counts_by_pdg(_t_sec(b))}
 
 
 def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
+    r_maps = {n: sum_merge([p["r"][n] for p in parts]) for n in names}
     t = sum_merge([p["t"] for p in parts])
-    keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[: len(ctx.top_pdgs)]
+
+    all_keys = set(t)
+    for m in r_maps.values():
+        all_keys |= set(m)
+
+    def _total(k: str) -> float:
+        return t.get(k, 0) + sum(m.get(k, 0) for m in r_maps.values())
+
+    keys = sorted(all_keys, key=lambda k: -_total(k))[: len(ctx.top_pdgs)]
     return Reduced(
         id="sec_count_per_species",
         family="secondaries",
@@ -714,39 +777,167 @@ def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
         xlabel="species",
         payload={
             "labels": [pdg_label(int(k)) for k in keys],
-            _ROLL: [float(r.get(k, 0)) for k in keys],
-            _REF: [float(t.get(k, 0)) for k in keys],
+            "series": {n: [float(r_maps[n].get(k, 0)) for k in keys] for n in names},
+            "reference": [float(t.get(k, 0)) for k in keys],
             "ylabel": "secondary count",
         },
     )
 
 
+# Per-step secondary multiplicity. Fixed integer edges (bin i == exactly i
+# secondaries, the top bin an overflow bucket) keep both plots sum-mergeable
+# across chunks — no shared-range pass needed. The species heatmap gets a
+# shorter row axis because a single step rarely emits many of *one* species.
+_N_SEC_STEP_CAP = 20
+_N_SEC_SPECIES_CAP = 10
+_OTHER_KEY = "other"
+
+
+def _n_sec_edges(cap: int) -> np.ndarray:
+    return np.arange(-0.5, cap + 1.5)
+
+
+def _sec_step_key_lf(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
+    """Secondaries with their emitting-step key.
+
+    The rollout side reads *all* rows, not just physical ones: a secondary
+    whose very first row is a synthetic termination row (born, then immediately
+    escaped or cut) was still produced by its parent step, and dropping it would
+    undercount that step's multiplicity.
+    """
+    return secondaries_by_step(lf, side)
+
+
+def _n_steps(lf: pl.LazyFrame) -> int:
+    """Number of (physical) step rows — the denominator the zero rows come from."""
+    return int(lf.select(pl.len()).collect(engine="streaming").item())
+
+
+def _sec_count_per_step_partial(b: Bundle) -> dict:
+    edges = _n_sec_edges(_N_SEC_STEP_CAP)
+
+    def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict:
+        per_step = sec_lf.group_by("step_key").agg(pl.len().alias("n"))
+        return {
+            "h": _partial_hist(per_step, pl.col("n").clip(0, _N_SEC_STEP_CAP), edges),
+            "n_steps": _n_steps(steps_lf),
+        }
+
+    return {
+        "r": _per_rollout(b, lambda rs: _side(_sec_step_key_lf(rs.all, Side.rollout), rs.phys)),
+        "t": _side(_sec_step_key_lf(b.t_all, Side.reference), b.t_phys),
+    }
+
+
+def _zero_filled(part_hists: list[dict], n_steps: int, key, nbins: int) -> list[int]:
+    """Merged counts for one series, with bin 0 (= steps that emitted none) filled in.
+
+    The reduction only ever sees steps that produced at least one secondary, so
+    the empty ones are recovered by subtraction from the total step count.
+    """
+    counts = _finalize_counts(sum_merge(part_hists), key, nbins)
+    counts[0] = max(n_steps - int(sum(counts)), 0)
+    return [int(c) for c in counts]
+
+
+def _sec_count_per_step_finalize(parts: list[dict], ctx: Context) -> Reduced:
+    edges = _n_sec_edges(_N_SEC_STEP_CAP)
+    nb = len(edges) - 1
+    names = list(parts[0]["r"])
+    series = {
+        name: _zero_filled([p["r"][name]["h"] for p in parts], sum(p["r"][name]["n_steps"] for p in parts), 0, nb)
+        for name in names
+    }
+    return Reduced(
+        id="sec_count_per_step",
+        family="secondaries",
+        kind="overlay_hist",
+        title="Number of secondaries per step",
+        xlabel="secondaries per step",
+        payload={
+            "edges": edges.tolist(),
+            "series": series,
+            "reference": _zero_filled([p["t"]["h"] for p in parts], sum(p["t"]["n_steps"] for p in parts), 0, nb),
+            "log_y": True,
+        },
+    )
+
+
+def _species_key_expr(top_pdgs: list[int]) -> pl.Expr:
+    """``pdg`` bucketed into the shared top-K columns plus one ``other`` bin."""
+    return pl.when(pl.col("pdg").is_in(list(top_pdgs))).then(pl.col("pdg").cast(pl.Utf8)).otherwise(pl.lit(_OTHER_KEY))
+
+
+def _sec_count_per_step_by_species_partial(b: Bundle) -> dict:
+    edges = _n_sec_edges(_N_SEC_SPECIES_CAP)
+    group = _species_key_expr(b.ctx.top_pdgs)
+
+    def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict:
+        per_step_species = sec_lf.group_by("step_key", "pdg").agg(pl.len().alias("n"))
+        return {
+            "h": _partial_hist(per_step_species, pl.col("n").clip(0, _N_SEC_SPECIES_CAP), edges, group=group),
+            "n_steps": _n_steps(steps_lf),
+        }
+
+    return {
+        "r": _per_rollout(b, lambda rs: _side(_sec_step_key_lf(rs.all, Side.rollout), rs.phys)),
+        "t": _side(_sec_step_key_lf(b.t_all, Side.reference), b.t_phys),
+    }
+
+
+def _sec_count_per_step_by_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
+    edges = _n_sec_edges(_N_SEC_SPECIES_CAP)
+    nb = len(edges) - 1
+    names = list(parts[0]["r"])
+    keys = [str(p) for p in ctx.top_pdgs] + [_OTHER_KEY]
+
+    def _matrix(hists: list[dict], n_steps: int) -> list[list[int]]:
+        # columns = species, rows = multiplicity; every species gets its own
+        # zero row (steps that produced none of *that* species).
+        cols = [_zero_filled(hists, n_steps, k, nb) for k in keys]
+        return [[cols[j][i] for j in range(len(keys))] for i in range(nb)]
+
+    return Reduced(
+        id="sec_count_per_step_by_species",
+        family="secondaries",
+        kind="heatmap",
+        title="Per-step secondary multiplicity by species",
+        xlabel="species",
+        payload={
+            "series": {
+                n: _matrix([p["r"][n]["h"] for p in parts], sum(p["r"][n]["n_steps"] for p in parts)) for n in names
+            },
+            "reference": _matrix([p["t"]["h"] for p in parts], sum(p["t"]["n_steps"] for p in parts)),
+            "row_labels": [str(i) for i in range(_N_SEC_SPECIES_CAP)] + [f"{_N_SEC_SPECIES_CAP}+"],
+            "col_labels": [pdg_label(k) for k in ctx.top_pdgs] + [_OTHER_KEY],
+            "ylabel": "secondaries of this species per step",
+            "cbar_label": "step count",
+            "log_color": True,
+        },
+    )
+
+
 def _sec_energy_partial(b: Bundle) -> dict:
-    r_sec, t_sec = _sec_frames(b)
     edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1)
     return {
-        "r": _partial_hist(r_sec, pl.col("energy"), edges),
-        "t": _partial_hist(t_sec, pl.col("energy"), edges),
+        "r": _per_rollout(b, lambda rs: _partial_hist(_r_sec(rs), pl.col("energy"), edges)),
+        "t": _partial_hist(_t_sec(b), pl.col("energy"), edges),
     }
 
 
 def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
     edges = np.linspace(*ctx.sec_energy_range, ctx.n_sec_bins + 1)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
     t = sum_merge([p["t"] for p in parts])
+    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
     return Reduced(
         id="sec_energy",
         family="secondaries",
         kind="overlay_hist",
         title="Secondary birth energy",
         xlabel="secondary energy [MeV]",
-        payload={
-            "edges": edges.tolist(),
-            _ROLL: _finalize_counts(r, 0, nb),
-            _REF: _finalize_counts(t, 0, nb),
-            "log_y": True,
-        },
+        payload={"edges": edges.tolist(), "series": series, "reference": _finalize_counts(t, 0, nb), "log_y": True},
     )
 
 
@@ -760,64 +951,25 @@ def _sec_cos_angle_partial(b: Bundle) -> dict:
         ea = entry_axis(steps_lf)
         return _partial_hist(attach_entry_axis(sec_lf, ea), cos, edges)
 
-    r_sec, t_sec = _sec_frames(b)
-    return {"r": _side(r_sec, b.r_phys), "t": _side(t_sec, b.t_all)}
+    return {
+        "r": _per_rollout(b, lambda rs: _side(_r_sec(rs), rs.phys)),
+        "t": _side(_t_sec(b), b.t_all),
+    }
 
 
 def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
     edges = np.linspace(-1.0, 1.0, ctx.n_sec_bins + 1)
     nb = len(edges) - 1
-    r = sum_merge([p["r"] for p in parts])
+    names = list(parts[0]["r"])
     t = sum_merge([p["t"] for p in parts])
+    series = {name: _finalize_counts(sum_merge([p["r"][name] for p in parts]), 0, nb) for name in names}
     return Reduced(
         id="sec_cos_angle",
         family="secondaries",
         kind="overlay_hist",
         title="Secondary emission angle relative to the shower axis",
         xlabel="cos of emission angle",
-        payload={
-            "edges": edges.tolist(),
-            _ROLL: _finalize_counts(r, 0, nb),
-            _REF: _finalize_counts(t, 0, nb),
-            "log_y": False,
-        },
-    )
-
-
-def _n_sec_confusion_partial(b: Bundle) -> dict:
-    r_sec, t_sec = _sec_frames(b)
-    r_ids, r_n = sec_count_by_event(b.r_phys, r_sec)
-    t_ids, t_n = sec_count_by_event(b.t_all, t_sec)
-    return {"r_ids": r_ids.tolist(), "r_n": r_n.tolist(), "t_ids": t_ids.tolist(), "t_n": t_n.tolist()}
-
-
-def _n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
-    r_ids = np.concatenate([np.asarray(p["r_ids"], dtype=np.int64) for p in parts])
-    r_n = np.concatenate([np.asarray(p["r_n"], dtype=np.int64) for p in parts])
-    t_ids = np.concatenate([np.asarray(p["t_ids"], dtype=np.int64) for p in parts])
-    t_n = np.concatenate([np.asarray(p["t_n"], dtype=np.int64) for p in parts])
-    # event-disjoint chunking (see Bundle.open) means each event_id appears in
-    # exactly one part on each side, so a plain dict build is a safe merge.
-    r_map = dict(zip(r_ids.tolist(), r_n.tolist()))
-    t_map = dict(zip(t_ids.tolist(), t_n.tolist()))
-    common = sorted(set(r_map) & set(t_map))
-    true_n = np.array([t_map[e] for e in common], dtype=np.int64)
-    pred_n = np.array([r_map[e] for e in common], dtype=np.int64)
-    labels, mat = _integer_confusion(true_n, pred_n)
-    return Reduced(
-        id="n_sec_confusion",
-        family="secondaries",
-        kind="heatmap",
-        title="Predicted vs true secondary count per event",
-        xlabel="predicted secondaries (rollout)",
-        payload={
-            "matrix": mat.tolist(),
-            "row_labels": labels,
-            "col_labels": labels,
-            "ylabel": "true secondaries (reference)",
-            "cbar_label": "event count",
-            "vmin": 0.0,
-        },
+        payload={"edges": edges.tolist(), "series": series, "reference": _finalize_counts(t, 0, nb), "log_y": False},
     )
 
 
@@ -825,20 +977,18 @@ def _n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
 # router diagnostics (not chunked — already bounded/subsampled)
 # ---------------------------------------------------------------------------
 
-_router_gating_partial, _router_gating_finalize = _unchunkable(
-    lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys)
-)
+_router_gating_partial, _router_gating_finalize = _unchunkable(lambda b: compute_router_gating(b.rollouts, b.t_phys))
 _router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
-    lambda b: compute_router_share_by_pdg(b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs)
+    lambda b: compute_router_share_by_pdg(b.rollouts, b.t_phys, b.ctx.top_pdgs)
 )
 _router_share_process_partial, _router_share_process_finalize = _unchunkable(
-    lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
+    lambda b: compute_router_share_by_process(b.rollouts, b.t_phys)
 )
 _router_specialization_partial, _router_specialization_finalize = _unchunkable(
-    lambda b: compute_router_specialization(b.checkpoint, b.r_phys, b.t_phys)
+    lambda b: compute_router_specialization(b.rollouts, b.t_phys)
 )
 _type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = _unchunkable(
-    lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist)
+    lambda b: compute_type_embedding_l1_distance(b.rollouts)
 )
 
 
@@ -986,6 +1136,18 @@ def build_catalog() -> list[PlotSpec]:
             compute_partial=_sec_count_per_species_partial,
             finalize=_sec_count_per_species_finalize,
         ),
+        PlotSpec(
+            "sec_count_per_step",
+            "secondaries",
+            compute_partial=_sec_count_per_step_partial,
+            finalize=_sec_count_per_step_finalize,
+        ),
+        PlotSpec(
+            "sec_count_per_step_by_species",
+            "secondaries",
+            compute_partial=_sec_count_per_step_by_species_partial,
+            finalize=_sec_count_per_step_by_species_finalize,
+        ),
         PlotSpec(
             "sec_energy",
             "secondaries",
@@ -998,12 +1160,6 @@ def build_catalog() -> list[PlotSpec]:
             compute_partial=_sec_cos_angle_partial,
             finalize=_sec_cos_angle_finalize,
         ),
-        PlotSpec(
-            "n_sec_confusion",
-            "secondaries",
-            compute_partial=_n_sec_confusion_partial,
-            finalize=_n_sec_confusion_finalize,
-        ),
         PlotSpec(
             "router_gating",
             "model",

giant/analysis/condor.py

@@ -1,4 +1,4 @@
-"""HTCondor orchestration driven by a ``giant rollout`` YAML sidecar.
+"""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
@@ -10,8 +10,11 @@ provenance:
 * ``checkpoint``, ``geometry_oracle``, ``energy_cutoff``, ``steps``, ... —
   metadata that flows straight into every plot's gallery ``metadata.yaml``.
 
-So the analysis takes that one YAML as input, derives its own **run directory**
-next to the rollout parquet, and lays everything out under it:
+The analysis takes N such YAMLs — one series per rollout, all required to
+share the same ``dataset`` (the premise is "N candidates vs one ground
+truth") — resolves each one's series name (``load_rollout_yamls``), derives
+its own **run directory** next to the first rollout's parquet, and lays
+everything out under it:
 
     <run_dir>/shared.json     fixed bin edges / group sets (prep)
     <run_dir>/run_meta.json   resolved rollout/reference paths + plot metadata
@@ -44,6 +47,7 @@ from __future__ import annotations
 import json
 import shutil
 import sys
+from collections.abc import Sequence
 from dataclasses import dataclass, field
 from pathlib import Path
 
@@ -54,7 +58,7 @@ from giant.analysis.catalog import Bundle, catalog_ids, get_spec
 from giant.analysis.context import Context, build_context
 from giant.analysis.reduced import Partial
 from giant.analysis.runtime_estimate import estimate_runtime_s
-from giant.analysis.sources import Side, open_side
+from giant.analysis.sources import RolloutSpec, Side, open_side
 
 # Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
 _PLOT_META_KEYS = (
@@ -111,8 +115,60 @@ def load_rollout_yaml(path: str | Path) -> dict:
     return d
 
 
+@dataclass
+class LoadedRollout:
+    """One rollout YAML plus its resolved series ``name`` (see ``load_rollout_yamls``)."""
+
+    name: str
+    yaml: dict
+
+
+def load_rollout_yamls(
+    paths: Sequence[str | Path], labels: Sequence[str] | None = None
+) -> tuple[list[LoadedRollout], str]:
+    """Load every rollout YAML, resolve each one's series name, and verify they
+    all share one reference (``dataset``) file — the premise is "N candidates
+    vs one ground truth", not N independent comparisons.
+
+    Names: an explicit ``labels[i]`` if given (``labels`` must be empty or
+    exactly ``len(paths)`` long); otherwise the YAML's stem for N>1, or
+    ``"rollout"`` for the single-YAML case — matching today's one-series
+    legend/payload key, so a single-rollout run renders identically to
+    before this feature existed. Raises ``ValueError`` if two rollouts
+    resolve to the same name, or if the YAMLs don't all name the same
+    ``dataset``.
+    """
+    if labels and len(labels) != len(paths):
+        raise ValueError(f"--label given {len(labels)} time(s) but {len(paths)} rollout YAML(s) were passed")
+    yamls = [load_rollout_yaml(p) for p in paths]
+    if labels:
+        names = list(labels)
+    elif len(paths) == 1:
+        names = ["rollout"]
+    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"rollout series names collide: {dupes} — pass --label to disambiguate")
+
+    references = {str(y["dataset"]) for y in yamls}
+    if len(references) > 1:
+        detail = "\n".join(f"  {p}: dataset={y['dataset']!r}" for p, y in zip(paths, yamls))
+        raise ValueError(
+            "all rollout YAMLs must be seeded from the same reference (dataset) "
+            f"file — got {len(references)} distinct ones:\n{detail}"
+        )
+
+    return [LoadedRollout(name=n, yaml=y) for n, y in zip(names, yamls)], yamls[0]["dataset"]
+
+
+def _run_tag(y: dict) -> str:
+    rollout = Path(y["output"])
+    return str(y.get("prediction_id") or rollout.stem)[:8]
+
+
 def derive_run_dir(
-    rollout_yaml: dict,
+    rollout_yamls: list[dict],
     run_dir: str | Path | None = None,
     default_base: str | Path | None = None,
 ) -> Path:
@@ -122,14 +178,24 @@ def derive_run_dir(
     ``default_base / analysis_<tag>`` if ``default_base`` is given (the CLI
     passes the repo's gitignored ``analysis_runs/``, so run directories don't
     pile up on ``/ceph`` next to the rollout parquet). Falls back to next to
-    the rollout parquet — the original convention — for callers that don't
-    care where the run directory lives.
+    the *first* rollout's parquet — the original convention — for callers
+    that don't care where the run directory lives.
+
+    ``tag`` is a single rollout's ``prediction_id``/output stem (matching
+    today's single-rollout convention exactly) when there's only one; for
+    N>1 it joins up to three tags with ``-``, then ``-plus<K>`` for any
+    beyond that, so a many-rollout run still gets a short, stable directory
+    name.
     """
     if run_dir is not None:
         return Path(run_dir)
-    rollout = Path(rollout_yaml["output"])
-    tag = str(rollout_yaml.get("prediction_id") or rollout.stem)[:8]
-    base = Path(default_base) if default_base is not None else rollout.parent
+    tags = [_run_tag(y) for y in rollout_yamls]
+    if len(tags) == 1:
+        tag = tags[0]
+    else:
+        shown, rest = tags[:3], tags[3:]
+        tag = "-".join(shown) + (f"-plus{len(rest)}" if rest else "")
+    base = Path(default_base) if default_base is not None else Path(rollout_yamls[0]["output"]).parent
     return base / f"analysis_{tag}"
 
 
@@ -139,17 +205,21 @@ def _plot_meta(rollout_yaml: dict) -> dict:
 
 @dataclass
 class RunMeta:
-    """Resolved paths + plot metadata for one analysis run (``run_meta.json``)."""
+    """Resolved paths + plot metadata for one analysis run (``run_meta.json``).
 
-    rollout: str
+    ``rollouts`` is ``[{"name", "path", "plot_meta"}, ...]``, insertion order
+    = the order rollouts were given on the CLI (and so the order every
+    ``Reduced.payload["series"]`` dict is built in — see ``catalog.py``).
+    """
+
+    rollouts: list[dict]
     reference: str
     run_dir: str
     title: str
-    plot_meta: dict
     n_chunks: int = 1
-    # rollout+reference row count of each event_id-disjoint chunk, and the
-    # dataset total — inputs to `runtime_estimate.estimate_runtime_s`. Empty/0
-    # on run directories written before this field existed.
+    # combined rollout+reference row count of each event_id-disjoint chunk,
+    # and the dataset total — inputs to `runtime_estimate.estimate_runtime_s`.
+    # Empty/0 on run directories written before this field existed.
     rows_per_chunk: list[int] = field(default_factory=list)
     total_rows: int = 0
 
@@ -161,8 +231,8 @@ class RunMeta:
         return cls(**json.loads(Path(path).read_text()))
 
 
-def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -> list[int]:
-    """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) -> list[int]:
+    """Combined rollout+reference 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
@@ -178,7 +248,8 @@ def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -
         )
 
     out = [0] * n_chunks
-    for lf in (open_side(rollout, Side.rollout), open_side(reference, Side.reference)):
+    sides = [open_side(reference, Side.reference)] + [open_side(r, Side.rollout) for r in rollouts]
+    for lf in sides:
         df = counts(lf)
         for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
             out[c] += n
@@ -186,20 +257,22 @@ def _rows_per_chunk(rollout: str | Path, reference: str | Path, n_chunks: int) -
 
 
 def prep(
-    rollout_yaml: str | Path,
+    rollout_yamls: Sequence[str | Path],
     run_dir: str | Path | None = None,
     n_chunks: int = 1,
     default_base: str | Path | None = None,
+    labels: Sequence[str] | None = None,
     **ctx_kwargs,
 ) -> Path:
-    """Read the rollout YAML, build the shared context, and lay out the run dir.
+    """Read the rollout 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.
+    resolve the actual directory, and ``load_rollout_yamls`` for how
+    ``labels``/YAML stems resolve each rollout's series name.
 
     Clears any existing ``reduced_partial/``/``reduced/`` from a prior prep of
     this same ``run_dir``: partial files carry no record of what context
@@ -208,8 +281,8 @@ def prep(
     rollout/reference files changed) would otherwise let ``merge_one`` silently
     merge stale partials against the new ``shared.json``.
     """
-    y = load_rollout_yaml(rollout_yaml)
-    run_path = derive_run_dir(y, run_dir, default_base=default_base)
+    loaded, reference = load_rollout_yamls(list(rollout_yamls), 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)
 
     for stale in ("reduced_partial", "reduced"):
@@ -217,19 +290,22 @@ def prep(
         if stale_dir.exists():
             shutil.rmtree(stale_dir)
 
-    rollout, reference = y["output"], y["dataset"]
-    ctx = build_context(rollout, reference, **ctx_kwargs)
+    rollout_specs = [RolloutSpec(name=lr.name, source=lr.yaml["output"]) for lr in loaded]
+    ctx = build_context(rollout_specs, reference, **ctx_kwargs)
     ctx.save(run_path / "shared.json")
 
-    rows_per_chunk = _rows_per_chunk(rollout, reference, n_chunks)
+    rows_per_chunk = _rows_per_chunk([lr.yaml["output"] for lr in loaded], reference, n_chunks)
+
+    rollouts_meta = [
+        {"name": lr.name, "path": str(lr.yaml["output"]), "plot_meta": _plot_meta(lr.yaml)} for lr in loaded
+    ]
+    ckpts = ", ".join(Path(lr.yaml.get("checkpoint", "")).name or "rollout" for lr in loaded)
 
-    ckpt = Path(y.get("checkpoint", "")).name or "rollout"
     RunMeta(
-        rollout=str(rollout),
+        rollouts=rollouts_meta,
         reference=str(reference),
         run_dir=str(run_path),
-        title=f"GIANT rollout analysis — {ckpt}",
-        plot_meta=_plot_meta(y),
+        title=f"GIANT rollout analysis — {ckpts}",
         n_chunks=n_chunks,
         rows_per_chunk=rows_per_chunk,
         total_rows=sum(rows_per_chunk),
@@ -244,17 +320,19 @@ def prep(
 
 def compute_reduced(
     spec_id: str,
-    rollout: str | Path,
+    rollouts: list[dict],
     reference: str | Path,
     shared: str | Path,
     out: str | Path,
-    checkpoint: str | None = None,
     chunk_index: int = 0,
     n_chunks: int = 1,
-    type_embedding_l1_dist: dict | None = None,
 ) -> Path:
     """Core: run one (plot, chunk)'s partial reduction against explicit paths.
 
+    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?},
+    ...]``, one per rollout series (insertion order preserved through to every
+    plot's ``Reduced.payload["series"]``).
+
     Writes a ``Partial`` JSON — the raw, not-yet-merged output of
     ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one``
     is what combines every chunk's ``Partial`` for a plot into the final
@@ -268,14 +346,16 @@ def compute_reduced(
         raise ValueError(
             f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})"
         )
-    bundle = Bundle.open(
-        rollout,
-        reference,
-        ctx,
-        checkpoint=checkpoint,
-        chunk=(chunk_index, effective_n),
-        type_embedding_l1_dist=type_embedding_l1_dist,
-    )
+    rollout_specs = [
+        RolloutSpec(
+            name=r["name"],
+            source=r["path"],
+            checkpoint=r.get("checkpoint"),
+            type_embedding_l1_dist=r.get("type_embedding_l1_dist"),
+        )
+        for r in rollouts
+    ]
+    bundle = Bundle.open(rollout_specs, reference, ctx, chunk=(chunk_index, effective_n))
     partial = Partial(
         id=spec_id,
         family=spec.family,
@@ -291,16 +371,23 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
     """Run one (plot, chunk)'s partial reduction from a prepped run directory."""
     run_path = Path(run_dir)
     meta = RunMeta.load(run_path / "run_meta.json")
+    rollouts = [
+        {
+            "name": ro["name"],
+            "path": ro["path"],
+            "checkpoint": ro["plot_meta"].get("checkpoint"),
+            "type_embedding_l1_dist": ro["plot_meta"].get("type_embedding_l1_dist"),
+        }
+        for ro in meta.rollouts
+    ]
     return compute_reduced(
         spec_id,
-        meta.rollout,
+        rollouts,
         meta.reference,
         run_path / "shared.json",
         run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json",
-        checkpoint=meta.plot_meta.get("checkpoint"),
         chunk_index=chunk_index,
         n_chunks=meta.n_chunks,
-        type_embedding_l1_dist=meta.plot_meta.get("type_embedding_l1_dist"),
     )
 
 

giant/analysis/context.py

@@ -26,7 +26,7 @@ from giant.analysis.reduce import (
     entry_axis,
     transverse_expr,
 )
-from giant.analysis.sources import Side, open_side, physical_steps, secondaries
+from giant.analysis.sources import RolloutSpec, Side, open_side, physical_steps, secondaries
 from giant.analysis.variables import RANGED_VARS
 
 
@@ -74,9 +74,9 @@ def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFram
     return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
 
 
-def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float) -> tuple[float, float]:
-    """Robust (lo_q, hi_q) range over the union of two value samples."""
-    both = np.concatenate([r_vals, t_vals])
+def _combined_quantiles(vals: list[np.ndarray], lo_q: float, hi_q: float) -> tuple[float, float]:
+    """Robust (lo_q, hi_q) range over the union of several value samples."""
+    both = np.concatenate(vals)
     lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
     if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
         lo, hi = lo - 0.5, hi + 0.5
@@ -84,7 +84,7 @@ def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_
 
 
 def build_context(
-    rollout: str | Path | pl.LazyFrame,
+    rollouts: list[RolloutSpec],
     reference: str | Path | pl.LazyFrame,
     *,
     n_energy_bins: int = 4,
@@ -94,41 +94,51 @@ def build_context(
     sample_rows: int = 1_000_000,
     seed: int = 0,
 ) -> Context:
-    """Resolve the shared context from the two files (the ``prep`` step)."""
-    r_all = open_side(rollout, Side.rollout)
+    """Resolve the shared context from the reference + every rollout (the ``prep`` step).
+
+    Every range/quantile below is the union of the reference and *all*
+    rollouts, so a single set of fixed bin edges/group sets is valid for
+    every series a compute job streams over.
+    """
     t_all = open_side(reference, Side.reference)
-    r_lf = physical_steps(r_all, Side.rollout)
     t_lf = physical_steps(t_all, Side.reference)
+    r_lfs = {rs.name: physical_steps(open_side(rs.source, Side.rollout), Side.rollout) for rs in rollouts}
 
     # Ranged marginal variables: robust ranges over a shared row subsample.
     exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
-    r_s = _row_subsample(r_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
     t_s = _row_subsample(t_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
+    r_s = {
+        name: _row_subsample(lf, sample_rows, seed).select(exprs).collect(engine="streaming")
+        for name, lf in r_lfs.items()
+    }
     var_ranges = {
-        name: _combined_quantiles(r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q) for name in RANGED_VARS
+        name: _combined_quantiles([t_s[name].to_numpy(), *(df[name].to_numpy() for df in r_s.values())], _LO_Q, _HI_Q)
+        for name in RANGED_VARS
     }
 
     # Energy-bin edges from exact per-event incident energies (cheap group_by).
     def _incident(lf: pl.LazyFrame) -> np.ndarray:
         return lf.group_by("event_id").agg(pl.col("pre_E").max()).collect(engine="streaming")["pre_E"].to_numpy()
 
-    r_inc, t_inc = _incident(r_lf), _incident(t_lf)
-    energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
+    t_inc = _incident(t_lf)
+    r_inc = {name: _incident(lf) for name, lf in r_lfs.items()}
+    energy_edges = energy_bin_edges(np.concatenate([t_inc, *r_inc.values()]), n_energy_bins)
 
     # Top PDG species and material list (cheap single-column group_bys).
     def _counts(lf: pl.LazyFrame, col: str) -> pl.DataFrame:
         return lf.group_by(col).agg(pl.len().alias("n")).collect(engine="streaming")
 
     pdg_counts = (
-        pl.concat([_counts(r_lf, "pdg"), _counts(t_lf, "pdg")])
+        pl.concat([_counts(t_lf, "pdg"), *(_counts(lf, "pdg") for lf in r_lfs.values())])
         .group_by("pdg")
         .agg(pl.col("n").sum())
         .sort("n", descending=True)
     )
     top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
-    materials = sorted(
-        set(_counts(r_lf, "material")["material"].to_list()) | set(_counts(t_lf, "material")["material"].to_list())
-    )
+    material_set: set[str] = set(_counts(t_lf, "material")["material"].to_list())
+    for lf in r_lfs.values():
+        material_set |= set(_counts(lf, "material")["material"].to_list())
+    materials = sorted(material_set)
 
     # Shower depth / transverse ranges from a subsampled proxy.
     def _proxy(lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
@@ -140,19 +150,20 @@ def build_context(
         )
         return sub["d"].to_numpy(), sub["t"].to_numpy()
 
-    r_d, r_t = _proxy(r_lf)
     t_d, t_t = _proxy(t_lf)
-    d_lo, d_hi = _combined_quantiles(r_d, t_d, _LO_Q, _HI_Q)
+    r_proxy = {name: _proxy(lf) for name, lf in r_lfs.items()}
+    d_lo, d_hi = _combined_quantiles([t_d, *(p[0] for p in r_proxy.values())], _LO_Q, _HI_Q)
     depth_edges = np.linspace(d_lo, d_hi, n_marginal_bins + 1)
-    t_hi = max(float(np.quantile(np.concatenate([r_t, t_t]), _HI_Q)), 1e-6)
+    t_hi = max(float(np.quantile(np.concatenate([t_t, *(p[1] for p in r_proxy.values())]), _HI_Q)), 1e-6)
     transverse_edges = np.linspace(0.0, t_hi, n_marginal_bins + 1)
 
     # Secondary energy range.
-    r_se = secondaries(r_lf, Side.rollout).select("energy")
-    t_se = secondaries(t_all, Side.reference).select("energy")
-    r_se = _row_sample_col(r_se, sample_rows, seed)
-    t_se = _row_sample_col(t_se, sample_rows, seed)
-    sec_energy_range = _combined_quantiles(r_se, t_se, _LO_Q, _HI_Q)
+    t_se = _row_sample_col(secondaries(t_all, Side.reference).select("energy"), sample_rows, seed)
+    r_se = {
+        name: _row_sample_col(secondaries(lf, Side.rollout).select("energy"), sample_rows, seed)
+        for name, lf in r_lfs.items()
+    }
+    sec_energy_range = _combined_quantiles([t_se, *r_se.values()], _LO_Q, _HI_Q)
 
     return Context(
         n_marginal_bins=n_marginal_bins,
@@ -165,8 +176,8 @@ def build_context(
         sec_energy_range=sec_energy_range,
         n_sec_bins=n_sec_bins,
         n_events={
-            "rollout": len(r_inc),
             "reference": len(t_inc),
+            **{name: len(arr) for name, arr in r_inc.items()},
         },
     )
 

giant/analysis/reduce.py

@@ -271,20 +271,3 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
     escaped = per_event["escaped"].fill_null(0.0).to_numpy()
     total = deposited + escaped
     return np.where(total > 0, escaped / total, 0.0)
-
-
-def sec_count_by_event(lf_all: pl.LazyFrame, sec_lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
-    """Per-event secondary count, zero-filled for events that produced none.
-
-    Two bounded per-event ``group_by``s — the full event set (from ``lf_all``)
-    and the secondary counts (from ``sec_lf``, see ``sources.secondaries``) —
-    merged in Python via a dict. Both results are event-granularity (not
-    per-row), so this stays in the same bounded-memory budget as
-    ``event_scalars``; a plain ``group_by`` on ``sec_lf`` alone would silently
-    drop zero-secondary events instead of zero-filling them.
-    """
-    ev = lf_all.select("event_id").unique().collect(engine="streaming")["event_id"].to_numpy()
-    cnt_df = sec_lf.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")
-    cnt = dict(zip(cnt_df["event_id"].to_list(), cnt_df["n"].to_list()))
-    counts = np.array([cnt.get(int(e), 0) for e in ev], dtype=np.int64)
-    return ev, counts

giant/analysis/reduced.py

@@ -11,19 +11,24 @@ import json
 from dataclasses import asdict, dataclass, field
 from pathlib import Path
 
-# Reduced.kind values:
-#   "overlay_hist"   rollout vs reference density histogram over shared edges
+# Reduced.kind values (payload keys a rollout series by name under
+# payload["series"], with the reference — where one exists — kept as one
+# distinguished payload["reference"] entry; see catalog.py's module
+# docstring for the full per-kind payload shape):
+#   "overlay_hist"   N-rollout-series vs reference density histogram over shared edges
 #   "grouped_hist"   one panel per group (energy/pdg/material), each an overlay
-#   "profile"        edep-weighted mean +/- event-RMS vs depth/radius, two series
-#   "bar"            per-category rollout vs reference bars (share / counts)
-#   "single_hist"    one series only (e.g. rollout leakage; reference has none)
-#   "router_gating"  stacked mean MoE gate weight vs energy, rollout + reference
-#   "router_share"   stacked bar of MoE top-1 dispatch share by category
-#   "router_specialization"  max gate weight vs energy, rollout + reference (one
-#                    scalar trend line summarizing "router_gating")
-#   "heatmap"        row x col matrix + colorbar (distance scorecard or a
-#                    predicted-vs-true confusion matrix)
-#   "unavailable"    plot not applicable to this run (e.g. non-MoE checkpoint)
+#   "profile"        edep-weighted mean +/- event-RMS vs depth/radius, N series + reference
+#   "bar"            per-category N-rollout-series vs reference bars (share / counts)
+#   "single_hist"    rollout-only series (e.g. leakage; reference has none)
+#   "router_gating"  stacked mean MoE gate weight vs energy, one rollout+reference
+#                    panel-pair per rollout with an enabled MoE router
+#   "router_share"   stacked bar of MoE top-1 dispatch share by category, one
+#                    panel per rollout with an enabled MoE router
+#   "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)
 
 
 @dataclass

giant/analysis/render.py

@@ -9,20 +9,55 @@ streaming compute.
 For each reduced artifact it writes ``<out>/<family>/<id>.pdf`` plus a sibling
 ``<id>.yaml`` (per-plot gallery metadata) and a per-family ``metadata.yaml``.
 Optionally runs ``gallery generate`` to build the static HTML site.
+
+Every rollout series gets a stable color via ``ps.get_color(i)``, ``i`` being
+its position in ``payload["series"]`` — that position is fixed by the run's
+YAML/``--label`` order (threaded unchanged from ``condor.RunMeta.rollouts``
+through every ``PlotSpec``), so a given rollout keeps the same color across
+every plot in a run. The reference, where a plot has one, always draws in one
+fixed, distinct style (dark ink, dashed) instead of taking a slot in that
+cycle.
 """
 
 from __future__ import annotations
 
+import dataclasses
 import subprocess
 from pathlib import Path
 
 import numpy as np
 import plotstyle as ps
+from matplotlib.colors import LogNorm
 import yaml
 
 from giant.analysis.reduced import Reduced
 
-_SERIES_LABELS = {"rollout": "rollout", "reference": "reference (Geant4)"}
+_REFERENCE_LABEL = "reference (Geant4)"
+
+_TEX_ESCAPE_MAP = {
+    "\\": r"\textbackslash{}",
+    "%": r"\%",
+    "&": r"\&",
+    "#": r"\#",
+    "$": r"\$",
+    "_": r"\_",
+    "{": r"\{",
+    "}": r"\}",
+}
+
+
+def _tex_escape(text: str) -> str:
+    """Escape characters LaTeX treats specially in catalog-authored title/xlabel
+    text (e.g. a literal ``%`` in a "90% of deposited energy" title, which
+    ``usetex`` otherwise reads as a comment marker and aborts the whole figure —
+    see gitea #81). A single pass over the *original* characters, so the
+    backslashes an escape itself introduces (e.g. ``\textbackslash{}``) are
+    never re-escaped."""
+    return "".join(_TEX_ESCAPE_MAP.get(c, c) for c in text)
+
+
+def _ref_color() -> str:
+    return ps.colors.INK["primary"]
 
 
 def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
@@ -33,10 +68,13 @@ def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
     return counts / (total * (edges[1] - edges[0]))
 
 
-def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> None:
-    for key in ("reference", "rollout"):
-        if key in series:
-            ax.stairs(_density(series[key], edges), edges, label=_SERIES_LABELS[key])
+def _overlay(ax, edges: np.ndarray, payload: dict, log_y: bool) -> None:
+    if "reference" in payload:
+        ax.stairs(
+            _density(payload["reference"], edges), edges, label=_REFERENCE_LABEL, color=_ref_color(), linestyle="--"
+        )
+    for i, (name, counts) in enumerate(payload.get("series", {}).items()):
+        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
     if log_y:
         ax.set_yscale("log")
 
@@ -47,8 +85,8 @@ def _router_summary(router_cfg: dict) -> str:
     return f"{router_cfg.get('type', '?')}×{router_cfg.get('n_experts', '?')}"
 
 
-def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
-    """`_figure_params` for a new-shape (nested) `model_config` — has a
+def _figure_params_v2(mc: dict, meta: dict) -> dict:
+    """`_figure_params_single` for a new-shape (nested) `model_config` — has a
     `stage1_model` key. Reports stage 1's architecture (the headline
     generator); stage 2's generator is only added (`mode_s2`) when it
     differs from stage 1's, since a mixed run (the `stage1=flow` +
@@ -71,23 +109,24 @@ def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
     if particle_type is not None:
         params["conditioning"] = particle_type
     params["router"] = _router_summary(s1.get("router") or {})
-    if run_meta.get("training_epoch") is not None:
-        params["epoch"] = run_meta["training_epoch"]
-    if run_meta.get("best_val_loss") is not None:
-        params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
+    if meta.get("training_epoch") is not None:
+        params["epoch"] = meta["training_epoch"]
+    if meta.get("best_val_loss") is not None:
+        params["best_val_loss"] = round(meta["best_val_loss"], 4)
     if mode == "wgan":
         noise_dim = (s1.get("wgan") or {}).get("noise_dim")
         if noise_dim is not None:
             params["noise_dim"] = noise_dim
-    elif run_meta.get("steps") is not None:
-        params["steps"] = run_meta["steps"]
+    elif meta.get("steps") is not None:
+        params["steps"] = meta["steps"]
     return params
 
 
-def _figure_params(run_meta: dict) -> dict:
-    """Curated run identity for the figure subtitle (``new_figure(params=...)``).
+def _figure_params_single(meta: dict) -> dict:
+    """Curated run identity for the figure subtitle (``new_figure(params=...)``),
+    for exactly one rollout's ``plot_meta``.
 
-    ``run_meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
+    ``meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
     carry every threaded model/training/rollout/dataset parameter for
     after-the-fact lookup — this picks only the handful that matter for
     telling figures apart at a glance while flipping through a gallery, since
@@ -99,9 +138,9 @@ def _figure_params(run_meta: dict) -> dict:
     Handles both a v0.2 checkpoint's flat ``model_config`` and a v0.3.0
     nested one (has a ``stage1_model`` key — see ``_figure_params_v2``).
     """
-    mc = run_meta.get("model_config") or {}
+    mc = meta.get("model_config") or {}
     if "stage1_model" in mc:
-        return _figure_params_v2(mc, run_meta)
+        return _figure_params_v2(mc, meta)
 
     mode = mc.get("mode")
     params: dict = {}
@@ -114,18 +153,35 @@ def _figure_params(run_meta: dict) -> dict:
     if mc.get("conditioning") is not None:
         params["conditioning"] = mc["conditioning"]
     params["router"] = _router_summary(mc.get("router") or {})
-    if run_meta.get("training_epoch") is not None:
-        params["epoch"] = run_meta["training_epoch"]
-    if run_meta.get("best_val_loss") is not None:
-        params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
+    if meta.get("training_epoch") is not None:
+        params["epoch"] = meta["training_epoch"]
+    if meta.get("best_val_loss") is not None:
+        params["best_val_loss"] = round(meta["best_val_loss"], 4)
     if mode == "wgan":
         if mc.get("noise_dim") is not None:
             params["noise_dim"] = mc["noise_dim"]
-    elif run_meta.get("steps") is not None:
-        params["steps"] = run_meta["steps"]
+    elif meta.get("steps") is not None:
+        params["steps"] = meta["steps"]
     return params
 
 
+def _figure_params(run_meta: dict) -> dict:
+    """Curated run identity for the figure subtitle.
+
+    A single-rollout run reuses that rollout's ``plot_meta`` (same curated
+    model/training/rollout subset as always — see ``_figure_params_single``);
+    a multi-rollout run instead names the series being compared, since no
+    single ``model_config`` applies to the figure as a whole (each plot's own
+    gallery YAML still carries every rollout's full ``plot_meta`` for
+    after-the-fact lookup, via ``_plot_metadata``).
+    """
+    rollouts = run_meta.get("rollouts") or {}
+    if len(rollouts) == 1:
+        ((_, meta),) = rollouts.items()
+        return _figure_params_single(meta)
+    return {"rollouts": ", ".join(rollouts)} if rollouts else {}
+
+
 def _render_overlay(r: Reduced, params: dict):
     edges = np.asarray(r.payload["edges"])
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
@@ -139,7 +195,8 @@ def _render_overlay(r: Reduced, params: dict):
 def _render_single(r: Reduced, params: dict):
     edges = np.asarray(r.payload["edges"])
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.stairs(_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"])
+    for i, (name, counts) in enumerate(r.payload.get("series", {}).items()):
+        ax.stairs(_density(counts, edges), edges, label=name, color=ps.get_color(i))
     if r.payload.get("log_y"):
         ax.set_yscale("log")
     if r.payload.get("log_x"):
@@ -181,11 +238,17 @@ def _render_profile(r: Reduced, params: dict):
     edges = np.asarray(r.payload["edges"])
     centers = 0.5 * (edges[:-1] + edges[1:])
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    for key in ("reference", "rollout"):
-        mean = np.asarray(r.payload[f"{key}_mean"])
-        std = np.asarray(r.payload[f"{key}_std"])
-        (line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
-        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=line.get_color())
+    if "reference" in r.payload:
+        ref = r.payload["reference"]
+        mean, std = np.asarray(ref["mean"]), np.asarray(ref["std"])
+        color = _ref_color()
+        ax.plot(centers, mean, label=_REFERENCE_LABEL, color=color, linestyle="--")
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
+    for i, (name, side) in enumerate(r.payload.get("series", {}).items()):
+        mean, std = np.asarray(side["mean"]), np.asarray(side["std"])
+        color = ps.get_color(i)
+        ax.plot(centers, mean, label=name, color=color)
+        ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=color)
     ax.set_xlabel(r.xlabel)
     ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
     ps.style_legend(ax, title="source")
@@ -195,10 +258,19 @@ def _render_profile(r: Reduced, params: dict):
 def _render_bar(r: Reduced, params: dict):
     labels = r.payload["labels"]
     x = np.arange(len(labels))
-    width = 0.4
+    series = r.payload.get("series", {})
+    has_ref = "reference" in r.payload
+    n_bars = len(series) + (1 if has_ref else 0)
+    width = 0.8 / max(n_bars, 1)
+    offsets = np.linspace(-0.4 + width / 2, 0.4 - width / 2, n_bars)
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    ax.bar(x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"])
-    ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
+    idx = 0
+    if has_ref:
+        ax.bar(x + offsets[idx], r.payload["reference"], width, label=_REFERENCE_LABEL, color=_ref_color())
+        idx += 1
+    for i, (name, vals) in enumerate(series.items()):
+        ax.bar(x + offsets[idx], vals, width, label=name, color=ps.get_color(i))
+        idx += 1
     ax.set_xticks(x)
     ax.set_xticklabels(labels, rotation=45, ha="right")
     ax.set_ylabel(r.payload.get("ylabel", "value"))
@@ -207,97 +279,144 @@ def _render_bar(r: Reduced, params: dict):
 
 
 def _render_router_gating(r: Reduced, params: dict):
-    n_experts = r.payload["n_experts"]
+    series = r.payload.get("series", {})
+    names = list(series)
     log_x = r.payload.get("log_x", False)
-    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False)
-    flat = axes.ravel()
-    for ax, key in zip(flat, ("rollout", "reference")):
-        side = r.payload.get(key, {})
-        centers = np.asarray(side.get("centers", []))
-        means = np.asarray(side.get("means", []))
-        if len(centers) and means.size:
-            cum = np.zeros(len(centers))
-            for i in range(n_experts):
-                ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}")
-                cum = cum + means[:, i]
-        if log_x:
-            ax.set_xscale("log")
-        ax.set_ylim(0, 1)
-        ax.set_title(_SERIES_LABELS[key], fontsize=8)
-        ax.set_xlabel(r.xlabel)
-    flat[0].set_ylabel("mean gate weight")
-    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=2, squeeze=False)
+    for row, name in enumerate(names):
+        entry = series[name]
+        n_experts = entry["n_experts"]
+        for col, key in enumerate(("rollout", "reference")):
+            ax = axes[row, col]
+            side = entry.get(key, {})
+            centers = np.asarray(side.get("centers", []))
+            means = np.asarray(side.get("means", []))
+            if len(centers) and means.size:
+                cum = np.zeros(len(centers))
+                for i in range(n_experts):
+                    ax.fill_between(centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}")
+                    cum = cum + means[:, i]
+            if log_x:
+                ax.set_xscale("log")
+            ax.set_ylim(0, 1)
+            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
+            ax.set_title(f"{name} — {panel_label}", fontsize=8)
+            if row == len(names) - 1:
+                ax.set_xlabel(r.xlabel)
+        axes[row, 0].set_ylabel("mean gate weight")
+    if names:
+        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
     return fig
 
 
 def _render_router_share(r: Reduced, params: dict):
-    categories = r.payload["categories"]
-    n_experts = r.payload["n_experts"]
-    x = np.arange(len(categories))
-    present = [k for k in ("rollout", "reference") if k in r.payload]
-    fig, axes = ps.new_figure(
-        "slide-16x9",
-        title=r.title,
-        params=params,
-        nrows=1,
-        ncols=len(present),
-        squeeze=False,
-    )
-    flat = axes.ravel()
-    for ax, key in zip(flat, present):
-        side = r.payload[key]
-        shares = np.array([side[c] for c in categories])  # (n_cat, n_experts)
-        bottom = np.zeros(len(categories))
-        for i in range(n_experts):
-            ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
-            bottom += shares[:, i]
-        ax.set_xticks(x)
-        ax.set_xticklabels(categories, rotation=45, ha="right")
-        ax.set_ylim(0, 1)
-        ax.set_title(_SERIES_LABELS[key], fontsize=8)
-    flat[0].set_ylabel("share of rows dispatched to expert")
-    ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
+    series = r.payload.get("series", {})
+    names = list(series)
+    present: tuple[str, ...] = ("rollout", "reference")
+    if names:
+        present = tuple(k for k in ("rollout", "reference") if k in series[names[0]])
+    ncols = max(len(present), 1)
+    fig, axes = ps.new_figure("slide-16x9", title=r.title, params=params, nrows=len(names), ncols=ncols, squeeze=False)
+    for row, name in enumerate(names):
+        entry = series[name]
+        n_experts = entry["n_experts"]
+        cats = entry["categories"]
+        x = np.arange(len(cats))
+        for col, key in enumerate(present):
+            ax = axes[row, col]
+            side = entry.get(key)
+            if side is not None:
+                shares = np.array([side[c] for c in cats])  # (n_cat, n_experts)
+                bottom = np.zeros(len(cats))
+                for i in range(n_experts):
+                    ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
+                    bottom += shares[:, i]
+            ax.set_xticks(x)
+            ax.set_xticklabels(cats, rotation=45, ha="right")
+            ax.set_ylim(0, 1)
+            panel_label = _REFERENCE_LABEL if key == "reference" else "rollout"
+            ax.set_title(f"{name} — {panel_label}", fontsize=8)
+        axes[row, 0].set_ylabel("share of rows dispatched to expert")
+    if names:
+        ps.style_legend(axes[0, 0], title=f"{series[names[0]]['router_type']} router")
     return fig
 
 
 def _render_router_specialization(r: Reduced, params: dict):
     fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    for key in ("reference", "rollout"):
-        side = r.payload.get(key)
-        if side and side["centers"]:
-            ax.plot(side["centers"], side["score"], label=_SERIES_LABELS[key], marker="o", markersize=3)
-    chance = r.payload.get("chance_level")
-    if chance is not None:
-        ax.axhline(chance, linestyle="--", color="gray", label="chance level (1/n_experts)")
+    series = r.payload.get("series", {})
+    chance_levels: set[float] = set()
+    for i, (name, entry) in enumerate(series.items()):
+        color = ps.get_color(i)
+        if entry.get("chance_level") is not None:
+            chance_levels.add(entry["chance_level"])
+        for key, linestyle, label in (
+            ("rollout", "-", name),
+            ("reference", "--", f"{name} ({_REFERENCE_LABEL})"),
+        ):
+            side = entry.get(key)
+            if side and side["centers"]:
+                ax.plot(
+                    side["centers"],
+                    side["score"],
+                    label=label,
+                    color=color,
+                    linestyle=linestyle,
+                    marker="o",
+                    markersize=3,
+                )
+    for lvl in sorted(chance_levels):
+        ax.axhline(lvl, linestyle=":", color="gray")
     if r.payload.get("log_x"):
         ax.set_xscale("log")
     ax.set_ylim(0, 1)
     ax.set_xlabel(r.xlabel)
     ax.set_ylabel("max gate weight")
-    ps.style_legend(ax, title=f"{r.payload.get('router_type', '')} router")
+    ps.style_legend(ax, title="router")
     return fig
 
 
 def _render_heatmap(r: Reduced, params: dict):
-    mat = np.asarray(r.payload["matrix"], dtype=float)
+    series = dict(r.payload["series"])
     row_labels = r.payload["row_labels"]
     col_labels = r.payload["col_labels"]
-    fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
-    im = ax.imshow(
-        mat,
-        origin="upper",
-        aspect="auto",
-        cmap=r.payload.get("cmap", "viridis"),
-        vmin=r.payload.get("vmin"),
-        vmax=r.payload.get("vmax"),
+    # A heatmap-shaped plot is one matrix per rollout, so the reference (when the
+    # comparison has one — the distance scorecard doesn't) becomes one more panel
+    # rather than another line.
+    if r.payload.get("reference") is not None:
+        series["reference"] = r.payload["reference"]
+    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,
     )
-    ax.set_xticks(range(len(col_labels)))
-    ax.set_xticklabels(col_labels, rotation=45, ha="right")
-    ax.set_yticks(range(len(row_labels)))
-    ax.set_yticklabels(row_labels)
-    ax.set_xlabel(r.xlabel)
-    ax.set_ylabel(r.payload.get("ylabel", ""))
-    fig.colorbar(im, ax=ax, label=r.payload.get("cbar_label", "value"))
+    flat = axes.ravel()
+    im = None
+    for ax, name in zip(flat, names):
+        mat = np.asarray(series[name], dtype=float)
+        im = ax.imshow(
+            mat,
+            origin="upper",
+            aspect="auto",
+            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"),
+        )
+        ax.set_xticks(range(len(col_labels)))
+        ax.set_xticklabels(col_labels, rotation=45, ha="right")
+        ax.set_yticks(range(len(row_labels)))
+        ax.set_yticklabels(row_labels)
+        ax.set_xlabel(r.xlabel)
+        if len(names) > 1:
+            ax.set_title(name, fontsize=8)
+    flat[0].set_ylabel(r.payload.get("ylabel", ""))
+    fig.colorbar(im, ax=list(flat), label=r.payload.get("cbar_label", "value"))
     return fig
 
 
@@ -332,8 +451,14 @@ _RENDERERS = {
 
 
 def render(r: Reduced, run_meta: dict | None = None):
-    """Build the matplotlib figure for one reduced artifact (dispatch on kind)."""
-    return _RENDERERS[r.kind](r, _figure_params(run_meta or {}))
+    """Build the matplotlib figure for one reduced artifact (dispatch on kind).
+
+    ``title``/``xlabel`` are LaTeX-escaped here, at the one point every kind's
+    renderer draws them from — ``_plot_metadata`` deliberately keeps using the
+    unescaped ``r`` for the gallery YAML, which isn't LaTeX.
+    """
+    escaped = dataclasses.replace(r, title=_tex_escape(r.title), xlabel=_tex_escape(r.xlabel))
+    return _RENDERERS[r.kind](escaped, _figure_params(run_meta or {}))
 
 
 def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
@@ -392,7 +517,7 @@ def render_all(
         yaml.safe_dump(
             {
                 "title": run_meta.get("title", "GIANT rollout analysis"),
-                "description": "Autoregressive rollout compared against held-out Geant4 reference steps.",
+                "description": "Autoregressive rollout(s) compared against a held-out Geant4 reference steps file.",
                 "experiment": "GIANT",
                 "parameters": {k: v for k, v in run_meta.items() if k != "title"},
             },
@@ -425,8 +550,7 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
     meta = RunMeta.load(run_dir / "run_meta.json")
     run_meta = {
         "title": meta.title,
-        "rollout": meta.rollout,
         "reference": meta.reference,
-        **meta.plot_meta,
+        "rollouts": {ro["name"]: ro["plot_meta"] for ro in meta.rollouts},
     }
     return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)

giant/analysis/router_gating.py

@@ -35,6 +35,7 @@ from giant.analysis.reduced import Reduced
 if TYPE_CHECKING:
     import torch
 
+    from giant.analysis.sources import RolloutSide
     from giant.data.transforms import Normalizer
 
 _SAMPLE_ROWS = 200_000
@@ -218,47 +219,46 @@ def _unavailable(spec_id: str) -> Reduced:
     )
 
 
-def compute_router_gating(
-    checkpoint: str | Path | None,
-    r_phys: pl.LazyFrame,
-    t_phys: pl.LazyFrame,
-    seed: int = 0,
-) -> Reduced:
-    """`Reduced` for the router-gating figure, or an explanatory note if n/a."""
+def _gating_entry(checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, seed: int) -> dict | None:
+    """One rollout's ``router_gating`` panel data, or ``None`` if not a MoE checkpoint."""
     handle = load_router(checkpoint) if checkpoint else None
     if handle is None:
-        return _unavailable("router_gating")
-
+        return None
     sides: dict[str, dict] = {}
     for name, lf in (("rollout", r_phys), ("reference", t_phys)):
         df = _subsample(lf, _SAMPLE_ROWS, seed)
         df, gate = _gate_for_df(handle, df)
         x = df["pre_E"].to_numpy()
         sides[name] = _quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
+    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:
+    """`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 = {}
+    for name, rs in rollouts.items():
+        entry = _gating_entry(rs.checkpoint, rs.phys, t_phys, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_gating")
     return Reduced(
         id="router_gating",
         family="model",
         kind="router_gating",
         title=_TITLES["router_gating"],
         xlabel="pre-step energy [MeV]",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "log_x": True,
-            **sides,
-        },
+        payload={"series": series, "log_x": True},
     )
 
 
-def compute_router_specialization(
-    checkpoint: str | Path | None,
-    r_phys: pl.LazyFrame,
-    t_phys: pl.LazyFrame,
-    seed: int = 0,
-) -> Reduced:
-    """Scalar specialization trend: max gate weight vs energy, per side.
+def _specialization_entry(
+    checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, seed: int
+) -> dict | None:
+    """One rollout's ``router_specialization`` curve data, or ``None`` if not a MoE checkpoint.
 
+    Scalar specialization trend: max gate weight vs energy, per side.
     Summarizes `router_gating`'s full per-expert stacked area into one curve —
     the routing plan's own "how sharp is the boundary here" number (1/n_experts
     = uniform/no specialization, 1.0 = one expert fully owns that energy). Same
@@ -268,8 +268,7 @@ def compute_router_specialization(
     """
     handle = load_router(checkpoint) if checkpoint else None
     if handle is None:
-        return _unavailable("router_specialization")
-
+        return None
     sides: dict[str, dict] = {}
     for name, lf in (("rollout", r_phys), ("reference", t_phys)):
         df = _subsample(lf, _SAMPLE_ROWS, seed)
@@ -282,35 +281,41 @@ def compute_router_specialization(
             sides[name] = {"centers": binned["centers"], "score": score}
         else:
             sides[name] = {"centers": [], "score": []}
+    return {
+        "router_type": handle.router_type,
+        "n_experts": handle.router.n_experts,
+        "chance_level": 1.0 / handle.router.n_experts,
+        **sides,
+    }
 
+
+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 = {}
+    for name, rs in rollouts.items():
+        entry = _specialization_entry(rs.checkpoint, rs.phys, t_phys, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_specialization")
     return Reduced(
         id="router_specialization",
         family="model",
         kind="router_specialization",
         title=_TITLES["router_specialization"],
         xlabel="pre-step energy [MeV]",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "log_x": True,
-            "chance_level": 1.0 / handle.router.n_experts,
-            **sides,
-        },
+        payload={"series": series, "log_x": True},
     )
 
 
-def compute_router_share_by_pdg(
-    checkpoint: str | Path | None,
-    r_phys: pl.LazyFrame,
-    t_phys: pl.LazyFrame,
-    top_pdgs: list[int],
-    seed: int = 0,
-) -> Reduced:
-    """Stacked-bar share of each particle species dispatched to each expert."""
+def _share_by_pdg_entry(
+    checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int
+) -> dict | None:
+    """One rollout's ``router_share_by_pdg`` panel-pair data, or ``None`` if not a MoE checkpoint."""
     handle = load_router(checkpoint) if checkpoint else None
     if handle is None:
-        return _unavailable("router_share_by_pdg")
-
+        return None
     labels = [pdg_label(p) for p in top_pdgs]
     sides: dict[str, dict] = {}
     for name, lf in (("rollout", r_phys), ("reference", t_phys)):
@@ -322,41 +327,36 @@ def compute_router_share_by_pdg(
         else:
             shares = {str(p): [0.0] * handle.router.n_experts for p in top_pdgs}
         sides[name] = {labels[i]: shares[str(p)] for i, p in enumerate(top_pdgs)}
+    return {"router_type": handle.router_type, "n_experts": handle.router.n_experts, "categories": labels, **sides}
 
+
+def compute_router_share_by_pdg(
+    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."""
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _share_by_pdg_entry(rs.checkpoint, rs.phys, t_phys, top_pdgs, seed)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_share_by_pdg")
     return Reduced(
         id="router_share_by_pdg",
         family="model",
         kind="router_share",
         title=_TITLES["router_share_by_pdg"],
         xlabel="particle species",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "categories": labels,
-            **sides,
-        },
+        payload={"series": series},
     )
 
 
-def compute_router_share_by_process(
-    checkpoint: str | Path | None,
-    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.
-
-    Reference-only: ``process`` is the true post-step physics process — a
-    label the rollout side has no equivalent of (see
-    `giant.model.network.ProcessRouter`, which predicts it from pre-step
-    conditioning alone, never observes it at eval time). This plot instead
-    checks *after the fact*, on real data, how well the router's conditioning
-    -based dispatch lines up with the true process.
-    """
+def _share_by_process_entry(checkpoint: str | Path | None, t_phys: pl.LazyFrame, seed: int, top_k: int) -> dict | None:
+    """One rollout checkpoint's ``router_share_by_process`` panel data (reference-only), or ``None`` if not MoE."""
     handle = load_router(checkpoint) if checkpoint else None
     if handle is None:
-        return _unavailable("router_share_by_process")
-
+        return None
     df = _subsample(t_phys, _SAMPLE_ROWS, seed, extra_cols=("process",))
     df, gate = _gate_for_df(handle, df)
     if len(df):
@@ -366,17 +366,39 @@ def compute_router_share_by_process(
         shares = _top1_shares(df["process"].to_numpy(), idx, order, handle.router.n_experts)
     else:
         order, shares = [], {}
+    return {
+        "router_type": handle.router_type,
+        "n_experts": handle.router.n_experts,
+        "categories": order,
+        "reference": {p: shares[p] for p in order},
+    }
 
+
+def compute_router_share_by_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.
+
+    Reference-only: ``process`` is the true post-step physics process — a
+    label the rollout side has no equivalent of (see
+    `giant.model.network.ProcessRouter`, which predicts it from pre-step
+    conditioning alone, never observes it at eval time). This plot instead
+    checks *after the fact*, on real data, how well each checkpoint's router
+    -based dispatch lines up with the true process.
+    """
+    series = {}
+    for name, rs in rollouts.items():
+        entry = _share_by_process_entry(rs.checkpoint, t_phys, seed, top_k)
+        if entry is not None:
+            series[name] = entry
+    if not series:
+        return _unavailable("router_share_by_process")
     return Reduced(
         id="router_share_by_process",
         family="model",
         kind="router_share",
         title=_TITLES["router_share_by_process"],
         xlabel="physics process",
-        payload={
-            "router_type": handle.router_type,
-            "n_experts": handle.router.n_experts,
-            "categories": order,
-            "reference": {p: shares[p] for p in order},
-        },
+        payload={"series": series},
     )

giant/analysis/sources.py

@@ -1,7 +1,12 @@
-"""Canonical world-frame LazyFrame builders for the two sides of a comparison.
+"""Canonical world-frame LazyFrame builders for the two kinds of comparison input.
 
-The analysis compares one autoregressive ``giant rollout`` (the *generated* side)
-against a raw miniCaloSim steps file (the *reference* / real side). Both carry a
+The analysis compares one or more autoregressive ``giant rollout`` runs (the
+*generated* side — one named series each, see ``RolloutSpec``) against a single
+raw miniCaloSim steps file shared by all of them (the *reference* / real side).
+Every rollout is the same *kind* of file regardless of how many there are, so
+``Side`` stays binary: it describes a file's schema (rollout column layout +
+synthetic-termination rows + per-track secondary view, vs. reference
+``sec_*_list`` columns), not series identity. Both kinds carry a
 **shared world-frame physical column subset** under identical names, so no
 renaming or coordinate decode is needed — everything is already in world-frame
 mm / MeV:
@@ -26,6 +31,7 @@ HTCondor workers that have no LaTeX toolchain.
 
 from __future__ import annotations
 
+from dataclasses import dataclass
 from enum import Enum
 from pathlib import Path
 
@@ -82,12 +88,44 @@ SYNTHETIC_TERMINATION_REASONS: frozenset[str] = frozenset(
 
 
 class Side(str, Enum):
-    """Which of the two comparison inputs a file is."""
+    """Which of the two comparison-input *kinds* a file is."""
 
     rollout = "rollout"
     reference = "reference"
 
 
+@dataclass
+class RolloutSpec:
+    """One named rollout input, as fed to ``build_context``/``Bundle.open``.
+
+    ``name`` is the series' identity throughout the rest of the pipeline (a
+    plot's ``payload["series"]`` key, a figure's legend label, its color) —
+    resolved once in ``condor.load_rollout_yamls`` from ``--label`` or the
+    YAML stem, then threaded through unchanged. ``checkpoint`` /
+    ``type_embedding_l1_dist`` are only used by the router/type-embedding
+    diagnostics (``catalog.py``'s ``chunkable=False`` specs).
+    """
+
+    name: str
+    source: str | Path | pl.LazyFrame
+    checkpoint: str | None = None
+    type_embedding_l1_dist: dict | None = None
+
+
+@dataclass
+class RolloutSide:
+    """One rollout's opened frames + per-checkpoint diagnostic inputs (``catalog.Bundle.rollouts`` value)."""
+
+    all: pl.LazyFrame  # rollout, all rows (incl. synthetic termination rows)
+    phys: pl.LazyFrame  # rollout, physical steps only
+    checkpoint: str | None = None  # from the rollout YAML; router_gating only
+    # Diagnostic pre-aggregated at rollout time (giant.rollout.
+    # L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
+    # only. Unlike checkpoint, this needs no live model: it's already a
+    # finished histogram, just passed through.
+    type_embedding_l1_dist: dict | None = None
+
+
 def _check_rollout_metadata(path: Path) -> None:
     """Raise if ``path`` carries coord metadata that isn't the rollout tag.
 
@@ -198,3 +236,34 @@ def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
             pl.col("sec_dz_list").alias("sdz"),
         )
     )
+
+
+def secondaries_by_step(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
+    """One row per produced secondary, tagged with the step that produced it.
+
+    Canonical columns: ``step_key`` (an opaque struct identifying the emitting
+    step) and ``pdg``. ``secondaries`` deliberately drops that link; the
+    per-step multiplicity plots need it, so this is a separate view rather than
+    extra columns every other consumer would pay for.
+
+    - rollout: a secondary's birth row carries ``parent_id`` and a birth
+      position copied verbatim from the parent step's ``post_pos``, so
+      ``(event_id, parent_id, pre_pos)`` identifies the emitting step exactly —
+      no join against the (large) step frame is needed.
+    - reference: secondaries already live on their parent step's row, so the
+      row index *is* the step key. It is only ever used as a group key inside
+      one chunk's own aggregation, so indices repeating across chunks is
+      harmless.
+    """
+    if side is Side.rollout:
+        return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
+            pl.struct("event_id", "parent_id", "pre_x", "pre_y", "pre_z").alias("step_key"),
+            "pdg",
+        )
+    return (
+        lf.select("sec_pdg_list")
+        .with_row_index("_row")
+        .explode("sec_pdg_list")
+        .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

@@ -20,26 +20,37 @@ redesign exists to fix.
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 from giant.analysis.reduced import Reduced
 
+if TYPE_CHECKING:
+    from giant.analysis.sources import RolloutSide
+
 _NOTE_NOT_APPLICABLE = (
-    "not applicable: this rollout's checkpoint doesn't use "
+    "not applicable: none of these rollouts' checkpoints use "
     "stage2_model.particle_type.target='embedding' (or generated no "
     "secondaries), so giant rollout recorded no type_embedding_l1_dist "
-    "diagnostic in its YAML sidecar"
+    "diagnostic in their YAML sidecar"
 )
 
 
-def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
-    """`Reduced` for the type-embedding-distance figure, or an explanatory
-    note if this checkpoint never populated the diagnostic.
+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.
 
-    `l1_dist`: `giant.rollout.L1DistCollector.summary()`'s dict, as recorded
-    in the rollout YAML's `type_embedding_l1_dist` key (`Bundle.
-    type_embedding_l1_dist`) — `{"n", "mean", "std", "min", "max",
-    "hist_edges", "hist_counts"}`.
+    Each rollout's `RolloutSide.type_embedding_l1_dist` is
+    `giant.rollout.L1DistCollector.summary()`'s dict, as recorded in that
+    rollout's YAML `type_embedding_l1_dist` key — `{"n", "mean", "std",
+    "min", "max", "hist_edges", "hist_counts"}`. Every collector uses the
+    same fixed log-spaced edges (`L1DistCollector.__init__`'s defaults, never
+    overridden — see `giant/cli.py`'s rollout command), so it's safe to plot
+    every rollout's counts against the first one's edges.
     """
-    if l1_dist is None:
+    entries = {
+        name: rs.type_embedding_l1_dist for name, rs in rollouts.items() if rs.type_embedding_l1_dist is not None
+    }
+    if not entries:
         return Reduced(
             id="type_embedding_l1_distance",
             family="model",
@@ -49,6 +60,11 @@ def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
             payload={"note": _NOTE_NOT_APPLICABLE},
         )
 
+    edges = next(iter(entries.values()))["hist_edges"]
+    notes = [
+        f"{name}: n={d['n']:,} mean={d['mean']:.4g} std={d['std']:.4g} min={d['min']:.4g} max={d['max']:.4g}"
+        for name, d in entries.items()
+    ]
     return Reduced(
         id="type_embedding_l1_distance",
         family="model",
@@ -56,15 +72,10 @@ def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
         title="Secondary-type embedding L1 distance (predicted vector -> nearest PDG row)",
         xlabel="L1 distance",
         payload={
-            "edges": l1_dist["hist_edges"],
-            "rollout": l1_dist["hist_counts"],
+            "edges": edges,
+            "series": {name: d["hist_counts"] for name, d in entries.items()},
             "log_y": True,
             "log_x": True,
-            "note": (
-                f"n={l1_dist['n']:,}  mean={l1_dist['mean']:.4g}  "
-                f"std={l1_dist['std']:.4g}  min={l1_dist['min']:.4g}  "
-                f"max={l1_dist['max']:.4g}; rollout only, no reference "
-                "concept for a raw pre-decode vector"
-            ),
+            "note": "; ".join(notes) + "; rollout only, no reference concept for a raw pre-decode vector",
         },
     )

giant/checkpoint_io.py

@@ -14,11 +14,18 @@ directly and imported from non-CLI code (`giant.analysis.router_gating`,
 lazily — see that module's docstring for why). Failures raise
 `CheckpointCompatibilityError` with the same wording the CLI has always
 shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
+
+`load_for_inference`'s `config_overrides` (gitea #87) lets a caller change a
+checkpoint's `model_config` at load time, restricted to
+`giant.config.INFERENCE_OVERRIDES` — the allowlist of keys that only affect
+sampling, never module construction/shapes or the preprocessing normalizers/
+vocab maps were fit under.
 """
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+import copy
+from dataclasses import dataclass, field
 from pathlib import Path
 
 import torch
@@ -29,13 +36,45 @@ from giant.constants import K_MAX
 from giant.data.loader import TopNMap
 from giant.data.setup_cache import topnmap_from_json
 from giant.data.transforms import Normalizer
-from giant.model.network import build_models
+from giant.model.network import _migrate_legacy_model_config, build_models
 
 
 class CheckpointCompatibilityError(Exception):
     """Checkpoint is missing something `load_for_inference` needs."""
 
 
+def apply_config_overrides(model_cfg: dict, overrides: dict[str, object] | None) -> dict:
+    """Deep-merge dotted-path *overrides* into a checkpoint's `model_config`,
+    validated against `giant.config.INFERENCE_OVERRIDES` — the allowlist of
+    keys that only affect sampling, not module construction/shapes or the
+    preprocessing normalizers/vocab maps were fit under (gitea #87).
+
+    Migrates a v0.2 flat `model_config` to the nested v0.3 shape first: a
+    dotted path like "stage1_model.ddpm.n_steps" would otherwise silently
+    write into a dict that `build_models` still reads as flat (it decides
+    v0.2-vs-v0.3 by `"stage1_model" in model_config`), suppressing migration.
+
+    Raises `CheckpointCompatibilityError` — never a bare `ValueError` or a
+    downstream `load_state_dict` size mismatch — for an unknown/disallowed
+    path or a value that fails its allowlisted check.
+    """
+    if not overrides:
+        return model_cfg
+    cfg = model_cfg if "stage1_model" in model_cfg else _migrate_legacy_model_config(model_cfg)
+    cfg = copy.deepcopy(cfg)
+    for path, value in overrides.items():
+        spec = gconfig.INFERENCE_OVERRIDES.get(path)
+        if spec is None:
+            allowed = ", ".join(sorted(gconfig.INFERENCE_OVERRIDES))
+            raise CheckpointCompatibilityError(f"{path!r} is not an inference-safe override — allowed paths: {allowed}")
+        try:
+            spec.check(path, value)
+        except ValueError as exc:
+            raise CheckpointCompatibilityError(str(exc)) from exc
+        gconfig._set_path(cfg, path, value)
+    return cfg
+
+
 def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
     """(particle_conditioning, material_conditioning) for
     `giant.data.transforms.build_cond_features`/`build_features` — from
@@ -128,6 +167,7 @@ class InferenceContext:
     model_config: dict
     epoch: int | None
     best_val_loss: float | None
+    config_overrides: dict[str, object] = field(default_factory=dict)
 
 
 def load_for_inference(
@@ -136,6 +176,7 @@ def load_for_inference(
     command_name: str,
     weights: str = "raw",
     require_stage2: bool = True,
+    config_overrides: dict[str, object] | None = None,
 ) -> InferenceContext:
     """Load *checkpoint* and reconstruct everything `predict`/`rollout` need
     to run it forward, on *device*, in `eval()` mode.
@@ -148,6 +189,13 @@ def load_for_inference(
     both stages) or an acceptable `stage2 = None` result — kept as a real
     parameter since `stage{1,2}_model.active` is a real, if currently
     stage1+stage2-only-in-practice, config option.
+
+    *config_overrides* deep-merges dotted `model_config` paths (e.g.
+    `{"stage2_model.n_sec.sampling": "sample"}`) before anything is
+    derived from `model_config` or built — see `apply_config_overrides` for
+    the allowlist and validation. Every derived `InferenceContext` field
+    (`other_policy`, `stage{1,2}_ddpm_steps`, the built modules, ...)
+    reflects the overridden config.
     """
     ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
     for key in ("model_config", "sec_decoder"):
@@ -159,7 +207,7 @@ def load_for_inference(
 
     gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
 
-    model_cfg = ckpt["model_config"]
+    model_cfg = apply_config_overrides(ckpt["model_config"], config_overrides)
     particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
     pdg_topn_map = load_pdg_topn_map(ckpt)
     mat_topn_map = load_mat_topn_map(ckpt)
@@ -231,4 +279,5 @@ def load_for_inference(
         model_config=model_cfg,
         epoch=ckpt.get("epoch"),
         best_val_loss=ckpt.get("best_val_loss"),
+        config_overrides=dict(config_overrides) if config_overrides else {},
     )

giant/cli.py

@@ -1,21 +1,19 @@
+from __future__ import annotations
+
 from collections import Counter
 from datetime import datetime, timezone
 from enum import Enum
 import math
 from pathlib import Path
 import re
-from typing import Optional
+from typing import TYPE_CHECKING, Optional
 import uuid as uuid_mod
 
-import numpy as np
-import yaml
-import torch
 import typer
 from typing_extensions import Annotated
 
-import pyarrow as pa
-import pyarrow.parquet as pq
-from tqdm import tqdm
+if TYPE_CHECKING:
+    import numpy as np
 
 from giant import config as gconfig
 from giant.constants import (
@@ -25,30 +23,11 @@ from giant.constants import (
     PREDICT_SCHEMA_VERSION_KEY,
     ROLLOUT_COORD_VALUE,
 )
-from giant.data.loader import (
-    event_id_offset,
-    find_parquet_files,
-    iter_file_chunks,
-    iter_cond_chunks,
-)
-from giant.data.transforms import (
-    build_features,
-    build_cond_features,
-    energy_simplex_decode,
-    inv_local_frame_rotation,
-    inv_log_transform,
-    reconstruct_post_pos,
-)
-from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
-from giant.geometry import GeometryOracle
+
+# giant.materials only pulls in numpy (no torch/pandas), and MATERIAL_PROPERTIES
+# is needed at decoration time below (a Typer option default), so it can't be
+# deferred into a command body like the rest of this module's heavy imports.
 from giant.materials import MATERIAL_PROPERTIES
-from giant.pipeline import run_train_job
-from giant.rollout import (
-    L1DistCollector,
-    decode_secondary_identity,
-    rollout as run_rollout,
-)
-from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -141,6 +120,23 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
     return out
 
 
+def _parse_set_flags(specs: Optional[list[str]]) -> dict[str, object]:
+    """Parse repeated `--set dotted.path=value` flags into a dict, typing
+    each value with `_coerce_scalar` the same way a TOML file's native types
+    would arrive. Validation against the inference-safe allowlist happens
+    downstream in `giant.checkpoint_io.apply_config_overrides` — this only
+    parses syntax.
+    """
+    out: dict[str, object] = {}
+    for spec in specs or []:
+        path, sep, val = spec.partition("=")
+        if not sep:
+            typer.echo(f"error: --set {spec!r} must be 'dotted.path=value'", err=True)
+            raise typer.Exit(1)
+        out[path] = _coerce_scalar(val)
+    return out
+
+
 def _router_cli_overrides(
     router: bool | None,
     router_type: str | None,
@@ -193,6 +189,8 @@ def _write_prediction_ref(
     comment: str | None = None,
 ) -> Path:
     """Write a YAML sidecar in the checkpoint directory and return its path."""
+    import yaml
+
     ref = {
         "prediction_id": pred_uuid,
         "output": str(out),
@@ -622,6 +620,10 @@ def train(
     ] = None,
 ) -> None:
     """Train the GIANT surrogate model."""
+    import torch
+
+    from giant.pipeline import run_train_job
+
     batch_size_auto = False
     batch_size_value: Optional[int] = None
     if batch_size is not None:
@@ -1020,8 +1022,36 @@ def predict(
             help="Free-text note recorded in the prediction's YAML sidecar",
         ),
     ] = None,
+    set_: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--set",
+            help="Override a sampling-only model_config key on this checkpoint, "
+            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
+            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
+        ),
+    ] = None,
 ) -> None:
     """Run trained model on a parquet file and save predictions."""
+    import numpy as np
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import torch
+    from tqdm import tqdm
+
+    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
+    from giant.data.loader import event_id_offset, find_parquet_files, iter_cond_chunks, iter_file_chunks
+    from giant.data.transforms import (
+        build_cond_features,
+        build_features,
+        energy_simplex_decode,
+        inv_local_frame_rotation,
+        inv_log_transform,
+        reconstruct_post_pos,
+    )
+    from giant.rollout import decode_secondary_identity
+    from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
+
     batch_size_auto = False
     batch_size_value: Optional[int] = None
     if batch_size.strip().lower() == "auto":
@@ -1040,8 +1070,11 @@ def predict(
     typer.echo(f"device: {_device}")
 
     # --- Load checkpoint ---
+    config_overrides = _parse_set_flags(set_)
     try:
-        ctx = load_for_inference(checkpoint, _device, "predict", weights=weights.value)
+        ctx = load_for_inference(
+            checkpoint, _device, "predict", weights=weights.value, config_overrides=config_overrides
+        )
     except CheckpointCompatibilityError as exc:
         typer.echo(f"error: {exc}", err=True)
         raise typer.Exit(1)
@@ -1314,6 +1347,10 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
     the codebase's convention for the primary (a secondary always carries less
     energy than its parent). See giant/analysis/reduce.py:entry_axis.
     """
+    import numpy as np
+
+    from giant.data.loader import event_id_offset, iter_cond_chunks
+
     best_E: dict[int, float] = {}
     best: dict[int, tuple] = {}
     for file_idx, path in enumerate(files):
@@ -1407,8 +1444,28 @@ def rollout(
         Optional[int],
         typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
     ] = None,
+    set_: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--set",
+            help="Override a sampling-only model_config key on this checkpoint, "
+            "'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
+            "for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
+        ),
+    ] = None,
 ) -> None:
     """Roll the surrogate forward into full showers (autoregressive)."""
+    import numpy as np
+    import pyarrow as pa
+    import pyarrow.parquet as pq
+    import torch
+    import yaml
+
+    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
+    from giant.data.loader import find_parquet_files
+    from giant.geometry import GeometryOracle
+    from giant.rollout import L1DistCollector, rollout as run_rollout
+
     if seed is not None:
         torch.manual_seed(seed)
         np.random.seed(seed)
@@ -1416,8 +1473,11 @@ def rollout(
     _device = torch.device(device) if device else gconfig.auto_device()
     typer.echo(f"device: {_device}")
 
+    config_overrides = _parse_set_flags(set_)
     try:
-        ctx = load_for_inference(checkpoint, _device, "rollout", weights=weights.value)
+        ctx = load_for_inference(
+            checkpoint, _device, "rollout", weights=weights.value, config_overrides=config_overrides
+        )
     except CheckpointCompatibilityError as exc:
         typer.echo(f"error: {exc}", err=True)
         raise typer.Exit(1)
@@ -1532,6 +1592,7 @@ def rollout(
             # model knob (router type/n_experts, noise_dim, vocab sizes, ...)
             # is available downstream without touching this command again.
             "model_config": dict(model_cfg),
+            "config_overrides": dict(ctx.config_overrides),
             "training_epoch": ctx.epoch,
             "best_val_loss": ctx.best_val_loss,
             # [train]/[meta] from the sibling config.toml (giant.config.save_config)
@@ -1556,10 +1617,23 @@ app.add_typer(analyze_app, name="analyze")
 
 @analyze_app.command("prep")
 def analyze_prep(
-    rollout_yaml: Annotated[
-        Path,
-        typer.Argument(help="giant rollout YAML sidecar (names the rollout + reference files)"),
+    rollout_yamls: Annotated[
+        list[Path],
+        typer.Argument(
+            help="giant rollout YAML sidecar(s) (names the rollout + reference files). "
+            "Multiple compare N rollouts against one shared reference — every YAML must "
+            "name the same `dataset`."
+        ),
     ],
+    label: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--label",
+            help="Series name for a rollout YAML, positionally matched to it — give none, "
+            'or exactly one per YAML. Defaults to the YAML stem (or "rollout" for a '
+            "single YAML).",
+        ),
+    ] = None,
     run_dir: Annotated[
         Optional[Path],
         typer.Option(
@@ -1576,14 +1650,15 @@ def analyze_prep(
         typer.Option("--chunks", help="Split each plot's data into this many event_id chunks"),
     ] = 1,
 ) -> None:
-    """Read the rollout YAML → shared.json + run_meta.json in the run directory."""
+    """Read the rollout YAML(s) → shared.json + run_meta.json in the run directory."""
     from giant.analysis import prep
 
     path = prep(
-        rollout_yaml,
+        rollout_yamls,
         run_dir,
         n_chunks=chunks,
         default_base=Path.cwd() / "analysis_runs",
+        labels=label,
         n_energy_bins=n_energy_bins,
         n_marginal_bins=n_marginal_bins,
         top_k_pdg=top_k_pdg,
@@ -1665,8 +1740,23 @@ def analyze_metrics(
 
 @analyze_app.command("submit")
 def analyze_submit(
-    rollout_yaml: Annotated[Path, typer.Argument(help="giant rollout YAML sidecar")],
+    rollout_yamls: Annotated[
+        list[Path],
+        typer.Argument(
+            help="giant rollout YAML sidecar(s). Multiple compare N rollouts against one "
+            "shared reference — every YAML must name the same `dataset`."
+        ),
+    ],
     accounting_group: Annotated[str, typer.Option("--accounting-group")],
+    label: Annotated[
+        Optional[list[str]],
+        typer.Option(
+            "--label",
+            help="Series name for a rollout YAML, positionally matched to it — give none, "
+            'or exactly one per YAML. Defaults to the YAML stem (or "rollout" for a '
+            "single YAML).",
+        ),
+    ] = None,
     run_dir: Annotated[
         Optional[Path],
         typer.Option(
@@ -1699,10 +1789,11 @@ def analyze_submit(
     from giant.analysis import SubmitConfig, prep, write_submit
 
     path = prep(
-        rollout_yaml,
+        rollout_yamls,
         run_dir,
         n_chunks=chunks,
         default_base=Path.cwd() / "analysis_runs",
+        labels=label,
         n_energy_bins=n_energy_bins,
         n_marginal_bins=n_marginal_bins,
         top_k_pdg=top_k_pdg,

giant/config.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import copy
 import difflib
 import hashlib
@@ -10,12 +12,12 @@ from dataclasses import dataclass, field
 from datetime import datetime, timezone
 from enum import Enum
 from pathlib import Path
-
-import numpy as np
-import torch
+from typing import TYPE_CHECKING
 
 from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
-from giant.model.history import HISTORY_REGISTRY
+
+if TYPE_CHECKING:
+    import torch
 
 
 class Conditioning(str, Enum):
@@ -404,6 +406,16 @@ class Stage2RouterConfig(RouterConfig):
         return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()}
 
 
+# stage2_model.n_sec.sampling choices — single source of truth for both
+# validate_config's train-time check and INFERENCE_OVERRIDES below.
+STOP_SAMPLING_CHOICES = ("greedy", "sample")
+
+# stage2_model.particle_type.other_policy choices — see ParticleTypeConfig's
+# docstring for what each means; only documented there until now, since
+# nothing validated it at train time.
+OTHER_POLICY_CHOICES = ("sample", "modal", "drop")
+
+
 @dataclass(frozen=True)
 class NSecConfig:
     # "head": a classifier over {0..k_max} on the condition encoding alone
@@ -424,11 +436,15 @@ class NSecConfig:
     # n_sec head was trained against Stage 1's own ConditionEncoder output and so has
     # to stay attached there, not just be labeled as such).
     owner: str = "stage2"
-    # mode="stop_token" only: how sample_secondaries_ar turns a slot's stop logit into a
-    # stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample":
-    # a Bernoulli draw at sigmoid(logit) (a real sample from the learned length
-    # distribution, at the cost of an extra RNG draw per slot).
-    stop_sampling: str = "greedy"
+    # How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an
+    # actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits
+    # (deterministic — the conditional mode, not a sample); "sample" is a categorical draw
+    # from softmax(logits) (a real sample from the learned count distribution). mode=
+    # "stop_token": "greedy" is sigmoid(stop_logit) >= 0.5 per slot (deterministic);
+    # "sample" is a Bernoulli draw at sigmoid(stop_logit) per slot. Renamed from
+    # "stop_sampling" (gitea #86), which is still accepted as a deprecated alias since it
+    # appears in existing checkpoints' model_config.
+    sampling: str = "greedy"
 
     @classmethod
     def from_dict(cls, d: dict | None) -> "NSecConfig":
@@ -437,7 +453,7 @@ class NSecConfig:
             mode=d.get("mode", "head"),
             lambda_weight=d.get("lambda", 0.1),
             owner=d.get("owner", "stage2"),
-            stop_sampling=d.get("stop_sampling", "greedy"),
+            sampling=d.get("sampling", d.get("stop_sampling", "greedy")),
         )
 
     def to_dict(self) -> dict:
@@ -445,7 +461,7 @@ class NSecConfig:
             "mode": self.mode,
             "lambda": self.lambda_weight,
             "owner": self.owner,
-            "stop_sampling": self.stop_sampling,
+            "sampling": self.sampling,
         }
 
 
@@ -927,6 +943,8 @@ def git_hash() -> str:
 
 
 def auto_device() -> torch.device:
+    import torch
+
     if torch.cuda.is_available():
         return torch.device("cuda")
     if torch.backends.mps.is_available():
@@ -972,6 +990,8 @@ def estimate_batch_size(
     inference (e.g. `predict`), which uses a much lower per-sample memory
     calibration since there's no backward graph or optimizer state.
     """
+    import torch
+
     if device.type != "cuda":
         raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}")
     device_index = device.index if device.index is not None else torch.cuda.current_device()
@@ -1075,6 +1095,27 @@ def _set_path(d: dict, dotted: str, value) -> None:
     cur[parts[-1]] = value
 
 
+def _pop_path(d: dict, dotted: str) -> None:
+    """Remove a dotted path from a nested dict, if present. No-op if any
+    component along the path is missing."""
+    parts = dotted.split(".")
+    cur = d
+    for part in parts[:-1]:
+        if not isinstance(cur, dict) or part not in cur:
+            return
+        cur = cur[part]
+    if isinstance(cur, dict):
+        cur.pop(parts[-1], None)
+
+
+# Config keys renamed within v0.3 itself (not part of the v0.2->v0.3 migration
+# above) — normalized by migrate_config so a config.toml still using an older
+# v0.3 key name keeps passing validate_config_keys.
+_RENAMED_KEYS = {
+    "stage2_model.n_sec.stop_sampling": "stage2_model.n_sec.sampling",  # gitea #86
+}
+
+
 def _deep_merge(base: dict, override: dict) -> dict:
     """Recursively merge `override` onto a copy of `base`.
 
@@ -1093,6 +1134,72 @@ def _deep_merge(base: dict, override: dict) -> dict:
     return result
 
 
+@dataclass(frozen=True)
+class InferenceOverride:
+    """One dotted `model_config` path that `giant.checkpoint_io.load_for_inference`
+    is allowed to change on an already-trained checkpoint, without retraining.
+
+    A path only belongs here if it affects neither module construction/tensor
+    shapes nor the data preprocessing the normalizers/vocab maps were fit
+    under — see the module docstring on `giant.model.summary` for the class
+    of key this targets (`_fingerprint`'s "plain scalar attribute" leaves),
+    and `giant.checkpoint_io.apply_config_overrides` for where this is used.
+    """
+
+    why: str
+    choices: tuple[str, ...] | None = None
+    minimum: float | None = None
+    numeric: bool = False  # int/float leaf (vs. str, the default)
+
+    def check(self, path: str, value: object) -> None:
+        if self.choices is not None:
+            if value not in self.choices:
+                raise ValueError(f"{path} = {value!r} — must be one of {self.choices}")
+            return
+        if self.numeric:
+            if isinstance(value, bool) or not isinstance(value, (int, float)):
+                raise ValueError(f"{path} = {value!r} — must be a number")
+            if self.minimum is not None and value < self.minimum:
+                raise ValueError(f"{path} = {value!r} — must be >= {self.minimum}")
+
+
+# Inference-safe dotted `model_config` paths — the allowlist gitea #87 asked
+# for, so a typo or a shape-bearing key (e.g. "stage1_model.hidden_dim")
+# raises a clear CheckpointCompatibilityError instead of surfacing as an
+# opaque load_state_dict size mismatch later. Extend this table, not a
+# per-call bypass, when a new inference-only key needs the capability.
+INFERENCE_OVERRIDES: dict[str, InferenceOverride] = {
+    "stage2_model.n_sec.sampling": InferenceOverride(
+        why="giant.sample's n_sec head/stop-token sampling reads this at sample time only (gitea #86)",
+        choices=STOP_SAMPLING_CHOICES,
+    ),
+    "stage1_model.ddpm.n_steps": InferenceOverride(
+        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
+        numeric=True,
+        minimum=1,
+    ),
+    "stage2_model.ddpm.n_steps": InferenceOverride(
+        why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
+        numeric=True,
+        minimum=1,
+    ),
+    "stage2_model.particle_type.other_policy": InferenceOverride(
+        why="giant.rollout resolves an 'other'-bucket secondary's PDG code with this at rollout time",
+        choices=OTHER_POLICY_CHOICES,
+    ),
+    "stage1_model.router.temperature": InferenceOverride(
+        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
+        numeric=True,
+        minimum=1e-6,
+    ),
+    "stage2_model.router.temperature": InferenceOverride(
+        why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
+        numeric=True,
+        minimum=1e-6,
+    ),
+}
+
+
 @dataclass(frozen=True)
 class FlagSpec:
     """One CLI flag's mapping into the config-overrides tree.
@@ -1271,11 +1378,21 @@ def migrate_config(cfg: dict) -> dict:
     (which additionally carries n_sec_head ownership and needs
     `network.build_models`'s cooperation) is a separate migration surface,
     deferred to the network.py refactor.
-    """
-    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
-        return copy.deepcopy(cfg)
 
+    Independently of the v0.2/v0.3 branch below, `_RENAMED_KEYS` normalizes
+    keys renamed within v0.3 itself (e.g. `stop_sampling` -> `sampling`,
+    gitea #86) so a config.toml written against an older v0.3 key name still
+    passes `validate_config_keys`.
+    """
     cfg = copy.deepcopy(cfg)
+    for old_path, new_path in _RENAMED_KEYS.items():
+        if _get_path(cfg, old_path) is not None and _get_path(cfg, new_path) is None:
+            _set_path(cfg, new_path, _get_path(cfg, old_path))
+        _pop_path(cfg, old_path)
+
+    if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
+        return cfg
+
     old_train = cfg.pop("train", {})
     old_model = cfg.pop("model", {})
     old_router = dict(old_model.pop("router", {}))
@@ -1512,9 +1629,9 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
                 "conditioning to hang an EOS decision off"
             )
 
-    stop_sampling = _get_path(cfg, "stage2_model.n_sec.stop_sampling")
-    if stop_sampling not in ("greedy", "sample"):
-        raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
+    n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
+    if n_sec_sampling not in STOP_SAMPLING_CHOICES:
+        raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")
 
     precision = _get_path(cfg, "train.precision")
     if precision not in ("fp32", "bf16"):
@@ -1569,6 +1686,8 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
                 "'energy_desc' (the only implemented ordering; see "
                 "AutoregressiveConfig.order's docstring)"
             )
+        from giant.model.history import HISTORY_REGISTRY
+
         history = _get_path(cfg, "stage2_model.autoregressive.history")
         if history not in HISTORY_REGISTRY:
             raise ValueError(
@@ -1770,6 +1889,9 @@ def resolve_default_out_dir(cfg: dict, base: Path = Path("checkpoints")) -> Path
 
 
 def seed_everything(seed: int) -> None:
+    import numpy as np
+    import torch
+
     random.seed(seed)
     np.random.seed(seed)
     torch.manual_seed(seed)
@@ -1823,6 +1945,8 @@ def build_run_meta(
     n_val_events: int,
     n_train_steps: int,
 ) -> dict:
+    import torch
+
     return {
         "config_version": CONFIG_VERSION,
         "git_hash": git_hash(),

giant/model/builders.py

@@ -138,7 +138,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
                 build_stop_head=stop_token,
-                stop_sampling=s2_spec.n_sec.stop_sampling,
+                n_sec_sampling=s2_spec.n_sec.sampling,
                 stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
             )
         else:
@@ -168,6 +168,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
                 cond_enc=shared_cond_enc,
                 n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
                 type_head_cfg=s2_spec.heads.type.to_dict(),
+                n_sec_sampling=s2_spec.n_sec.sampling,
             )
 
     return result

giant/model/models.py

@@ -370,6 +370,7 @@ class Stage2OneShot(StageModel):
         cond_enc: ConditionEncoder | None = None,
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
+        n_sec_sampling: str = "greedy",
     ) -> None:
         super().__init__(
             pdg_vocab,
@@ -383,6 +384,7 @@ class Stage2OneShot(StageModel):
             particle_type_cfg=particle_type_cfg,
             cond_enc=cond_enc,
         )
+        self.n_sec_sampling = n_sec_sampling
         self._build_context_fusion(x_dim, context_dim, cond_out_dim)
         target = self.particle_type_cfg.target
         type_head_out_dim = None if target == "physical" else k_max * self.type_dim
@@ -498,7 +500,7 @@ class Stage2Autoregressive(StageModel):
         n_sec_head_cfg: dict | None = None,
         type_head_cfg: dict | None = None,
         build_stop_head: bool = False,
-        stop_sampling: str = "greedy",
+        n_sec_sampling: str = "greedy",
         stop_head_cfg: dict | None = None,
     ) -> None:
         super().__init__(
@@ -514,7 +516,7 @@ class Stage2Autoregressive(StageModel):
             cond_enc=cond_enc,
         )
         self.history_kind = history
-        self.stop_sampling = stop_sampling
+        self.n_sec_sampling = n_sec_sampling
         self.context_adapter = ContextAdapter(x_dim, context_dim)
         self.base_fuse = nn.Sequential(
             nn.Linear(cond_out_dim + context_dim, cond_out_dim),

giant/model/summary.py

@@ -37,7 +37,7 @@ from dataclasses import dataclass, field
 
 import torch.nn as nn
 
-from giant.config import _get_path, _set_path, leaf_paths
+from giant.config import INFERENCE_OVERRIDES, _get_path, _set_path, leaf_paths
 from giant.model.builders import build_critics, build_models
 from giant.model.trunks import RoutedTrunk
 
@@ -111,6 +111,7 @@ class ModelSummary:
     pdg_vocab: int
     mat_vocab: int
     vocab_caveats: list[str] = field(default_factory=list)
+    overridable: list[str] = field(default_factory=list)
 
 
 def _build_model_config(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict:
@@ -140,7 +141,7 @@ def _fingerprint(modules: dict[str, nn.Module]) -> list:
     exist, every parameter's/buffer's shape+dtype (never values — those are
     randomly initialized and irrelevant to *structure*), and every plain
     scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
-    .stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
+    .n_sec_sampling`, `EnergyRouter.temperature`) — this is what makes a
     non-parametric key's effect on construction observable."""
     sig = []
     for stage_name, module in modules.items():
@@ -234,6 +235,8 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
         else:
             inert.append(path)
 
+    overridable = sorted(p for p in in_scope if p in INFERENCE_OVERRIDES)
+
     return ModelSummary(
         modules=modules,
         consumed=sorted(consumed),
@@ -242,6 +245,7 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
         pdg_vocab=pdg_vocab,
         mat_vocab=mat_vocab,
         vocab_caveats=_vocab_caveats(cfg),
+        overridable=overridable,
     )
 
 
@@ -309,6 +313,12 @@ def render_summary(summary: ModelSummary) -> str:
     else:
         lines.append("  (none)")
 
+    if summary.overridable:
+        lines.append("")
+        lines.append("inference-overridable without retraining (giant predict/rollout --set):")
+        for path in summary.overridable:
+            lines.append(f"  {path}  ({INFERENCE_OVERRIDES[path].why})")
+
     if summary.vocab_caveats:
         lines.append("")
         lines.append("vocab placeholder caveats:")

giant/sample.py

@@ -261,7 +261,7 @@ def sample_secondaries_ar(
     slot's own stop logit (`predict_stop`, evaluated on the same prefix
     conditioning as the token itself — see `predict_type`'s docstring for
     why this needs no extra state) decides whether generation should have
-    already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
+    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
@@ -345,7 +345,7 @@ def sample_secondaries_ar(
                 slot_idx,
                 hist=hist,
             ).squeeze(1)
-            if sec_decoder.stop_sampling == "sample":
+            if sec_decoder.n_sec_sampling == "sample":
                 stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
             else:
                 stop_now = stop_logit >= 0.0
@@ -496,7 +496,12 @@ def resolve_n_sec(
 
     Raises if neither stage owns any n_sec mechanism at all — the only way
     that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
-    rollout-/predict-capable checkpoint."""
+    rollout-/predict-capable checkpoint.
+
+    `n_sec.mode = "head"` resolves the classifier logits per
+    `sec_decoder.n_sec_sampling`: "greedy" (default) takes the conditional
+    mode via argmax; "sample" draws a real sample from the learned count
+    distribution via `torch.multinomial` on the softmax — see gitea #86."""
     if n_sec_pred is not None:
         return n_sec_pred
     if getattr(sec_decoder, "stop_head", None) is not None:
@@ -508,4 +513,6 @@ def resolve_n_sec(
             "'truth' is standalone-evaluation-only"
         )
     logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
+    if sec_decoder.n_sec_sampling == "sample":
+        return torch.multinomial(logits.softmax(dim=-1), 1).squeeze(-1)
     return logits.argmax(dim=-1)

giant/tools/dwarf.py

@@ -5,6 +5,8 @@ simulation-fanout tools into one Typer app so there's a single command name
 (and `--help`) to remember instead of five differently-hyphenated ones.
 """
 
+from __future__ import annotations
+
 import os
 from enum import Enum
 from pathlib import Path
@@ -14,20 +16,15 @@ import typer
 from typing_extensions import Annotated
 
 from giant.config import Conditioning
-from giant.tools.bump_dataset_version import (
-    run_bump_gen,
-    run_bump_schema,
-    run_create_manifest,
-    run_status,
-    run_update_manifest,
-)
-from giant.tools.create_root_files import run_make_root
-from giant.tools.geometry_oracle import run_build_geometry_oracle
-from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
-from giant.tools.migrate_geant_steps import run_migration
-from giant.tools.steps_to_parquet import convert_steps_to_parquet
-from giant.tools.steps_to_parquet_parallel import run_parallel_job
-from giant.tools.warm_setup_cache import run_warm_setup_cache
+
+# DATA_DEFAULT/SCAN_DIR_DEFAULT are Typer option defaults (evaluated at
+# decoration time below), so that one name has to stay eager — the module
+# itself is stdlib-only, so it costs nothing. Every other giant.tools.*
+# import here is deferred into the one command body that uses it, since
+# several (steps_to_parquet: uproot/awkward/polars; warm_setup_cache:
+# giant.pipeline -> torch; geometry_oracle: pandas) are expensive and
+# `dwarf --help`/tab-completion shouldn't pay for all of them upfront.
+from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT
 
 app = typer.Typer(no_args_is_help=True)
 
@@ -121,6 +118,9 @@ def convert(
     ] = None,
 ) -> None:
     """Convert ROOT Steps tree(s) to Parquet."""
+    from giant.tools.steps_to_parquet import convert_steps_to_parquet
+    from giant.tools.steps_to_parquet_parallel import run_parallel_job
+
     if jobs < 1:
         typer.echo("error: --jobs must be >= 1", err=True)
         raise typer.Exit(1)
@@ -183,6 +183,8 @@ def migrate(
     ] = False,
 ) -> None:
     """One-time migration into the versioned raw/processed/pools/derived layout."""
+    from giant.tools.migrate_geant_steps import run_migration
+
     run_migration(str(root), execute=execute, copy=copy)
 
 
@@ -207,6 +209,8 @@ def bump_gen(
     root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new raw generation."""
+    from giant.tools.bump_dataset_version import run_bump_gen
+
     run_bump_gen(
         kind=kind,
         reason=reason,
@@ -240,6 +244,8 @@ def bump_schema(
     root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """Cut a new schema within a gen."""
+    from giant.tools.bump_dataset_version import run_bump_schema
+
     run_bump_schema(
         kind=kind,
         gen=gen,
@@ -257,6 +263,8 @@ def status(
     root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
 ) -> None:
     """List existing gens/schemas per kind."""
+    from giant.tools.bump_dataset_version import run_status
+
     run_status(str(root))
 
 
@@ -281,6 +289,8 @@ def update_manifest(
     ] = False,
 ) -> None:
     """Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
+    from giant.tools.bump_dataset_version import run_update_manifest
+
     run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)
 
 
@@ -311,6 +321,8 @@ def create_manifest(
     ] = False,
 ) -> None:
     """Create a new manifest from a list of parquet files."""
+    from giant.tools.bump_dataset_version import run_create_manifest
+
     run_create_manifest(
         [str(f) for f in files],
         execute=execute,
@@ -358,6 +370,8 @@ def make_root(
     ] = False,
 ) -> None:
     """Generate new ROOT shards via a minicalosim executable."""
+    from giant.tools.create_root_files import run_make_root
+
     _warn_if_exceeds_shared_quota(jobs, "--jobs")
     run_make_root(
         executable=executable,
@@ -423,6 +437,8 @@ def build_geometry_oracle(
     ] = 2000,
 ) -> None:
     """Fit a position -> (material, layer_id) oracle for `giant rollout`."""
+    from giant.tools.geometry_oracle import run_build_geometry_oracle
+
     run_build_geometry_oracle(
         data=data,
         out=out,
@@ -515,6 +531,8 @@ def warm_cache(
     such entry across every run) skips straight to training. See
     giant/data/setup_cache.py.
     """
+    from giant.tools.warm_setup_cache import run_warm_setup_cache
+
     flag_overrides = {
         "--val-fraction": val_fraction,
         "--seed": seed,
@@ -557,6 +575,8 @@ def hparam_scan(
     dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
 ) -> None:
     """Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
+    from giant.tools.hparam_scan import run_hparam_scan
+
     run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)
 
 

giant/tools/profile_analysis_costs.py

@@ -28,6 +28,7 @@ import polars as pl
 from giant.analysis.catalog import catalog_ids, get_spec
 from giant.analysis.condor import compute_reduced
 from giant.analysis.context import build_context
+from giant.analysis.sources import RolloutSpec
 
 # Row counts (per side) to benchmark at. Kept in local memory/CPU range so the
 # whole sweep finishes in about a minute; the fit is linear so it extrapolates
@@ -165,11 +166,10 @@ def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path)
     t0 = time.perf_counter()
     compute_reduced(
         spec_id,
-        rollout,
+        [{"name": "rollout", "path": str(rollout)}],
         reference,
         shared,
         out,
-        checkpoint=None,
         chunk_index=0,
         n_chunks=1,
     )
@@ -191,7 +191,7 @@ def main() -> None:
 
             shared = tmp_path / f"shared_{n_side}.json"
             ctx = build_context(
-                rollout,
+                [RolloutSpec(name="rollout", source=rollout)],
                 reference,
                 n_energy_bins=4,
                 n_marginal_bins=50,

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.8"
+version = "0.3.15"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"

tests/test_analysis_reduce.py

@@ -13,6 +13,7 @@ from giant.analysis.sources import (
     open_side,
     physical_steps,
     secondaries,
+    secondaries_by_step,
 )
 from giant.data.loader import EVENT_ID_FILE_STRIDE
 
@@ -159,18 +160,17 @@ def test_secondaries_rollout_vs_reference_align():
     assert t["pdg"].to_list() == [22, 22]
 
 
-def test_sec_count_by_event_zero_fills_events_with_no_secondaries():
-    r_phys = physical_steps(_rollout_frame(), Side.rollout)
-    r_sec = secondaries(_rollout_frame(), Side.rollout)
-    ev, n = R.sec_count_by_event(r_phys, r_sec)
-    # event 1 has one secondary track; event 2 has none and must still appear (as 0),
-    # not silently drop out of a plain group_by on the secondaries frame alone.
-    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}
+def test_secondaries_by_step_keys_each_secondary_to_its_emitting_step():
+    r = secondaries_by_step(_rollout_frame(), Side.rollout).collect()
+    assert r["pdg"].to_list() == [22]
+    # the rollout key is (event_id, parent_id, birth position) — the parent
+    # step's post_pos, copied verbatim onto the child's birth row.
+    assert r["step_key"][0] == {"event_id": 1, "parent_id": 0, "pre_x": 0.0, "pre_y": 0.0, "pre_z": 1.0}
 
-    t_all = _reference_frame()
-    t_sec = secondaries(t_all, Side.reference)
-    ev, n = R.sec_count_by_event(t_all, t_sec)
-    assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1}
+    t = secondaries_by_step(_reference_frame(), Side.reference).collect()
+    assert t["pdg"].to_list() == [22, 22]
+    # one row per emitting step; the empty-list step drops out entirely
+    assert [k["_row"] for k in t["step_key"]] == [0, 2]
 
 
 def test_leakage_fraction():

tests/test_catalog.py

@@ -10,16 +10,25 @@ from giant.analysis.catalog import (
     Bundle,
     PlotSpec,
     _containment_depths,
-    _integer_confusion,
     _ks_statistic,
 )
 from giant.analysis.context import Context, build_context
+from giant.analysis.grouping import pdg_label
+from giant.analysis.sources import RolloutSpec
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 
 def _build_ctx() -> Context:
     r, t = _rollout_frame(), _reference_frame()
-    return build_context(r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
+    return build_context(
+        [RolloutSpec("rollout", r)], t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
+    )
+
+
+def _two_rollout_specs() -> list[RolloutSpec]:
+    # Two distinct rollout sources so multi-series merging/finalize code is
+    # exercised even though the underlying frame is the same fixture.
+    return [RolloutSpec("flow", _rollout_frame()), RolloutSpec("wgan", _rollout_frame())]
 
 
 @pytest.fixture(scope="module")
@@ -27,9 +36,20 @@ def ctx() -> Context:
     return _build_ctx()
 
 
+@pytest.fixture(scope="module")
+def two_ctx() -> Context:
+    t = _reference_frame()
+    return build_context(_two_rollout_specs(), t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
+
+
 @pytest.fixture(scope="module")
 def bundle(ctx: Context) -> Bundle:
-    return Bundle.open(_rollout_frame(), _reference_frame(), ctx)
+    return Bundle.open([RolloutSpec("rollout", _rollout_frame())], _reference_frame(), ctx)
+
+
+@pytest.fixture(scope="module")
+def two_bundle(two_ctx: Context) -> Bundle:
+    return Bundle.open(_two_rollout_specs(), _reference_frame(), two_ctx)
 
 
 def test_catalog_ids_unique_and_nonempty():
@@ -64,46 +84,71 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle):
             "unavailable",
         }
         assert r.title and r.xlabel
-        _validate_payload(r)
+        _validate_payload(r, ["rollout"])
 
 
-def _validate_payload(r) -> None:
+def test_every_spec_computes_valid_reduced_with_two_rollouts(two_bundle: Bundle):
+    for spec in build_catalog():
+        r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
+        assert r.id == spec.id
+        _validate_payload(r, ["flow", "wgan"])
+
+
+def _validate_payload(r, names: list[str]) -> None:
     p = r.payload
     if r.kind == "overlay_hist":
         n = len(p["edges"]) - 1
-        assert len(p["rollout"]) == n and len(p["reference"]) == n
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(v) == n
+        assert len(p["reference"]) == n
     elif r.kind == "single_hist":
-        assert len(p["rollout"]) == len(p["edges"]) - 1
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(v) == len(p["edges"]) - 1
     elif r.kind == "grouped_hist":
         n = len(p["edges"]) - 1
         assert p["groups"], "grouped hist must have at least one group"
         for g in p["groups"].values():
-            assert len(g["rollout"]) == n and len(g["reference"]) == n
+            assert list(g["series"]) == names
+            for v in g["series"].values():
+                assert len(v) == n
+            assert len(g["reference"]) == n
     elif r.kind == "profile":
         n = len(p["edges"]) - 1
-        for k in ("rollout_mean", "rollout_std", "reference_mean", "reference_std"):
-            assert len(p[k]) == n
+        assert list(p["series"]) == names
+        for side in p["series"].values():
+            assert len(side["mean"]) == n and len(side["std"]) == n
+        assert len(p["reference"]["mean"]) == n and len(p["reference"]["std"]) == n
     elif r.kind == "bar":
-        assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
+        assert list(p["series"]) == names
+        for v in p["series"].values():
+            assert len(p["labels"]) == len(v)
+        assert len(p["labels"]) == len(p["reference"])
     elif r.kind == "unavailable":
         assert p["note"]
     elif r.kind == "router_gating":
-        for side in ("rollout", "reference"):
-            if side in p:
-                assert len(p[side]["centers"]) == len(p[side]["means"])
-    elif r.kind == "router_share":
-        for cat in p["categories"]:
+        for entry in p["series"].values():
             for side in ("rollout", "reference"):
-                if side in p:
-                    assert cat in p[side]
+                if side in entry:
+                    assert len(entry[side]["centers"]) == len(entry[side]["means"])
+    elif r.kind == "router_share":
+        for entry in p["series"].values():
+            for cat in entry["categories"]:
+                for side in ("rollout", "reference"):
+                    if side in entry:
+                        assert cat in entry[side]
     elif r.kind == "router_specialization":
-        for side in ("rollout", "reference"):
-            if side in p:
-                assert len(p[side]["centers"]) == len(p[side]["score"])
+        for entry in p["series"].values():
+            for side in ("rollout", "reference"):
+                if side in entry:
+                    assert len(entry[side]["centers"]) == len(entry[side]["score"])
     elif r.kind == "heatmap":
-        assert len(p["matrix"]) == len(p["row_labels"])
-        for row in p["matrix"]:
-            assert len(row) == len(p["col_labels"])
+        assert list(p["series"]) == names
+        for mat in p["series"].values():
+            assert len(mat) == len(p["row_labels"])
+            for row in mat:
+                assert len(row) == len(p["col_labels"])
 
 
 # ---------------------------------------------------------------------------
@@ -115,8 +160,9 @@ def _validate_payload(r) -> None:
 # data-dependent edges (event_total_edep), concat-then-mean/std (shower_
 # longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
 # ratio (species_edep_share), a chunkable=False passthrough (router_gating),
-# nested sum-merge into a scorecard (marginal_distance_summary), concat-then-
-# event-id-join (n_sec_confusion), and concat-then-per-event-derived-quantity
+# sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
+# nested sum-merge into a scorecard (marginal_distance_summary), and
+# concat-then-per-event-derived-quantity
 # (shower_containment_depth_90, reusing the profile matrix's own merge shape).
 _CHUNK_EQUIVALENCE_IDS = [
     "marginal_edep",
@@ -125,9 +171,10 @@ _CHUNK_EQUIVALENCE_IDS = [
     "shower_longitudinal",
     "leakage_fraction",
     "sec_count_per_species",
+    "sec_count_per_step",
+    "sec_count_per_step_by_species",
     "router_gating",
     "marginal_distance_summary",
-    "n_sec_confusion",
     "shower_containment_depth_90",
 ]
 
@@ -150,20 +197,21 @@ def _assert_payload_close(a, b, path: str = "payload") -> None:
 
 
 @pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS)
-def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
+def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
     """A plot computed over N event-disjoint chunks then merged must equal the
     same plot computed in one unchunked pass — the core chunking correctness
-    guarantee (see the analysis-rollout-plots chunking plan)."""
+    guarantee (see the analysis-rollout-plots chunking plan). Exercised with
+    two rollout series so the per-rollout merge path is covered too."""
     spec: PlotSpec = get_spec(spec_id)
-    r, t = _rollout_frame(), _reference_frame()
+    rollouts, t = _two_rollout_specs(), _reference_frame()
 
-    unchunked_bundle = Bundle.open(r, t, ctx)
-    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)
+    unchunked_bundle = Bundle.open(rollouts, t, two_ctx)
+    unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], two_ctx)
 
     # 4 chunks over only 2 distinct event_ids also exercises empty chunks.
     n_chunks = 4 if spec.chunkable else 1
-    parts = [spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
-    chunked = spec.finalize(parts, ctx)
+    parts = [spec.compute_partial(Bundle.open(rollouts, t, two_ctx, chunk=(k, n_chunks))) for k in range(n_chunks)]
+    chunked = spec.finalize(parts, two_ctx)
 
     assert chunked.id == unchunked.id
     assert chunked.kind == unchunked.kind
@@ -171,7 +219,7 @@ def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
 
 
 # ---------------------------------------------------------------------------
-# new (gitea #76) reductions: KS distance, confusion matrix, containment depth
+# new (gitea #76) reductions: KS distance and containment depth
 # ---------------------------------------------------------------------------
 
 
@@ -182,20 +230,6 @@ def test_ks_statistic():
     assert _ks_statistic([10, 0], [0, 0]) == 1.0  # one side empty, other isn't -> maximal mismatch
 
 
-def test_integer_confusion_matches_event_pairing():
-    # true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
-    labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
-    assert labels == ["0", "1+"]
-    assert mat.tolist() == [[0, 0], [1, 1]]  # row=true, col=pred
-
-
-def test_integer_confusion_caps_pathological_outliers():
-    labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
-    assert labels[-1] == "4+"
-    assert mat.shape == (5, 5)
-    assert mat.sum() == 2
-
-
 def test_containment_depths_simple_ramp():
     # one event, edep concentrated in the first bin -> 90%/95% containment
     # depth is the first bin's right edge; a zero-energy event is dropped.
@@ -205,8 +239,25 @@ def test_containment_depths_simple_ramp():
     assert depths.tolist() == [1.0]
 
 
-def test_n_sec_confusion_spec(bundle):
-    spec = get_spec("n_sec_confusion")
+def test_sec_count_per_step_counts_empty_steps(bundle):
+    spec = get_spec("sec_count_per_step")
     r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
-    assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
-    assert r.payload["matrix"] == [[0, 0], [1, 1]]
+    # reference: 3 steps, two of which emit exactly one secondary
+    assert r.payload["reference"][:2] == [1, 2]
+    # rollout: 4 physical steps, one of which emits a single secondary
+    assert r.payload["series"]["rollout"][:2] == [3, 1]
+    assert sum(r.payload["reference"]) == 3
+
+
+def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
+    spec = get_spec("sec_count_per_step_by_species")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    cols = r.payload["col_labels"]
+    ref = r.payload["reference"]
+    g = cols.index(pdg_label(22))
+    # two reference steps emit one photon each; the third emits none
+    assert [row[g] for row in ref][:2] == [1, 2]
+    # every other species column is "no such secondary" on all 3 steps
+    for j, _ in enumerate(cols):
+        if j != g:
+            assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0

tests/test_checkpoint_io.py

@@ -14,6 +14,7 @@ from giant import config as gconfig
 from giant.checkpoint_io import (
     CheckpointCompatibilityError,
     InferenceContext,
+    apply_config_overrides,
     conditioning_axes,
     load_for_inference,
     stage_cfg,
@@ -278,3 +279,122 @@ def test_stage_cfg_new_shape_returns_subdict():
 def test_stage_cfg_v02_flat_shape_returns_empty_dict():
     model_cfg = {"hidden_dim": 32, "n_blocks": 4}
     assert stage_cfg(model_cfg, "stage2") == {}
+
+
+# ---------------------------------------------------------------------------
+# config_overrides (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def _router_model_cfg() -> dict:
+    cfg = _model_cfg()
+    cfg["stage1_model"]["router"] = {"enabled": True, "type": "energy", "n_experts": 2}
+    return cfg
+
+
+def test_config_override_n_sec_sampling_changes_stage2_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.n_sec.sampling": "sample"},
+    )
+    assert ctx.stage2 is not None
+    assert ctx.stage2.n_sec_sampling == "sample"
+    assert ctx.config_overrides == {"stage2_model.n_sec.sampling": "sample"}
+
+
+def test_config_override_ddpm_n_steps_changes_context_fields(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.ddpm.n_steps": 42, "stage2_model.ddpm.n_steps": 7},
+    )
+    assert ctx.stage1_ddpm_steps == 42
+    assert ctx.stage2_ddpm_steps == 7
+
+
+def test_config_override_other_policy_changes_context_field(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage2_model.particle_type.other_policy": "modal"},
+    )
+    assert ctx.other_policy == "modal"
+
+
+def test_config_override_router_temperature_changes_router_attribute(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path, model_cfg=_router_model_cfg())
+    ctx = load_for_inference(
+        checkpoint,
+        torch.device("cpu"),
+        "predict",
+        config_overrides={"stage1_model.router.temperature": 1.5},
+    )
+    assert ctx.stage1 is not None
+    assert ctx.stage1.trunk.router.temperature == pytest.approx(1.5)
+
+
+def test_config_override_no_overrides_defaults_to_empty_dict(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
+    assert ctx.config_overrides == {}
+
+
+def test_config_override_unknown_path_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.typo": "sample"},
+        )
+
+
+def test_config_override_shape_bearing_key_raises_up_front(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage1_model.hidden_dim": 999},
+        )
+
+
+def test_config_override_bad_value_raises(tmp_path):
+    checkpoint = _write_checkpoint(tmp_path)
+    with pytest.raises(CheckpointCompatibilityError, match="must be one of"):
+        load_for_inference(
+            checkpoint,
+            torch.device("cpu"),
+            "predict",
+            config_overrides={"stage2_model.n_sec.sampling": "maybe"},
+        )
+
+
+def test_apply_config_overrides_no_overrides_returns_same_object():
+    cfg = _model_cfg()
+    assert apply_config_overrides(cfg, None) is cfg
+    assert apply_config_overrides(cfg, {}) is cfg
+
+
+def test_apply_config_overrides_migrates_legacy_flat_model_config_first():
+    legacy_cfg = {
+        "pdg_vocab": len(PDG_MAP),
+        "mat_vocab": len(MAT_MAP),
+        "hidden_dim": 32,
+        "n_blocks": 4,
+        "emb_dim": 8,
+        "dropout": 0.1,
+        "k_max": 5,
+    }
+    merged = apply_config_overrides(legacy_cfg, {"stage1_model.ddpm.n_steps": 10})
+    assert merged["stage1_model"]["ddpm"]["n_steps"] == 10
+    assert merged["stage1_model"]["hidden_dim"] == 32

tests/test_cli_predict.py

@@ -172,3 +172,43 @@ def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
 
     assert result.exit_code == 1
     assert "checkpoint has no model_config" in result.output
+
+
+# ---------------------------------------------------------------------------
+# --set (gitea #87)
+# ---------------------------------------------------------------------------
+
+
+def test_predict_set_flag_without_equals_exits_1(tmp_path):
+    checkpoint = tmp_path / "missing.pt"
+
+    result = runner.invoke(
+        app,
+        ["predict", "dummy.parquet", "--checkpoint", str(checkpoint), "--set", "sampling"],
+    )
+
+    assert result.exit_code == 1
+    assert "must be 'dotted.path=value'" in result.output
+
+
+def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "predict",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--set",
+            "stage1_model.hidden_dim=999",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output

tests/test_cli_rollout.py

@@ -31,3 +31,28 @@ def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
 
     assert result.exit_code == 1
     assert "checkpoint has no model_config" in result.output
+
+
+def test_rollout_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
+    checkpoint = tmp_path / "ckpt.pt"
+    torch.save(
+        {"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
+        checkpoint,
+    )
+
+    result = runner.invoke(
+        app,
+        [
+            "rollout",
+            "dummy.parquet",
+            "--checkpoint",
+            str(checkpoint),
+            "--geometry",
+            "dummy_geometry.pkl",
+            "--set",
+            "stage2_model.n_sec.typo=sample",
+        ],
+    )
+
+    assert result.exit_code == 1
+    assert "not an inference-safe override" in result.output

tests/test_cli_train_overrides.py

@@ -20,7 +20,7 @@ def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dic
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["cfg"] = cfg
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     result = runner.invoke(
         cli.app,
@@ -125,7 +125,7 @@ def test_stage2_init_from_and_freeze_flags_land_in_cfg_and_dont_touch_stage1(mon
 
 
 def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
-    monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
+    monkeypatch.setattr("giant.pipeline.run_train_job", lambda *a, **kw: None)
     result = runner.invoke(
         cli.app,
         ["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
@@ -140,7 +140,7 @@ def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_pa
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     resume_dir = tmp_path / "resumed_run"
     resume_dir.mkdir()
@@ -161,7 +161,7 @@ def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
 
     resume_dir = tmp_path / "resumed_run"
     resume_dir.mkdir()
@@ -178,7 +178,7 @@ def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypat
     def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
         captured["out_dir"] = out_dir
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
     monkeypatch.chdir(tmp_path)
 
     result = runner.invoke(cli.app, ["train", "dummy.parquet"])
@@ -192,7 +192,7 @@ def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
     def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
         captured["batch_size"] = cfg["train"]["batch_size"]
 
-    monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
+    monkeypatch.setattr("giant.pipeline.run_train_job", _fake_run_train_job)
     monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
 
     result = runner.invoke(

tests/test_condor.py

@@ -1,4 +1,4 @@
-"""Tests for the rollout-YAML → run-directory flow, compute, and submit."""
+"""Tests for the rollout-YAML(s) → run-directory flow, compute, and submit."""
 
 from __future__ import annotations
 
@@ -17,6 +17,7 @@ from giant.analysis import (
     compute_reduced,
     derive_run_dir,
     load_rollout_yaml,
+    load_rollout_yamls,
     merge_one,
     prep,
     write_submit,
@@ -28,13 +29,17 @@ from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
 from tests.test_analysis_reduce import _reference_frame, _rollout_frame
 
 
+def _write_rollout(path: Path) -> None:
+    tbl = _rollout_frame().collect().to_arrow()
+    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
+    pq.write_table(tbl, path)
+
+
 def _write_inputs(tmp_path: Path) -> Path:
     """Materialize rollout+reference parquet and a rollout YAML; return the YAML path."""
     rollout = tmp_path / "rollout.parquet"
     reference = tmp_path / "reference.parquet"
-    tbl = _rollout_frame().collect().to_arrow()
-    tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
-    pq.write_table(tbl, rollout)
+    _write_rollout(rollout)
     _reference_frame().collect().write_parquet(reference)
 
     yaml_path = tmp_path / "run.yaml"
@@ -54,6 +59,33 @@ def _write_inputs(tmp_path: Path) -> Path:
     return yaml_path
 
 
+def _write_two_inputs(tmp_path: Path) -> tuple[Path, Path]:
+    """Two rollout YAMLs (distinct output files) sharing one reference file."""
+    reference = tmp_path / "reference.parquet"
+    _reference_frame().collect().write_parquet(reference)
+
+    paths = []
+    for tag, pred_id in (("a", "aaaa1111ef"), ("b", "bbbb2222ef")):
+        rollout = tmp_path / f"rollout_{tag}.parquet"
+        _write_rollout(rollout)
+        yaml_path = tmp_path / f"run_{tag}.yaml"
+        yaml_path.write_text(
+            yaml.safe_dump(
+                {
+                    "prediction_id": pred_id,
+                    "output": str(rollout),
+                    "dataset": str(reference),
+                    "checkpoint": f"/ckpt/{tag}.pt",
+                    "kind": "rollout",
+                    "energy_cutoff": 0.1,
+                    "steps": 10,
+                }
+            )
+        )
+        paths.append(yaml_path)
+    return paths[0], paths[1]
+
+
 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"
@@ -62,12 +94,13 @@ def _fake_venv(repo_dir: Path) -> None:
     giant.chmod(0o755)
 
 
-def _prep(rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1) -> Path:
+def _prep(rollout_yamls, run_dir: str | Path | None = None, chunks: int = 1, labels=None) -> Path:
     """``prep`` with small test-sized context bins/sampling."""
     return prep(
-        rollout_yaml,
+        rollout_yamls,
         run_dir,
         n_chunks=chunks,
+        labels=labels,
         n_energy_bins=2,
         n_marginal_bins=8,
         top_k_pdg=3,
@@ -82,39 +115,108 @@ def test_load_rollout_yaml_requires_paths(tmp_path: Path):
         load_rollout_yaml(bad)
 
 
+def test_load_rollout_yamls_single_defaults_to_rollout_name(tmp_path: Path):
+    yaml_path = _write_inputs(tmp_path)
+    loaded, reference = load_rollout_yamls([yaml_path])
+    assert [lr.name for lr in loaded] == ["rollout"]
+    assert reference.endswith("reference.parquet")
+
+
+def test_load_rollout_yamls_multi_defaults_to_stem(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    loaded, _ = load_rollout_yamls([a, b])
+    assert [lr.name for lr in loaded] == ["run_a", "run_b"]
+
+
+def test_load_rollout_yamls_explicit_labels(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    loaded, _ = load_rollout_yamls([a, b], labels=["flow", "wgan"])
+    assert [lr.name for lr in loaded] == ["flow", "wgan"]
+
+
+def test_load_rollout_yamls_label_count_mismatch(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    with pytest.raises(ValueError, match="--label"):
+        load_rollout_yamls([a, b], labels=["only-one"])
+
+
+def test_load_rollout_yamls_rejects_duplicate_names(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    with pytest.raises(ValueError, match="collide"):
+        load_rollout_yamls([a, b], labels=["same", "same"])
+
+
+def test_load_rollout_yamls_rejects_mismatched_reference(tmp_path: Path):
+    a, _ = _write_two_inputs(tmp_path)
+    other_ref = tmp_path / "other_reference.parquet"
+    _reference_frame().collect().write_parquet(other_ref)
+    c = tmp_path / "run_c.yaml"
+    c.write_text(
+        yaml.safe_dump(
+            {"prediction_id": "cccc3333ef", "output": str(tmp_path / "rollout_c.parquet"), "dataset": str(other_ref)}
+        )
+    )
+    _write_rollout(tmp_path / "rollout_c.parquet")
+    with pytest.raises(ValueError, match="same reference"):
+        load_rollout_yamls([a, c])
+
+
 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")
-    assert derive_run_dir(y, "/somewhere") == Path("/somewhere")
+    assert derive_run_dir([y]) == Path("/data/analysis_abcd1234")
+    assert derive_run_dir([y], "/somewhere") == Path("/somewhere")
 
 
 def test_derive_run_dir_default_base():
     y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
-    assert derive_run_dir(y, default_base="/work/lbogner/giant2/analysis_runs") == Path(
+    assert derive_run_dir([y], default_base="/work/lbogner/giant2/analysis_runs") == Path(
         "/work/lbogner/giant2/analysis_runs/analysis_abcd1234"
     )
     # an explicit run_dir still wins over default_base
-    assert derive_run_dir(y, "/somewhere", default_base="/other") == Path("/somewhere")
+    assert derive_run_dir([y], "/somewhere", default_base="/other") == Path("/somewhere")
+
+
+def test_derive_run_dir_multi_rollout_joins_tags():
+    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(2)]
+    assert derive_run_dir(ys, default_base="/base") == Path("/base/analysis_tag0xxxx-tag1xxxx")
+
+
+def test_derive_run_dir_many_rollouts_truncates_with_plus_count():
+    ys = [{"output": f"/data/roll_{i}.parquet", "prediction_id": f"tag{i}xxxx", "dataset": "d"} for i in range(5)]
+    run_dir = derive_run_dir(ys, default_base="/base")
+    assert run_dir == Path("/base/analysis_tag0xxxx-tag1xxxx-tag2xxxx-plus2")
 
 
 def test_prep_lays_out_run_dir(tmp_path: Path):
     yaml_path = _write_inputs(tmp_path)
-    run_dir = _prep(yaml_path)
+    run_dir = _prep([yaml_path])
     assert run_dir == tmp_path / "analysis_abcd1234"
     assert (run_dir / "shared.json").exists()
     ctx = Context.load(run_dir / "shared.json")
     assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
     meta = RunMeta.load(run_dir / "run_meta.json")
     assert meta.reference.endswith("reference.parquet")
-    assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt"
+    assert [ro["name"] for ro in meta.rollouts] == ["rollout"]
+    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/best.pt"
     assert "best.pt" in meta.title
     assert meta.n_chunks == 1
     assert meta.rows_per_chunk == [meta.total_rows]  # single chunk holds everything
     assert meta.total_rows == 8  # 5 rollout rows + 3 reference rows
 
 
+def test_prep_multi_rollout_lays_out_run_dir(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    run_dir = _prep([a, b], labels=["flow", "wgan"])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert [ro["name"] for ro in meta.rollouts] == ["flow", "wgan"]
+    assert meta.rollouts[0]["plot_meta"]["checkpoint"] == "/ckpt/a.pt"
+    assert meta.rollouts[1]["plot_meta"]["checkpoint"] == "/ckpt/b.pt"
+    # 5 rows from each rollout + 3 from the shared reference
+    assert meta.total_rows == 13
+
+
 def test_prep_splits_rows_per_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
     meta = RunMeta.load(run_dir / "run_meta.json")
     assert len(meta.rows_per_chunk) == 2
     assert sum(meta.rows_per_chunk) == meta.total_rows == 8
@@ -125,7 +227,7 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
     partials on disk for merge_one to silently merge against the new
     context (they'd be keyed/sized for the old n_chunks)."""
     yaml_path = _write_inputs(tmp_path)
-    run_dir = _prep(yaml_path, chunks=2)
+    run_dir = _prep([yaml_path], chunks=2)
     compute_one("marginal_edep", run_dir, chunk_index=0)
     compute_one("marginal_edep", run_dir, chunk_index=1)
     stale = run_dir / "reduced_partial" / "marginal_edep__0.json"
@@ -133,7 +235,7 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
     (run_dir / "reduced").mkdir(exist_ok=True)
     (run_dir / "reduced" / "marginal_edep.json").write_text("{}")
 
-    _prep(yaml_path, run_dir, chunks=1)
+    _prep([yaml_path], run_dir, chunks=1)
 
     assert not stale.exists()
     assert not (run_dir / "reduced" / "marginal_edep.json").exists()
@@ -141,20 +243,22 @@ def test_reprep_clears_stale_partials_from_a_different_chunk_count(tmp_path: Pat
 
 
 def test_compute_one_from_run_dir(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
     out = compute_one("marginal_edep", run_dir)
     assert out == run_dir / "reduced_partial" / "marginal_edep__0.json"
     partial = Partial.load(out)
     assert partial.id == "marginal_edep" and partial.chunk == 0
     assert "r" in partial.data and "t" in partial.data
+    assert list(partial.data["r"]) == ["rollout"]
 
 
 def test_compute_reduced_explicit_paths(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
     meta = RunMeta.load(run_dir / "run_meta.json")
+    rollouts = [{"name": ro["name"], "path": ro["path"]} for ro in meta.rollouts]
     out = compute_reduced(
         "marginal_step_length",
-        meta.rollout,
+        rollouts,
         meta.reference,
         run_dir / "shared.json",
         tmp_path / "r.json",
@@ -163,17 +267,17 @@ def test_compute_reduced_explicit_paths(tmp_path: Path):
 
 
 def test_merge_one_produces_reduced(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))
+    run_dir = _prep([_write_inputs(tmp_path)])
     compute_one("marginal_edep", run_dir)
     out = merge_one("marginal_edep", run_dir)
     assert out == run_dir / "reduced" / "marginal_edep.json"
     reduced = Reduced.load(out)
     assert reduced.id == "marginal_edep"
-    assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
+    assert len(reduced.payload["series"]["rollout"]) == len(reduced.payload["edges"]) - 1
 
 
 def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path), chunks=2)
+    run_dir = _prep([_write_inputs(tmp_path)], chunks=2)
     compute_one("marginal_edep", run_dir, chunk_index=0)  # chunk 1 never computed
     with pytest.raises(FileNotFoundError, match="missing chunk"):
         merge_one("marginal_edep", run_dir)
@@ -182,11 +286,11 @@ def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
 def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
     (tmp_path / "a").mkdir()
     (tmp_path / "b").mkdir()
-    unchunked_dir = _prep(_write_inputs(tmp_path / "a"))
+    unchunked_dir = _prep([_write_inputs(tmp_path / "a")])
     compute_one("marginal_step_length", unchunked_dir)
     unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir))
 
-    chunked_dir = _prep(_write_inputs(tmp_path / "b"), chunks=2)
+    chunked_dir = _prep([_write_inputs(tmp_path / "b")], chunks=2)
     for k in range(2):
         compute_one("marginal_step_length", chunked_dir, chunk_index=k)
     chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir))
@@ -194,14 +298,23 @@ def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
     assert chunked.payload == unchunked.payload
 
 
+def test_two_rollout_compute_and_merge_produces_both_series(tmp_path: Path):
+    a, b = _write_two_inputs(tmp_path)
+    run_dir = _prep([a, b], labels=["flow", "wgan"])
+    compute_one("marginal_edep", run_dir)
+    reduced = Reduced.load(merge_one("marginal_edep", run_dir))
+    assert list(reduced.payload["series"]) == ["flow", "wgan"]
+    assert "reference" in reduced.payload
+
+
 def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_path))  # n_chunks=1 (default)
+    run_dir = _prep([_write_inputs(tmp_path)])  # n_chunks=1 (default)
     with pytest.raises(ValueError, match="out of range"):
         compute_one("marginal_edep", run_dir, chunk_index=1)
 
 
 def test_write_submit_description(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_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()
@@ -223,7 +336,7 @@ def test_write_submit_description(tmp_path: Path):
 
 
 def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch):
-    run_dir = _prep(_write_inputs(tmp_path))
+    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.
@@ -233,7 +346,7 @@ def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.M
 
 
 def test_write_submit_remote_flag(tmp_path: Path):
-    run_dir = _prep(_write_inputs(tmp_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()
@@ -243,7 +356,7 @@ def test_write_submit_remote_flag(tmp_path: Path):
 
 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)
+    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)
@@ -260,7 +373,7 @@ def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
     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)
+    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"):
@@ -282,7 +395,7 @@ 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)
+    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)

tests/test_config.py

@@ -171,17 +171,40 @@ def test_n_sec_config_owner_defaults_to_stage2():
 def test_n_sec_config_owner_round_trips():
     n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
     assert n_sec.owner == "stage1"
-    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
+    assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "sampling": "greedy"}
 
 
-def test_n_sec_config_stop_sampling_defaults_to_greedy():
-    assert gconfig.NSecConfig().stop_sampling == "greedy"
+def test_n_sec_config_sampling_defaults_to_greedy():
+    assert gconfig.NSecConfig().sampling == "greedy"
 
 
-def test_n_sec_config_stop_sampling_round_trips():
+def test_n_sec_config_sampling_round_trips():
+    n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "sampling": "sample"})
+    assert n_sec.sampling == "sample"
+    assert n_sec.to_dict()["sampling"] == "sample"
+
+
+def test_n_sec_config_stop_sampling_alias_still_honored():
+    """gitea #86: stop_sampling was renamed to sampling; old checkpoints'
+    model_config still carries the old key and must keep working."""
     n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
-    assert n_sec.stop_sampling == "sample"
-    assert n_sec.to_dict()["stop_sampling"] == "sample"
+    assert n_sec.sampling == "sample"
+    assert "stop_sampling" not in n_sec.to_dict()
+
+
+def test_n_sec_config_sampling_key_wins_over_stop_sampling_alias():
+    n_sec = gconfig.NSecConfig.from_dict({"sampling": "sample", "stop_sampling": "greedy"})
+    assert n_sec.sampling == "sample"
+
+
+def test_migrate_config_renames_stop_sampling_key():
+    cfg = {
+        "meta": {"config_version": gconfig.CONFIG_VERSION},
+        "stage2_model": {"n_sec": {"stop_sampling": "sample"}},
+    }
+    migrated = gconfig.migrate_config(cfg)
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.sampling") == "sample"
+    assert gconfig._get_path(migrated, "stage2_model.n_sec.stop_sampling") is None
 
 
 # ---------------------------------------------------------------------------
@@ -851,13 +874,13 @@ def test_validate_config_stop_token_rejected_for_stage1_owner():
         assert "stop_token" in str(e) and "owner" in str(e)
 
 
-def test_validate_config_bad_stop_sampling_rejected():
-    cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
+def test_validate_config_bad_n_sec_sampling_rejected():
+    cfg = _cfg_with(**{"stage2_model.n_sec.sampling": "bogus"})
     try:
         gconfig.validate_config(cfg)
         assert False, "expected ValueError"
     except ValueError as e:
-        assert "stop_sampling" in str(e)
+        assert "sampling" in str(e)
 
 
 def test_validate_config_default_precision_is_fp32():

tests/test_model_summary.py

@@ -12,6 +12,8 @@ from giant.cli import app
 from giant.materials import MATERIAL_PROPERTIES
 from giant.model.summary import _NOT_BUILD_TIME, _built_modules, _vocab_caveats, summarize_model
 
+INFERENCE_OVERRIDES = gconfig.INFERENCE_OVERRIDES
+
 runner = CliRunner()
 
 _PDG_VOCAB = 300
@@ -53,6 +55,16 @@ def test_not_build_time_allow_list_has_no_stale_entries():
     assert not stale, f"_NOT_BUILD_TIME entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
 
 
+def test_inference_overrides_allow_list_has_no_stale_entries():
+    in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
+    stale = set(INFERENCE_OVERRIDES) - in_scope
+    assert not stale, f"INFERENCE_OVERRIDES entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
+
+
+def test_default_config_overridable_lists_every_allowlisted_path(default_summary):
+    assert set(default_summary.overridable) == set(INFERENCE_OVERRIDES)
+
+
 def test_router_disabled_by_default_so_its_fields_are_inert(default_summary):
     assert "stage1_model.router.n_experts" in default_summary.inert
     assert "stage1_model.router.temperature" in default_summary.inert
@@ -135,3 +147,5 @@ def test_cli_default_smoke():
     assert "parameters" in result.output
     assert "trunk" in result.output
     assert "inert under this config" in result.output
+    assert "inference-overridable without retraining" in result.output
+    assert "stage2_model.n_sec.sampling" in result.output

tests/test_render.py

@@ -25,41 +25,96 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
     reduced = [
         Reduced(
             "rg",
-            "router",
+            "model",
             "router_gating",
             "Router gating",
             "pre-step energy [MeV]",
             {
-                "n_experts": 2,
                 "log_x": True,
-                "router_type": "energy",
-                "rollout": {
-                    "centers": [1.0, 10.0, 100.0],
-                    "means": [[0.6, 0.4], [0.5, 0.5], [0.4, 0.6]],
-                },
-                "reference": {
-                    "centers": [1.0, 10.0, 100.0],
-                    "means": [[0.55, 0.45], [0.5, 0.5], [0.45, 0.55]],
+                "series": {
+                    "flow": {
+                        "n_experts": 2,
+                        "router_type": "energy",
+                        "rollout": {
+                            "centers": [1.0, 10.0, 100.0],
+                            "means": [[0.6, 0.4], [0.5, 0.5], [0.4, 0.6]],
+                        },
+                        "reference": {
+                            "centers": [1.0, 10.0, 100.0],
+                            "means": [[0.55, 0.45], [0.5, 0.5], [0.45, 0.55]],
+                        },
+                    },
+                    "wgan": {
+                        "n_experts": 2,
+                        "router_type": "energy",
+                        "rollout": {"centers": [1.0], "means": [[0.5, 0.5]]},
+                        "reference": {"centers": [1.0], "means": [[0.5, 0.5]]},
+                    },
                 },
             },
         ),
         Reduced(
             "rs",
-            "router",
+            "model",
             "router_share",
             "Router share",
             "species",
             {
-                "categories": ["e-", "gamma"],
-                "n_experts": 2,
-                "router_type": "energy",
-                "rollout": {"e-": [0.7, 0.3], "gamma": [0.2, 0.8]},
-                "reference": {"e-": [0.6, 0.4], "gamma": [0.3, 0.7]},
+                "series": {
+                    "flow": {
+                        "categories": ["e-", "gamma"],
+                        "n_experts": 2,
+                        "router_type": "energy",
+                        "rollout": {"e-": [0.7, 0.3], "gamma": [0.2, 0.8]},
+                        "reference": {"e-": [0.6, 0.4], "gamma": [0.3, 0.7]},
+                    },
+                },
+            },
+        ),
+        Reduced(
+            "rp",
+            "model",
+            "router_share",
+            "Router share by process (reference-only)",
+            "process",
+            {
+                "series": {
+                    "flow": {
+                        "categories": ["compt", "phot"],
+                        "n_experts": 2,
+                        "router_type": "energy",
+                        "reference": {"compt": [0.4, 0.6], "phot": [0.9, 0.1]},
+                    },
+                },
+            },
+        ),
+        Reduced(
+            "rz",
+            "model",
+            "router_specialization",
+            "Router specialization",
+            "pre-step energy [MeV]",
+            {
+                "log_x": True,
+                "series": {
+                    "flow": {
+                        "n_experts": 2,
+                        "chance_level": 0.5,
+                        "rollout": {"centers": [1.0, 10.0], "score": [0.6, 0.7]},
+                        "reference": {"centers": [1.0, 10.0], "score": [0.55, 0.65]},
+                    },
+                    "wgan": {
+                        "n_experts": 4,
+                        "chance_level": 0.25,
+                        "rollout": {"centers": [1.0, 10.0], "score": [0.3, 0.4]},
+                        "reference": {"centers": [], "score": []},
+                    },
+                },
             },
         ),
         Reduced(
             "ru",
-            "router",
+            "model",
             "unavailable",
             "Router unavailable",
             "x",
@@ -73,7 +128,10 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
             "x",
             {
                 "edges": [0, 1, 2],
-                "groups": {lbl: {"rollout": [1, 2], "reference": [2, 1]} for lbl in ("a", "b", "c", "d")},
+                "groups": {
+                    lbl: {"series": {"flow": [1, 2], "wgan": [2, 1]}, "reference": [2, 1]}
+                    for lbl in ("a", "b", "c", "d")
+                },
                 "log_y": True,
             },
         ),
@@ -83,7 +141,22 @@ def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
             "single_hist",
             "Single (log-x)",
             "x",
-            {"edges": [1, 10, 100], "rollout": [5, 1], "log_x": True, "log_y": True},
+            {"edges": [1, 10, 100], "series": {"flow": [5, 1], "wgan": [3, 2]}, "log_x": True, "log_y": True},
+        ),
+        Reduced(
+            "hm",
+            "quality",
+            "heatmap",
+            "Distance summary (2 rollouts)",
+            "grouping axis",
+            {
+                "series": {"flow": [[0.1, 0.2], [0.3, 0.4]], "wgan": [[0.5, 0.6], [0.7, 0.8]]},
+                "row_labels": ["step_length", "edep"],
+                "col_labels": ["overall", "energy"],
+                "cbar_label": "KS statistic",
+                "vmin": 0.0,
+                "vmax": 1.0,
+            },
         ),
     ]
     try:
@@ -117,7 +190,7 @@ def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
             "single_hist",
             "Single",
             "x",
-            {"edges": [0, 1, 2], "rollout": [5, 1]},
+            {"edges": [0, 1, 2], "series": {"rollout": [5, 1]}},
         )
     ]
     for r in reduced:
@@ -142,15 +215,14 @@ def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkey
     merge_calls = []
     monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
     meta = condor_mod.RunMeta(
-        rollout="rollout.parquet",
+        rollouts=[{"name": "rollout", "path": "rollout.parquet", "plot_meta": {"checkpoint": "ckpt/best.pt"}}],
         reference="reference.parquet",
         run_dir=str(run_dir),
         title="my-run",
-        plot_meta={"checkpoint": "ckpt/best.pt"},
     )
     monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))
 
-    Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]}).save(
+    Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "series": {"rollout": [1]}}).save(
         run_dir / "reduced" / "s.json"
     )
 
@@ -177,7 +249,7 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "x",
             {
                 "edges": [0, 1, 2, 3],
-                "rollout": [1, 2, 3],
+                "series": {"flow": [1, 2, 3], "wgan": [2, 2, 2]},
                 "reference": [3, 2, 1],
                 "log_y": False,
             },
@@ -190,7 +262,7 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "x",
             {
                 "edges": [0, 1, 2],
-                "groups": {"a": {"rollout": [1, 2], "reference": [2, 1]}},
+                "groups": {"a": {"series": {"flow": [1, 2]}, "reference": [2, 1]}},
                 "log_y": False,
             },
         ),
@@ -202,10 +274,8 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "depth",
             {
                 "edges": [0, 1, 2],
-                "rollout_mean": [1, 2],
-                "rollout_std": [0.1, 0.2],
-                "reference_mean": [1.1, 1.9],
-                "reference_std": [0.1, 0.1],
+                "series": {"flow": {"mean": [1, 2], "std": [0.1, 0.2]}},
+                "reference": {"mean": [1.1, 1.9], "std": [0.1, 0.1]},
                 "ylabel": "e",
             },
         ),
@@ -217,7 +287,7 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "species",
             {
                 "labels": ["e-", "gamma"],
-                "rollout": [0.6, 0.4],
+                "series": {"flow": [0.6, 0.4], "wgan": [0.55, 0.45]},
                 "reference": [0.5, 0.5],
                 "ylabel": "frac",
             },
@@ -228,7 +298,22 @@ def test_render_one_of_each_kind(tmp_path: Path):
             "single_hist",
             "Single",
             "x",
-            {"edges": [0, 1, 2], "rollout": [5, 1], "log_y": True},
+            {"edges": [0, 1, 2], "series": {"flow": [5, 1]}, "log_y": True},
+        ),
+        Reduced(
+            "hm1",
+            "secondaries",
+            "heatmap",
+            "Heatmap (single rollout)",
+            "predicted",
+            {
+                "series": {"flow": [[1, 0], [0, 1]]},
+                "reference": [[2, 0], [0, 1]],
+                "row_labels": ["0", "1+"],
+                "col_labels": ["0", "1+"],
+                "cbar_label": "count",
+                "log_color": True,
+            },
         ),
     ]
     try:
@@ -276,8 +361,8 @@ def test_figure_params_v2_basics_and_router_and_epoch():
         },
         "conditioning": {"particle": {"type": "physical"}},
     }
-    run_meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10}
-    params = render_mod._figure_params(run_meta | {"model_config": mc})
+    meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10, "model_config": mc}
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params == {
         "hidden_dim": 256,
         "n_res_blocks": 4,
@@ -297,8 +382,8 @@ def test_figure_params_v2_wgan_reports_noise_dim_not_steps():
             "wgan": {"noise_dim": 32},
         },
     }
-    run_meta = {"model_config": mc, "steps": 10}
-    params = render_mod._figure_params(run_meta)
+    meta = {"model_config": mc, "steps": 10}
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params["mode"] == "wgan"
     assert params["noise_dim"] == 32
     assert "steps" not in params
@@ -309,18 +394,18 @@ def test_figure_params_v2_reports_mode_s2_only_when_it_differs():
         "stage1_model": {"generator": "flow"},
         "stage2_model": {"generator": "flow"},
     }
-    assert "mode_s2" not in render_mod._figure_params({"model_config": same})
+    assert "mode_s2" not in render_mod._figure_params({"rollouts": {"rollout": {"model_config": same}}})
 
     mixed = {
         "stage1_model": {"generator": "flow"},
         "stage2_model": {"generator": "wgan"},
     }
-    params = render_mod._figure_params({"model_config": mixed})
+    params = render_mod._figure_params({"rollouts": {"rollout": {"model_config": mixed}}})
     assert params["mode_s2"] == "wgan"
 
 
 def test_figure_params_old_shape_basics():
-    run_meta = {
+    meta = {
         "model_config": {
             "hidden_dim": 128,
             "n_blocks": 3,
@@ -332,7 +417,7 @@ def test_figure_params_old_shape_basics():
         "best_val_loss": 0.5,
         "steps": 20,
     }
-    params = render_mod._figure_params(run_meta)
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params == {
         "hidden_dim": 128,
         "n_blocks": 3,
@@ -346,20 +431,30 @@ def test_figure_params_old_shape_basics():
 
 
 def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps():
-    run_meta = {
+    meta = {
         "model_config": {"mode": "wgan", "noise_dim": 16},
         "steps": 20,
     }
-    params = render_mod._figure_params(run_meta)
+    params = render_mod._figure_params({"rollouts": {"rollout": meta}})
     assert params["noise_dim"] == 16
     assert "steps" not in params
 
 
+def test_figure_params_multi_rollout_names_the_series():
+    run_meta = {"rollouts": {"flow": {"model_config": {"mode": "flow"}}, "wgan": {"model_config": {"mode": "wgan"}}}}
+    assert render_mod._figure_params(run_meta) == {"rollouts": "flow, wgan"}
+
+
+def test_figure_params_empty_rollouts_is_empty():
+    assert render_mod._figure_params({}) == {}
+    assert render_mod._figure_params({"rollouts": {}}) == {}
+
+
 def test_plot_metadata_includes_note_and_run_meta_parameters():
     r = Reduced("u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"})
-    meta = render_mod._plot_metadata(r, {"title": "run-1", "checkpoint": "ckpt.pt"})
+    meta = render_mod._plot_metadata(r, {"title": "run-1", "reference": "ref.parquet", "rollouts": {"rollout": {}}})
     assert meta["note"] == "no router data"
-    assert meta["parameters"] == {"checkpoint": "ckpt.pt"}
+    assert meta["parameters"] == {"reference": "ref.parquet", "rollouts": {"rollout": {}}}
     assert "title" not in meta["parameters"]
 
 
@@ -368,3 +463,30 @@ def test_plot_metadata_omits_parameters_when_run_meta_empty():
     meta = render_mod._plot_metadata(r, {})
     assert "parameters" not in meta
     assert "note" not in meta
+
+
+def test_tex_escape_handles_percent_and_other_special_chars():
+    assert render_mod._tex_escape("90% of deposited energy") == r"90\% of deposited energy"
+    assert render_mod._tex_escape(r"a_b & c#d $e {f} \bar") == r"a\_b \& c\#d \$e \{f\} \textbackslash{}bar"
+
+
+def test_render_survives_title_and_xlabel_with_literal_percent(tmp_path: Path):
+    # Regression test for gitea #81: a literal "%" in a catalog title (e.g.
+    # "Shower containment depth (90% of deposited energy)") crashed the whole
+    # LaTeX render, since usetex treats an unescaped "%" as a comment marker.
+    reduced = [
+        Reduced(
+            "shower_containment_depth_90",
+            "shower",
+            "single_hist",
+            "Shower containment depth (90% of deposited energy)",
+            "depth containing 90% of deposited energy [mm]",
+            {"edges": [0, 1, 2], "series": {"flow": [5, 1]}},
+        ),
+    ]
+    try:
+        pdfs = _try_render(reduced, tmp_path)
+    except RuntimeError as e:  # LaTeX missing at render time
+        pytest.skip(f"LaTeX rendering unavailable: {e}")
+    assert len(pdfs) == 1
+    assert pdfs[0].exists()

tests/test_router_gating.py

@@ -10,7 +10,9 @@ from giant.analysis.router_gating import (
     compute_router_gating,
     compute_router_share_by_pdg,
     compute_router_share_by_process,
+    compute_router_specialization,
 )
+from giant.analysis.sources import RolloutSide
 from giant.data.transforms import Normalizer
 from giant.model.network import build_models
 
@@ -34,7 +36,7 @@ def _model_cfg() -> dict:
     }
 
 
-def _write_checkpoint(tmp_path) -> str:
+def _write_checkpoint(tmp_path, name: str = "ckpt.pt") -> str:
     cfg = _model_cfg()
     stage1 = build_models(cfg)["stage1"]
     assert stage1 is not None
@@ -48,7 +50,7 @@ def _write_checkpoint(tmp_path) -> str:
         "mat_map": _MAT_MAP,
         "normalizer": {"cond": norm.to_dict()},
     }
-    path = tmp_path / "ckpt.pt"
+    path = tmp_path / name
     torch.save(ckpt, path)
     return str(path)
 
@@ -86,42 +88,80 @@ def _steps_frame(process: bool = False) -> pl.LazyFrame:
     return pl.DataFrame(data).lazy()
 
 
+def _side(checkpoint: str | None, lf: pl.LazyFrame) -> RolloutSide:
+    return RolloutSide(all=lf, phys=lf, checkpoint=checkpoint)
+
+
 def test_compute_router_gating_shapes(tmp_path):
     checkpoint = _write_checkpoint(tmp_path)
     lf = _steps_frame()
-    r = compute_router_gating(checkpoint, lf, lf)
+    r = compute_router_gating({"rollout": _side(checkpoint, lf)}, lf)
     assert r.kind == "router_gating"
-    assert r.payload["n_experts"] == 2
+    assert list(r.payload["series"]) == ["rollout"]
+    entry = r.payload["series"]["rollout"]
+    assert entry["n_experts"] == 2
     for side in ("rollout", "reference"):
-        means = r.payload[side]["means"]
+        means = entry[side]["means"]
         assert means, f"{side} produced no bins"
         assert all(abs(sum(row) - 1.0) < 1e-5 for row in means)
 
 
 def test_compute_router_gating_missing_checkpoint_is_unavailable():
     lf = _steps_frame()
-    r = compute_router_gating(None, lf, lf)
+    r = compute_router_gating({"rollout": _side(None, lf)}, lf)
     assert r.kind == "unavailable"
     assert "note" in r.payload
     assert r.title
 
 
+def test_compute_router_gating_two_rollouts_only_moe_ones_included(tmp_path):
+    lf = _steps_frame()
+    ckpt = _write_checkpoint(tmp_path)
+    rollouts = {"flow": _side(None, lf), "moe": _side(ckpt, lf)}
+    r = compute_router_gating(rollouts, lf)
+    assert list(r.payload["series"]) == ["moe"]
+
+
+def test_compute_router_specialization_two_rollouts(tmp_path):
+    lf = _steps_frame()
+    ckpt_a = _write_checkpoint(tmp_path, "a.pt")
+    ckpt_b = _write_checkpoint(tmp_path, "b.pt")
+    rollouts = {"a": _side(ckpt_a, lf), "b": _side(ckpt_b, lf)}
+    r = compute_router_specialization(rollouts, lf)
+    assert r.kind == "router_specialization"
+    assert list(r.payload["series"]) == ["a", "b"]
+    for entry in r.payload["series"].values():
+        assert entry["chance_level"] == 0.5
+        assert len(entry["rollout"]["centers"]) == len(entry["rollout"]["score"])
+
+
 def test_compute_router_share_by_pdg(tmp_path):
     checkpoint = _write_checkpoint(tmp_path)
     lf = _steps_frame()
-    r = compute_router_share_by_pdg(checkpoint, lf, lf, top_pdgs=[11, 22])
+    r = compute_router_share_by_pdg({"rollout": _side(checkpoint, lf)}, lf, top_pdgs=[11, 22])
     assert r.kind == "router_share"
+    entry = r.payload["series"]["rollout"]
     for side in ("rollout", "reference"):
-        assert set(r.payload[side]) == {"e-", "gamma"}
-        for shares in r.payload[side].values():
+        assert set(entry[side]) == {"e-", "gamma"}
+        for shares in entry[side].values():
             assert abs(sum(shares) - 1.0) < 1e-5
 
 
 def test_compute_router_share_by_process(tmp_path):
     checkpoint = _write_checkpoint(tmp_path)
     lf = _steps_frame(process=True)
-    r = compute_router_share_by_process(checkpoint, lf)
+    r = compute_router_share_by_process({"rollout": _side(checkpoint, lf)}, lf)
     assert r.kind == "router_share"
-    assert set(r.payload["categories"]) <= {"eIoni", "compt"}
-    for shares in r.payload["reference"].values():
+    entry = r.payload["series"]["rollout"]
+    assert set(entry["categories"]) <= {"eIoni", "compt"}
+    for shares in entry["reference"].values():
         assert abs(sum(shares) - 1.0) < 1e-5
+
+
+def test_no_moe_rollouts_are_unavailable(tmp_path):
+    lf = _steps_frame()
+    rollouts = {"flow": _side(None, lf), "wgan": _side(None, lf)}
+    assert compute_router_gating(rollouts, lf).kind == "unavailable"
+    assert compute_router_share_by_pdg(rollouts, lf, top_pdgs=[11, 22]).kind == "unavailable"
+    assert compute_router_share_by_process(rollouts, lf).kind == "unavailable"
+    assert compute_router_specialization(rollouts, lf).kind == "unavailable"

tests/test_sample.py

@@ -14,6 +14,7 @@ from giant.model.network import (
     stage2_trunk_sec_dim,
 )
 from giant.sample import (
+    resolve_n_sec,
     sample_flow,
     sample_secondaries,
     sample_secondaries_ar,
@@ -72,6 +73,7 @@ def _stage2_ar(
     mat: int = 2,
     k_max: int = 5,
     history: str = "markov",
+    n_sec_sampling: str = "greedy",
 ) -> Stage2Autoregressive:
     particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
     return Stage2Autoregressive(
@@ -89,6 +91,7 @@ def _stage2_ar(
         history=history,
         attn_n_heads=2,
         attn_n_layers=1,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -99,7 +102,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
 def _stage2_ar_stop_token(
     target: str,
     generator: str,
-    stop_sampling: str = "greedy",
+    n_sec_sampling: str = "greedy",
     emb_dim: int = 6,
     pdg: int = 3,
     mat: int = 2,
@@ -120,7 +123,7 @@ def _stage2_ar_stop_token(
         particle_type_cfg=ParticleTypeConfig(target=target),
         build_n_sec_head=False,
         build_stop_head=True,
-        stop_sampling=stop_sampling,
+        n_sec_sampling=n_sec_sampling,
     ).eval()
 
 
@@ -267,14 +270,14 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
 # ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_sampling):
     """A stop_head pinned to a large positive logit fires at slot 0 for
     every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
     a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
     generating any token."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, 50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -283,13 +286,13 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(sto
     assert not sec_valid.any()
 
 
-@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
-def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_sampling):
     """A stop_head pinned to a large negative logit never fires under either
     policy, so every row is capped at k_max (the safety cap, not a modeling
     ceiling)."""
     B, k_max = 4, 5
-    decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
+    decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
     _force_stop_head_logit(decoder, -50.0)
     cond_cont, cond_cat = _cond(B)
     stage1_out = torch.randn(B, X_DIM)
@@ -336,3 +339,60 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
     stage1_out = torch.randn(3, X_DIM)
     with pytest.raises(AssertionError):
         sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+
+
+# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
+
+
+def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
+    """Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
+    returns `bias` (broadcast over the batch) as logits regardless of
+    conditioning — mirrors `_force_stop_head_logit`."""
+    assert decoder.n_sec_head is not None
+    last_linear = decoder.n_sec_head[-1]
+    with torch.no_grad():
+        last_linear.weight.zero_()
+        last_linear.bias.copy_(bias)
+
+
+@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
+def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
+    """A logit vector overwhelmingly favoring one class gives the same
+    answer under both policies — greedy because it's the argmax, sample
+    because softmax puts ~all mass on it."""
+    B, k_max = 8, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
+    bias = torch.full((k_max + 1,), -50.0)
+    bias[2] = 50.0
+    _force_n_sec_head_bias(decoder, bias)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))
+
+
+def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
+    B, k_max = 32, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() == 1
+
+
+def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
+    """Under a flat logit vector, a categorical draw across a large batch
+    should hit more than one class — the whole point of gitea #86: greedy
+    always collapses to one, sample should not."""
+    torch.manual_seed(0)
+    B, k_max = 256, 5
+    decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
+    _force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
+    assert n_sec is not None
+    assert n_sec.unique().numel() > 1

tests/test_type_embedding_distance.py

@@ -3,6 +3,9 @@
 
 from __future__ import annotations
 
+import polars as pl
+
+from giant.analysis.sources import RolloutSide
 from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
 
 
@@ -18,26 +21,47 @@ def _summary(n=100):
     }
 
 
+def _side(l1_dist: dict | None) -> RolloutSide:
+    empty = pl.LazyFrame()
+    return RolloutSide(all=empty, phys=empty, type_embedding_l1_dist=l1_dist)
+
+
 def test_none_is_unavailable():
-    r = compute_type_embedding_l1_distance(None)
+    r = compute_type_embedding_l1_distance({"rollout": _side(None)})
     assert r.kind == "unavailable"
     assert r.id == "type_embedding_l1_distance"
     assert r.payload["note"]
 
 
 def test_summary_produces_single_hist():
-    r = compute_type_embedding_l1_distance(_summary())
+    r = compute_type_embedding_l1_distance({"rollout": _side(_summary())})
     assert r.kind == "single_hist"
     assert r.id == "type_embedding_l1_distance"
     assert r.payload["edges"] == [0.0, 1.0, 2.0, 3.0]
-    assert r.payload["rollout"] == [30, 40, 30]
+    assert r.payload["series"]["rollout"] == [30, 40, 30]
     assert r.payload["log_x"] is True
     assert r.payload["log_y"] is True
     assert "n=100" in r.payload["note"]
 
 
 def test_single_hist_payload_shape_matches_render_contract():
-    """_render_single (giant.analysis.render) requires len(rollout) ==
+    """_render_single (giant.analysis.render) requires each series' length ==
     len(edges) - 1."""
-    r = compute_type_embedding_l1_distance(_summary())
-    assert len(r.payload["rollout"]) == len(r.payload["edges"]) - 1
+    r = compute_type_embedding_l1_distance({"rollout": _side(_summary())})
+    assert len(r.payload["series"]["rollout"]) == len(r.payload["edges"]) - 1
+
+
+def test_two_rollouts_both_populated():
+    r = compute_type_embedding_l1_distance({"flow": _side(_summary(50)), "wgan": _side(_summary(80))})
+    assert list(r.payload["series"]) == ["flow", "wgan"]
+    assert "n=50" in r.payload["note"] and "n=80" in r.payload["note"]
+
+
+def test_one_of_two_rollouts_populated_only_that_one_appears():
+    r = compute_type_embedding_l1_distance({"flow": _side(None), "wgan": _side(_summary())})
+    assert list(r.payload["series"]) == ["wgan"]
+
+
+def test_none_populated_across_rollouts_is_unavailable():
+    r = compute_type_embedding_l1_distance({"flow": _side(None), "wgan": _side(None)})
+    assert r.kind == "unavailable"

uv.lock

@@ -675,7 +675,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.8"
+version = "0.3.15"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },