chore: update changelog for v0.3.17

Reviewed-on: #93

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.13"
+current_version = "0.3.17"
 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:
+  pull_request: {}
-    branches: [master]
 
 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,35 @@
 # Changelog
 
+## [0.3.17] - 2026-09-02
+
+### Changed
+
+- Feat: add WGAN + AR stop-token config variant
+
+- Perf: replace pandas with polars in the setup-stage scan
+
+## [0.3.16] - 2026-08-31
+
+### Changed
+
+- Docs: record analysis_341dfb14 baseline rollout benchmark results
+
+- Feat: add eval-cost benchmark — Geant4 reference vs surrogate rollout timing
+
+## [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

CLAUDE.md

@@ -113,6 +113,8 @@ Secondary energies are a **stick-breaking partition of the `e_sec` budget** from
 
 **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.
 
+**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.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.
 
 **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.

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
+# Known weak spots, now measured against this exact config rather than
-# reason for the comparisons, not a reason to retune this file): every model
+# extrapolated from the pre-v0.3 field (analysis_341dfb14, best.pt @ epoch
-# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
+# 50/50, full writeup: knowledge-base/experiments/
-# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
+# giant-baseline-flow-ar-rollout-validation.md). Unlike every pre-v0.3
-# head accuracy sits at 0.863-0.867 regardless of size or objective.
+# 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

configs/wgan_ar_stoptoken.toml

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

giant/analysis/catalog.py

@@ -45,6 +45,7 @@ import numpy as np
 import polars as pl
 
 from giant.analysis.context import Context
+from giant.analysis.geant4_reference import GEANT4_REFERENCE, geant4_per_step_us
 from giant.analysis.grouping import (
     energy_bin_labels,
     event_energy_bins,
@@ -125,6 +126,7 @@ class Bundle:
                 phys=physical_steps(r_all, Side.rollout),
                 checkpoint=rs.checkpoint,
                 type_embedding_l1_dist=rs.type_embedding_l1_dist,
+                timing=rs.timing,
             )
         return cls(ctx=ctx, rollouts=sides, t_all=t_all, t_phys=physical_steps(t_all, Side.reference))
 
@@ -973,6 +975,80 @@ def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
     )
 
 
+# ---------------------------------------------------------------------------
+# eval cost (not chunked — metadata-only, no row scan)
+# ---------------------------------------------------------------------------
+
+_EVAL_COST_LABELS = ["sampling / simulation", "parquet write / convert", "total"]
+_EVAL_COST_NOTE = (
+    "no rollout in this run carries a `timing` block — re-run `giant rollout` "
+    "(timing instrumentation added after this checkpoint's rollout run) to "
+    "populate this plot"
+)
+
+
+def _eval_cost_per_step(b: Bundle) -> Reduced:
+    """Per-rollout µs/physical-step vs the measured Geant4 reference.
+
+    ``timing`` (``giant.cli``'s ``rollout`` command) is metadata carried on
+    the rollout YAML, not derived from the row data, so this needs no chunked
+    scan — same shape as the router diagnostics above.
+    """
+    series: dict[str, list[float]] = {}
+    speedup: dict[str, float] = {}
+    for name, rs in b.rollouts.items():
+        t = rs.timing
+        if not t or t.get("us_per_step") is None:
+            continue
+        sample_us = t["us_per_step"]
+        write_us = t.get("write_us_per_step") or 0.0
+        series[name] = [sample_us, write_us, sample_us + write_us]
+
+    if not series:
+        return Reduced(
+            id="eval_cost_per_step",
+            family="cost",
+            kind="unavailable",
+            title="Eval cost per step: surrogate vs Geant4",
+            xlabel="n/a",
+            payload={"note": _EVAL_COST_NOTE},
+        )
+
+    g4 = geant4_per_step_us()
+    reference = [g4["sim_us_per_step"], g4["convert_us_per_step"], g4["total_us_per_step"]]
+    for name, vals in series.items():
+        speedup[name] = reference[-1] / vals[-1] if vals[-1] else float("inf")
+
+    return Reduced(
+        id="eval_cost_per_step",
+        family="cost",
+        kind="bar",
+        title="Eval cost per step: surrogate vs Geant4",
+        xlabel="phase",
+        payload={
+            "labels": _EVAL_COST_LABELS,
+            "series": series,
+            "reference": reference,
+            "ylabel": "µs per physical step",
+            "log_y": True,
+        },
+        meta={
+            "speedup_vs_geant4_total": speedup,
+            "geant4_provenance": GEANT4_REFERENCE["provenance"],
+            "caveat": (
+                "The Geant4 reference is measured single-threaded on one CPU core "
+                "(see giant.analysis.geant4_reference); a rollout's timing is "
+                "whatever device it actually ran on (see each series' device in "
+                "run_meta.json's plot_meta). This is a deployment-speedup ratio, "
+                "not a same-hardware or per-FLOP comparison."
+            ),
+        },
+    )
+
+
+_eval_cost_per_step_partial, _eval_cost_per_step_finalize = _unchunkable(_eval_cost_per_step)
+
+
 # ---------------------------------------------------------------------------
 # router diagnostics (not chunked — already bounded/subsampled)
 # ---------------------------------------------------------------------------
