25 Commits

Author SHA1 Message Date
gitea-actions f3aa28eac7 chore: update changelog for v0.3.17
CI / Tests (push) Successful in 5m54s
CI / Lint (ruff check) (push) Successful in 2m47s
CI / Format (ruff format) (push) Successful in 2m45s
CI / Type check (ty) (push) Successful in 1m56s
CI / Sync project version with tag (push) Has been skipped
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 7s
2026-09-02 08:13:02 +00:00
gitea-actions a1df1faf51 chore: bump version 0.3.16 -> 0.3.17 2026-09-02 08:12:56 +00:00
lars 674f7254cd Merge pull request 'perf: replace pandas with polars in the setup-stage scan' (#93) from warm-cache-polars-scan into master
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 1m6s
CI / Lint (ruff check) (push) Successful in 1m10s
CI / Format (ruff format) (push) Successful in 1m10s
CI / Tests (push) Successful in 3m52s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 1m49s
Reviewed-on: #93
2026-09-02 10:07:04 +02:00
lars 7df1945384 perf: replace pandas with polars in the setup-stage scan
CI / Sync project version with tag (pull_request) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 1m6s
CI / Format (ruff format) (pull_request) Successful in 1m11s
CI / Type check (ty) (pull_request) Successful in 1m12s
CI / Tests (pull_request) Successful in 4m16s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
giant.pipeline.run_setup_stage (used by both giant train and dwarf
warm-cache) previously opened and fully read each parquet file 4-6
separate times via pandas, with per-row Python loops padding the
secondary list columns on every chunk of the normalizer-fitting pass.

- giant/data/loader.py: pandas -> polars throughout; ragged sec_*_list
  padding is now a single vectorized polars expression instead of a
  per-row Python loop (including a .iloc[i] loop for directions).
- giant/data/scan.py (new): a fused metadata scan answering the event
  index, pdg/material vocab, process counts, and pooled-pdg counts in
  one pass per file instead of one pass per section. Frequency-ranking
  ties are now an explicit (-count, first_seen) contract instead of an
  accident of pandas' value_counts iteration order.
- giant/pipeline.py: run_setup_stage restructured to consult the cache
  for every section first, then issue one combined scan request for
  whatever's missing.
- giant/geometry.py: ported the one remaining pandas groupby to polars.
- pyproject.toml: polars promoted to a core dependency, pandas moved
  to dev (only test fixtures still use it).
- giant/tools/profile_setup_scan.py (new): synthetic-data benchmark
  for this scan, mirroring profile_analysis_costs.py's pattern.

Also fixes a real deadlock this surfaced: DataLoader worker
subprocesses fork() on Linux, and polars' native thread pool doesn't
survive a fork — a worker touching polars after the parent already had
hangs instantly. giant/pipeline.py's train/val DataLoaders now use
multiprocessing_context="spawn" whenever num_workers>0.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01DdT32YWNEwnVLZUHsgdeSC
2026-09-02 09:59:37 +02:00
lars deb9e8e7de feat: add WGAN + AR stop-token config variant
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m23s
CI / Format (ruff format) (push) Successful in 1m23s
CI / Type check (ty) (push) Successful in 1m24s
CI / Tests (push) Successful in 3m32s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 8s
Branches off configs/baseline.toml with both stages on WGAN-GP and
stage-2 n_sec.mode = stop_token, 30 epochs — combines two unbenchmarked
roadmap axes (post-v0.3.0 WGAN, and the AR stop-token multiplicity mode)
into one variant that stays a single edit away from baseline for
attribution.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01KzPghrmFcAJYrWUvHApY9N
2026-08-31 14:32:26 +02:00
gitea-actions 292bf3d29f chore: update changelog for v0.3.16
CI / Tests (push) Successful in 5m2s
CI / Lint (ruff check) (push) Successful in 1m26s
CI / Format (ruff format) (push) Successful in 1m11s
CI / Type check (ty) (push) Successful in 27s
CI / Sync project version with tag (push) Successful in 7s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
CI / Publish package to Gitea package registry (push) Successful in 5s
2026-08-31 10:07:24 +00:00
gitea-actions 96748d1c5a chore: bump version 0.3.15 -> 0.3.16 2026-08-31 10:07:23 +00:00
lars 461fa33878 Merge pull request 'feat: add eval-cost benchmark — Geant4 reference vs surrogate rollout timing' (#92) from eval-cost-benchmark into master
CI / Format (ruff format) (push) Successful in 34s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 35s
CI / Lint (ruff check) (push) Successful in 54s
CI / Tests (push) Successful in 4m47s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 24s
Reviewed-on: #92
2026-08-31 12:01:15 +02:00
lars 2358a75ee1 feat: add eval-cost benchmark — Geant4 reference vs surrogate rollout timing
CI / Sync project version with tag (pull_request) Has been skipped
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Lint (ruff check) (pull_request) Successful in 54s
CI / Type check (ty) (pull_request) Successful in 53s
CI / Tests (pull_request) Successful in 2m33s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
Closes the roadmap's long-standing "no eval-latency number exists for any
configuration" gap. Instruments `giant rollout` to record per-physical-step
wall-clock cost in its YAML sidecar, adds a measured Geant4/miniCaloSim
per-step reference (giant/analysis/geant4_reference.py, from a 3-energy,
4-event-count-per-energy local benchmark), and wires both into a new
eval_cost_per_step PlotSpec in the giant analyze gallery.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-31 11:52:20 +02:00
lars 50d8368415 docs: record analysis_341dfb14 baseline rollout benchmark results
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 26s
CI / Format (ruff format) (push) Successful in 33s
CI / Type check (ty) (push) Successful in 33s
CI / Tests (push) Successful in 3m4s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 34s
Replaces the extrapolated pre-v0.3 weak-spot claims in baseline.toml's
header (which had the wrong sign on step-count error) with measured
numbers from the first full rollout validation of this exact config,
and adds a matching Roadmap entry in CLAUDE.md.
2026-08-28 15:01:35 +02:00
gitea-actions 95d5fc6d89 chore: update changelog for v0.3.15
CI / Publish package to Gitea package registry (push) Successful in 52s
CI / Lint (ruff check) (push) Successful in 24s
CI / Format (ruff format) (push) Successful in 1m8s
CI / Type check (ty) (push) Successful in 32s
CI / Tests (push) Successful in 3m2s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
CI / Sync project version with tag (push) Successful in 15s
2026-08-28 12:28:02 +00:00
gitea-actions b8bd1ec982 chore: bump version 0.3.14 -> 0.3.15 2026-08-28 12:28:01 +00:00
lars 70d018982b Merge pull request 'perf: defer heavy imports in giant/dwarf CLIs until commands run' (#91) from cli-lazy-imports into master
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m33s
CI / Type check (ty) (push) Successful in 1m36s
CI / Format (ruff format) (push) Successful in 1m37s
CI / Tests (push) Successful in 3m40s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 2m58s
Reviewed-on: #91
2026-08-28 14:21:08 +02:00
lars 8d1c29efdd Merge branch 'master' into cli-lazy-imports
CI / Format (ruff format) (pull_request) Successful in 29s
CI / Lint (ruff check) (pull_request) Successful in 34s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 32s
CI / Tests (pull_request) Successful in 2m51s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
2026-08-28 14:11:11 +02:00
lars 516a8a9ee1 perf: defer heavy imports in giant/dwarf CLIs until commands run
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 19s
CI / Format (ruff format) (push) Successful in 21s
CI / Lint (ruff check) (push) Successful in 25s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 4m5s
CI / Format (ruff format) (pull_request) Successful in 4m5s
CI / Lint (ruff check) (pull_request) Successful in 4m15s
CI / Tests (pull_request) Successful in 5m15s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Tests (push) Successful in 9m55s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
torch/pandas/pyarrow/polars/uproot/awkward/particle were all imported
at module scope in giant/cli.py and giant/tools/dwarf.py, so even
`--help` paid ~1.6-1.9s of import cost. Move those imports into the
command bodies that actually need them (following the deferred-import
pattern already used for analysis/render/plots/sklearn/wandb), cutting
`giant --help` to ~0.3s and `dwarf --help` to ~0.2s with no change to
any command's actual behavior.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-28 14:10:19 +02:00
gitea-actions c12acfdade chore: update changelog for v0.3.14
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 31s
CI / Lint (ruff check) (push) Successful in 2m29s
CI / Format (ruff format) (push) Successful in 2m27s
CI / Type check (ty) (push) Successful in 2m34s
CI / Tests (push) Successful in 6m1s
CI / Sync project version with tag (push) Has been skipped
CI / Publish package to Gitea package registry (push) Has been skipped
2026-08-28 11:19:14 +00:00
gitea-actions e06d9e9581 chore: bump version 0.3.13 -> 0.3.14 2026-08-28 11:18:49 +00:00
lars b0998a7d86 Merge pull request 'ci: give automated commits visible checks, scope CI triggers, publish releases' (#90) from ci/pr-scoped-checks-and-package-publish into master
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 25s
CI / Lint (ruff check) (push) Successful in 38s
CI / Format (ruff format) (push) Successful in 36s
CI / Tests (push) Successful in 2m45s
CI / Publish package to Gitea package registry (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 2m36s
Reviewed-on: #90
2026-08-28 13:14:18 +02:00
lars cc11efb3ae ci: fix pull_request trigger not registering
CI / Sync project version with tag (pull_request) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 39s
CI / Format (ruff format) (pull_request) Successful in 43s
CI / Type check (ty) (pull_request) Successful in 1m6s
CI / Tests (pull_request) Successful in 3m29s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Publish package to Gitea package registry (pull_request) Has been skipped
A bare "pull_request:" key parses to null in YAML, which the runner
apparently doesn't treat as "trigger with defaults" the way an empty
mapping does — the open PR for this branch got no CI run at all.
2026-08-28 13:08:25 +02:00
lars 7de3e92871 ci: give automated commits visible checks, scope CI triggers, publish releases
- Drop [skip ci] from the bump-version/changelog/tag-sync commits so
  master's tip always has a check run instead of only the merge commit.
- Filter those chore commits out of the changelog via message pattern
  instead of the now-removed [skip ci] tag.
- Only run CI on push to master (plus tags); pull requests to any branch
  still run the full suite.
- Add a publish-package job that builds and publishes to the Gitea PyPI
  registry on every tag push, after tests and version sync pass.
2026-08-28 13:06:26 +02:00
gitea-actions f80fc90758 chore: update changelog for v0.3.13 [skip ci] 2026-08-28 09:50:33 +00:00
gitea-actions 1cf16526c9 chore: bump version 0.3.12 -> 0.3.13 [skip ci] 2026-08-28 09:50:33 +00:00
lars 5c93457081 Merge pull request 'Fix/issue 87' (#89) from fix/issue-87 into master
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 25s
CI / Format (ruff format) (push) Successful in 1m8s
CI / Type check (ty) (push) Successful in 1m29s
CI / Tests (push) Successful in 3m21s
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 1m32s
Reviewed-on: #89
2026-08-28 11:45:00 +02:00
lars 1ec333ff6d Merge remote-tracking branch 'origin/master' into fix/issue-87
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m36s
CI / Format (ruff format) (push) Successful in 1m41s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (push) Successful in 2m17s
CI / Lint (ruff check) (pull_request) Successful in 1m19s
CI / Format (ruff format) (pull_request) Successful in 2m15s
CI / Type check (ty) (pull_request) Successful in 4m13s
CI / Tests (pull_request) Successful in 7m18s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CI / Tests (push) Successful in 11m5s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
# Conflicts:
#	giant/config.py
2026-08-28 11:31:05 +02:00
lars bd255419e1 Add inference-time model_config overrides with a sampling-key allowlist (gitea #87)
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 1m45s
CI / Format (ruff format) (push) Successful in 3m0s
CI / Type check (ty) (push) Successful in 3m16s
CI / Tests (push) Successful in 3m30s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
giant predict/rollout rebuilt models straight from ckpt["model_config"] with
no way to change sampling-only keys (e.g. stage2_model.n_sec.stop_sampling)
without retraining. Adds config_overrides to load_for_inference, validated
against giant.config.INFERENCE_OVERRIDES so a typo or shape-bearing key
raises CheckpointCompatibilityError up front instead of an opaque
load_state_dict mismatch. Wired as a repeatable --set dotted.path=value on
both CLI commands, recorded in the rollout YAML sidecar, and surfaced in
`giant model summary`'s output.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-28 11:21:36 +02:00
36 changed files with 1646 additions and 274 deletions
+2 -2
View File
@@ -1,5 +1,5 @@
[tool.bumpversion]
current_version = "0.3.12"
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]]
+34 -5
View File
@@ -2,10 +2,9 @@ name: CI
"on":
push:
branches: ["**"]
branches: ["master"]
tags: ["**"]
pull_request:
branches: [master]
pull_request: {}
env:
UV_CACHE_DIR: /uv-cache
@@ -156,7 +155,7 @@ jobs:
uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
git add CHANGELOG.md
if ! git diff --cached --quiet -- CHANGELOG.md; then
git commit -m "chore: update changelog for $TAG [skip ci]"
git commit -m "chore: update changelog for $TAG"
else
git restore --staged CHANGELOG.md
fi
@@ -193,7 +192,7 @@ jobs:
git config user.name "gitea-actions"
git config user.email "actions@git.larsbogner.de"
git add pyproject.toml uv.lock
git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME} [skip ci]"
git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME}"
git push origin HEAD:master
git push origin ":refs/tags/${GITHUB_REF_NAME}"
git tag -f "${GITHUB_REF_NAME}" HEAD
@@ -201,3 +200,33 @@ jobs:
else
echo "Tag version matches project version ($CURRENT_VERSION)"
fi
publish-package:
name: Publish package to Gitea package registry
needs: [ruff-check, ruff-format, type-check, test, sync-version-on-tag]
if: startsWith(github.ref, 'refs/tags/')
runs-on: ubuntu-latest
container:
image: docker.gitea.com/runner-images:ubuntu-latest
volumes:
- /srv/act-runner-cache/uv:/uv-cache
steps:
# Check out by tag name (not the triggering SHA) since sync-version-on-tag
# may have force-moved the tag to a version-corrected commit.
- uses: actions/checkout@v4
with:
ref: ${{ github.ref_name }}
- uses: astral-sh/setup-uv@v5
with:
enable-cache: false
- run: |
echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
- run: uv build
# CI_TOKEN needs write:package scope (in addition to write:repository,
# used elsewhere) for this upload to authenticate.
- run: |
uv publish \
--publish-url "https://git.larsbogner.de/api/packages/lars/pypi" \
--username gitea-actions \
--password "${{ secrets.CI_TOKEN }}"
+36
View File
@@ -1,5 +1,41 @@
# 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
- Add inference-time model_config overrides with a sampling-key allowlist [gitea #87](https://git.larsbogner.de/lars/giant/issues/87)
## [0.3.12] - 2026-08-28
### Added
+2
View File
@@ -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.
+1 -1
View File
@@ -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" },
+15 -5
View File
@@ -29,11 +29,21 @@
# capacity overfitting is not the binding constraint, and every recent
# run used 0.0.
#
# Known weak spots this baseline is expected to *exhibit* (they are the
# reason for the comparisons, not a reason to retune this file): every model
# on record under-produces steps per event by ~2x (rollout ~7e4 vs Geant4
# ~1.4e5) and secondaries per event by 2-3.5x (~2-3e4 vs 7.2e4), and n_sec
# head accuracy sits at 0.863-0.867 regardless of size or objective.
# Known weak spots, now measured against this exact config rather than
# extrapolated from the pre-v0.3 field (analysis_341dfb14, best.pt @ epoch
# 50/50, full writeup: knowledge-base/experiments/
# giant-baseline-flow-ar-rollout-validation.md). Unlike every pre-v0.3
# checkpoint (which under-produced steps/event by 1.6-5x), this baseline
# OVER-produces steps/event by 1.32x (1.86e5 vs Geant4 1.41e5) and
# under-produces secondaries/event by 0.84x (5.97e4 vs 7.14e4) — the sign on
# steps flipped with the v0.3 autoregressive pivot, so don't assume it still
# undershoots. Secondary-species hallucination (zero photons, hallucinated
# `-14` muon antineutrinos) that broke every prior checkpoint is gone; the
# remaining species gap is a total absence of hadronic/nuclear secondaries
# (protons, neutrons, ion recoils), not miscalibration of the ones produced.
# Total deposited energy/event is +1.9% high but its event-to-event spread is
# ~16x too narrow (31 MeV vs Geant4's 491 MeV). Per-step deposited energy is
# the worst per-step marginal (KS 0.179 vs 0.004-0.071 for the others).
[meta]
# REQUIRED. Without it config.migrate_config reads this file as v0.2 and
+122
View File
@@ -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"
+83
View File
@@ -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",
+6 -3
View File
@@ -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"?},
...]``, one per rollout series (insertion order preserved through to every
plot's ``Reduced.payload["series"]``).
``rollouts``: ``[{"name", "path", "checkpoint"?, "type_embedding_l1_dist"?,
"timing"?}, ...]``, one per rollout series (insertion order preserved
through to every plot's ``Reduced.payload["series"]``).
Writes a ``Partial`` JSON the raw, not-yet-merged output of
``PlotSpec.compute_partial`` never a finished ``Reduced``; ``merge_one``
@@ -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
]
+76
View File
@@ -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,
}
+2
View File
@@ -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
+4
View File
@@ -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),
}
+5
View File
@@ -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:
+52 -3
View File
@@ -14,11 +14,18 @@ directly and imported from non-CLI code (`giant.analysis.router_gating`,
lazily see that module's docstring for why). Failures raise
`CheckpointCompatibilityError` with the same wording the CLI has always
shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
`load_for_inference`'s `config_overrides` (gitea #87) lets a caller change a
checkpoint's `model_config` at load time, restricted to
`giant.config.INFERENCE_OVERRIDES` the allowlist of keys that only affect
sampling, never module construction/shapes or the preprocessing normalizers/
vocab maps were fit under.
"""
from __future__ import annotations
from dataclasses import dataclass
import copy
from dataclasses import dataclass, field
from pathlib import Path
import torch
@@ -29,13 +36,45 @@ from giant.constants import K_MAX
from giant.data.loader import TopNMap
from giant.data.setup_cache import topnmap_from_json
from giant.data.transforms import Normalizer
from giant.model.network import build_models
from giant.model.network import _migrate_legacy_model_config, build_models
class CheckpointCompatibilityError(Exception):
"""Checkpoint is missing something `load_for_inference` needs."""
def apply_config_overrides(model_cfg: dict, overrides: dict[str, object] | None) -> dict:
"""Deep-merge dotted-path *overrides* into a checkpoint's `model_config`,
validated against `giant.config.INFERENCE_OVERRIDES` the allowlist of
keys that only affect sampling, not module construction/shapes or the
preprocessing normalizers/vocab maps were fit under (gitea #87).
Migrates a v0.2 flat `model_config` to the nested v0.3 shape first: a
dotted path like "stage1_model.ddpm.n_steps" would otherwise silently
write into a dict that `build_models` still reads as flat (it decides
v0.2-vs-v0.3 by `"stage1_model" in model_config`), suppressing migration.
Raises `CheckpointCompatibilityError` never a bare `ValueError` or a
downstream `load_state_dict` size mismatch for an unknown/disallowed
path or a value that fails its allowlisted check.
"""
if not overrides:
return model_cfg
cfg = model_cfg if "stage1_model" in model_cfg else _migrate_legacy_model_config(model_cfg)
cfg = copy.deepcopy(cfg)
for path, value in overrides.items():
spec = gconfig.INFERENCE_OVERRIDES.get(path)
if spec is None:
allowed = ", ".join(sorted(gconfig.INFERENCE_OVERRIDES))
raise CheckpointCompatibilityError(f"{path!r} is not an inference-safe override — allowed paths: {allowed}")
try:
spec.check(path, value)
except ValueError as exc:
raise CheckpointCompatibilityError(str(exc)) from exc
gconfig._set_path(cfg, path, value)
return cfg
def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
"""(particle_conditioning, material_conditioning) for
`giant.data.transforms.build_cond_features`/`build_features` from
@@ -128,6 +167,7 @@ class InferenceContext:
model_config: dict
epoch: int | None
best_val_loss: float | None
config_overrides: dict[str, object] = field(default_factory=dict)
def load_for_inference(
@@ -136,6 +176,7 @@ def load_for_inference(
command_name: str,
weights: str = "raw",
require_stage2: bool = True,
config_overrides: dict[str, object] | None = None,
) -> InferenceContext:
"""Load *checkpoint* and reconstruct everything `predict`/`rollout` need
to run it forward, on *device*, in `eval()` mode.
@@ -148,6 +189,13 @@ def load_for_inference(
both stages) or an acceptable `stage2 = None` result kept as a real
parameter since `stage{1,2}_model.active` is a real, if currently
stage1+stage2-only-in-practice, config option.
*config_overrides* deep-merges dotted `model_config` paths (e.g.
`{"stage2_model.n_sec.sampling": "sample"}`) before anything is
derived from `model_config` or built see `apply_config_overrides` for
the allowlist and validation. Every derived `InferenceContext` field
(`other_policy`, `stage{1,2}_ddpm_steps`, the built modules, ...)
reflects the overridden config.
"""
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
for key in ("model_config", "sec_decoder"):
@@ -159,7 +207,7 @@ def load_for_inference(
gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
model_cfg = ckpt["model_config"]
model_cfg = apply_config_overrides(ckpt["model_config"], config_overrides)
particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
pdg_topn_map = load_pdg_topn_map(ckpt)
mat_topn_map = load_mat_topn_map(ckpt)
@@ -231,4 +279,5 @@ def load_for_inference(
model_config=model_cfg,
epoch=ckpt.get("epoch"),
best_val_loss=ckpt.get("best_val_loss"),
config_overrides=dict(config_overrides) if config_overrides else {},
)
+168 -33
View File
@@ -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
import pyarrow.parquet as pq
from tqdm import tqdm
if TYPE_CHECKING:
import numpy as np
from giant import config as gconfig
from giant.constants import (
@@ -25,30 +23,11 @@ from giant.constants import (
PREDICT_SCHEMA_VERSION_KEY,
ROLLOUT_COORD_VALUE,
)
from giant.data.loader import (
event_id_offset,
find_parquet_files,
iter_file_chunks,
iter_cond_chunks,
)
from giant.data.transforms import (
build_features,
build_cond_features,
energy_simplex_decode,
inv_local_frame_rotation,
inv_log_transform,
reconstruct_post_pos,
)
from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
from giant.geometry import GeometryOracle
# giant.materials only pulls in numpy (no torch/pandas), and MATERIAL_PROPERTIES
# is needed at decoration time below (a Typer option default), so it can't be
# deferred into a command body like the rest of this module's heavy imports.
from giant.materials import MATERIAL_PROPERTIES
from giant.pipeline import run_train_job
from giant.rollout import (
L1DistCollector,
decode_secondary_identity,
rollout as run_rollout,
)
from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
app = typer.Typer(no_args_is_help=True)
@@ -141,6 +120,23 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
return out
def _parse_set_flags(specs: Optional[list[str]]) -> dict[str, object]:
"""Parse repeated `--set dotted.path=value` flags into a dict, typing
each value with `_coerce_scalar` the same way a TOML file's native types
would arrive. Validation against the inference-safe allowlist happens
downstream in `giant.checkpoint_io.apply_config_overrides` this only
parses syntax.
"""
out: dict[str, object] = {}
for spec in specs or []:
path, sep, val = spec.partition("=")
if not sep:
typer.echo(f"error: --set {spec!r} must be 'dotted.path=value'", err=True)
raise typer.Exit(1)
out[path] = _coerce_scalar(val)
return out
def _router_cli_overrides(
router: bool | None,
router_type: str | None,
@@ -193,6 +189,8 @@ def _write_prediction_ref(
comment: str | None = None,
) -> Path:
"""Write a YAML sidecar in the checkpoint directory and return its path."""
import yaml
ref = {
"prediction_id": pred_uuid,
"output": str(out),
@@ -207,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."""
@@ -622,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:
@@ -1020,8 +1061,36 @@ def predict(
help="Free-text note recorded in the prediction's YAML sidecar",
),
] = None,
set_: Annotated[
Optional[list[str]],
typer.Option(
"--set",
help="Override a sampling-only model_config key on this checkpoint, "
"'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
"for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
),
] = None,
) -> None:
"""Run trained model on a parquet file and save predictions."""
import numpy as np
import pyarrow as pa
import pyarrow.parquet as pq
import torch
from tqdm import tqdm
from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
from giant.data.loader import event_id_offset, find_parquet_files, iter_cond_chunks, iter_file_chunks
from giant.data.transforms import (
build_cond_features,
build_features,
energy_simplex_decode,
inv_local_frame_rotation,
inv_log_transform,
reconstruct_post_pos,
)
from giant.rollout import decode_secondary_identity
from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
batch_size_auto = False
batch_size_value: Optional[int] = None
if batch_size.strip().lower() == "auto":
@@ -1040,8 +1109,11 @@ def predict(
typer.echo(f"device: {_device}")
# --- Load checkpoint ---
config_overrides = _parse_set_flags(set_)
try:
ctx = load_for_inference(checkpoint, _device, "predict", weights=weights.value)
ctx = load_for_inference(
checkpoint, _device, "predict", weights=weights.value, config_overrides=config_overrides
)
except CheckpointCompatibilityError as exc:
typer.echo(f"error: {exc}", err=True)
raise typer.Exit(1)
@@ -1314,6 +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):
@@ -1407,8 +1483,32 @@ def rollout(
Optional[int],
typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
] = None,
set_: Annotated[
Optional[list[str]],
typer.Option(
"--set",
help="Override a sampling-only model_config key on this checkpoint, "
"'dotted.path=value' (repeatable) — see giant.config.INFERENCE_OVERRIDES "
"for the allowlist, e.g. --set stage2_model.n_sec.sampling=sample",
),
] = None,
) -> None:
"""Roll the surrogate forward into full showers (autoregressive)."""
import 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)
@@ -1416,8 +1516,11 @@ def rollout(
_device = torch.device(device) if device else gconfig.auto_device()
typer.echo(f"device: {_device}")
config_overrides = _parse_set_flags(set_)
try:
ctx = load_for_inference(checkpoint, _device, "rollout", weights=weights.value)
ctx = load_for_inference(
checkpoint, _device, "rollout", weights=weights.value, config_overrides=config_overrides
)
except CheckpointCompatibilityError as exc:
typer.echo(f"error: {exc}", err=True)
raise typer.Exit(1)
@@ -1451,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(
@@ -1464,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,
@@ -1500,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()
@@ -1522,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
@@ -1532,6 +1661,7 @@ def rollout(
# model knob (router type/n_experts, noise_dim, vocab sizes, ...)
# is available downstream without touching this command again.
"model_config": dict(model_cfg),
"config_overrides": dict(ctx.config_overrides),
"training_epoch": ctx.epoch,
"best_val_loss": ctx.best_val_loss,
# [train]/[meta] from the sibling config.toml (giant.config.save_config)
@@ -1544,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}")
+94 -5
View File
@@ -1,3 +1,5 @@
from __future__ import annotations
import copy
import difflib
import hashlib
@@ -10,12 +12,12 @@ from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from pathlib import Path
import numpy as np
import torch
from typing import TYPE_CHECKING
from giant._migration import V02_FIXED_FACTS, V02_MODEL_KEY_TO_STAGES, reject_legacy_router_expert_sizing
from giant.model.history import HISTORY_REGISTRY
if TYPE_CHECKING:
import torch
class Conditioning(str, Enum):
@@ -404,6 +406,16 @@ class Stage2RouterConfig(RouterConfig):
return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()}
# stage2_model.n_sec.sampling choices — single source of truth for both
# validate_config's train-time check and INFERENCE_OVERRIDES below.
STOP_SAMPLING_CHOICES = ("greedy", "sample")
# stage2_model.particle_type.other_policy choices — see ParticleTypeConfig's
# docstring for what each means; only documented there until now, since
# nothing validated it at train time.
OTHER_POLICY_CHOICES = ("sample", "modal", "drop")
@dataclass(frozen=True)
class NSecConfig:
# "head": a classifier over {0..k_max} on the condition encoding alone
@@ -931,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():
@@ -976,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()
@@ -1118,6 +1134,72 @@ def _deep_merge(base: dict, override: dict) -> dict:
return result
@dataclass(frozen=True)
class InferenceOverride:
"""One dotted `model_config` path that `giant.checkpoint_io.load_for_inference`
is allowed to change on an already-trained checkpoint, without retraining.
A path only belongs here if it affects neither module construction/tensor
shapes nor the data preprocessing the normalizers/vocab maps were fit
under see the module docstring on `giant.model.summary` for the class
of key this targets (`_fingerprint`'s "plain scalar attribute" leaves),
and `giant.checkpoint_io.apply_config_overrides` for where this is used.
"""
why: str
choices: tuple[str, ...] | None = None
minimum: float | None = None
numeric: bool = False # int/float leaf (vs. str, the default)
def check(self, path: str, value: object) -> None:
if self.choices is not None:
if value not in self.choices:
raise ValueError(f"{path} = {value!r} — must be one of {self.choices}")
return
if self.numeric:
if isinstance(value, bool) or not isinstance(value, (int, float)):
raise ValueError(f"{path} = {value!r} — must be a number")
if self.minimum is not None and value < self.minimum:
raise ValueError(f"{path} = {value!r} — must be >= {self.minimum}")
# Inference-safe dotted `model_config` paths — the allowlist gitea #87 asked
# for, so a typo or a shape-bearing key (e.g. "stage1_model.hidden_dim")
# raises a clear CheckpointCompatibilityError instead of surfacing as an
# opaque load_state_dict size mismatch later. Extend this table, not a
# per-call bypass, when a new inference-only key needs the capability.
INFERENCE_OVERRIDES: dict[str, InferenceOverride] = {
"stage2_model.n_sec.sampling": InferenceOverride(
why="giant.sample's n_sec head/stop-token sampling reads this at sample time only (gitea #86)",
choices=STOP_SAMPLING_CHOICES,
),
"stage1_model.ddpm.n_steps": InferenceOverride(
why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
numeric=True,
minimum=1,
),
"stage2_model.ddpm.n_steps": InferenceOverride(
why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
numeric=True,
minimum=1,
),
"stage2_model.particle_type.other_policy": InferenceOverride(
why="giant.rollout resolves an 'other'-bucket secondary's PDG code with this at rollout time",
choices=OTHER_POLICY_CHOICES,
),
"stage1_model.router.temperature": InferenceOverride(
why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
numeric=True,
minimum=1e-6,
),
"stage2_model.router.temperature": InferenceOverride(
why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
numeric=True,
minimum=1e-6,
),
}
@dataclass(frozen=True)
class FlagSpec:
"""One CLI flag's mapping into the config-overrides tree.
@@ -1548,7 +1630,7 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
)
n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
if n_sec_sampling not in ("greedy", "sample"):
if n_sec_sampling not in STOP_SAMPLING_CHOICES:
raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")
precision = _get_path(cfg, "train.precision")
@@ -1604,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(
@@ -1805,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)
@@ -1858,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(),
+92 -115
View File
@@ -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:
"""Pad / truncate a list-valued Series to fixed width K → (N, K) float32."""
out = np.full((len(series), K), fill, dtype=np.float32)
for i, lst in enumerate(series):
if lst is not None and len(lst) > 0:
n = min(len(lst), K)
out[i, :n] = lst[:n]
return out
def _pad_list_column(df: pl.DataFrame, col: str, k: int, fill, dtype: type[pl.DataType] | pl.DataType) -> np.ndarray:
"""Pad / truncate a list-valued column to fixed width `k` → (N, k) numpy array.
def _pad_list_col_int(series: pd.Series, K: int, fill: int = 0) -> np.ndarray:
"""Pad / truncate a list-valued integer Series to fixed width K → (N, K) int64."""
out = np.full((len(series), K), fill, dtype=np.int64)
for i, lst in enumerate(series):
if lst is not None and len(lst) > 0:
n = min(len(lst), K)
out[i, :n] = lst[:n]
return out
def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndarray:
"""Pad three list-valued direction columns → (N, K, 3) float32.
Padding direction defaults to (0,0,1) (forward) so it is a valid unit vector.
Concatenating `k` fill values before truncating to `k` guarantees every
row ends up with exactly `k` non-null elements regardless of how short
(including empty) or long the original list was, so `list.to_array(k)`
(a fixed-size-array dtype) converts to a plain 2D numpy array with a
single vectorized expression no per-row Python loop.
"""
N = len(dx)
out = np.zeros((N, K, 3), dtype=np.float32)
out[:, :, 2] = 1.0
for i in range(N):
lx, ly, lz = dx.iloc[i], dy.iloc[i], dz.iloc[i]
if lx is not None and len(lx) > 0:
n = min(len(lx), K)
out[i, :n, 0] = lx[:n]
out[i, :n, 1] = ly[:n]
out[i, :n, 2] = lz[:n]
return out
fill_tail = pl.lit([fill] * k, dtype=pl.List(dtype))
out = df.select(pl.col(col).cast(pl.List(dtype)).list.concat(fill_tail).list.head(k).list.to_array(k).alias("_p"))
return out["_p"].to_numpy()
def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
def _pad_dir_col(df: pl.DataFrame, dx: str, dy: str, dz: str, k: int) -> np.ndarray:
"""Pad three list-valued direction columns → (N, k, 3) float32.
Padding direction defaults to (0, 0, 1) (forward) so it is a valid unit vector.
"""
px = _pad_list_column(df, dx, k, 0.0, pl.Float64)
py = _pad_list_column(df, dy, k, 0.0, pl.Float64)
pz = _pad_list_column(df, dz, k, 1.0, pl.Float64)
return np.stack([px, py, pz], axis=-1).astype(np.float32)
def _df_to_dict(df: pl.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
has_sec_lists = "sec_E_list" in df.columns
d: dict[str, np.ndarray] = {
"event_id": _offset_event_id(df["event_id"].to_numpy(), offset),
"pdg": df["pdg"].to_numpy(dtype=np.int32),
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_E": df["pre_E"].to_numpy(dtype=np.float32),
"pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
"material": df["material"].to_numpy(dtype=object),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
"pdg": df["pdg"].to_numpy().astype(np.int32),
"pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
"pre_E": df["pre_E"].to_numpy().astype(np.float32),
"pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
"material": df["material"].to_numpy().astype(object),
"layer_id": df["layer_id"].to_numpy().astype(np.int32),
"n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
"e_sec": df["e_sec"].to_numpy().astype(np.float32),
# The physics process that ended the step (e.g. "compt", "phot",
# "eBrem") — a post-step outcome, so it's a router/classifier
# supervision label only, never conditioning (see build_process_map*
# / ProcessRouter). Guarded like has_sec_lists: older parquet
# conversions predating this column still load fine.
"process": (
df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
df["process"].to_numpy().astype(object) if "process" in df.columns else np.full(len(df), "", dtype=object)
),
"step_length": df["step_length"].to_numpy(dtype=np.float32),
"post_E": df["post_E"].to_numpy(dtype=np.float32),
"delta_e": (df["pre_E"] - df["post_E"]).to_numpy(dtype=np.float32),
"edep": df["edep"].to_numpy(dtype=np.float32),
"post_dir": df[["post_dx", "post_dy", "post_dz"]].to_numpy(dtype=np.float32),
"post_pos": df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32),
"step_length": df["step_length"].to_numpy().astype(np.float32),
"post_E": df["post_E"].to_numpy().astype(np.float32),
"delta_e": (df["pre_E"] - df["post_E"]).to_numpy().astype(np.float32),
"edep": df["edep"].to_numpy().astype(np.float32),
"post_dir": df.select(["post_dx", "post_dy", "post_dz"]).to_numpy().astype(np.float32),
"post_pos": df.select(["post_x", "post_y", "post_z"]).to_numpy().astype(np.float32),
}
if has_sec_lists:
d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)
d["sec_E_list"] = _pad_list_column(df, "sec_E_list", k_max, 0.0, pl.Float64).astype(np.float32)
d["sec_pdg_list"] = _pad_list_column(df, "sec_pdg_list", k_max, 0, pl.Int64).astype(np.int64)
d["sec_dir_list"] = _pad_dir_col(df, "sec_dx_list", "sec_dy_list", "sec_dz_list", k_max)
return d
def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
return _df_to_dict(pl.read_parquet(path), offset=offset, k_max=k_max)
def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
"""Read only the event_id column — cheap scan for split assignment."""
ids = pd.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
ids = pl.read_parquet(path, columns=["event_id"])["event_id"].to_numpy()
return _offset_event_id(ids, offset)
@@ -170,7 +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),
"pre_pos": df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32),
"pre_E": df["pre_E"].to_numpy(dtype=np.float32),
"pre_dir": df[["pre_dx", "pre_dy", "pre_dz"]].to_numpy(dtype=np.float32),
"material": df["material"].to_numpy(dtype=object),
"layer_id": df["layer_id"].to_numpy(dtype=np.int32),
"n_sec": df["child_track_ids"].apply(len).to_numpy(dtype=np.int32),
"e_sec": df["e_sec"].to_numpy(dtype=np.float32),
"pdg": df["pdg"].to_numpy().astype(np.int32),
"pre_pos": df.select(["pre_x", "pre_y", "pre_z"]).to_numpy().astype(np.float32),
"pre_E": df["pre_E"].to_numpy().astype(np.float32),
"pre_dir": df.select(["pre_dx", "pre_dy", "pre_dz"]).to_numpy().astype(np.float32),
"material": df["material"].to_numpy().astype(object),
"layer_id": df["layer_id"].to_numpy().astype(np.int32),
"n_sec": df["child_track_ids"].list.len().to_numpy().astype(np.int32),
"e_sec": df["e_sec"].to_numpy().astype(np.float32),
}
@@ -208,7 +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)."""
pdg_vals: set[int] = set()
mat_vals: set[str] = set()
for path in files:
df = pd.read_parquet(path, columns=["pdg", "material"])
pdg_vals.update(int(v) for v in df["pdg"].unique())
mat_vals.update(str(v) for v in df["material"].unique())
"""Scan only pdg and material columns across all files (fused single-pass scan)."""
from giant.data.scan import ScanRequest, scan_metadata
result = scan_metadata(files, ScanRequest(pdg=True, material=True))
assert result.pdg is not None and result.material is not None
return (
{v: i for i, v in enumerate(sorted(pdg_vals))},
{v: i for i, v in enumerate(sorted(mat_vals))},
{v: i for i, v in enumerate(sorted(result.pdg))},
{v: i for i, v in enumerate(sorted(result.material))},
)
def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None:
for name, count in series.value_counts().items():
name = cast(name)
counts[name] = counts.get(name, 0) + int(count)
def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
"""Scan `column` across `files` and return `{cast(value): total_count}`,
accumulated in file order (see `fingerprint_files`'s docstring on why
scan order not a normalized/sorted order is preserved: it drives
tie-breaking in the frequency ranking below)."""
counts: dict = {}
for path in files:
df = pd.read_parquet(path, columns=[column])
_accumulate_value_counts(counts, df[column], cast)
return counts
def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict, dict]:
def _topn_plus_other_map(counts: "Mapping[Any, ValueStat]", n_classes: int) -> tuple[dict, dict, dict]:
"""Frequency-capped value->index map: the `n_classes - 1` most frequent
keys get their own index; every rarer key is bucketed into a shared
"other" index (`n_classes - 1`).
`counts` maps each key to something with `.count` and `.first_seen`
attributes (`giant.data.scan.ValueStat`) ties in `.count` are broken by
`.first_seen` (whichever value was scanned first: file order, then row
order within a file see `giant.data.scan`'s module docstring). This is
an explicit, documented contract, not an accident of iteration order.
Returns `(class_map, other_members, class_counts)` `other_members` is
`{key: count}` for every key bucketed into "other" (the empirical
within-bucket distribution, for `other_policy = "sample"` at rollout);
@@ -270,15 +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)
class_map, _, _ = _topn_plus_other_map(counts, n_experts)
from giant.data.scan import ScanRequest, scan_metadata
result = scan_metadata(files, ScanRequest(process=True))
assert result.process is not None
class_map, _, _ = _topn_plus_other_map(result.process, n_experts)
return class_map
@@ -327,8 +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)
class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
from giant.data.scan import ScanRequest, scan_metadata
if column != "material":
raise ValueError(f"build_topn_map_from_files only supports column='material', got {column!r}")
result = scan_metadata(files, ScanRequest(material=True))
assert result.material is not None
class_map, other_members, class_counts = _topn_plus_other_map(result.material, n_classes)
return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
@@ -349,16 +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 = {}
for path in files:
columns = ["pdg"]
has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names
if has_sec:
columns.append("sec_pdg_list")
df = pd.read_parquet(path, columns=columns)
_accumulate_value_counts(counts, df["pdg"], int)
if has_sec:
exploded = df["sec_pdg_list"].explode().dropna()
_accumulate_value_counts(counts, exploded, int)
class_map, other_members, class_counts = _topn_plus_other_map(counts, n_classes)
from giant.data.scan import ScanRequest, scan_metadata
result = scan_metadata(files, ScanRequest(pooled_pdg=True))
assert result.pooled_pdg is not None
class_map, other_members, class_counts = _topn_plus_other_map(result.pooled_pdg, n_classes)
return TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
+171
View File
@@ -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,
)
+14 -5
View File
@@ -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)])
unique_ids, counts = np.unique(all_ids, return_counts=True)
return unique_ids, counts
result = scan_metadata(files, ScanRequest(event_index=True))
assert result.event_index is not None
return result.event_index
def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:
+12 -8
View File
@@ -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})
.groupby(["bin", "material", "layer_id"])
.size()
.to_frame("n")
.reset_index()
.sort_values("n", ascending=False)
.drop_duplicates("bin")
pl.DataFrame({"bin": bin_idx, "material": mat, "layer_id": lay})
.group_by(["bin", "material", "layer_id"])
.agg(pl.len().alias("n"))
.sort(["bin", "material", "layer_id"])
.sort("n", descending=True, maintain_order=True)
.unique(subset="bin", keep="first", maintain_order=True)
)
bin_material = np.full(n_bins, "", dtype=object)
+11 -1
View File
@@ -37,7 +37,7 @@ from dataclasses import dataclass, field
import torch.nn as nn
from giant.config import _get_path, _set_path, leaf_paths
from giant.config import INFERENCE_OVERRIDES, _get_path, _set_path, leaf_paths
from giant.model.builders import build_critics, build_models
from giant.model.trunks import RoutedTrunk
@@ -111,6 +111,7 @@ class ModelSummary:
pdg_vocab: int
mat_vocab: int
vocab_caveats: list[str] = field(default_factory=list)
overridable: list[str] = field(default_factory=list)
def _build_model_config(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict:
@@ -234,6 +235,8 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
else:
inert.append(path)
overridable = sorted(p for p in in_scope if p in INFERENCE_OVERRIDES)
return ModelSummary(
modules=modules,
consumed=sorted(consumed),
@@ -242,6 +245,7 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
pdg_vocab=pdg_vocab,
mat_vocab=mat_vocab,
vocab_caveats=_vocab_caveats(cfg),
overridable=overridable,
)
@@ -309,6 +313,12 @@ def render_summary(summary: ModelSummary) -> str:
else:
lines.append(" (none)")
if summary.overridable:
lines.append("")
lines.append("inference-overridable without retraining (giant predict/rollout --set):")
for path in summary.overridable:
lines.append(f" {path} ({INFERENCE_OVERRIDES[path].why})")
if summary.vocab_caveats:
lines.append("")
lines.append("vocab placeholder caveats:")
+106 -54
View File
@@ -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"]
cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
if cached_proc_map is not None:
assert process_n_experts is not None
if not need_process:
assert cache is not None
cached_proc_map = cache.proc_maps.get(process_n_experts)
assert cached_proc_map is not None
proc_map = cached_proc_map
echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {process_n_experts} experts)")
else:
echo("building process vocabulary …")
proc_map = build_process_map_from_files(files, n_experts=n_experts)
echo(f" {len(proc_map)} process labels mapped to {n_experts} experts")
assert scan.process is not None
proc_map, _, _ = _topn_plus_other_map(scan.process, process_n_experts)
echo(f" {len(proc_map)} process labels mapped to {process_n_experts} experts")
if cache is not None:
cache.proc_maps[n_experts] = proc_map
# Top-N-plus-other maps for onehot conditioning/type axes.
# The PDG axis is used independently by conditioning.particle.type="onehot"
# (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
# (secondary-species decode) — their class counts can now differ (gitea
# #29: stage2_model.particle_type.n_classes, 0 = inherit
# conditioning.particle.emb_dim), so each is resolved and built
# independently via _pdg_topn below. cache.topn_maps is keyed by
# (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
# the same N the second call is a cache hit against the first — no extra
# scan in the common case where they still match. The material axis is
# independent of both.
particle_cfg = cfg["conditioning"]["particle"]
material_cfg = cfg["conditioning"]["material"]
particle_type_cfg = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type"))
particle_type_target = particle_type_cfg.target
cache.proc_maps[process_n_experts] = proc_map
def _pdg_topn(n_classes: int) -> TopNMap:
cache_key = setup_cache.topn_key("pdg", n_classes)
@@ -198,33 +234,36 @@ def run_setup_stage(
echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
return cached
echo("building pdg top-N map …")
topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
assert scan.pooled_pdg is not None
class_map, other_members, class_counts = _topn_plus_other_map(scan.pooled_pdg, n_classes)
topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
echo(f" {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
if cache is not None:
cache.topn_maps[cache_key] = topn_map
return topn_map
pdg_topn_map: TopNMap | None = None
if particle_cfg["type"] == "onehot":
pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
pdg_topn_map: TopNMap | None = _pdg_topn(particle_cfg["emb_dim"]) if need_pdg_onehot else None
sec_type_topn_map: TopNMap | None = None
if particle_type_target == "onehot":
sec_type_n_classes = resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
if need_sec_type_onehot:
assert sec_type_n_classes is not None
sec_type_topn_map = _pdg_topn(sec_type_n_classes)
mat_topn_map: TopNMap | None = None
if material_cfg["type"] == "onehot":
n_classes = material_cfg["emb_dim"]
cache_key = setup_cache.topn_key("material", n_classes)
if need_material_onehot:
assert material_n_classes is not None
cache_key = setup_cache.topn_key("material", material_n_classes)
cached = cache.topn_maps.get(cache_key) if cache is not None else None
if cached is not None:
mat_topn_map = cached
echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
echo(
f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {material_n_classes} classes)"
)
else:
echo("building material top-N map …")
mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
echo(f" {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
assert scan.material is not None
class_map, other_members, class_counts = _topn_plus_other_map(scan.material, material_n_classes)
mat_topn_map = TopNMap(class_map=class_map, other_members=other_members, class_counts=class_counts)
echo(f" {len(mat_topn_map.class_map)} materials mapped to {material_n_classes} classes")
if cache is not None:
cache.topn_maps[cache_key] = mat_topn_map
@@ -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 = {
+34 -14
View File
@@ -5,6 +5,8 @@ simulation-fanout tools into one Typer app so there's a single command name
(and `--help`) to remember instead of five differently-hyphenated ones.
"""
from __future__ import annotations
import os
from enum import Enum
from pathlib import Path
@@ -14,20 +16,15 @@ import typer
from typing_extensions import Annotated
from giant.config import Conditioning
from giant.tools.bump_dataset_version import (
run_bump_gen,
run_bump_schema,
run_create_manifest,
run_status,
run_update_manifest,
)
from giant.tools.create_root_files import run_make_root
from giant.tools.geometry_oracle import run_build_geometry_oracle
from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
from giant.tools.migrate_geant_steps import run_migration
from giant.tools.steps_to_parquet import convert_steps_to_parquet
from giant.tools.steps_to_parquet_parallel import run_parallel_job
from giant.tools.warm_setup_cache import run_warm_setup_cache
# DATA_DEFAULT/SCAN_DIR_DEFAULT are Typer option defaults (evaluated at
# decoration time below), so that one name has to stay eager — the module
# itself is stdlib-only, so it costs nothing. Every other giant.tools.*
# import here is deferred into the one command body that uses it, since
# several (steps_to_parquet: uproot/awkward/polars; warm_setup_cache:
# giant.pipeline -> torch; geometry_oracle: pandas) are expensive and
# `dwarf --help`/tab-completion shouldn't pay for all of them upfront.
from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT
app = typer.Typer(no_args_is_help=True)
@@ -121,6 +118,9 @@ def convert(
] = None,
) -> None:
"""Convert ROOT Steps tree(s) to Parquet."""
from giant.tools.steps_to_parquet import convert_steps_to_parquet
from giant.tools.steps_to_parquet_parallel import run_parallel_job
if jobs < 1:
typer.echo("error: --jobs must be >= 1", err=True)
raise typer.Exit(1)
@@ -183,6 +183,8 @@ def migrate(
] = False,
) -> None:
"""One-time migration into the versioned raw/processed/pools/derived layout."""
from giant.tools.migrate_geant_steps import run_migration
run_migration(str(root), execute=execute, copy=copy)
@@ -207,6 +209,8 @@ def bump_gen(
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new raw generation."""
from giant.tools.bump_dataset_version import run_bump_gen
run_bump_gen(
kind=kind,
reason=reason,
@@ -240,6 +244,8 @@ def bump_schema(
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""Cut a new schema within a gen."""
from giant.tools.bump_dataset_version import run_bump_schema
run_bump_schema(
kind=kind,
gen=gen,
@@ -257,6 +263,8 @@ def status(
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
) -> None:
"""List existing gens/schemas per kind."""
from giant.tools.bump_dataset_version import run_status
run_status(str(root))
@@ -281,6 +289,8 @@ def update_manifest(
] = False,
) -> None:
"""Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
from giant.tools.bump_dataset_version import run_update_manifest
run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)
@@ -311,6 +321,8 @@ def create_manifest(
] = False,
) -> None:
"""Create a new manifest from a list of parquet files."""
from giant.tools.bump_dataset_version import run_create_manifest
run_create_manifest(
[str(f) for f in files],
execute=execute,
@@ -358,6 +370,8 @@ def make_root(
] = False,
) -> None:
"""Generate new ROOT shards via a minicalosim executable."""
from giant.tools.create_root_files import run_make_root
_warn_if_exceeds_shared_quota(jobs, "--jobs")
run_make_root(
executable=executable,
@@ -423,6 +437,8 @@ def build_geometry_oracle(
] = 2000,
) -> None:
"""Fit a position -> (material, layer_id) oracle for `giant rollout`."""
from giant.tools.geometry_oracle import run_build_geometry_oracle
run_build_geometry_oracle(
data=data,
out=out,
@@ -515,6 +531,8 @@ def warm_cache(
such entry across every run) skips straight to training. See
giant/data/setup_cache.py.
"""
from giant.tools.warm_setup_cache import run_warm_setup_cache
flag_overrides = {
"--val-fraction": val_fraction,
"--seed": seed,
@@ -557,6 +575,8 @@ def hparam_scan(
dry_run: Annotated[bool, typer.Option("--dry-run")] = False,
) -> None:
"""Grid-scan dropout x n_blocks x hidden_dim via sequential `giant train` runs."""
from giant.tools.hparam_scan import run_hparam_scan
run_hparam_scan(data=data, scan_dir=scan_dir, seed=seed, dry_run=dry_run)
+162
View File
@@ -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()
+6 -2
View File
@@ -1,12 +1,12 @@
[project]
name = "giant"
version = "0.3.12"
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 = [
+29
View File
@@ -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"]
+120
View File
@@ -14,6 +14,7 @@ from giant import config as gconfig
from giant.checkpoint_io import (
CheckpointCompatibilityError,
InferenceContext,
apply_config_overrides,
conditioning_axes,
load_for_inference,
stage_cfg,
@@ -278,3 +279,122 @@ def test_stage_cfg_new_shape_returns_subdict():
def test_stage_cfg_v02_flat_shape_returns_empty_dict():
model_cfg = {"hidden_dim": 32, "n_blocks": 4}
assert stage_cfg(model_cfg, "stage2") == {}
# ---------------------------------------------------------------------------
# config_overrides (gitea #87)
# ---------------------------------------------------------------------------
def _router_model_cfg() -> dict:
cfg = _model_cfg()
cfg["stage1_model"]["router"] = {"enabled": True, "type": "energy", "n_experts": 2}
return cfg
def test_config_override_n_sec_sampling_changes_stage2_attribute(tmp_path):
checkpoint = _write_checkpoint(tmp_path)
ctx = load_for_inference(
checkpoint,
torch.device("cpu"),
"predict",
config_overrides={"stage2_model.n_sec.sampling": "sample"},
)
assert ctx.stage2 is not None
assert ctx.stage2.n_sec_sampling == "sample"
assert ctx.config_overrides == {"stage2_model.n_sec.sampling": "sample"}
def test_config_override_ddpm_n_steps_changes_context_fields(tmp_path):
checkpoint = _write_checkpoint(tmp_path)
ctx = load_for_inference(
checkpoint,
torch.device("cpu"),
"predict",
config_overrides={"stage1_model.ddpm.n_steps": 42, "stage2_model.ddpm.n_steps": 7},
)
assert ctx.stage1_ddpm_steps == 42
assert ctx.stage2_ddpm_steps == 7
def test_config_override_other_policy_changes_context_field(tmp_path):
checkpoint = _write_checkpoint(tmp_path)
ctx = load_for_inference(
checkpoint,
torch.device("cpu"),
"predict",
config_overrides={"stage2_model.particle_type.other_policy": "modal"},
)
assert ctx.other_policy == "modal"
def test_config_override_router_temperature_changes_router_attribute(tmp_path):
checkpoint = _write_checkpoint(tmp_path, model_cfg=_router_model_cfg())
ctx = load_for_inference(
checkpoint,
torch.device("cpu"),
"predict",
config_overrides={"stage1_model.router.temperature": 1.5},
)
assert ctx.stage1 is not None
assert ctx.stage1.trunk.router.temperature == pytest.approx(1.5)
def test_config_override_no_overrides_defaults_to_empty_dict(tmp_path):
checkpoint = _write_checkpoint(tmp_path)
ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
assert ctx.config_overrides == {}
def test_config_override_unknown_path_raises(tmp_path):
checkpoint = _write_checkpoint(tmp_path)
with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
load_for_inference(
checkpoint,
torch.device("cpu"),
"predict",
config_overrides={"stage2_model.n_sec.typo": "sample"},
)
def test_config_override_shape_bearing_key_raises_up_front(tmp_path):
checkpoint = _write_checkpoint(tmp_path)
with pytest.raises(CheckpointCompatibilityError, match="not an inference-safe override"):
load_for_inference(
checkpoint,
torch.device("cpu"),
"predict",
config_overrides={"stage1_model.hidden_dim": 999},
)
def test_config_override_bad_value_raises(tmp_path):
checkpoint = _write_checkpoint(tmp_path)
with pytest.raises(CheckpointCompatibilityError, match="must be one of"):
load_for_inference(
checkpoint,
torch.device("cpu"),
"predict",
config_overrides={"stage2_model.n_sec.sampling": "maybe"},
)
def test_apply_config_overrides_no_overrides_returns_same_object():
cfg = _model_cfg()
assert apply_config_overrides(cfg, None) is cfg
assert apply_config_overrides(cfg, {}) is cfg
def test_apply_config_overrides_migrates_legacy_flat_model_config_first():
legacy_cfg = {
"pdg_vocab": len(PDG_MAP),
"mat_vocab": len(MAT_MAP),
"hidden_dim": 32,
"n_blocks": 4,
"emb_dim": 8,
"dropout": 0.1,
"k_max": 5,
}
merged = apply_config_overrides(legacy_cfg, {"stage1_model.ddpm.n_steps": 10})
assert merged["stage1_model"]["ddpm"]["n_steps"] == 10
assert merged["stage1_model"]["hidden_dim"] == 32
+40
View File
@@ -172,3 +172,43 @@ def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
assert result.exit_code == 1
assert "checkpoint has no model_config" in result.output
# ---------------------------------------------------------------------------
# --set (gitea #87)
# ---------------------------------------------------------------------------
def test_predict_set_flag_without_equals_exits_1(tmp_path):
checkpoint = tmp_path / "missing.pt"
result = runner.invoke(
app,
["predict", "dummy.parquet", "--checkpoint", str(checkpoint), "--set", "sampling"],
)
assert result.exit_code == 1
assert "must be 'dotted.path=value'" in result.output
def test_predict_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
checkpoint = tmp_path / "ckpt.pt"
torch.save(
{"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
checkpoint,
)
result = runner.invoke(
app,
[
"predict",
"dummy.parquet",
"--checkpoint",
str(checkpoint),
"--set",
"stage1_model.hidden_dim=999",
],
)
assert result.exit_code == 1
assert "not an inference-safe override" in result.output
+67 -1
View File
@@ -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)
@@ -31,3 +72,28 @@ def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
assert result.exit_code == 1
assert "checkpoint has no model_config" in result.output
def test_rollout_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
checkpoint = tmp_path / "ckpt.pt"
torch.save(
{"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
checkpoint,
)
result = runner.invoke(
app,
[
"rollout",
"dummy.parquet",
"--checkpoint",
str(checkpoint),
"--geometry",
"dummy_geometry.pkl",
"--set",
"stage2_model.n_sec.typo=sample",
],
)
assert result.exit_code == 1
assert "not an inference-safe override" in result.output
+6 -6
View File
@@ -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(
+16
View File
@@ -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")
+23 -6
View File
@@ -108,12 +108,12 @@ def test_build_process_map_from_files_spans_multiple_files(tmp_path):
def test_build_process_map_from_files_tie_breaking_pins_first_seen_order(tmp_path):
"""When two processes end up with equal total counts, ranking falls back
to whichever was accumulated first (`sorted(..., reverse=True)` is stable,
and `counts` is built in file/row-scan order) this is implementation-
defined, not a documented contract, so pin it explicitly: a future
rewrite (e.g. a polars-based single-scan) that ties differently would
silently reshuffle which processes get their own expert slot across a
retrain, and this test is what should catch that."""
to whichever was scanned first file order, then row order within a
file (`giant.data.scan`'s `first_seen` ordinal, ranked by
`giant.data.loader._topn_plus_other_map`'s `(-count, first_seen)` key).
This is an explicit, documented contract (not an accident of iteration
order), pinned here so a future change to the ranking can't silently
reshuffle which processes get their own expert slot across a retrain."""
path = tmp_path / "a.parquet"
pd.DataFrame({"process": ["compt", "phot", "compt", "phot"]}).to_parquet(path)
@@ -233,6 +233,23 @@ def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path)
assert m.class_counts == {0: 11, 1: 5}
def test_build_pdg_topn_map_from_files_pooled_tie_breaks_by_row_position(tmp_path):
"""Pooled pdg counting merges the primary `pdg` column and the exploded
`sec_pdg_list` column via one `group_by` over both (see
`giant.data.scan._pooled_pdg_lazy`), keyed by row position regardless of
which role (primary or secondary) a code was seen in not "all
primaries before all secondaries" the way a two-pass accumulation would.
11 (primary, row 0), 33 (primary, row 1) and 22 (secondary, row 1) all
end up with count 1; 11's strictly earlier row wins the tie over both,
whatever order 33/22 (tied with each other, same row) land in."""
path = tmp_path / "a.parquet"
pd.DataFrame({"pdg": [11, 33], "sec_pdg_list": [[], [22]]}).to_parquet(path)
m = build_pdg_topn_map_from_files([path], n_classes=4)
assert m.class_map[11] == 0
def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):
"""Files predating the parent->child join have no sec_pdg_list column —
must not raise, just count the primary pdg column alone."""
+14
View File
@@ -12,6 +12,8 @@ from giant.cli import app
from giant.materials import MATERIAL_PROPERTIES
from giant.model.summary import _NOT_BUILD_TIME, _built_modules, _vocab_caveats, summarize_model
INFERENCE_OVERRIDES = gconfig.INFERENCE_OVERRIDES
runner = CliRunner()
_PDG_VOCAB = 300
@@ -53,6 +55,16 @@ def test_not_build_time_allow_list_has_no_stale_entries():
assert not stale, f"_NOT_BUILD_TIME entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
def test_inference_overrides_allow_list_has_no_stale_entries():
in_scope = set(gconfig.leaf_paths(gconfig.DEFAULT_CONFIG))
stale = set(INFERENCE_OVERRIDES) - in_scope
assert not stale, f"INFERENCE_OVERRIDES entries no longer in DEFAULT_CONFIG: {sorted(stale)}"
def test_default_config_overridable_lists_every_allowlisted_path(default_summary):
assert set(default_summary.overridable) == set(INFERENCE_OVERRIDES)
def test_router_disabled_by_default_so_its_fields_are_inert(default_summary):
assert "stage1_model.router.n_experts" in default_summary.inert
assert "stage1_model.router.temperature" in default_summary.inert
@@ -135,3 +147,5 @@ def test_cli_default_smoke():
assert "parameters" in result.output
assert "trunk" in result.output
assert "inert under this config" in result.output
assert "inference-overridable without retraining" in result.output
assert "stage2_model.n_sec.sampling" in result.output
+2 -2
View File
@@ -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)
+14
View File
@@ -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",
Generated
+5 -3
View File
@@ -675,12 +675,12 @@ wheels = [
[[package]]
name = "giant"
version = "0.3.12"
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" },