@@ -1160,6 +1236,13 @@ def build_catalog() -> list[PlotSpec]:
             compute_partial=_sec_cos_angle_partial,
             finalize=_sec_cos_angle_finalize,
         ),
+        PlotSpec(
+            "eval_cost_per_step",
+            "cost",
+            compute_partial=_eval_cost_per_step_partial,
+            finalize=_eval_cost_per_step_finalize,
+            chunkable=False,
+        ),
         PlotSpec(
             "router_gating",
             "model",

giant/analysis/condor.py

@@ -82,6 +82,7 @@ _PLOT_META_KEYS = (
     "rollout_seed",
     "n_rows",
     "termination_reason_counts",
+    "timing",
     "model_config",
     "training_epoch",
     "best_val_loss",
@@ -329,9 +330,9 @@ def compute_reduced(
 ) -> Path:
     """Core: run one (plot, chunk)'s partial reduction against explicit paths.
 
-    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?},
+    ``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?,
-    ...]``, one per rollout series (insertion order preserved through to every
+    "timing"?}, ...]``, one per rollout series (insertion order preserved
-    plot's ``Reduced.payload["series"]``).
+    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``
@@ -352,6 +353,7 @@ def compute_reduced(
             source=r["path"],
             checkpoint=r.get("checkpoint"),
             type_embedding_l1_dist=r.get("type_embedding_l1_dist"),
+            timing=r.get("timing"),
         )
         for r in rollouts
     ]
@@ -377,6 +379,7 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
             "path": ro["path"],
             "checkpoint": ro["plot_meta"].get("checkpoint"),
             "type_embedding_l1_dist": ro["plot_meta"].get("type_embedding_l1_dist"),
+            "timing": ro["plot_meta"].get("timing"),
         }
         for ro in meta.rollouts
     ]

giant/analysis/geant4_reference.py

@@ -0,0 +1,76 @@
+"""Measured Geant4 (miniCaloSim) per-step eval cost — the reference line for
+``eval_cost_per_step`` in ``catalog.py``.
+
+Mirrors the precedent set by ``runtime_estimate.py``'s ``_COST_MODEL``: a
+constant table measured once on a specific machine and pasted in, with the
+methodology and provenance recorded in this docstring rather than derived at
+runtime (there is no live Geant4 install on the machines that run
+``giant analyze``, and re-measuring per invocation would be both slow and
+noisy — see the module docstring precedent).
+
+**Methodology** (``scratchpad/bench_geant4.py``, a one-off, not a `dwarf`
+subcommand): ``run_pbwo4`` (the default homogeneous-PbWO4 miniCaloSim
+executable, see ``~/Programming/minicalosim``) was timed at 3 beam energies
+(1/10/50 GeV) and **4 event counts each**, converting each run's ROOT output
+to Parquet with ``giant.tools.steps_to_parquet.convert_steps_to_parquet``
+immediately after. Event counts were scaled down as energy rose (100/400/
+1000/2000 at 1 GeV, 30/100/200/300 at 10 GeV, 10/25/45/60 at 50 GeV) to keep
+every run's row count under ~8.1M — a naive 50/200 pair at 50 GeV produces
+~27M steps and OOM'd the conversion step on a 14GB laptop. Per-energy linear
+fits (``t = intercept + slope * n``) separate Geant4's one-time init (physics
+tables, geometry construction) from its true marginal per-event cost — the
+slope, not a naive ``t / n_events`` from a single run, is what feeds
+``sim_us_per_step`` below. The per-step denominator is the produced
+``Steps``-tree/Parquet row count, matching the "physical step" unit
+``giant rollout``'s ``timing.n_physical_rows`` uses on the surrogate side.
+Both stages ran single-threaded (default Geant4 threading), pinned to one
+CPU core.
+
+``sim_us_per_step``/``convert_us_per_step``/``sim_ms_per_event`` below are
+the mean across the 3 energies. With 4 event-count points per energy (up
+from an initial 2-point pass, which had ~80% spread and nonsensical negative
+fitted intercepts at 10/50 GeV — an artifact of extrapolating a 2-point
+line), both quantities are now energy-flat as physically expected:
+``sim_us_per_step`` spread ~5%, ``convert_us_per_step`` spread ~13.5%. Treat
+these as reliable to about that precision.
+
+**Caveat — hardware asymmetry**: this reference is single-core CPU. A
+surrogate rollout's ``timing`` block will typically be measured on a batched
+GPU. The resulting ratio in ``eval_cost_per_step`` is a *deployment* speedup
+(what you'd actually see swapping Geant4 for the surrogate in a production
+pipeline), not a same-hardware or per-FLOP comparison — state this whenever
+quoting the number.
+
+**Staleness**: re-run ``scratchpad/bench_geant4.py`` (and update this file)
+if measured on different hardware, after a miniCaloSim/Geant4 version bump,
+or if this reference is more than a year or two stale.
+"""
+
+from __future__ import annotations
+
+GEANT4_REFERENCE: dict = {
+    "sim_us_per_step": 11.2903,
+    "convert_us_per_step": 11.1014,
+    "sim_ms_per_event": 609.6848,
+    "provenance": {
+        "cpu": "AMD Ryzen 7 PRO 4750U with Radeon Graphics",
+        "geant4_version": "11.4.1",
+        "minicalosim_sha": "ea917da",
+        "measured": "2026-08-31",
+        "energies_gev": [1.0, 10.0, 50.0],
+        "spread_pct_sim": 4.96,
+        "spread_pct_convert": 13.52,
+        "threads": 1,
+    },
+}
+
+
+def geant4_per_step_us() -> dict[str, float]:
+    """Sim / convert / total microseconds per physical step, from ``GEANT4_REFERENCE``."""
+    sim = GEANT4_REFERENCE["sim_us_per_step"]
+    convert = GEANT4_REFERENCE["convert_us_per_step"]
+    return {
+        "sim_us_per_step": sim,
+        "convert_us_per_step": convert,
+        "total_us_per_step": sim + convert,
+    }

giant/analysis/render.py

@@ -274,6 +274,8 @@ def _render_bar(r: Reduced, params: dict):
     ax.set_xticks(x)
     ax.set_xticklabels(labels, rotation=45, ha="right")
     ax.set_ylabel(r.payload.get("ylabel", "value"))
+    if r.payload.get("log_y"):
+        ax.set_yscale("log")
     ps.style_legend(ax, title="source")
     return fig
 

giant/analysis/runtime_estimate.py

@@ -96,6 +96,10 @@ _COST_MODEL: dict[str, tuple[float, float]] = {
     "sec_count_per_species": (0.0, 4.963e-07),
     "sec_energy": (0.0, 4.727e-07),
     "sec_cos_angle": (0.0, 2.749e-06),
+    # Metadata-only (YAML-carried `timing`, no row scan) — same shape as the
+    # router diagnostics' fixed cost, just cheaper since there's no live
+    # torch checkpoint to load.
+    "eval_cost_per_step": (0.0, 0.0),
 }
 
 

giant/analysis/sources.py

@@ -110,6 +110,7 @@ class RolloutSpec:
     source: str | Path | pl.LazyFrame
     checkpoint: str | None = None
     type_embedding_l1_dist: dict | None = None
+    timing: dict | None = None
 
 
 @dataclass
@@ -124,6 +125,10 @@ class RolloutSide:
     # only. Unlike checkpoint, this needs no live model: it's already a
     # finished histogram, just passed through.
     type_embedding_l1_dist: dict | None = None
+    # Wall-clock cost of this rollout run (giant.cli's rollout command),
+    # from the rollout YAML — eval_cost_per_step only. None on rollout runs
+    # that predate timing instrumentation.
+    timing: dict | None = None
 
 
 def _check_rollout_metadata(path: Path) -> None:

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, cast
 import uuid as uuid_mod
 
-import numpy as np
-import yaml
-import torch
 import typer
 from typing_extensions import Annotated
 
-import pyarrow as pa
+if TYPE_CHECKING:
-import pyarrow.parquet as pq
+    import numpy as np
-from tqdm import tqdm
 
 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,
+# giant.materials only pulls in numpy (no torch/pandas), and MATERIAL_PROPERTIES
-    find_parquet_files,
+# is needed at decoration time below (a Typer option default), so it can't be
-    iter_file_chunks,
+# deferred into a command body like the rest of this module's heavy imports.
-    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
 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)
 
@@ -210,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),
@@ -224,6 +205,45 @@ def _write_prediction_ref(
     return ref_path
 
 
+def _build_rollout_timing(
+    *,
+    setup_s: float,
+    rollout_s: float,
+    write_s: float,
+    n_rows: int,
+    termination_reason_counts: dict[str, int],
+    n_seed_events: int,
+    device: str,
+    torch_threads: int,
+) -> dict:
+    """Assemble ``giant rollout``'s ``timing`` sidecar block.
+
+    ``n_physical_rows`` excludes the synthetic termination rows (escape/
+    unknown-pdg/energy-cutoff/max-steps markers `giant.rollout` emits but
+    Geant4 never does) so ``us_per_step`` is comparable to
+    ``giant.analysis.geant4_reference``'s per-step Geant4 measurement — see
+    ``giant/analysis/catalog.py``'s ``eval_cost_per_step`` spec.
+    """
+    from giant.analysis.sources import SYNTHETIC_TERMINATION_REASONS
+
+    sample_s = rollout_s - write_s
+    n_synthetic_rows = sum(termination_reason_counts.get(reason, 0) for reason in SYNTHETIC_TERMINATION_REASONS)
+    n_physical_rows = n_rows - n_synthetic_rows
+    return {
+        "setup_s": setup_s,
+        "rollout_s": rollout_s,
+        "write_s": write_s,
+        "sample_s": sample_s,
+        "n_rows": n_rows,
+        "n_physical_rows": n_physical_rows,
+        "us_per_step": (sample_s / n_physical_rows * 1e6) if n_physical_rows else None,
+        "write_us_per_step": (write_s / n_physical_rows * 1e6) if n_physical_rows else None,
+        "ms_per_event": (rollout_s / n_seed_events * 1e3) if n_seed_events else None,
+        "device": device,
+        "torch_threads": torch_threads,
+    }
+
+
 @app.callback()
 def _main() -> None:
     """GIANT — Geant4 step-function surrogate."""
@@ -639,6 +659,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:
@@ -1048,6 +1072,25 @@ def predict(
     ] = 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":
@@ -1343,6 +1386,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):
@@ -1447,6 +1494,21 @@ def rollout(
     ] = None,
 ) -> None:
     """Roll the surrogate forward into full showers (autoregressive)."""
+    import time
+
+    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, RolloutSummary, rollout as run_rollout
+
+    _t_setup_start = time.perf_counter()
+
     if seed is not None:
         torch.manual_seed(seed)
         np.random.seed(seed)
@@ -1492,9 +1554,11 @@ def rollout(
     # avg_tracks_per_event) — mirrors the row-group streaming `giant predict`
     # already does on its input side.
     writer: pq.ParquetWriter | None = None
+    _write_s = 0.0
 
     def _write_chunk(row: dict[str, np.ndarray]) -> None:
-        nonlocal writer
+        nonlocal writer, _write_s
+        _t0 = time.perf_counter()
         table = pa.table(row)
         if writer is None:
             table = table.replace_schema_metadata(
@@ -1505,11 +1569,14 @@ def rollout(
             )
             writer = pq.ParquetWriter(out, table.schema)
         writer.write_table(table)
+        _write_s += time.perf_counter() - _t0
 
     # Only meaningful under particle_type.target="embedding" — a
     # no-op collector otherwise, cheaper than branching the call itself.
     l1_dist_collector = L1DistCollector()
 
+    _setup_s = time.perf_counter() - _t_setup_start
+    _t_rollout_start = time.perf_counter()
     summary = run_rollout(
         model,
         sec_decoder,
@@ -1541,6 +1608,23 @@ def rollout(
     )
     if writer is not None:
         writer.close()
+    # on_chunk=_write_chunk is always passed above, so rollout() always
+    # returns the streaming-summary shape (RolloutSummary), never the
+    # materialized dict[str, np.ndarray] alternative its return type allows.
+    summary = cast(RolloutSummary, summary)
+    _rollout_s = time.perf_counter() - _t_rollout_start
+    timing = _build_rollout_timing(
+        setup_s=_setup_s,
+        rollout_s=_rollout_s,
+        write_s=_write_s,
+        n_rows=summary["n_rows"],
+        termination_reason_counts=summary["termination_reason_counts"],
+        n_seed_events=len(seeds["event_id"]),
+        device=str(_device),
+        torch_threads=torch.get_num_threads(),
+    )
+    _sample_s = timing["sample_s"]
+    n_physical_rows = timing["n_physical_rows"]
 
     l1_summary = l1_dist_collector.summary()
 
@@ -1563,6 +1647,10 @@ def rollout(
             "rollout_seed": seed,
             "n_rows": summary["n_rows"],
             "termination_reason_counts": summary["termination_reason_counts"],
+            # Wall-clock cost of this run, normalized per physical step (the
+            # comparable unit against giant.analysis.geant4_reference) — see
+            # eval_cost_per_step in giant/analysis/catalog.py.
+            "timing": timing,
             # Diagnostic — only present under
             # stage2_model.particle_type.target="embedding"; omitted (not
             # written as null) otherwise, so giant.analysis can tell "not
@@ -1586,6 +1674,11 @@ def rollout(
 
     typer.echo(f"wrote {summary['n_rows']:,} step rows → {out}")
     typer.echo(f"terminations: {summary['termination_reason_counts']}")
+    if timing["us_per_step"] is not None:
+        typer.echo(
+            f"timing: {_rollout_s:.1f}s total ({_sample_s:.1f}s sample + {_write_s:.1f}s write), "
+            f"{timing['us_per_step']:.1f} us/step over {n_physical_rows:,} physical steps"
+        )
     typer.echo(f"reference: {ref_path}")
 
 

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
-
+from typing import TYPE_CHECKING
-import numpy as np
-import torch
 
 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):
@@ -941,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():
@@ -986,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()
@@ -1680,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(
@@ -1881,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)
@@ -1934,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/data/loader.py

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

giant/data/scan.py

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

giant/data/setup_cache.py

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

giant/geometry.py

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

giant/pipeline.py

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

giant/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,
+# DATA_DEFAULT/SCAN_DIR_DEFAULT are Typer option defaults (evaluated at
-    run_bump_schema,
+# decoration time below), so that one name has to stay eager — the module
-    run_create_manifest,
+# itself is stdlib-only, so it costs nothing. Every other giant.tools.*
-    run_status,
+# import here is deferred into the one command body that uses it, since
-    run_update_manifest,
+# several (steps_to_parquet: uproot/awkward/polars; warm_setup_cache:
-)
+# giant.pipeline -> torch; geometry_oracle: pandas) are expensive and
-from giant.tools.create_root_files import run_make_root
+# `dwarf --help`/tab-completion shouldn't pay for all of them upfront.
-from giant.tools.geometry_oracle import run_build_geometry_oracle
+from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT
-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
 
 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_setup_scan.py

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

pyproject.toml

@@ -1,12 +1,12 @@
 [project]
 name = "giant"
-version = "0.3.13"
+version = "0.3.17"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "numpy>=1.26,<3",
-    "pandas>=2.2,<4",
+    "polars>=1.0,<2",
     "pyarrow>=16,<25",
     "tqdm>=4.60,<5",
     "typer>=0.12,<1",
@@ -28,6 +28,10 @@ dev = [
     "ty>=0.0.50,<0.1",
     "bump-my-version>=1.2,<2",
     "git-cliff>=2,<3",
+    # Only used by test fixtures (writing small parquet files) — not a
+    # runtime dependency of giant itself since the pandas -> polars
+    # data-loading rewrite.
+    "pandas>=2.2,<4",
     "giant[convert,analysis,geometry,wandb]",
 ]
 geometry = [

tests/test_catalog.py

@@ -261,3 +261,32 @@ def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
     for j, _ in enumerate(cols):
         if j != g:
             assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0
+
+
+# ---------------------------------------------------------------------------
+# eval_cost_per_step
+# ---------------------------------------------------------------------------
+
+
+def test_eval_cost_per_step_unavailable_without_timing(bundle: Bundle):
+    # `bundle`'s RolloutSpec carries no `timing` -> no rollout to compare.
+    spec = get_spec("eval_cost_per_step")
+    r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
+    assert r.kind == "unavailable"
+    assert r.payload["note"]
+
+
+def test_eval_cost_per_step_bar_with_timing(ctx: Context):
+    spec = get_spec("eval_cost_per_step")
+    rs = RolloutSpec(
+        "rollout",
+        _rollout_frame(),
+        timing={"us_per_step": 12.5, "write_us_per_step": 2.5},
+    )
+    b = Bundle.open([rs], _reference_frame(), ctx)
+    r = spec.finalize([spec.compute_partial(b)], ctx)
+    assert r.kind == "bar"
+    assert r.payload["series"]["rollout"] == [12.5, 2.5, 15.0]
+    assert len(r.payload["reference"]) == 3
+    assert r.payload["log_y"] is True
+    assert "rollout" in r.meta["speedup_vs_geant4_total"]

tests/test_cli_rollout.py

@@ -8,11 +8,52 @@ from __future__ import annotations
 import torch
 from typer.testing import CliRunner
 
-from giant.cli import app
+from giant.cli import _build_rollout_timing, app
 
 runner = CliRunner()
 
 
+def test_build_rollout_timing_excludes_synthetic_rows_from_per_step_cost():
+    # 100 rows total, 30 of them synthetic termination markers (escape) ->
+    # us_per_step should be normalized over the 70 physical rows only, the
+    # same unit giant.analysis.geant4_reference measures Geant4 in.
+    timing = _build_rollout_timing(
+        setup_s=1.0,
+        rollout_s=10.0,
+        write_s=2.0,
+        n_rows=100,
+        termination_reason_counts={"escaped": 30, "natural_end": 70},
+        n_seed_events=5,
+        device="cpu",
+        torch_threads=4,
+    )
+    assert timing["n_rows"] == 100
+    assert timing["n_physical_rows"] == 70
+    assert timing["n_physical_rows"] < timing["n_rows"]
+    assert timing["sample_s"] == 8.0  # rollout_s - write_s
+    assert timing["us_per_step"] == 8.0 / 70 * 1e6
+    assert timing["write_us_per_step"] == 2.0 / 70 * 1e6
+    assert timing["ms_per_event"] == 10.0 / 5 * 1e3
+    assert timing["device"] == "cpu" and timing["torch_threads"] == 4
+
+
+def test_build_rollout_timing_handles_zero_physical_rows_and_events():
+    timing = _build_rollout_timing(
+        setup_s=1.0,
+        rollout_s=1.0,
+        write_s=0.0,
+        n_rows=5,
+        termination_reason_counts={"escaped": 5},
+        n_seed_events=0,
+        device="cpu",
+        torch_threads=1,
+    )
+    assert timing["n_physical_rows"] == 0
+    assert timing["us_per_step"] is None
+    assert timing["write_us_per_step"] is None
+    assert timing["ms_per_event"] is None
+
+
 def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
     checkpoint = tmp_path / "bad.pt"
     torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)

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

@@ -252,6 +252,22 @@ def test_compute_one_from_run_dir(tmp_path: Path):
     assert list(partial.data["r"]) == ["rollout"]
 
 
+def test_timing_survives_plot_meta_to_compute_one(tmp_path: Path):
+    yaml_path = _write_inputs(tmp_path)
+    d = yaml.safe_load(yaml_path.read_text())
+    d["timing"] = {"us_per_step": 7.0, "write_us_per_step": 1.0}
+    yaml_path.write_text(yaml.safe_dump(d))
+
+    run_dir = _prep([yaml_path])
+    meta = RunMeta.load(run_dir / "run_meta.json")
+    assert meta.rollouts[0]["plot_meta"]["timing"] == {"us_per_step": 7.0, "write_us_per_step": 1.0}
+
+    out = compute_one("eval_cost_per_step", run_dir)
+    reduced = Reduced(**Partial.load(out).data["reduced"])
+    assert reduced.kind == "bar"
+    assert reduced.payload["series"]["rollout"] == [7.0, 1.0, 8.0]
+
+
 def test_compute_reduced_explicit_paths(tmp_path: Path):
     run_dir = _prep([_write_inputs(tmp_path)])
     meta = RunMeta.load(run_dir / "run_meta.json")

tests/test_loader.py

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

tests/test_pipeline.py

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

tests/test_render.py

@@ -292,6 +292,20 @@ def test_render_one_of_each_kind(tmp_path: Path):
                 "ylabel": "frac",
             },
         ),
+        Reduced(
+            "cost",
+            "cost",
+            "bar",
+            "Cost",
+            "phase",
+            {
+                "labels": ["sample", "write", "total"],
+                "series": {"flow": [10.0, 1.0, 11.0]},
+                "reference": [5.0, 0.5, 5.5],
+                "ylabel": "us/step",
+                "log_y": True,
+            },
+        ),
         Reduced(
             "s",
             "species",

uv.lock

@@ -675,12 +675,12 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.13"
+version = "0.3.17"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },
-    { name = "pandas" },
     { name = "particle" },
+    { name = "polars" },
     { name = "pyarrow" },
     { name = "pyyaml" },
     { name = "tqdm" },
@@ -712,6 +712,7 @@ dev = [
     { name = "git-cliff" },
     { name = "ipykernel" },
     { name = "matplotlib" },
+    { name = "pandas" },
     { name = "plotstyle" },
     { name = "polars" },
     { name = "pytest" },
@@ -738,9 +739,10 @@ requires-dist = [
     { name = "ipykernel", marker = "extra == 'analysis'", specifier = ">=7.3.0" },
     { name = "matplotlib", marker = "extra == 'analysis'", specifier = ">=3.8,<4" },
     { name = "numpy", specifier = ">=1.26,<3" },
-    { name = "pandas", specifier = ">=2.2,<4" },
+    { name = "pandas", marker = "extra == 'dev'", specifier = ">=2.2,<4" },
     { name = "particle", specifier = ">=1.0,<2" },
     { name = "plotstyle", marker = "extra == 'analysis'", specifier = ">=1.0.0", index = "https://git.larsbogner.de/api/packages/lars/pypi/simple/" },
+    { name = "polars", specifier = ">=1.0,<2" },
     { name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0,<2" },
     { name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" },
     { name = "pyarrow", specifier = ">=16,<25" },