Compare commits
18 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| f80fc90758 | |||
| 1cf16526c9 | |||
| 5c93457081 | |||
| 1ec333ff6d | |||
| 0654fa3f12 | |||
| 36fe9bd66d | |||
| bd255419e1 | |||
| 73975a4587 | |||
| fcd77c2f4b | |||
| bb8d16caba | |||
| c8a1b4f25d | |||
| 23efd6d9ff | |||
| 9fa6420183 | |||
| b66574877b | |||
| 9b77e04731 | |||
| 8dee2feab7 | |||
| f2da0642b2 | |||
| 1e92902c8d |
+1
-1
@@ -1,5 +1,5 @@
|
||||
[tool.bumpversion]
|
||||
current_version = "0.3.9"
|
||||
current_version = "0.3.13"
|
||||
parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
|
||||
serialize = ["{major}.{minor}.{patch}"]
|
||||
search = "{current_version}"
|
||||
|
||||
+502
-1
@@ -1,5 +1,31 @@
|
||||
# Changelog
|
||||
|
||||
## [0.3.13] - 2026-08-28
|
||||
|
||||
### Added
|
||||
|
||||
- Add inference-time model_config overrides with a sampling-key allowlist [gitea #87](https://git.larsbogner.de/lars/giant/issues/87)
|
||||
|
||||
## [0.3.12] - 2026-08-28
|
||||
|
||||
### Added
|
||||
|
||||
- Add sampled n_sec under n_sec.mode = 'head' [gitea #86](https://git.larsbogner.de/lars/giant/issues/86)
|
||||
|
||||
## [0.3.11] - 2026-08-26
|
||||
|
||||
### Changed
|
||||
|
||||
- Feat(analysis): per-step secondary multiplicity plots
|
||||
|
||||
## [0.3.10] - 2026-08-26
|
||||
|
||||
### Changed
|
||||
|
||||
- Backfill CHANGELOG.md for v0.2.0-v0.3.2
|
||||
|
||||
- Docs: bring README and CLAUDE.md in line with v0.3.9
|
||||
|
||||
## [0.3.9] - 2026-08-24
|
||||
|
||||
### Added
|
||||
@@ -65,4 +91,479 @@
|
||||
|
||||
- Document CI_TOKEN's write:repository scope requirement [gitea #50](https://git.larsbogner.de/lars/giant/issues/50)
|
||||
|
||||
# Changelog
|
||||
## [0.3.2] - 2026-08-17
|
||||
|
||||
### Added
|
||||
|
||||
- Add configs/baseline.toml as the kept reference model
|
||||
|
||||
|
||||
### Fixed
|
||||
|
||||
- Clamp analysis histogram bins before the i32 cast, not after [gitea #61](https://git.larsbogner.de/lars/giant/issues/61)
|
||||
|
||||
- Clip raw predicted log_mass in decode_secondaries [gitea #54](https://git.larsbogner.de/lars/giant/issues/54)
|
||||
|
||||
|
||||
### Changed
|
||||
|
||||
- Let dwarf warm-cache take --config so it can't under-warm a config's cache keys [gitea #59](https://git.larsbogner.de/lars/giant/issues/59)
|
||||
|
||||
- Implement n_sec.mode = "stop_token" for the AR secondary decoder [gitea #40](https://git.larsbogner.de/lars/giant/issues/40)
|
||||
|
||||
- Bump patch version to 0.3.2
|
||||
|
||||
## [0.3.1] - 2026-08-14
|
||||
|
||||
### Added
|
||||
|
||||
- Add an Objective registry for the flow/ddpm/wgan generator choice [gitea #32](https://git.larsbogner.de/lars/giant/issues/32)
|
||||
|
||||
|
||||
### Changed
|
||||
|
||||
- Make trunk architecture selectable via a registry [gitea #33](https://git.larsbogner.de/lars/giant/issues/33)
|
||||
|
||||
- Make ResBlock's conditioning-injection mechanism selectable [gitea #34](https://git.larsbogner.de/lars/giant/issues/34)
|
||||
|
||||
- Make HistoryEncoder a pluggable registry, like Router/Objective [gitea #35](https://git.larsbogner.de/lars/giant/issues/35)
|
||||
|
||||
- Deduplicate n_sec_head/type_head MLPs into build_mlp_head [gitea #36](https://git.larsbogner.de/lars/giant/issues/36)
|
||||
|
||||
- Give the cond_cat/cond_cont column layout one owner [gitea #37](https://git.larsbogner.de/lars/giant/issues/37)
|
||||
|
||||
- Give Stage1Model/Stage2OneShot/Stage2Autoregressive a shared StageModel base [gitea #39](https://git.larsbogner.de/lars/giant/issues/39)
|
||||
|
||||
- Pass ConditioningAxisConfig/ParticleTypeConfig themselves instead of raw dicts [gitea #38](https://git.larsbogner.de/lars/giant/issues/38)
|
||||
|
||||
- Bump patch version to 0.3.1
|
||||
|
||||
## [0.3.0] - 2026-08-13
|
||||
|
||||
### Added
|
||||
|
||||
- Add v0.3.0 design doc: Stage-2 autoregressive redesign
|
||||
|
||||
- Add pytest-cov to dev deps and run coverage in CI
|
||||
|
||||
- Add coverage for router-center seeding, geometry batch reader, material topN cache, and setup-cache corruption paths
|
||||
|
||||
- Add render.py coverage: figure params, router diagnostics plots, gallery/condor glue
|
||||
|
||||
- Add unknown-key validation to config.toml merge (issues.md Issue 2)
|
||||
|
||||
- Add consumed-keys audit test (issues.md Issue 5)
|
||||
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fix test_render_all_run_gallery_invokes_subprocess clobbering LaTeX's own subprocess.run
|
||||
|
||||
|
||||
### Removed
|
||||
|
||||
- Remove issues.md
|
||||
|
||||
|
||||
### Changed
|
||||
|
||||
- Refine v0.3.0 design: defaults, deferred scope, open questions
|
||||
|
||||
- Document the differentiability position and its validation obligation
|
||||
|
||||
- V0.3.0 step 1: new nested config schema, v0.2 migration shim
|
||||
|
||||
- V0.3.0 step 2: network.py refactor to composable stage models
|
||||
|
||||
- V0.3.0 step 3: per-stage train.py trainers + pipeline.py/cli.py rewrite
|
||||
|
||||
- V0.3.0 step 4: type map + particle_type.target = "onehot"/"embedding"
|
||||
|
||||
- V0.3.0 step 5: Stage2Autoregressive (history=markov) + §11.4 grad instrumentation
|
||||
|
||||
- V0.3.0 step 6: sample.py/rollout.py AR generation + class->PDG decode
|
||||
|
||||
- V0.3.0 step 7: AttentionHistory (KV-cached) + scheduled/never teacher forcing
|
||||
|
||||
- V0.3.0 post-implementation audit: resolve all 9 tracked discrepancies
|
||||
|
||||
- Refactor train.py into giant/training/ around a metrics collector
|
||||
|
||||
- Silence the fork-safety warning from num_workers>0 pipeline tests
|
||||
|
||||
- Deduplicate giant/training/trainers.py shared per-stage logic
|
||||
|
||||
- Rewrite README for v0.3.0 architecture, quick start, and data columns
|
||||
|
||||
- Bump version to 0.3.0
|
||||
|
||||
- Delete docs/v0.3.0-design.md and strip all references to it
|
||||
|
||||
- Apply ruff format
|
||||
|
||||
- Downgrade coverage-report upload to actions/upload-artifact@v3
|
||||
|
||||
- Bump ruff line-length to 120 and reformat
|
||||
|
||||
- Make config dataclasses the single source of truth for DEFAULT_CONFIG
|
||||
|
||||
- Extract giant train/new-run's CLI override mapping into a table-driven function (issues.md Issues 3 & 4)
|
||||
|
||||
- Mark issues.md Issues 3 & 4 as fixed
|
||||
|
||||
- Extract predict/rollout's duplicated inference bootstrap into giant.checkpoint_io (issues.md Issue 5)
|
||||
|
||||
- Mark issues.md Issue 5 as fixed
|
||||
|
||||
- Unify the two v0.2->v0.3 migration surfaces (issues.md Issue 6)
|
||||
|
||||
- Type the data/model/training batch contracts with NamedTuples (issues.md Issue 7)
|
||||
|
||||
- Split giant/model/network.py into giant/model/ (issues.md Issue 8)
|
||||
|
||||
- Move scripts/ to giant/tools/ (issues.md Issue 9)
|
||||
|
||||
- Reject stage2_model.stage1_context = 'sampled' as unimplemented (issues.md Issue 1)
|
||||
|
||||
- Honour wgan.critic_hidden_dim/critic_n_res_blocks in build_critics [gitea #28](https://git.larsbogner.de/lars/giant/issues/28)
|
||||
|
||||
- Validate stage2_model.autoregressive.order in validate_config [gitea #30](https://git.larsbogner.de/lars/giant/issues/30)
|
||||
|
||||
- Decouple secondary-species vocabulary from conditioning.particle.emb_dim [gitea #29](https://git.larsbogner.de/lars/giant/issues/29)
|
||||
|
||||
- Skip router auxiliary loss compute when their lambda is 0 [gitea #31](https://git.larsbogner.de/lars/giant/issues/31)
|
||||
|
||||
## [0.2.0] - 2026-08-04
|
||||
|
||||
### Added
|
||||
|
||||
- Add CLAUDE.md with architecture overview and dev commands
|
||||
|
||||
- Add streaming data pipeline and giant CLI entry point
|
||||
|
||||
- Add giant predict command
|
||||
|
||||
- Add ROOT-to-parquet conversion script with convert dependency group
|
||||
|
||||
- Add post_pos as a model target via travel_dir decomposition
|
||||
|
||||
- Add --coord local mode to predict for raw-space prediction debugging
|
||||
|
||||
- Add KL divergence to marginal validation and hook it into the training loop
|
||||
|
||||
- Add graceful shutdown on SIGINT/SIGTERM
|
||||
|
||||
- Add configurable dropout to ResBlocks
|
||||
|
||||
- Add giant.analysis module for notebook-based model quality diagnostics
|
||||
|
||||
- Add lazy polars I/O and duplicate KL/constraint checks for giant.analysis
|
||||
|
||||
- Add ruff and ty as dev dependencies, fix lint/type findings
|
||||
|
||||
- Add linear warmup before cosine LR decay
|
||||
|
||||
- Add --batch-size auto to estimate batch size from free GPU memory
|
||||
|
||||
- Add hyperparameter scan
|
||||
|
||||
- Add --batch-size auto to predict, matching train
|
||||
|
||||
- Add tqdm progress bar to predict
|
||||
|
||||
- Add KL bar plots and sample_frac to load_predicted_local; ignore root parquet scratch files
|
||||
|
||||
- Add event-level shower observables to giant.analysis
|
||||
|
||||
- Add total length traveled per event to event observables
|
||||
|
||||
- Add pdg energy/length contribution pie plots
|
||||
|
||||
- Add export script for Tier 4 event-level/pdg-share plots
|
||||
|
||||
- Add mean/median deposited energy and step length plots per event
|
||||
|
||||
- Add export script for ETP group-update presentation plots
|
||||
|
||||
- Add photon edep export scripts and per-step presentation plots
|
||||
|
||||
- Add tooling for a versioned geant_steps dataset layout
|
||||
|
||||
- Add --copy mode to migrate_geant_steps.py
|
||||
|
||||
- Add update-manifest and create-manifest subcommands to bump_dataset_version
|
||||
|
||||
- Add --to flag for bump-gen/bump-schema and --gen flag for update-manifest
|
||||
|
||||
- Add disk usage summary to dwarf status
|
||||
|
||||
- Add file counts and reference tracking to dwarf status
|
||||
|
||||
- Add --comment option to predict, recorded in YAML sidecar
|
||||
|
||||
- Add energy-conservation PoC ODE-step comparison scripts
|
||||
|
||||
- Add autoregressive shower rollout driver
|
||||
|
||||
- Add fast slab lookup for the GeometryOracle, replacing knn as the default
|
||||
|
||||
- Add load_rollout_vs_truth to compare rollouts against held-out truth data
|
||||
|
||||
- Add mixture-of-experts routing prototype for Stage 1 and Stage 2
|
||||
|
||||
- Add ProcessRouter for physics-process-based expert gating
|
||||
|
||||
- Add PdgRouter for particle-type-based expert gating
|
||||
|
||||
- Add ComposedRouter for multi-axis MoE gating
|
||||
|
||||
- Add EMA weights, weight decay, step-based LR schedule, and grad-norm logging to training
|
||||
|
||||
- Add WGAN-GP mode as a throwaway fast-eval experiment
|
||||
|
||||
- Add router gating diagnostic for MoE checkpoints
|
||||
|
||||
- Add Gitea Actions CI pipeline
|
||||
|
||||
- Add configs for router energy (embedding/physical) and WGAN baseline runs
|
||||
|
||||
- Add opt-in Weights & Biases logging for the training loop
|
||||
|
||||
- Add test coverage for resolve_expert_dims
|
||||
|
||||
- Add regression coverage for vocab/process index-map builders
|
||||
|
||||
- Add dwarf warm-cache to precompute the setup-stage sidecar
|
||||
|
||||
- Add giant new-run to scaffold a config.toml + run dir ahead of training
|
||||
|
||||
- Add learnable per-expert width and shared temperature to EnergyRouter
|
||||
|
||||
- Add opt-in straight-through Gumbel-softmax combine weights to MoE router
|
||||
|
||||
- Add gumbel router configs sweeping learnable-knob combinations
|
||||
|
||||
- Add gumbel/learn_centers/learn_width/learn_temperature to out-dir naming
|
||||
|
||||
- Add bigger WGAN config (hidden_dim=512, n_blocks=6)
|
||||
|
||||
- Add data-integrity guards against silent NaN/Inf propagation and races
|
||||
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fix column names to match actual parquet schema
|
||||
|
||||
- Fix installed torch version to be compatible with cuda drivers
|
||||
|
||||
- Fix miniCaloSim link in README
|
||||
|
||||
- Fix giant.analysis import after Phase 2 dataset API changes
|
||||
|
||||
- Fix silent failure modes surfaced by extensive code review
|
||||
|
||||
- Fix ruff, ty, and pytest failures; apply ruff format
|
||||
|
||||
- Clamp n_sec classification label to K_MAX
|
||||
|
||||
- Fix rollout edep mismatch and add truth overlay to Tier 4 observables
|
||||
|
||||
- Fix crashes in physical-property conditioning edge cases
|
||||
|
||||
- Fix router experts silently ignoring --hidden-dim/--n-blocks
|
||||
|
||||
- Fix conditioning="physical" so it can actually generalize past training vocab
|
||||
|
||||
- Fix training-loop checkpoint/resume and WGAN bugs
|
||||
|
||||
- Fix stale-partial reuse and n_chunks mismatch in analysis condor pipeline
|
||||
|
||||
- Fix CLI/tooling robustness gaps and dedupe the Conditioning enum
|
||||
|
||||
- Fix test_write_submit_requires_synced_venv for active-venv resolution
|
||||
|
||||
|
||||
### Removed
|
||||
|
||||
- Remove scripts/train.py in favor of the giant train CLI
|
||||
|
||||
- Drop orphaned child tracks instead of nulling secondary targets
|
||||
|
||||
|
||||
### Changed
|
||||
|
||||
- Initial commit: giant surrogate model with two-phase roadmap in README
|
||||
|
||||
- Implement Phase 1: full data pipeline, model, training, and config support
|
||||
|
||||
- Handle material column as string type
|
||||
|
||||
- Rename pre_energy/post_energy columns to pre_E/post_E
|
||||
|
||||
- Rename direction columns from pre_dir_x/y/z to pre_dx/dy/dz
|
||||
|
||||
- Batch StreamingStepsDataset internally instead of per-row collate
|
||||
|
||||
- Dedup training pipeline, add seeding/resume and per-epoch metrics logging
|
||||
|
||||
- Split torch into cpu/cuda extras and pin dependency version bounds
|
||||
|
||||
- Apply ruff format and document lint/type tooling in CLAUDE.md
|
||||
|
||||
- Update README to match current architecture and tooling
|
||||
|
||||
- Make sampler step count configurable for validation
|
||||
|
||||
- Calibrate auto batch size separately for inference vs training
|
||||
|
||||
- Skip rows with unknown PDG codes during predict
|
||||
|
||||
- Buffer predict rows across row-group boundaries before inference
|
||||
|
||||
- Export plots for knowledge base
|
||||
|
||||
- Rework validation notebook with markdown sections and Tier 4 plots
|
||||
|
||||
- Allow steps_to_parquet.py to accept multiple ROOT input files
|
||||
|
||||
- Encode edep/secondary/post energy as a conservation-constrained simplex
|
||||
|
||||
- Expose dataset/conversion scripts as uv entry points
|
||||
|
||||
- Restrict holdout overlap check to holdout vs dev/full only
|
||||
|
||||
- Route predict output to UUID-named parquet with YAML reference sidecar
|
||||
|
||||
- Implement Phase 2: secondary particle prediction
|
||||
|
||||
- Unify dataset/tooling scripts into a single `dwarf` Typer CLI
|
||||
|
||||
- Fold --to/--gen dataset-versioning flags into the dwarf CLI
|
||||
|
||||
- Prefix default train output dir with current date
|
||||
|
||||
- Color-code dwarf status output by tree level
|
||||
|
||||
- Show VERSIONS.md reason extracts in dwarf status
|
||||
|
||||
- Wire up predict CLI to load and run the Stage-2 sec_decoder
|
||||
|
||||
- Wire up n_sec/species/energy-fraction validation for Stage 2
|
||||
|
||||
- Detach Stage-2 type-embedding target to stop self-referential collapse
|
||||
|
||||
- Weight Stage-2 secondary loss equally between direction and type-embedding dims
|
||||
|
||||
- Recalibrate batch-size estimate for the post-Phase-2 model size
|
||||
|
||||
- Update CLAUDE.md and README for the implemented Phase 2 model
|
||||
|
||||
- Error on missing secondary lists instead of silently zeroing Stage-2 targets
|
||||
|
||||
- Derive a unique per-job seed for minicalosim shard generation
|
||||
|
||||
- Rescale secondary energies to exactly consume the e_sec budget
|
||||
|
||||
- Support --energy-gev in dwarf make-root for the new minicalosim energy arg
|
||||
|
||||
- Stream giant rollout output instead of buffering the whole run
|
||||
|
||||
- Scale auto batch-size estimate by MoE expert count during training
|
||||
|
||||
- Rewrite analysis module as a lean, fully-streaming pipeline
|
||||
|
||||
- Reimplement rollout-vs-truth comparison on the streaming analysis module
|
||||
|
||||
- Condition on material/particle physical properties instead of learned embeddings
|
||||
|
||||
- Ignore the scratchpad working directory
|
||||
|
||||
- Quote the on: key in the CI workflow
|
||||
|
||||
- Split CI lint stage into parallel jobs
|
||||
|
||||
- Rewrite analysis as streaming rollout-vs-reference plotting pipeline
|
||||
|
||||
- Analyze: drive prep/submit from the rollout YAML sidecar
|
||||
|
||||
- Analyze: show model/training params on rendered figures
|
||||
|
||||
- Deps: install plotstyle from git.larsbogner.de package index
|
||||
|
||||
- Analyze: drop stale ty:ignore on plotstyle import
|
||||
|
||||
- Test: replace prep(**_CTX) splat with a typed _prep helper
|
||||
|
||||
- Analyze: add MoE router gating/share diagnostic plots
|
||||
|
||||
- Chore: remove stray CUDA sanity script and stale Phase 2 planning doc
|
||||
|
||||
- Docs: document compute environment, WGAN/MoE status, and condor-gpu-train-rollout
|
||||
|
||||
- Analyze: normalize pdg dtype in open_side to fix rollout/reference concat
|
||||
|
||||
- Analyze: chunk per-plot aggregation across HTCondor jobs
|
||||
|
||||
- Analyze: expose bin/pdg options on `analyze submit`
|
||||
|
||||
- Analyze: estimate per-job HTCondor walltime from chunk row count
|
||||
|
||||
- Analyze: run condor compute jobs via .venv/bin/giant, not uv run
|
||||
|
||||
- Analyze: default condor docker image to alma9-gridjob
|
||||
|
||||
- Analyze: raise default condor job memory request to 8192 MB
|
||||
|
||||
- Analyze: recalibrate condor walltime model from real cluster timings
|
||||
|
||||
- Transforms: pad legacy cond normalizers for pre-physical-conditioning checkpoints
|
||||
|
||||
- Analyze: default run directory to <repo>/analysis_runs, gitignored
|
||||
|
||||
- Docs: record first MoE router rollout benchmark result in the roadmap
|
||||
|
||||
- Router: seed EnergyRouter centers from data quantiles instead of a fixed linspace
|
||||
|
||||
- Docs: note the EnergyRouter centers_init fix in the roadmap
|
||||
|
||||
- Analyze: thread full model/training/rollout/dataset params to plots
|
||||
|
||||
- Ci: share one uv sync across jobs, gate tests on lint+type-check, sync tag/version on release tags
|
||||
|
||||
- Ci: replace unsupported artifact sharing with a bind-mounted uv cache
|
||||
|
||||
- Ci: stop setup-uv from overriding UV_CACHE_DIR
|
||||
|
||||
- Ci: re-pin UV_CACHE_DIR after setup-uv, which exports its own value regardless of enable-cache
|
||||
|
||||
- Ci: set UV_LINK_MODE=copy to silence the cross-filesystem hardlink warning
|
||||
|
||||
- Log batch-level metrics to W&B, not just per-epoch summaries
|
||||
|
||||
- Log router health, WGAN grad-norm split, n_sec accuracy, GPU/throughput to W&B
|
||||
|
||||
- Persist global_step across --resume so W&B step stays monotonic
|
||||
|
||||
- Timestamp default checkpoint dir to avoid W&B run-id collisions
|
||||
|
||||
- Skip empty-slice mean/std in sec phys validation print
|
||||
|
||||
- Speed up giant train's setup stage
|
||||
|
||||
- Speed up _WelfordAccumulator's per-chunk update
|
||||
|
||||
- Make default checkpoint out_dir name reflect only non-default hyperparams
|
||||
|
||||
- Cache giant train's setup stage in a sidecar file
|
||||
|
||||
- Pass --seed through to the train/val event split
|
||||
|
||||
- Offset event_id per file to avoid cross-file collisions
|
||||
|
||||
- Store a quantile grid instead of a raw reservoir sample in the setup cache
|
||||
|
||||
- Scope wandb run config to only-active hyperparameters
|
||||
|
||||
- Resolve giant condor wrapper from the active venv, not a hardcoded path
|
||||
|
||||
- Bump version to 0.2.0
|
||||
|
||||
@@ -7,18 +7,20 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
|
||||
```bash
|
||||
uv sync --extra cpu # install dependencies with CPU-only torch (standard/default)
|
||||
uv sync --extra cuda # install dependencies with CUDA 11.8 torch
|
||||
uv sync --extra cpu --extra dev # add dev extras (pytest, etc.)
|
||||
uv sync --extra cpu --extra dev # add dev extras (pytest, ruff, ty, bump-my-version, git-cliff, + all runtime extras)
|
||||
uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle (giant rollout)
|
||||
pytest # run tests
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 # scaffold a config.toml + run dir ahead of training
|
||||
giant train path/to/steps.parquet --mode flow # train (flow matching)
|
||||
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
|
||||
giant train path/to/steps.parquet --mode wgan # train (WGAN-GP, single-pass eval; implemented, not yet tested)
|
||||
giant train path/to/steps.parquet --router --router-type energy # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap)
|
||||
giant train path/to/steps.parquet # train (defaults: stage 1 flow, stage 2 wgan + autoregressive)
|
||||
giant train path/to/steps.parquet --mode flow # set both stages' generative objective at once
|
||||
giant train path/to/steps.parquet --stage1-generator flow --stage2-generator wgan # per-stage override
|
||||
giant train path/to/steps.parquet --router --router-type energy # MoE routing trunk (see Roadmap for status)
|
||||
giant model summary --config config.toml # build-only: parameter counts + which config keys actually bite
|
||||
giant predict path/to/steps.parquet --checkpoint ckpt/best.pt # per-step predictions
|
||||
giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl # full showers
|
||||
giant analyze submit rollout.yaml --accounting-group cms # parallel rollout-vs-reference analysis on HTCondor
|
||||
giant analyze render <run_dir> --gallery # render PDFs + HTML gallery (run_dir from prep/submit)
|
||||
giant analyze metrics <train_run_dir> # training-progress plots from metrics.csv
|
||||
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
|
||||
# bump-schema, status, update-manifest, create-manifest,
|
||||
# make-root, build-geometry-oracle, warm-cache, hparam-scan
|
||||
@@ -27,6 +29,8 @@ dwarf --help # dataset/tooling CLI: convert,
|
||||
|
||||
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.
|
||||
|
||||
`configs/` holds kept reference configs (`baseline.toml`, `default.toml`, the router/WGAN scan configs) — pass them with `--config`.
|
||||
|
||||
### Lint and type checking
|
||||
|
||||
```bash
|
||||
@@ -37,6 +41,10 @@ uv run ty check . # type check
|
||||
|
||||
Part of the `dev` extra. Run these periodically (not just at commit time) to catch drift early.
|
||||
|
||||
### Release tooling
|
||||
|
||||
Merges to `master` auto-bump the patch version, tag, and update `CHANGELOG.md` via the Gitea workflow in `.gitea/workflows/ci.yml` (bump-my-version + git-cliff). Don't hand-edit the version in `pyproject.toml` or write changelog entries by hand.
|
||||
|
||||
## Compute environment
|
||||
|
||||
Work on this repo happens across three kinds of machine:
|
||||
@@ -47,50 +55,70 @@ Work on this repo happens across three kinds of machine:
|
||||
|
||||
## Architecture
|
||||
|
||||
GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — now including the variable-length list of secondary particles the step produces (Phase 2, see Roadmap).
|
||||
GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — including the variable-length list of secondary particles the step produces.
|
||||
|
||||
**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`). Train/val split is by `event_id` to avoid leaking correlated steps from the same shower.
|
||||
**Data pipeline** (`giant/data/`): parquet files from miniCaloSim are loaded into numpy arrays (`loader.py`), then log-transformed and rotated into a local coordinate frame where `pre_dir = ẑ` (`transforms.py`), before being wrapped in a PyTorch `Dataset` (`dataset.py`, streaming variant included). Train/val split is by `event_id` (`--seed`-controlled) to avoid leaking correlated steps from the same shower. Loading a directory or `.manifest` of several parquet files offsets each file's `event_id`s by a per-file stride so ids stay globally unique. `setup_cache.py` persists the pre-epoch setup scan (vocab maps, event split, process maps, normalizer stats) as a sidecar so repeated runs over the same `data` path don't rescan (`--cache-setup`/`--rebuild-setup-cache`, precomputable with `dwarf warm-cache --config ...`).
|
||||
|
||||
**Stage-1 output space (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):** `log_step_length`, two additive-log-ratio (ALR) coordinates `edep_logit`/`sec_logit` of a **deposit / secondary / post-energy simplex**, `post_dir` (post-scattering momentum direction, unit vector in the local frame), and `travel_dir` (direction of `post_pos - pre_pos`, unit vector in the local frame). The energy simplex decodes via softmax over `[edep_logit, sec_logit, 0]` × `pre_E` so `edep + e_sec + post_E == pre_E` holds by construction — energy conservation is architectural, not learned (see `energy_simplex_decode`). `post_pos` is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, since `step_length` already encodes that displacement's magnitude and duplicating it would let the two become inconsistent.
|
||||
|
||||
**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus, since particle/material physical-property conditioning (`model.conditioning`, see below), 7 more columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** (that was Phase 1 / the energy-conservation PoC); the model predicts them.
|
||||
**Conditioning vector (15D continuous, `COND_DIM`):** pre-step position, log(pre-energy), pre-step direction, layer ID (`COND_DIM_BASE=8`) — plus 7 physical-property columns: particle `log(mass)`/`charge` (`PARTICLE_PHYS_DIM=2`, `giant/particles.py`) and material `Z_eff`/`A_eff`/`log(density)`/`log(X0)`/`log(λ_int)` (`MATERIAL_PHYS_DIM=5`, `giant/materials.py`). `n_sec` and `e_sec` are **not conditioning inputs** — the model predicts them. `giant/cond_layout.py` is the single source of truth for the `cond_cont`/`cond_cat` column layout shared by `giant.data.transforms`, `giant.model.encoders`, and `giant.model.routers`.
|
||||
|
||||
`ConditionEncoder`/`SecondaryConditionEncoder` (`giant/model/network.py`) support two mutually exclusive `conditioning` modes, selected per-checkpoint (`model_config["conditioning"]`, defaulting to `"embedding"` for old checkpoints without the key, `"physical"` for new `giant train` runs — see `--conditioning`):
|
||||
- **`"embedding"`** (original Phase 2 design): a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab (`pdg_map`/`mat_map`). Memorizes the training menu.
|
||||
- **`"physical"`** (default): the 7 physical-property columns above are each routed through a small MLP (`particle_mlp`/`material_mlp`) to the same `emb_dim` width the embedding tables would have produced — a drop-in replacement computable for any PDG code / material name, not just ones seen in training, which is what lets the surrogate generalize to a held-out material or species. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships real Geant4-11.4.1-derived `z_eff`/`a_eff`/`density`/`x0`/`lambda_int` values for every material the detector geometry actually produces; the sole exception is `G4_LYSO` (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays `MaterialProperties(None, ...)` and raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting if it's ever requested.
|
||||
`ConditionEncoder` (`giant/model/encoders.py`) configures the particle and material identity axes **independently** (`conditioning.particle` / `conditioning.material`, each a `ConditioningAxisConfig` with `type`/`emb_dim`/`n_layers`), so they may mix freely. Three per-axis modes:
|
||||
- **`"physical"`** (default): the axis's raw physical properties routed through a small MLP — computable for any PDG code / material name, which is what lets the surrogate generalize beyond the training menu. `giant/particles.py` decodes nuclear/ion PDG codes (the `10LZZZAAAI` scheme) via the scikit-HEP `particle` package with a Z/A-digit-decode fallback for isomer codes the package's ground-state-only table misses. `giant/materials.py` ships real Geant4-11.4.1-derived values for every material the detector geometry actually produces; the sole exception is `G4_LYSO` (not a stock Geant4 NIST material, never actually constructed by the geometry — see the module docstring), which stays `MaterialProperties(None, ...)` and raises loudly (`MaterialPropertiesNotFilledError`) rather than silently defaulting.
|
||||
- **`"embedding"`**: a learned `nn.Embedding` per PDG code / material name, indexed by a dataset-scoped dense vocab. Memorizes the training menu; the generalization-comparison baseline, and the only mode compatible with `stage2_model.particle_type.target = "embedding"`.
|
||||
- **`"onehot"`**: a fixed, unlearned vector over the top `emb_dim - 1` codes by training-set count plus one "other" bin. Not a reparameterization of `"embedding"` — the vocabulary cap is the real difference.
|
||||
|
||||
**Model** (`giant/model/network.py`): a two-stage model, both checkpointed together.
|
||||
- **Stage 1 — `DenoisingMLP`:** `ResBlock` stack with a `SinusoidalEmbedding` for the flow/diffusion time variable and a `ConditionEncoder` fusing the conditioning. Predicts the 9D primary vector field, plus an `n_sec_head` classifier over `{0..K_MAX}` (`K_MAX=15`) that runs on the condition encoding alone (no diffusion noise), callable via `predict_n_sec`.
|
||||
- **Stage 2 — `SecondaryDecoder`:** a second flow-matching net (`SecondaryConditionEncoder` fuses the pre-step conditioning with the Stage-1 outcome) that generates all `K_MAX` secondary slots at once. Each slot is `(stick-breaking energy logit, local-frame direction 3D, log-mass, charge)` = `SEC_SLOT_DIM=6`, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. A secondary's mass/charge are regressed directly against a fixed physics-derived target (its ground-truth PDG code's `giant.particles.particle_mass_charge`) — not a learned/moving embedding target, so nothing needs detaching. **No snapping at inference**: the predicted (mass, charge) are used as-is as the secondary's physical identity, including for its own future conditioning if it goes on to take further steps in a rollout. A separate, reporting-only nearest-known-PDG lookup (`giant.particles.nearest_known_pdg`) is used purely to populate a nominal `pdg` label for output rows / `"embedding"`-mode fallback conditioning — it never feeds back into the model.
|
||||
`conditioning.share_stages` decides whether the two stages get one shared encoder instance or two identically-configured independent ones.
|
||||
|
||||
`schedule.py` provides both a `CosineSchedule` for DDPM and the flow matching loss utilities (Lipman et al. 2022 conditional flow matching).
|
||||
**Model** (`giant/model/`, both stages checkpointed together). `network.py` is only a re-export shim now; the real code is split by concern:
|
||||
- `layers.py` — `ResBlock`/`AdaLNResBlock` + `BLOCK_REGISTRY` (conditioning-injection mechanism is selectable), `SinusoidalEmbedding`, `ContextAdapter`, `build_mlp_head`.
|
||||
- `encoders.py` — `ConditionEncoder` (above).
|
||||
- `trunks.py` — `TRUNK_REGISTRY`/`build_trunk`: everything downstream of the fused conditioning vector, as a registrable expert *body* (`resmlp` default, plus a `none` variant). `RoutedTrunk` builds `router.n_experts` instances of whichever body is named, so mixing is orthogonal to which body is mixed.
|
||||
- `routers.py` — `Router` base + `ROUTER_REGISTRY`: `energy`/`pdg`/`process`/`composed`/`none`. Soft-mixed at train time, **top-1 dispatched at eval time** (each row runs exactly one small expert), which is the actual inference-speed win. `EnergyRouter`/`PdgRouter` gate on a quantity known at inference; `ProcessRouter` runs its own small classifier (process isn't known upfront); `ComposedRouter` gates jointly over outer-product expert cells via repeated `--router-axis "type:key=val,..."`. The `--router*`/`--n-experts` CLI flags target `stage1_model.router` only; stage 2's router is config-file-only (`stage2_model.router`). `EnergyRouter` accepts `centers_init`, which `giant/pipeline.py` auto-populates from real data quantiles via a reservoir sample collected during the normalizer-fitting pass.
|
||||
- `history.py` — `HISTORY_REGISTRY`/`build_history`: `markov` (previous token only), `attention` (causal self-attention, KV-cached at inference via `init_cache`/`step`), `none`. Stage-2 autoregressive only.
|
||||
- `objectives.py` — `Objective` base + registry for `flow`/`ddpm`/`wgan`: answers in one place whether a stage needs a time embedding, is adversarial, folds the secondary type slice into its trunk output, what its trunk input is, and which loss it trains against.
|
||||
- `models.py` — the composed stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel`, all on a shared `StageModel` base.
|
||||
- `builders.py` — `build_models`/`build_critics`, assembling the above from a config dict.
|
||||
- `schedule.py` (`CosineSchedule` for DDPM + conditional-flow-matching losses), `wgan.py` (gradient penalty / critic / generator losses, Gulrajani et al. 2017), `summary.py` (`giant model summary`), `_legacy.py` (v0.2 checkpoint migration).
|
||||
|
||||
**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
|
||||
**Stage 1 — primary step.** Trunk (routed or not) over the fused conditioning, plus a `SinusoidalEmbedding` of the flow/diffusion time for non-adversarial objectives, predicting the 9D vector field. An `n_sec` classifier head over `{0..k_max}` runs on the condition encoding alone; `stage2_model.n_sec.owner` decides whether it lives on stage 1 (v0.2 checkpoints) or stage 2 (default).
|
||||
|
||||
**WGAN-GP mode (`--mode wgan`, implemented, not yet tested):** a throwaway fast-eval alternative to the flow/DDPM samplers above — single forward pass instead of ~10 ODE steps. Dedicated noise-conditioned generators (`WGANGenerator`/`WGANSecondaryGenerator`, `giant/model/network.py`) stand in for `DenoisingMLP`/`SecondaryDecoder`, trained against `Critic`/`SecondaryCritic` discriminators with the gradient-penalty loss in `giant/model/wgan.py` (Gulrajani et al. 2017); `sample_wgan` (`giant/sample.py`) does the single-pass draw at inference. Not yet validated against the flow-matching baseline.
|
||||
**Stage 2 — secondaries.** Conditioned on the pre-step state plus a projected stage-1 outcome (`stage2_model.context_dim`; `stage1_context` selects ground-truth vs sampled context, annealable via `ctx_p_start`/`ctx_p_end`). Two decoders (`stage2_model.decoder`):
|
||||
- **`autoregressive`** (default): one secondary at a time in descending-energy order, each token conditioned on a `HistoryEncoder` summary of prior tokens, with teacher forcing (`always`/`scheduled`/`never`, `tf_p_start`/`tf_p_end`). `n_sec.mode = "stop_token"` lets the length be emitted by the sequence itself instead of the classifier head.
|
||||
- **`one_shot`**: all `k_max` slots in one pass, masked past the predicted `n_sec` (the v0.2 behaviour).
|
||||
|
||||
**MoE routing trunk (`--router`, implemented; first rollout benchmark shows the experts don't specialize — see Roadmap):** an alternative to `DenoisingMLP`'s monolithic `ResBlock` trunk — a `Router` (`giant/model/network.py`, `ROUTER_REGISTRY`/`build_router`) gates between small per-expert `ResBlock` stacks (`Expert`), soft-mixed over all experts at train time but **top-1 dispatched at eval time** (each row runs exactly one small expert), which is the actual inference-speed win. Router types gate on different conditioning axes: `EnergyRouter`/`PdgRouter` read a quantity already known at inference time, `ProcessRouter` runs its own small classifier over pre-step conditioning (since process isn't known upfront); `ComposedRouter` gates jointly over multiple axes (outer-product expert cells) via repeated `--router-axis "type:key=val,..."` flags. Config lives under `model.router` (`giant/config.py`), deep-merged one level so `router.enabled` alone doesn't drop the rest of the defaults.
|
||||
Secondary energies are a **stick-breaking partition of the `e_sec` budget** from Stage 1 (they sum to it), so the whole chain conserves energy. Particle identity is set by `stage2_model.particle_type.target`: `"onehot"` (default — categorical over the top `n_classes - 1` PDG codes by training count plus "other", with configurable `other_policy` and `class_weighting`), `"physical"` (continuous `(log-mass, charge)` regressed against `giant.particles.particle_mass_charge`), or `"embedding"` (nearest-row snap into the conditioning embedding table; requires `conditioning.particle.type = "embedding"`).
|
||||
|
||||
**Validation** (`giant/validate.py`): step-level marginal comparisons.
|
||||
**Samplers** (`giant/sample.py`): DDPM, DDIM, flow matching (ODE integration, ~10 steps), and single-pass WGAN, plus the stage-2 secondary sampling loop (one-shot and autoregressive).
|
||||
|
||||
**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive `giant rollout` runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + `RolloutSpec`/`RolloutSide` — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json` over the union of the reference and every rollout, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries; `Bundle.rollouts` is a name-keyed dict of `RolloutSide`, and every `compute_partial`/`finalize` builds a `Reduced.payload["series"]` dict keyed the same way, with `payload["reference"]` as the one distinguished non-rollout entry), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`; each rollout gets a stable `ps.get_color(i)` slot by its position in `series`, the reference always draws in one fixed dashed-ink style). The two heatmap-shaped specs (`marginal_distance_summary`, `n_sec_confusion`) and the router/type-embedding diagnostics (`router_gating.py`, `type_embedding_distance.py`) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout. **Input is one or more `giant rollout` YAML sidecars** (`condor.py:load_rollout_yamls`, wrapping the single-YAML `load_rollout_yaml`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML — matching pre-multi-rollout output exactly). `prep` derives its own **run directory** next to the *first* rollout's parquet (`<...>/analysis_<tag(s)>/`) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit a.yaml [b.yaml ...] --chunks N` runs `prep` (recording the run's chunk count `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` (`catalog.py`) splits into a `compute_partial`/`finalize` pair so a plot's chunks can be summed/concatenated back together correctly per rollout (`chunkable=False` specs — the router diagnostics, already bounded/subsampled — always run as a single chunk regardless of `N`). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`), then turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
|
||||
**Training** (`giant/training/`): `loop.py` (epoch loop, graceful shutdown, best-checkpoint selection), `trainers.py` (`StageSpec` + per-stage flow/ddpm and WGAN-GP trainers, and the `MetricSpec` declarations that define `metrics.csv`'s columns), `stage2_inputs.py` (ground-truth stage-2 targets + teacher-forcing inputs), `metrics.py` (`MetricsCollector`: `metrics.csv`, W&B logging, progress/summary), `checkpoint.py`, `amp.py` (`train.precision = fp32|bf16` autocast), `plots.py` (`giant analyze metrics`). Per-stage `init_from`/`freeze` lets one stage be retrained against a fixed, known-good other stage while still producing a complete rollout-capable checkpoint.
|
||||
|
||||
**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
|
||||
**Config** (`giant/config.py`): frozen dataclasses are the single source of truth; `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather than hand-maintained. Blocks: `[conditioning]`, `[stage1_model]`, `[stage2_model]`, `[train]`, `[meta]`. Unknown keys are rejected on merge (with a did-you-mean suggestion), and `tests/test_config_consumed_keys.py` audits that every key is actually read somewhere.
|
||||
|
||||
**Validation** (`giant/validate.py`): step-level marginal + KL-divergence comparisons during training (`--validate-every`).
|
||||
|
||||
**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one or more autoregressive `giant rollout` runs against a single held-out miniCaloSim reference steps file shared by all of them, and produces publication-styled PDFs assembled into an HTML gallery — one distinctly colored series per rollout, one reference line/panel. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + `RolloutSpec`/`Side` — a rollout's opened frames + per-checkpoint diagnostic inputs — + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `variables.py` (the per-step value expressions shared by range sizing and the plot registry), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json` over the union of the reference and every rollout, so every compute job is one pass with no range scan), `reduced.py` (`Partial`/`Reduced` — the compact self-describing JSON a compute job emits), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles/containment, species/leakage, secondaries, distance summaries, router and type-embedding diagnostics; `giant analyze list` prints every id), `runtime_estimate.py` (per-(plot, chunk) walltime estimates for the submit description), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`; each rollout gets a stable `ps.get_color(i)` slot by its position in `series`, the reference always draws in one fixed dashed-ink style). `Bundle.rollouts` is a name-keyed dict of `Side`, and every `compute_partial`/`finalize` builds a `Reduced.payload["series"]` dict keyed the same way, with `payload["reference"]` as the one distinguished non-rollout entry. The heatmap-shaped specs (`marginal_distance_summary`, `sec_count_per_step_by_species` — the latter also drawing the reference as its own panel) and the checkpoint-bound diagnostics (`router_gating.py`, `type_embedding_distance.py`) are inherently one-matrix/one-checkpoint per rollout, so they render as one panel per rollout instead of one line/bar per rollout.
|
||||
|
||||
**Input is one or more `giant rollout` YAML sidecars** (`condor.py:load_rollout_yamls`): each YAML's `output`/`dataset` keys name its rollout parquet and seed file (= the reference truth); every supplied YAML must resolve to the same `dataset`, checked up front with a clear error otherwise (the premise is "N candidates vs one ground truth"). Each rollout's series name comes from a repeated `--label` CLI flag, else the YAML stem (N>1), else `"rollout"` (a single YAML). `prep` creates a **run directory** (`<cwd>/analysis_runs/analysis_<id>/` by default, `--run-dir` to override) holding `shared.json`, `run_meta.json` (`RunMeta.rollouts: list[{name,path,plot_meta}]`, insertion order = CLI order = every plot's series order), `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit a.yaml [b.yaml ...] --chunks N` runs `prep` (recording `N` in `run_meta.json`) then submits one HTCondor job per (plot, chunk) pair (`compute-one --id --chunk --run-dir`, polars/numpy only — no LaTeX on workers), each streaming over an `event_id`-disjoint slice (`event_id % N == chunk`) of the reference **and every rollout** and writing a small `reduced_partial/<id>__<chunk>.json`; every `PlotSpec` splits into a `compute_partial`/`finalize` pair so chunks can be summed/concatenated back per rollout (`chunkable=False` specs — the checkpoint-bound diagnostics, already bounded/subsampled — always run as a single chunk). The local `giant analyze render <run_dir>` first joins every plot's chunk partials into `reduced/<id>.json` (`merge_all`, a no-op join when `N=1`; `merge-one` does a single plot for debugging), then turns those into the styled PDF/gallery tree. `giant analyze metrics <train_run_dir>` is a separate, unrelated entry point: training-progress plots straight from a run's `metrics.csv`.
|
||||
|
||||
**Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower, advancing tracks breadth-first (every sweep steps all active tracks once, in `batch_size` chunks, so many tracks share each forward pass). Each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on one of the `TERM_*` reasons in `constants.py` (energy cutoff, max steps, escape, natural end, unknown pdg, max tracks); energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. `giant/checkpoint_io.py` is the shared checkpoint → ready-to-run-models path used by both `predict` and `rollout`.
|
||||
|
||||
## Roadmap
|
||||
|
||||
**Phase 1 (done):** number of secondaries and their total energy were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
|
||||
|
||||
**Phase 2 (implemented — baseline):** the two-stage model above jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected). This is the "get a baseline out" track agreed with Jan & Tobias (2026-07-07).
|
||||
**Phase 2 (done):** the two-stage model jointly predicts `n_sec`, the energy simplex (`e_sec` falls out of it), and each secondary's energy/direction/species, so a rollout is self-contained (no ground-truth secondary counts injected).
|
||||
|
||||
**Physical-property conditioning (implemented):** `model.conditioning = "physical" | "embedding"` (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. `"embedding"` stays available as the generalization-comparison baseline. `giant/materials.py`'s table is already filled with real values for every material the geometry produces. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.
|
||||
**Physical-property conditioning (implemented, default):** `conditioning.particle.type` / `conditioning.material.type` = `physical | embedding | onehot`. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment.
|
||||
|
||||
**Faster-eval architectures (implemented, validation in progress):** both tracks below target a ~10× native-Geant4 eval budget and are now wired into `giant train`/`giant/model/network.py`, but neither has a validated result yet — treat both as unproven until the corresponding analysis run says otherwise:
|
||||
- **WGAN-GP** (`--mode wgan`, see Architecture above): implemented, **not yet tested** — no rollout-vs-reference analysis run against it yet.
|
||||
- **MoE routing trunk** (`--router`, see Architecture above): implemented, **first rollout benchmark done (2026-07-22), result: needs retraining with a different router config, not abandoned.** A 10-expert `EnergyRouter` run (`n_experts=10`, `temperature=0.5`, `learn_centers=true`, **`lambda_balance=0.0`**, only 20 fine-tuning epochs resumed from a non-routed checkpoint) diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. The `router_gating` diagnostic plot points at the likely cause: the ten experts overlap heavily across ~5 decades of pre-step energy instead of partitioning it — even the top-energy expert only reaches ~60–65% gate weight at the highest energies plotted — so eval-time top-1 (Voronoi) dispatch is choosing among near-ties rather than real specialists. Two contributors were identified: the missing load-balancing loss (`lambda_balance=0.0`), and `EnergyRouter`'s center init (`torch.linspace(-2, 2, n_experts)`) assuming a roughly uniform z-normalized energy distribution, which real energy spectra don't match. **Fixed (2026-07-27):** `EnergyRouter` now accepts an optional `centers_init` (backward compatible — omitting it keeps the old linspace), and `giant train` auto-populates it from real data quantiles via a reservoir sample collected during the existing normalizer-fitting pass in `giant/pipeline.py` (no extra file scan), for `--router-type energy` only. The routing *strategy* itself may still be sound, but the specific benchmarked config wasn't. **Next step before further evaluation: retrain with `lambda_balance > 0` and the new quantile-seeded centers (and consider more epochs / a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Full writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
|
||||
**v0.3.0 — Stage-2 autoregressive redesign (implemented, released; on `master` since 2026-08-13):** motivated by the 2026-08-03 WGAN rollout benchmark, which failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). Stage 2 became autoregressive in descending-energy order with teacher forcing, and the particle-type representation went back to **categorical** (`particle_type.target = "onehot"`), reversing the 2026-07-17 continuous `(log-mass, charge)` target. The config break (`[conditioning]`/`[stage1_model]`/`[stage2_model]`/`[train]` replacing the flat `train.mode` + `[model]`) makes per-stage generators, stage-2-only training, and one-shot-vs-autoregressive comparison all expressible, and the `network.py` refactor into composable parts (encoder × trunk × objective) also makes routed WGAN work for the first time.
|
||||
|
||||
A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
|
||||
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.
|
||||
|
||||
**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`) — v0.2 checkpoint-loading support has **no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
|
||||
**Faster-eval architectures — both implemented, neither validated.** Target is a ~10× native-Geant4 eval budget; no eval-latency number exists for any configuration yet, so that budget is unverified across the board.
|
||||
- **WGAN-GP** (`--stage2-generator wgan`, now the stage-2 default): first rollout benchmark 2026-08-03 failed with secondary-species mode collapse — the failure v0.3.0 was designed to address. **No post-v0.3.0 benchmark has been run.** Writeup: `/home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md`.
|
||||
- **MoE routing trunk** (`--router`): first rollout benchmark 2026-07-22 diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. Cause identified as a bad config, not a bad idea: `lambda_balance=0.0` (no load-balancing loss) plus `EnergyRouter`'s `torch.linspace(-2, 2, n_experts)` center init assuming a roughly uniform z-normalized energy distribution — so the ten experts overlapped across ~5 decades of energy instead of partitioning it, and eval-time top-1 dispatch chose among near-ties rather than real specialists. Both prerequisites are fixed in code (quantile-seeded `centers_init` from `pipeline.py`, `lambda_balance` exposed). **Next step: retrain with `lambda_balance > 0` and quantile-seeded centers (consider a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
|
||||
|
||||
**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.
|
||||
A sampling-calorimeter (multi-material) dataset track is still open and unblocked, not yet started. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
|
||||
|
||||
**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2. Partway between "needs major features" and feature-complete — not ready to merge yet.
|
||||
|
||||
@@ -10,6 +10,7 @@ A conditional generative model that replaces the Geant4 step function: given a p
|
||||
uv sync --extra cpu # install deps (CPU torch; use --extra cuda for GPU)
|
||||
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 # scaffold config.toml + run dir
|
||||
giant model summary --config config.toml # parameter counts + which config keys actually bite
|
||||
giant train path/to/steps.parquet # train (flow + wgan by default)
|
||||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||||
|
||||
@@ -42,9 +43,9 @@ A **two-stage model**, checkpointed together. Either stage's outcome can be prod
|
||||
|
||||
Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).
|
||||
|
||||
**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above (`--conditioning`) — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
|
||||
**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above. The particle and material axes are configured independently (`conditioning.particle.type` / `conditioning.material.type`; `--conditioning` sets both at once) and may mix — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
|
||||
|
||||
**MoE routing** (`--router`, either stage): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk.
|
||||
**MoE routing** (`--router`): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk. The CLI flags configure Stage 1's router; Stage 2 has its own `stage2_model.router` block, config-file only.
|
||||
|
||||
## Data
|
||||
|
||||
@@ -64,12 +65,24 @@ giant/
|
||||
│ ├── data/
|
||||
│ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads)
|
||||
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
|
||||
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
||||
│ │ ├── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
||||
│ │ └── setup_cache.py # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
|
||||
│ ├── model/
|
||||
│ │ ├── network.py # ConditionEncoder, Stage1Model, Stage2OneShot/Stage2Autoregressive, Router/MoE, CriticModel
|
||||
│ │ ├── models.py # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
|
||||
│ │ ├── builders.py # build_models / build_critics — config dict → assembled stage models
|
||||
│ │ ├── encoders.py # ConditionEncoder (physical / embedding / onehot, per axis)
|
||||
│ │ ├── layers.py # ResBlock/AdaLNResBlock registry, SinusoidalEmbedding, MLP heads
|
||||
│ │ ├── trunks.py # trunk registry (resmlp, none) + RoutedTrunk (MoE expert bodies)
|
||||
│ │ ├── routers.py # Router registry: energy / pdg / process / composed / none
|
||||
│ │ ├── history.py # stage-2 AR history encoders: markov / attention (KV-cached) / none
|
||||
│ │ ├── objectives.py # flow / ddpm / wgan objective registry
|
||||
│ │ ├── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||||
│ │ └── wgan.py # WGAN-GP gradient penalty / critic / generator losses
|
||||
│ │ ├── wgan.py # WGAN-GP gradient penalty / critic / generator losses
|
||||
│ │ ├── summary.py # build-only introspection behind `giant model summary`
|
||||
│ │ ├── _legacy.py # v0.2 checkpoint model_config/state-dict migration
|
||||
│ │ └── network.py # re-export shim over all of the above
|
||||
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
|
||||
│ ├── cond_layout.py # single source of truth for the cond_cont/cond_cat column layout
|
||||
│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
|
||||
│ ├── materials.py # material name → (Z_eff, A_eff, density, X0, λ_int)
|
||||
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
|
||||
@@ -79,20 +92,28 @@ giant/
|
||||
│ │ ├── trainers.py # StageSpec + flow/ddpm and WGAN-GP per-stage trainers
|
||||
│ │ ├── stage2_inputs.py# ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
|
||||
│ │ ├── metrics.py # MetricsCollector: metrics.csv columns, W&B logging, progress/summary
|
||||
│ │ ├── amp.py # bf16 autocast (`train.precision`)
|
||||
│ │ ├── plots.py # training-progress plots (`giant analyze metrics`)
|
||||
│ │ └── checkpoint.py # checkpoint assembly/restore (format unchanged since v0.2)
|
||||
│ ├── sample.py # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
|
||||
│ ├── checkpoint_io.py # checkpoint → ready-to-run models/normalizers (predict + rollout)
|
||||
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
|
||||
│ ├── rollout.py # autoregressive shower rollout driver
|
||||
│ ├── validate.py # step-level marginal + KL-divergence validation
|
||||
│ ├── _migration.py # shared v0.2 → v0.3 facts used by both migration surfaces
|
||||
│ ├── analysis/ # rollout-vs-reference analysis pipeline (see `giant analyze` below)
|
||||
│ │ ├── sources.py # canonical LazyFrames + secondary view
|
||||
│ │ ├── variables.py # per-step value expressions shared by range sizing and the catalog
|
||||
│ │ ├── reduce.py # streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
|
||||
│ │ ├── grouping.py # fixed bin edges + energy/pdg/material group sets
|
||||
│ │ ├── context.py # resolves grouping into `shared.json` once per run
|
||||
│ │ ├── catalog.py # declarative PlotSpec registry
|
||||
│ │ ├── condor.py # prep / compute-one / submit-description plumbing
|
||||
│ │ ├── reduced.py # Partial/Reduced — the compact JSON a compute job emits
|
||||
│ │ ├── catalog.py # declarative PlotSpec registry (`giant analyze list`)
|
||||
│ │ ├── router_gating.py / type_embedding_distance.py # checkpoint-bound diagnostics
|
||||
│ │ ├── runtime_estimate.py # per-(plot, chunk) walltime estimates for submit
|
||||
│ │ ├── condor.py # prep / compute-one / merge / submit-description plumbing
|
||||
│ │ └── render.py # PDFs + HTML gallery (only module importing plotstyle/LaTeX)
|
||||
│ └── cli.py # `giant train` / `new-run` / `predict` / `rollout` / `analyze` Typer app
|
||||
│ └── cli.py # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze`
|
||||
├── giant/tools/ # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
|
||||
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
|
||||
│ │ # update-manifest, create-manifest, make-root,
|
||||
@@ -116,8 +137,13 @@ giant/
|
||||
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
|
||||
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
|
||||
uv sync --extra cpu --extra geometry # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
|
||||
uv sync --extra cpu --extra analysis # matplotlib/polars/plotstyle, for `giant analyze render`
|
||||
uv sync --extra cpu --extra convert # uproot/awkward/polars, for `dwarf convert`
|
||||
uv sync --extra cpu --extra wandb # W&B logging (`giant train --wandb`)
|
||||
```
|
||||
|
||||
The `dev` extra pulls in `convert`, `analysis`, `geometry` and `wandb` as well.
|
||||
|
||||
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||||
|
||||
## Training, prediction, rollout
|
||||
@@ -137,11 +163,13 @@ Useful flags on `giant train`:
|
||||
- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
|
||||
- `--conditioning {physical,embedding,onehot}` — conditioning representation
|
||||
- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
|
||||
- `--stage2-stage1-context {truth,sampled}` — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via `stage2_model.ctx_p_start`/`ctx_p_end`)
|
||||
- `--precision {fp32,bf16}` — bf16 autocast in the training loop
|
||||
- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
|
||||
- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
|
||||
- `--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
|
||||
|
||||
Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
|
||||
Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`/`.class_weighting`, `stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`, `stage*_model.trunk.*` and the finer `router` knobs (`lambda_balance`, `gumbel`, `learn_width`, …). `configs/` holds kept reference configs. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
|
||||
|
||||
`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.
|
||||
|
||||
@@ -151,12 +179,19 @@ Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_mod
|
||||
- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
|
||||
|
||||
```bash
|
||||
giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot (compute only)
|
||||
giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot × chunk (compute only)
|
||||
giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan # N rollouts vs one shared reference
|
||||
giant analyze render <run_dir> --gallery # local: styled PDFs + HTML gallery (needs LaTeX)
|
||||
giant analyze render <run_dir> --gallery # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
|
||||
|
||||
giant analyze list # every catalog plot id
|
||||
giant analyze prep rollout.yaml --chunks 8 # just the run directory, no submission
|
||||
giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0 # what a condor job runs
|
||||
giant analyze merge-one --id marginal_edep --run-dir <run_dir> # merge one plot's chunks (debugging)
|
||||
```
|
||||
|
||||
`<run_dir>` is derived next to the first rollout's parquet (`analyze prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
|
||||
`<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
|
||||
|
||||
Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.
|
||||
|
||||
## Development
|
||||
|
||||
|
||||
+157
-92
@@ -22,9 +22,9 @@ which is the order rollouts were given on the CLI) plus the single reference.
|
||||
``finalize`` merges each rollout's chunks independently and assembles a
|
||||
``Reduced.payload`` keyed the same way: ``"series": {name: ...}`` for the
|
||||
rollouts, ``"reference": ...`` as one distinguished entry (omitted on
|
||||
rollout-only plots like ``leakage_fraction``). The two heatmap-shaped specs
|
||||
(``marginal_distance_summary``, ``n_sec_confusion``) and the router
|
||||
diagnostics are inherently one-matrix/one-checkpoint per rollout, so their
|
||||
rollout-only plots like ``leakage_fraction``). The heatmap-shaped specs
|
||||
(``marginal_distance_summary``, ``sec_count_per_step_by_species``) and the
|
||||
router diagnostics are inherently one-matrix/one-checkpoint per rollout, so their
|
||||
``"series"`` entries are whole per-rollout artifacts (a matrix, a gating
|
||||
dict) rather than a single number/array — ``render.py`` draws those as one
|
||||
panel per rollout instead of one line/bar per rollout.
|
||||
@@ -60,7 +60,6 @@ from giant.analysis.reduce import (
|
||||
leakage_fraction,
|
||||
profile_finalize,
|
||||
profile_partial,
|
||||
sec_count_by_event,
|
||||
species_share,
|
||||
sum_merge,
|
||||
transverse_expr,
|
||||
@@ -72,7 +71,15 @@ from giant.analysis.router_gating import (
|
||||
compute_router_share_by_process,
|
||||
compute_router_specialization,
|
||||
)
|
||||
from giant.analysis.sources import RolloutSide, RolloutSpec, Side, open_side, physical_steps, secondaries
|
||||
from giant.analysis.sources import (
|
||||
RolloutSide,
|
||||
RolloutSpec,
|
||||
Side,
|
||||
open_side,
|
||||
physical_steps,
|
||||
secondaries,
|
||||
secondaries_by_step,
|
||||
)
|
||||
from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
|
||||
from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
|
||||
|
||||
@@ -211,31 +218,6 @@ def _ks_statistic(r_counts, t_counts) -> float:
|
||||
return float(np.max(np.abs(r_cdf - t_cdf)))
|
||||
|
||||
|
||||
def _integer_confusion(
|
||||
t: np.ndarray, r: np.ndarray, max_bins: int = 21, cap: int | None = None
|
||||
) -> tuple[list[str], np.ndarray]:
|
||||
"""Confusion matrix of two paired small-integer arrays (e.g. secondary counts).
|
||||
|
||||
Bins are consecutive integers ``0..cap``, with the last bin an overflow
|
||||
``"cap+"`` bucket, so an occasional pathological count doesn't blow up the
|
||||
heatmap. Returns ``(labels, matrix)`` with ``matrix[i, j]`` counting pairs
|
||||
with ``t == i`` and ``r == j`` (both clipped into ``[0, cap]``).
|
||||
|
||||
``cap``, if given, is used as-is instead of being derived from ``t``/``r``
|
||||
— lets a multi-rollout caller fix one shared cap (and so one shared label
|
||||
set) across every rollout's matrix rather than each panel picking its own.
|
||||
"""
|
||||
if cap is None:
|
||||
cap = min(max(int(t.max()) if len(t) else 0, int(r.max()) if len(r) else 0, 1), max_bins - 1)
|
||||
t_c = np.clip(t.astype(np.int64), 0, cap)
|
||||
r_c = np.clip(r.astype(np.int64), 0, cap)
|
||||
n = cap + 1
|
||||
mat = np.zeros((n, n), dtype=np.int64)
|
||||
np.add.at(mat, (t_c, r_c), 1)
|
||||
labels = [str(i) for i in range(cap)] + [f"{cap}+"]
|
||||
return labels, mat
|
||||
|
||||
|
||||
def _containment_depths(mat: np.ndarray, edges: np.ndarray, quantile: float) -> np.ndarray:
|
||||
"""Per-event depth containing ``quantile`` of that event's deposited energy.
|
||||
|
||||
@@ -802,6 +784,139 @@ def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
)
|
||||
|
||||
|
||||
# Per-step secondary multiplicity. Fixed integer edges (bin i == exactly i
|
||||
# secondaries, the top bin an overflow bucket) keep both plots sum-mergeable
|
||||
# across chunks — no shared-range pass needed. The species heatmap gets a
|
||||
# shorter row axis because a single step rarely emits many of *one* species.
|
||||
_N_SEC_STEP_CAP = 20
|
||||
_N_SEC_SPECIES_CAP = 10
|
||||
_OTHER_KEY = "other"
|
||||
|
||||
|
||||
def _n_sec_edges(cap: int) -> np.ndarray:
|
||||
return np.arange(-0.5, cap + 1.5)
|
||||
|
||||
|
||||
def _sec_step_key_lf(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
"""Secondaries with their emitting-step key.
|
||||
|
||||
The rollout side reads *all* rows, not just physical ones: a secondary
|
||||
whose very first row is a synthetic termination row (born, then immediately
|
||||
escaped or cut) was still produced by its parent step, and dropping it would
|
||||
undercount that step's multiplicity.
|
||||
"""
|
||||
return secondaries_by_step(lf, side)
|
||||
|
||||
|
||||
def _n_steps(lf: pl.LazyFrame) -> int:
|
||||
"""Number of (physical) step rows — the denominator the zero rows come from."""
|
||||
return int(lf.select(pl.len()).collect(engine="streaming").item())
|
||||
|
||||
|
||||
def _sec_count_per_step_partial(b: Bundle) -> dict:
|
||||
edges = _n_sec_edges(_N_SEC_STEP_CAP)
|
||||
|
||||
def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict:
|
||||
per_step = sec_lf.group_by("step_key").agg(pl.len().alias("n"))
|
||||
return {
|
||||
"h": _partial_hist(per_step, pl.col("n").clip(0, _N_SEC_STEP_CAP), edges),
|
||||
"n_steps": _n_steps(steps_lf),
|
||||
}
|
||||
|
||||
return {
|
||||
"r": _per_rollout(b, lambda rs: _side(_sec_step_key_lf(rs.all, Side.rollout), rs.phys)),
|
||||
"t": _side(_sec_step_key_lf(b.t_all, Side.reference), b.t_phys),
|
||||
}
|
||||
|
||||
|
||||
def _zero_filled(part_hists: list[dict], n_steps: int, key, nbins: int) -> list[int]:
|
||||
"""Merged counts for one series, with bin 0 (= steps that emitted none) filled in.
|
||||
|
||||
The reduction only ever sees steps that produced at least one secondary, so
|
||||
the empty ones are recovered by subtraction from the total step count.
|
||||
"""
|
||||
counts = _finalize_counts(sum_merge(part_hists), key, nbins)
|
||||
counts[0] = max(n_steps - int(sum(counts)), 0)
|
||||
return [int(c) for c in counts]
|
||||
|
||||
|
||||
def _sec_count_per_step_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
edges = _n_sec_edges(_N_SEC_STEP_CAP)
|
||||
nb = len(edges) - 1
|
||||
names = list(parts[0]["r"])
|
||||
series = {
|
||||
name: _zero_filled([p["r"][name]["h"] for p in parts], sum(p["r"][name]["n_steps"] for p in parts), 0, nb)
|
||||
for name in names
|
||||
}
|
||||
return Reduced(
|
||||
id="sec_count_per_step",
|
||||
family="secondaries",
|
||||
kind="overlay_hist",
|
||||
title="Number of secondaries per step",
|
||||
xlabel="secondaries per step",
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
"series": series,
|
||||
"reference": _zero_filled([p["t"]["h"] for p in parts], sum(p["t"]["n_steps"] for p in parts), 0, nb),
|
||||
"log_y": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _species_key_expr(top_pdgs: list[int]) -> pl.Expr:
|
||||
"""``pdg`` bucketed into the shared top-K columns plus one ``other`` bin."""
|
||||
return pl.when(pl.col("pdg").is_in(list(top_pdgs))).then(pl.col("pdg").cast(pl.Utf8)).otherwise(pl.lit(_OTHER_KEY))
|
||||
|
||||
|
||||
def _sec_count_per_step_by_species_partial(b: Bundle) -> dict:
|
||||
edges = _n_sec_edges(_N_SEC_SPECIES_CAP)
|
||||
group = _species_key_expr(b.ctx.top_pdgs)
|
||||
|
||||
def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict:
|
||||
per_step_species = sec_lf.group_by("step_key", "pdg").agg(pl.len().alias("n"))
|
||||
return {
|
||||
"h": _partial_hist(per_step_species, pl.col("n").clip(0, _N_SEC_SPECIES_CAP), edges, group=group),
|
||||
"n_steps": _n_steps(steps_lf),
|
||||
}
|
||||
|
||||
return {
|
||||
"r": _per_rollout(b, lambda rs: _side(_sec_step_key_lf(rs.all, Side.rollout), rs.phys)),
|
||||
"t": _side(_sec_step_key_lf(b.t_all, Side.reference), b.t_phys),
|
||||
}
|
||||
|
||||
|
||||
def _sec_count_per_step_by_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
edges = _n_sec_edges(_N_SEC_SPECIES_CAP)
|
||||
nb = len(edges) - 1
|
||||
names = list(parts[0]["r"])
|
||||
keys = [str(p) for p in ctx.top_pdgs] + [_OTHER_KEY]
|
||||
|
||||
def _matrix(hists: list[dict], n_steps: int) -> list[list[int]]:
|
||||
# columns = species, rows = multiplicity; every species gets its own
|
||||
# zero row (steps that produced none of *that* species).
|
||||
cols = [_zero_filled(hists, n_steps, k, nb) for k in keys]
|
||||
return [[cols[j][i] for j in range(len(keys))] for i in range(nb)]
|
||||
|
||||
return Reduced(
|
||||
id="sec_count_per_step_by_species",
|
||||
family="secondaries",
|
||||
kind="heatmap",
|
||||
title="Per-step secondary multiplicity by species",
|
||||
xlabel="species",
|
||||
payload={
|
||||
"series": {
|
||||
n: _matrix([p["r"][n]["h"] for p in parts], sum(p["r"][n]["n_steps"] for p in parts)) for n in names
|
||||
},
|
||||
"reference": _matrix([p["t"]["h"] for p in parts], sum(p["t"]["n_steps"] for p in parts)),
|
||||
"row_labels": [str(i) for i in range(_N_SEC_SPECIES_CAP)] + [f"{_N_SEC_SPECIES_CAP}+"],
|
||||
"col_labels": [pdg_label(k) for k in ctx.top_pdgs] + [_OTHER_KEY],
|
||||
"ylabel": "secondaries of this species per step",
|
||||
"cbar_label": "step count",
|
||||
"log_color": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _sec_energy_partial(b: Bundle) -> dict:
|
||||
edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1)
|
||||
return {
|
||||
@@ -858,62 +973,6 @@ def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
)
|
||||
|
||||
|
||||
def _n_sec_confusion_partial(b: Bundle) -> dict:
|
||||
t_ids, t_n = sec_count_by_event(b.t_all, _t_sec(b))
|
||||
|
||||
def _r(rs: RolloutSide) -> dict:
|
||||
ids, n = sec_count_by_event(rs.phys, _r_sec(rs))
|
||||
return {"ids": ids.tolist(), "n": n.tolist()}
|
||||
|
||||
return {"r": _per_rollout(b, _r), "t": {"ids": t_ids.tolist(), "n": t_n.tolist()}}
|
||||
|
||||
|
||||
def _n_sec_confusion_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
names = list(parts[0]["r"])
|
||||
# event-disjoint chunking (see Bundle.open) means each event_id appears in
|
||||
# exactly one part on each side, so a plain dict build is a safe merge.
|
||||
t_ids = np.concatenate([np.asarray(p["t"]["ids"], dtype=np.int64) for p in parts])
|
||||
t_n = np.concatenate([np.asarray(p["t"]["n"], dtype=np.int64) for p in parts])
|
||||
t_map = dict(zip(t_ids.tolist(), t_n.tolist()))
|
||||
|
||||
pairs: dict[str, tuple[np.ndarray, np.ndarray]] = {}
|
||||
max_val = 0
|
||||
for name in names:
|
||||
r_ids = np.concatenate([np.asarray(p["r"][name]["ids"], dtype=np.int64) for p in parts])
|
||||
r_n = np.concatenate([np.asarray(p["r"][name]["n"], dtype=np.int64) for p in parts])
|
||||
r_map = dict(zip(r_ids.tolist(), r_n.tolist()))
|
||||
common = sorted(set(r_map) & set(t_map))
|
||||
true_n = np.array([t_map[e] for e in common], dtype=np.int64)
|
||||
pred_n = np.array([r_map[e] for e in common], dtype=np.int64)
|
||||
pairs[name] = (true_n, pred_n)
|
||||
if len(true_n):
|
||||
max_val = max(max_val, int(true_n.max()), int(pred_n.max()))
|
||||
|
||||
cap = min(max(max_val, 1), 20)
|
||||
matrices: dict[str, list[list[int]]] = {}
|
||||
labels: list[str] = []
|
||||
for name in names:
|
||||
true_n, pred_n = pairs[name]
|
||||
labels, mat = _integer_confusion(true_n, pred_n, cap=cap)
|
||||
matrices[name] = mat.tolist()
|
||||
|
||||
return Reduced(
|
||||
id="n_sec_confusion",
|
||||
family="secondaries",
|
||||
kind="heatmap",
|
||||
title="Predicted vs true secondary count per event",
|
||||
xlabel="predicted secondaries (rollout)",
|
||||
payload={
|
||||
"series": matrices,
|
||||
"row_labels": labels,
|
||||
"col_labels": labels,
|
||||
"ylabel": "true secondaries (reference)",
|
||||
"cbar_label": "event count",
|
||||
"vmin": 0.0,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# router diagnostics (not chunked — already bounded/subsampled)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -1077,6 +1136,18 @@ def build_catalog() -> list[PlotSpec]:
|
||||
compute_partial=_sec_count_per_species_partial,
|
||||
finalize=_sec_count_per_species_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"sec_count_per_step",
|
||||
"secondaries",
|
||||
compute_partial=_sec_count_per_step_partial,
|
||||
finalize=_sec_count_per_step_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"sec_count_per_step_by_species",
|
||||
"secondaries",
|
||||
compute_partial=_sec_count_per_step_by_species_partial,
|
||||
finalize=_sec_count_per_step_by_species_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"sec_energy",
|
||||
"secondaries",
|
||||
@@ -1089,12 +1160,6 @@ def build_catalog() -> list[PlotSpec]:
|
||||
compute_partial=_sec_cos_angle_partial,
|
||||
finalize=_sec_cos_angle_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"n_sec_confusion",
|
||||
"secondaries",
|
||||
compute_partial=_n_sec_confusion_partial,
|
||||
finalize=_n_sec_confusion_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"router_gating",
|
||||
"model",
|
||||
|
||||
@@ -271,20 +271,3 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
|
||||
escaped = per_event["escaped"].fill_null(0.0).to_numpy()
|
||||
total = deposited + escaped
|
||||
return np.where(total > 0, escaped / total, 0.0)
|
||||
|
||||
|
||||
def sec_count_by_event(lf_all: pl.LazyFrame, sec_lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Per-event secondary count, zero-filled for events that produced none.
|
||||
|
||||
Two bounded per-event ``group_by``s — the full event set (from ``lf_all``)
|
||||
and the secondary counts (from ``sec_lf``, see ``sources.secondaries``) —
|
||||
merged in Python via a dict. Both results are event-granularity (not
|
||||
per-row), so this stays in the same bounded-memory budget as
|
||||
``event_scalars``; a plain ``group_by`` on ``sec_lf`` alone would silently
|
||||
drop zero-secondary events instead of zero-filling them.
|
||||
"""
|
||||
ev = lf_all.select("event_id").unique().collect(engine="streaming")["event_id"].to_numpy()
|
||||
cnt_df = sec_lf.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")
|
||||
cnt = dict(zip(cnt_df["event_id"].to_list(), cnt_df["n"].to_list()))
|
||||
counts = np.array([cnt.get(int(e), 0) for e in ev], dtype=np.int64)
|
||||
return ev, counts
|
||||
|
||||
@@ -27,7 +27,7 @@ from pathlib import Path
|
||||
# "router_specialization" max gate weight vs energy (one scalar trend line
|
||||
# summarizing "router_gating"), per rollout with an enabled router
|
||||
# "heatmap" row x col matrix + colorbar, one panel per rollout (a
|
||||
# distance scorecard or a predicted-vs-true confusion matrix)
|
||||
# distance scorecard)
|
||||
# "unavailable" plot not applicable to this run (e.g. no MoE checkpoint)
|
||||
|
||||
|
||||
|
||||
@@ -27,6 +27,7 @@ from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import plotstyle as ps
|
||||
from matplotlib.colors import LogNorm
|
||||
import yaml
|
||||
|
||||
from giant.analysis.reduced import Reduced
|
||||
@@ -376,10 +377,16 @@ def _render_router_specialization(r: Reduced, params: dict):
|
||||
|
||||
|
||||
def _render_heatmap(r: Reduced, params: dict):
|
||||
series = r.payload["series"]
|
||||
series = dict(r.payload["series"])
|
||||
row_labels = r.payload["row_labels"]
|
||||
col_labels = r.payload["col_labels"]
|
||||
# A heatmap-shaped plot is one matrix per rollout, so the reference (when the
|
||||
# comparison has one — the distance scorecard doesn't) becomes one more panel
|
||||
# rather than another line.
|
||||
if r.payload.get("reference") is not None:
|
||||
series["reference"] = r.payload["reference"]
|
||||
names = list(series)
|
||||
norm = LogNorm(vmin=1) if r.payload.get("log_color") else None
|
||||
fig, axes = ps.new_figure(
|
||||
"slide-16x9" if len(names) > 1 else "thesis-single",
|
||||
title=r.title,
|
||||
@@ -397,8 +404,9 @@ def _render_heatmap(r: Reduced, params: dict):
|
||||
origin="upper",
|
||||
aspect="auto",
|
||||
cmap=r.payload.get("cmap", "viridis"),
|
||||
vmin=r.payload.get("vmin"),
|
||||
vmax=r.payload.get("vmax"),
|
||||
norm=norm,
|
||||
vmin=None if norm else r.payload.get("vmin"),
|
||||
vmax=None if norm else r.payload.get("vmax"),
|
||||
)
|
||||
ax.set_xticks(range(len(col_labels)))
|
||||
ax.set_xticklabels(col_labels, rotation=45, ha="right")
|
||||
|
||||
@@ -236,3 +236,34 @@ def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
pl.col("sec_dz_list").alias("sdz"),
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def secondaries_by_step(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
"""One row per produced secondary, tagged with the step that produced it.
|
||||
|
||||
Canonical columns: ``step_key`` (an opaque struct identifying the emitting
|
||||
step) and ``pdg``. ``secondaries`` deliberately drops that link; the
|
||||
per-step multiplicity plots need it, so this is a separate view rather than
|
||||
extra columns every other consumer would pay for.
|
||||
|
||||
- rollout: a secondary's birth row carries ``parent_id`` and a birth
|
||||
position copied verbatim from the parent step's ``post_pos``, so
|
||||
``(event_id, parent_id, pre_pos)`` identifies the emitting step exactly —
|
||||
no join against the (large) step frame is needed.
|
||||
- reference: secondaries already live on their parent step's row, so the
|
||||
row index *is* the step key. It is only ever used as a group key inside
|
||||
one chunk's own aggregation, so indices repeating across chunks is
|
||||
harmless.
|
||||
"""
|
||||
if side is Side.rollout:
|
||||
return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
|
||||
pl.struct("event_id", "parent_id", "pre_x", "pre_y", "pre_z").alias("step_key"),
|
||||
"pdg",
|
||||
)
|
||||
return (
|
||||
lf.select("sec_pdg_list")
|
||||
.with_row_index("_row")
|
||||
.explode("sec_pdg_list")
|
||||
.drop_nulls("sec_pdg_list")
|
||||
.select(pl.struct("_row").alias("step_key"), pl.col("sec_pdg_list").cast(pl.Int64).alias("pdg"))
|
||||
)
|
||||
|
||||
+52
-3
@@ -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 {},
|
||||
)
|
||||
|
||||
+44
-2
@@ -141,6 +141,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,
|
||||
@@ -1020,6 +1037,15 @@ 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."""
|
||||
batch_size_auto = False
|
||||
@@ -1040,8 +1066,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)
|
||||
@@ -1407,6 +1436,15 @@ 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)."""
|
||||
if seed is not None:
|
||||
@@ -1416,8 +1454,11 @@ def rollout(
|
||||
_device = torch.device(device) if device else gconfig.auto_device()
|
||||
typer.echo(f"device: {_device}")
|
||||
|
||||
config_overrides = _parse_set_flags(set_)
|
||||
try:
|
||||
ctx = load_for_inference(checkpoint, _device, "rollout", weights=weights.value)
|
||||
ctx = load_for_inference(
|
||||
checkpoint, _device, "rollout", weights=weights.value, config_overrides=config_overrides
|
||||
)
|
||||
except CheckpointCompatibilityError as exc:
|
||||
typer.echo(f"error: {exc}", err=True)
|
||||
raise typer.Exit(1)
|
||||
@@ -1532,6 +1573,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)
|
||||
|
||||
+124
-13
@@ -404,6 +404,16 @@ class Stage2RouterConfig(RouterConfig):
|
||||
return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()}
|
||||
|
||||
|
||||
# stage2_model.n_sec.sampling choices — single source of truth for both
|
||||
# validate_config's train-time check and INFERENCE_OVERRIDES below.
|
||||
STOP_SAMPLING_CHOICES = ("greedy", "sample")
|
||||
|
||||
# stage2_model.particle_type.other_policy choices — see ParticleTypeConfig's
|
||||
# docstring for what each means; only documented there until now, since
|
||||
# nothing validated it at train time.
|
||||
OTHER_POLICY_CHOICES = ("sample", "modal", "drop")
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class NSecConfig:
|
||||
# "head": a classifier over {0..k_max} on the condition encoding alone
|
||||
@@ -424,11 +434,15 @@ class NSecConfig:
|
||||
# n_sec head was trained against Stage 1's own ConditionEncoder output and so has
|
||||
# to stay attached there, not just be labeled as such).
|
||||
owner: str = "stage2"
|
||||
# mode="stop_token" only: how sample_secondaries_ar turns a slot's stop logit into a
|
||||
# stop/continue decision. "greedy": sigmoid(logit) >= 0.5 (deterministic). "sample":
|
||||
# a Bernoulli draw at sigmoid(logit) (a real sample from the learned length
|
||||
# distribution, at the cost of an extra RNG draw per slot).
|
||||
stop_sampling: str = "greedy"
|
||||
# How resolve_n_sec/sample_secondaries_ar turn a count-bearing head's output into an
|
||||
# actual n_sec decision. mode="head": "greedy" is argmax over the classifier logits
|
||||
# (deterministic — the conditional mode, not a sample); "sample" is a categorical draw
|
||||
# from softmax(logits) (a real sample from the learned count distribution). mode=
|
||||
# "stop_token": "greedy" is sigmoid(stop_logit) >= 0.5 per slot (deterministic);
|
||||
# "sample" is a Bernoulli draw at sigmoid(stop_logit) per slot. Renamed from
|
||||
# "stop_sampling" (gitea #86), which is still accepted as a deprecated alias since it
|
||||
# appears in existing checkpoints' model_config.
|
||||
sampling: str = "greedy"
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, d: dict | None) -> "NSecConfig":
|
||||
@@ -437,7 +451,7 @@ class NSecConfig:
|
||||
mode=d.get("mode", "head"),
|
||||
lambda_weight=d.get("lambda", 0.1),
|
||||
owner=d.get("owner", "stage2"),
|
||||
stop_sampling=d.get("stop_sampling", "greedy"),
|
||||
sampling=d.get("sampling", d.get("stop_sampling", "greedy")),
|
||||
)
|
||||
|
||||
def to_dict(self) -> dict:
|
||||
@@ -445,7 +459,7 @@ class NSecConfig:
|
||||
"mode": self.mode,
|
||||
"lambda": self.lambda_weight,
|
||||
"owner": self.owner,
|
||||
"stop_sampling": self.stop_sampling,
|
||||
"sampling": self.sampling,
|
||||
}
|
||||
|
||||
|
||||
@@ -1075,6 +1089,27 @@ def _set_path(d: dict, dotted: str, value) -> None:
|
||||
cur[parts[-1]] = value
|
||||
|
||||
|
||||
def _pop_path(d: dict, dotted: str) -> None:
|
||||
"""Remove a dotted path from a nested dict, if present. No-op if any
|
||||
component along the path is missing."""
|
||||
parts = dotted.split(".")
|
||||
cur = d
|
||||
for part in parts[:-1]:
|
||||
if not isinstance(cur, dict) or part not in cur:
|
||||
return
|
||||
cur = cur[part]
|
||||
if isinstance(cur, dict):
|
||||
cur.pop(parts[-1], None)
|
||||
|
||||
|
||||
# Config keys renamed within v0.3 itself (not part of the v0.2->v0.3 migration
|
||||
# above) — normalized by migrate_config so a config.toml still using an older
|
||||
# v0.3 key name keeps passing validate_config_keys.
|
||||
_RENAMED_KEYS = {
|
||||
"stage2_model.n_sec.stop_sampling": "stage2_model.n_sec.sampling", # gitea #86
|
||||
}
|
||||
|
||||
|
||||
def _deep_merge(base: dict, override: dict) -> dict:
|
||||
"""Recursively merge `override` onto a copy of `base`.
|
||||
|
||||
@@ -1093,6 +1128,72 @@ def _deep_merge(base: dict, override: dict) -> dict:
|
||||
return result
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InferenceOverride:
|
||||
"""One dotted `model_config` path that `giant.checkpoint_io.load_for_inference`
|
||||
is allowed to change on an already-trained checkpoint, without retraining.
|
||||
|
||||
A path only belongs here if it affects neither module construction/tensor
|
||||
shapes nor the data preprocessing the normalizers/vocab maps were fit
|
||||
under — see the module docstring on `giant.model.summary` for the class
|
||||
of key this targets (`_fingerprint`'s "plain scalar attribute" leaves),
|
||||
and `giant.checkpoint_io.apply_config_overrides` for where this is used.
|
||||
"""
|
||||
|
||||
why: str
|
||||
choices: tuple[str, ...] | None = None
|
||||
minimum: float | None = None
|
||||
numeric: bool = False # int/float leaf (vs. str, the default)
|
||||
|
||||
def check(self, path: str, value: object) -> None:
|
||||
if self.choices is not None:
|
||||
if value not in self.choices:
|
||||
raise ValueError(f"{path} = {value!r} — must be one of {self.choices}")
|
||||
return
|
||||
if self.numeric:
|
||||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||||
raise ValueError(f"{path} = {value!r} — must be a number")
|
||||
if self.minimum is not None and value < self.minimum:
|
||||
raise ValueError(f"{path} = {value!r} — must be >= {self.minimum}")
|
||||
|
||||
|
||||
# Inference-safe dotted `model_config` paths — the allowlist gitea #87 asked
|
||||
# for, so a typo or a shape-bearing key (e.g. "stage1_model.hidden_dim")
|
||||
# raises a clear CheckpointCompatibilityError instead of surfacing as an
|
||||
# opaque load_state_dict size mismatch later. Extend this table, not a
|
||||
# per-call bypass, when a new inference-only key needs the capability.
|
||||
INFERENCE_OVERRIDES: dict[str, InferenceOverride] = {
|
||||
"stage2_model.n_sec.sampling": InferenceOverride(
|
||||
why="giant.sample's n_sec head/stop-token sampling reads this at sample time only (gitea #86)",
|
||||
choices=STOP_SAMPLING_CHOICES,
|
||||
),
|
||||
"stage1_model.ddpm.n_steps": InferenceOverride(
|
||||
why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
|
||||
numeric=True,
|
||||
minimum=1,
|
||||
),
|
||||
"stage2_model.ddpm.n_steps": InferenceOverride(
|
||||
why="giant.model.schedule.CosineSchedule's step count, resolved at sample time",
|
||||
numeric=True,
|
||||
minimum=1,
|
||||
),
|
||||
"stage2_model.particle_type.other_policy": InferenceOverride(
|
||||
why="giant.rollout resolves an 'other'-bucket secondary's PDG code with this at rollout time",
|
||||
choices=OTHER_POLICY_CHOICES,
|
||||
),
|
||||
"stage1_model.router.temperature": InferenceOverride(
|
||||
why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
|
||||
numeric=True,
|
||||
minimum=1e-6,
|
||||
),
|
||||
"stage2_model.router.temperature": InferenceOverride(
|
||||
why="giant.model.routers.EnergyRouter.temperature, a plain constructor attribute",
|
||||
numeric=True,
|
||||
minimum=1e-6,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class FlagSpec:
|
||||
"""One CLI flag's mapping into the config-overrides tree.
|
||||
@@ -1271,11 +1372,21 @@ def migrate_config(cfg: dict) -> dict:
|
||||
(which additionally carries n_sec_head ownership and needs
|
||||
`network.build_models`'s cooperation) is a separate migration surface,
|
||||
deferred to the network.py refactor.
|
||||
"""
|
||||
if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
|
||||
return copy.deepcopy(cfg)
|
||||
|
||||
Independently of the v0.2/v0.3 branch below, `_RENAMED_KEYS` normalizes
|
||||
keys renamed within v0.3 itself (e.g. `stop_sampling` -> `sampling`,
|
||||
gitea #86) so a config.toml written against an older v0.3 key name still
|
||||
passes `validate_config_keys`.
|
||||
"""
|
||||
cfg = copy.deepcopy(cfg)
|
||||
for old_path, new_path in _RENAMED_KEYS.items():
|
||||
if _get_path(cfg, old_path) is not None and _get_path(cfg, new_path) is None:
|
||||
_set_path(cfg, new_path, _get_path(cfg, old_path))
|
||||
_pop_path(cfg, old_path)
|
||||
|
||||
if _get_path(cfg, "meta.config_version") == CONFIG_VERSION:
|
||||
return cfg
|
||||
|
||||
old_train = cfg.pop("train", {})
|
||||
old_model = cfg.pop("model", {})
|
||||
old_router = dict(old_model.pop("router", {}))
|
||||
@@ -1512,9 +1623,9 @@ def validate_config(cfg: dict, *, resume: bool = False) -> None:
|
||||
"conditioning to hang an EOS decision off"
|
||||
)
|
||||
|
||||
stop_sampling = _get_path(cfg, "stage2_model.n_sec.stop_sampling")
|
||||
if stop_sampling not in ("greedy", "sample"):
|
||||
raise ValueError(f"stage2_model.n_sec.stop_sampling = {stop_sampling!r} — must be 'greedy' or 'sample'")
|
||||
n_sec_sampling = _get_path(cfg, "stage2_model.n_sec.sampling")
|
||||
if n_sec_sampling not in STOP_SAMPLING_CHOICES:
|
||||
raise ValueError(f"stage2_model.n_sec.sampling = {n_sec_sampling!r} — must be 'greedy' or 'sample'")
|
||||
|
||||
precision = _get_path(cfg, "train.precision")
|
||||
if precision not in ("fp32", "bf16"):
|
||||
|
||||
@@ -138,7 +138,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
|
||||
type_head_cfg=s2_spec.heads.type.to_dict(),
|
||||
build_stop_head=stop_token,
|
||||
stop_sampling=s2_spec.n_sec.stop_sampling,
|
||||
n_sec_sampling=s2_spec.n_sec.sampling,
|
||||
stop_head_cfg=s2_spec.heads.n_sec.to_dict(),
|
||||
)
|
||||
else:
|
||||
@@ -168,6 +168,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
cond_enc=shared_cond_enc,
|
||||
n_sec_head_cfg=s2_spec.heads.n_sec.to_dict(),
|
||||
type_head_cfg=s2_spec.heads.type.to_dict(),
|
||||
n_sec_sampling=s2_spec.n_sec.sampling,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
@@ -370,6 +370,7 @@ class Stage2OneShot(StageModel):
|
||||
cond_enc: ConditionEncoder | None = None,
|
||||
n_sec_head_cfg: dict | None = None,
|
||||
type_head_cfg: dict | None = None,
|
||||
n_sec_sampling: str = "greedy",
|
||||
) -> None:
|
||||
super().__init__(
|
||||
pdg_vocab,
|
||||
@@ -383,6 +384,7 @@ class Stage2OneShot(StageModel):
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
cond_enc=cond_enc,
|
||||
)
|
||||
self.n_sec_sampling = n_sec_sampling
|
||||
self._build_context_fusion(x_dim, context_dim, cond_out_dim)
|
||||
target = self.particle_type_cfg.target
|
||||
type_head_out_dim = None if target == "physical" else k_max * self.type_dim
|
||||
@@ -498,7 +500,7 @@ class Stage2Autoregressive(StageModel):
|
||||
n_sec_head_cfg: dict | None = None,
|
||||
type_head_cfg: dict | None = None,
|
||||
build_stop_head: bool = False,
|
||||
stop_sampling: str = "greedy",
|
||||
n_sec_sampling: str = "greedy",
|
||||
stop_head_cfg: dict | None = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
@@ -514,7 +516,7 @@ class Stage2Autoregressive(StageModel):
|
||||
cond_enc=cond_enc,
|
||||
)
|
||||
self.history_kind = history
|
||||
self.stop_sampling = stop_sampling
|
||||
self.n_sec_sampling = n_sec_sampling
|
||||
self.context_adapter = ContextAdapter(x_dim, context_dim)
|
||||
self.base_fuse = nn.Sequential(
|
||||
nn.Linear(cond_out_dim + context_dim, cond_out_dim),
|
||||
|
||||
+12
-2
@@ -37,7 +37,7 @@ from dataclasses import dataclass, field
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from giant.config import _get_path, _set_path, leaf_paths
|
||||
from giant.config import INFERENCE_OVERRIDES, _get_path, _set_path, leaf_paths
|
||||
from giant.model.builders import build_critics, build_models
|
||||
from giant.model.trunks import RoutedTrunk
|
||||
|
||||
@@ -111,6 +111,7 @@ class ModelSummary:
|
||||
pdg_vocab: int
|
||||
mat_vocab: int
|
||||
vocab_caveats: list[str] = field(default_factory=list)
|
||||
overridable: list[str] = field(default_factory=list)
|
||||
|
||||
|
||||
def _build_model_config(cfg: dict, pdg_vocab: int, mat_vocab: int) -> dict:
|
||||
@@ -140,7 +141,7 @@ def _fingerprint(modules: dict[str, nn.Module]) -> list:
|
||||
exist, every parameter's/buffer's shape+dtype (never values — those are
|
||||
randomly initialized and irrelevant to *structure*), and every plain
|
||||
scalar attribute any module stores on itself (e.g. `Stage2Autoregressive
|
||||
.stop_sampling`, `EnergyRouter.temperature`) — this is what makes a
|
||||
.n_sec_sampling`, `EnergyRouter.temperature`) — this is what makes a
|
||||
non-parametric key's effect on construction observable."""
|
||||
sig = []
|
||||
for stage_name, module in modules.items():
|
||||
@@ -234,6 +235,8 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
|
||||
else:
|
||||
inert.append(path)
|
||||
|
||||
overridable = sorted(p for p in in_scope if p in INFERENCE_OVERRIDES)
|
||||
|
||||
return ModelSummary(
|
||||
modules=modules,
|
||||
consumed=sorted(consumed),
|
||||
@@ -242,6 +245,7 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
vocab_caveats=_vocab_caveats(cfg),
|
||||
overridable=overridable,
|
||||
)
|
||||
|
||||
|
||||
@@ -309,6 +313,12 @@ def render_summary(summary: ModelSummary) -> str:
|
||||
else:
|
||||
lines.append(" (none)")
|
||||
|
||||
if summary.overridable:
|
||||
lines.append("")
|
||||
lines.append("inference-overridable without retraining (giant predict/rollout --set):")
|
||||
for path in summary.overridable:
|
||||
lines.append(f" {path} ({INFERENCE_OVERRIDES[path].why})")
|
||||
|
||||
if summary.vocab_caveats:
|
||||
lines.append("")
|
||||
lines.append("vocab placeholder caveats:")
|
||||
|
||||
+10
-3
@@ -261,7 +261,7 @@ def sample_secondaries_ar(
|
||||
slot's own stop logit (`predict_stop`, evaluated on the same prefix
|
||||
conditioning as the token itself — see `predict_type`'s docstring for
|
||||
why this needs no extra state) decides whether generation should have
|
||||
already stopped, per `sec_decoder.stop_sampling` ("greedy": threshold at
|
||||
already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
|
||||
0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
|
||||
`n_sec_pred` is the first slot index where this fires; once every row in
|
||||
the batch has fired, the loop breaks before spending a model call on the
|
||||
@@ -345,7 +345,7 @@ def sample_secondaries_ar(
|
||||
slot_idx,
|
||||
hist=hist,
|
||||
).squeeze(1)
|
||||
if sec_decoder.stop_sampling == "sample":
|
||||
if sec_decoder.n_sec_sampling == "sample":
|
||||
stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
|
||||
else:
|
||||
stop_now = stop_logit >= 0.0
|
||||
@@ -496,7 +496,12 @@ def resolve_n_sec(
|
||||
|
||||
Raises if neither stage owns any n_sec mechanism at all — the only way
|
||||
that happens is `stage2_model.n_sec.mode = "truth"`, which is not a valid
|
||||
rollout-/predict-capable checkpoint."""
|
||||
rollout-/predict-capable checkpoint.
|
||||
|
||||
`n_sec.mode = "head"` resolves the classifier logits per
|
||||
`sec_decoder.n_sec_sampling`: "greedy" (default) takes the conditional
|
||||
mode via argmax; "sample" draws a real sample from the learned count
|
||||
distribution via `torch.multinomial` on the softmax — see gitea #86."""
|
||||
if n_sec_pred is not None:
|
||||
return n_sec_pred
|
||||
if getattr(sec_decoder, "stop_head", None) is not None:
|
||||
@@ -508,4 +513,6 @@ def resolve_n_sec(
|
||||
"'truth' is standalone-evaluation-only"
|
||||
)
|
||||
logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
|
||||
if sec_decoder.n_sec_sampling == "sample":
|
||||
return torch.multinomial(logits.softmax(dim=-1), 1).squeeze(-1)
|
||||
return logits.argmax(dim=-1)
|
||||
|
||||
+1
-1
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "giant"
|
||||
version = "0.3.9"
|
||||
version = "0.3.13"
|
||||
description = "Geant4 step-function surrogate via conditional flow matching"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
|
||||
@@ -13,6 +13,7 @@ from giant.analysis.sources import (
|
||||
open_side,
|
||||
physical_steps,
|
||||
secondaries,
|
||||
secondaries_by_step,
|
||||
)
|
||||
from giant.data.loader import EVENT_ID_FILE_STRIDE
|
||||
|
||||
@@ -159,18 +160,17 @@ def test_secondaries_rollout_vs_reference_align():
|
||||
assert t["pdg"].to_list() == [22, 22]
|
||||
|
||||
|
||||
def test_sec_count_by_event_zero_fills_events_with_no_secondaries():
|
||||
r_phys = physical_steps(_rollout_frame(), Side.rollout)
|
||||
r_sec = secondaries(_rollout_frame(), Side.rollout)
|
||||
ev, n = R.sec_count_by_event(r_phys, r_sec)
|
||||
# event 1 has one secondary track; event 2 has none and must still appear (as 0),
|
||||
# not silently drop out of a plain group_by on the secondaries frame alone.
|
||||
assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}
|
||||
def test_secondaries_by_step_keys_each_secondary_to_its_emitting_step():
|
||||
r = secondaries_by_step(_rollout_frame(), Side.rollout).collect()
|
||||
assert r["pdg"].to_list() == [22]
|
||||
# the rollout key is (event_id, parent_id, birth position) — the parent
|
||||
# step's post_pos, copied verbatim onto the child's birth row.
|
||||
assert r["step_key"][0] == {"event_id": 1, "parent_id": 0, "pre_x": 0.0, "pre_y": 0.0, "pre_z": 1.0}
|
||||
|
||||
t_all = _reference_frame()
|
||||
t_sec = secondaries(t_all, Side.reference)
|
||||
ev, n = R.sec_count_by_event(t_all, t_sec)
|
||||
assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1}
|
||||
t = secondaries_by_step(_reference_frame(), Side.reference).collect()
|
||||
assert t["pdg"].to_list() == [22, 22]
|
||||
# one row per emitting step; the empty-list step drops out entirely
|
||||
assert [k["_row"] for k in t["step_key"]] == [0, 2]
|
||||
|
||||
|
||||
def test_leakage_fraction():
|
||||
|
||||
+26
-38
@@ -10,10 +10,10 @@ from giant.analysis.catalog import (
|
||||
Bundle,
|
||||
PlotSpec,
|
||||
_containment_depths,
|
||||
_integer_confusion,
|
||||
_ks_statistic,
|
||||
)
|
||||
from giant.analysis.context import Context, build_context
|
||||
from giant.analysis.grouping import pdg_label
|
||||
from giant.analysis.sources import RolloutSpec
|
||||
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
||||
|
||||
@@ -160,8 +160,9 @@ def _validate_payload(r, names: list[str]) -> None:
|
||||
# data-dependent edges (event_total_edep), concat-then-mean/std (shower_
|
||||
# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
|
||||
# ratio (species_edep_share), a chunkable=False passthrough (router_gating),
|
||||
# nested sum-merge into a scorecard (marginal_distance_summary), concat-then-
|
||||
# event-id-join (n_sec_confusion), and concat-then-per-event-derived-quantity
|
||||
# sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
|
||||
# nested sum-merge into a scorecard (marginal_distance_summary), and
|
||||
# concat-then-per-event-derived-quantity
|
||||
# (shower_containment_depth_90, reusing the profile matrix's own merge shape).
|
||||
_CHUNK_EQUIVALENCE_IDS = [
|
||||
"marginal_edep",
|
||||
@@ -170,9 +171,10 @@ _CHUNK_EQUIVALENCE_IDS = [
|
||||
"shower_longitudinal",
|
||||
"leakage_fraction",
|
||||
"sec_count_per_species",
|
||||
"sec_count_per_step",
|
||||
"sec_count_per_step_by_species",
|
||||
"router_gating",
|
||||
"marginal_distance_summary",
|
||||
"n_sec_confusion",
|
||||
"shower_containment_depth_90",
|
||||
]
|
||||
|
||||
@@ -217,7 +219,7 @@ def test_chunked_matches_unchunked(two_ctx: Context, spec_id: str):
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# new (gitea #76) reductions: KS distance, confusion matrix, containment depth
|
||||
# new (gitea #76) reductions: KS distance and containment depth
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@@ -228,28 +230,6 @@ def test_ks_statistic():
|
||||
assert _ks_statistic([10, 0], [0, 0]) == 1.0 # one side empty, other isn't -> maximal mismatch
|
||||
|
||||
|
||||
def test_integer_confusion_matches_event_pairing():
|
||||
# true (reference) n_sec = [1, 1]; predicted (rollout) n_sec = [1, 0]
|
||||
labels, mat = _integer_confusion(np.array([1, 1]), np.array([1, 0]))
|
||||
assert labels == ["0", "1+"]
|
||||
assert mat.tolist() == [[0, 0], [1, 1]] # row=true, col=pred
|
||||
|
||||
|
||||
def test_integer_confusion_caps_pathological_outliers():
|
||||
labels, mat = _integer_confusion(np.array([0, 500]), np.array([0, 0]), max_bins=5)
|
||||
assert labels[-1] == "4+"
|
||||
assert mat.shape == (5, 5)
|
||||
assert mat.sum() == 2
|
||||
|
||||
|
||||
def test_integer_confusion_explicit_cap_overrides_local_range():
|
||||
# Even though this pair's own max is 1, an explicit shared cap forces a
|
||||
# wider (and so cross-rollout-consistent) label set.
|
||||
labels, mat = _integer_confusion(np.array([1, 1]), np.array([0, 1]), cap=3)
|
||||
assert labels == ["0", "1", "2", "3+"]
|
||||
assert mat.shape == (4, 4)
|
||||
|
||||
|
||||
def test_containment_depths_simple_ramp():
|
||||
# one event, edep concentrated in the first bin -> 90%/95% containment
|
||||
# depth is the first bin's right edge; a zero-energy event is dropped.
|
||||
@@ -259,17 +239,25 @@ def test_containment_depths_simple_ramp():
|
||||
assert depths.tolist() == [1.0]
|
||||
|
||||
|
||||
def test_n_sec_confusion_spec(bundle):
|
||||
spec = get_spec("n_sec_confusion")
|
||||
def test_sec_count_per_step_counts_empty_steps(bundle):
|
||||
spec = get_spec("sec_count_per_step")
|
||||
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
||||
assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"]
|
||||
assert r.payload["series"]["rollout"] == [[0, 0], [1, 1]]
|
||||
# reference: 3 steps, two of which emit exactly one secondary
|
||||
assert r.payload["reference"][:2] == [1, 2]
|
||||
# rollout: 4 physical steps, one of which emits a single secondary
|
||||
assert r.payload["series"]["rollout"][:2] == [3, 1]
|
||||
assert sum(r.payload["reference"]) == 3
|
||||
|
||||
|
||||
def test_n_sec_confusion_shares_one_cap_across_rollouts(two_bundle):
|
||||
spec = get_spec("n_sec_confusion")
|
||||
r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx)
|
||||
assert list(r.payload["series"]) == ["flow", "wgan"]
|
||||
# both rollouts share the same fixture data here, so their matrices (and
|
||||
# the shared label set) must be identical.
|
||||
assert r.payload["series"]["flow"] == r.payload["series"]["wgan"]
|
||||
def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
|
||||
spec = get_spec("sec_count_per_step_by_species")
|
||||
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
||||
cols = r.payload["col_labels"]
|
||||
ref = r.payload["reference"]
|
||||
g = cols.index(pdg_label(22))
|
||||
# two reference steps emit one photon each; the third emits none
|
||||
assert [row[g] for row in ref][:2] == [1, 2]
|
||||
# every other species column is "no such secondary" on all 3 steps
|
||||
for j, _ in enumerate(cols):
|
||||
if j != g:
|
||||
assert ref[0][j] == 3 and sum(row[j] for row in ref[1:]) == 0
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -31,3 +31,28 @@ def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "checkpoint has no model_config" in result.output
|
||||
|
||||
|
||||
def test_rollout_set_flag_disallowed_path_surfaces_compat_error(tmp_path):
|
||||
checkpoint = tmp_path / "ckpt.pt"
|
||||
torch.save(
|
||||
{"model_config": {"stage1_model": {}, "stage2_model": {}}, "sec_decoder": {}, "normalizer": {"sec_phys": {}}},
|
||||
checkpoint,
|
||||
)
|
||||
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"rollout",
|
||||
"dummy.parquet",
|
||||
"--checkpoint",
|
||||
str(checkpoint),
|
||||
"--geometry",
|
||||
"dummy_geometry.pkl",
|
||||
"--set",
|
||||
"stage2_model.n_sec.typo=sample",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "not an inference-safe override" in result.output
|
||||
|
||||
+32
-9
@@ -171,17 +171,40 @@ def test_n_sec_config_owner_defaults_to_stage2():
|
||||
def test_n_sec_config_owner_round_trips():
|
||||
n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "owner": "stage1"})
|
||||
assert n_sec.owner == "stage1"
|
||||
assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "stop_sampling": "greedy"}
|
||||
assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "owner": "stage1", "sampling": "greedy"}
|
||||
|
||||
|
||||
def test_n_sec_config_stop_sampling_defaults_to_greedy():
|
||||
assert gconfig.NSecConfig().stop_sampling == "greedy"
|
||||
def test_n_sec_config_sampling_defaults_to_greedy():
|
||||
assert gconfig.NSecConfig().sampling == "greedy"
|
||||
|
||||
|
||||
def test_n_sec_config_stop_sampling_round_trips():
|
||||
def test_n_sec_config_sampling_round_trips():
|
||||
n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "sampling": "sample"})
|
||||
assert n_sec.sampling == "sample"
|
||||
assert n_sec.to_dict()["sampling"] == "sample"
|
||||
|
||||
|
||||
def test_n_sec_config_stop_sampling_alias_still_honored():
|
||||
"""gitea #86: stop_sampling was renamed to sampling; old checkpoints'
|
||||
model_config still carries the old key and must keep working."""
|
||||
n_sec = gconfig.NSecConfig.from_dict({"mode": "stop_token", "stop_sampling": "sample"})
|
||||
assert n_sec.stop_sampling == "sample"
|
||||
assert n_sec.to_dict()["stop_sampling"] == "sample"
|
||||
assert n_sec.sampling == "sample"
|
||||
assert "stop_sampling" not in n_sec.to_dict()
|
||||
|
||||
|
||||
def test_n_sec_config_sampling_key_wins_over_stop_sampling_alias():
|
||||
n_sec = gconfig.NSecConfig.from_dict({"sampling": "sample", "stop_sampling": "greedy"})
|
||||
assert n_sec.sampling == "sample"
|
||||
|
||||
|
||||
def test_migrate_config_renames_stop_sampling_key():
|
||||
cfg = {
|
||||
"meta": {"config_version": gconfig.CONFIG_VERSION},
|
||||
"stage2_model": {"n_sec": {"stop_sampling": "sample"}},
|
||||
}
|
||||
migrated = gconfig.migrate_config(cfg)
|
||||
assert gconfig._get_path(migrated, "stage2_model.n_sec.sampling") == "sample"
|
||||
assert gconfig._get_path(migrated, "stage2_model.n_sec.stop_sampling") is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -851,13 +874,13 @@ def test_validate_config_stop_token_rejected_for_stage1_owner():
|
||||
assert "stop_token" in str(e) and "owner" in str(e)
|
||||
|
||||
|
||||
def test_validate_config_bad_stop_sampling_rejected():
|
||||
cfg = _cfg_with(**{"stage2_model.n_sec.stop_sampling": "bogus"})
|
||||
def test_validate_config_bad_n_sec_sampling_rejected():
|
||||
cfg = _cfg_with(**{"stage2_model.n_sec.sampling": "bogus"})
|
||||
try:
|
||||
gconfig.validate_config(cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "stop_sampling" in str(e)
|
||||
assert "sampling" in str(e)
|
||||
|
||||
|
||||
def test_validate_config_default_precision_is_fp32():
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -304,13 +304,15 @@ def test_render_one_of_each_kind(tmp_path: Path):
|
||||
"hm1",
|
||||
"secondaries",
|
||||
"heatmap",
|
||||
"Confusion (single rollout)",
|
||||
"Heatmap (single rollout)",
|
||||
"predicted",
|
||||
{
|
||||
"series": {"flow": [[1, 0], [0, 1]]},
|
||||
"reference": [[2, 0], [0, 1]],
|
||||
"row_labels": ["0", "1+"],
|
||||
"col_labels": ["0", "1+"],
|
||||
"cbar_label": "count",
|
||||
"log_color": True,
|
||||
},
|
||||
),
|
||||
]
|
||||
|
||||
+68
-8
@@ -14,6 +14,7 @@ from giant.model.network import (
|
||||
stage2_trunk_sec_dim,
|
||||
)
|
||||
from giant.sample import (
|
||||
resolve_n_sec,
|
||||
sample_flow,
|
||||
sample_secondaries,
|
||||
sample_secondaries_ar,
|
||||
@@ -72,6 +73,7 @@ def _stage2_ar(
|
||||
mat: int = 2,
|
||||
k_max: int = 5,
|
||||
history: str = "markov",
|
||||
n_sec_sampling: str = "greedy",
|
||||
) -> Stage2Autoregressive:
|
||||
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
|
||||
return Stage2Autoregressive(
|
||||
@@ -89,6 +91,7 @@ def _stage2_ar(
|
||||
history=history,
|
||||
attn_n_heads=2,
|
||||
attn_n_layers=1,
|
||||
n_sec_sampling=n_sec_sampling,
|
||||
).eval()
|
||||
|
||||
|
||||
@@ -99,7 +102,7 @@ def _expected_type_dim(target: str, emb_dim: int) -> int:
|
||||
def _stage2_ar_stop_token(
|
||||
target: str,
|
||||
generator: str,
|
||||
stop_sampling: str = "greedy",
|
||||
n_sec_sampling: str = "greedy",
|
||||
emb_dim: int = 6,
|
||||
pdg: int = 3,
|
||||
mat: int = 2,
|
||||
@@ -120,7 +123,7 @@ def _stage2_ar_stop_token(
|
||||
particle_type_cfg=ParticleTypeConfig(target=target),
|
||||
build_n_sec_head=False,
|
||||
build_stop_head=True,
|
||||
stop_sampling=stop_sampling,
|
||||
n_sec_sampling=n_sec_sampling,
|
||||
).eval()
|
||||
|
||||
|
||||
@@ -267,14 +270,14 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
|
||||
# ── Stage2Autoregressive: n_sec.mode = "stop_token" ─────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
|
||||
def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(stop_sampling):
|
||||
@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
|
||||
def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_sec_sampling):
|
||||
"""A stop_head pinned to a large positive logit fires at slot 0 for
|
||||
every row under both policies (greedy: sigmoid(logit) >= 0.5; sample:
|
||||
a Bernoulli draw at sigmoid(logit) ~= 1) — the loop should break before
|
||||
generating any token."""
|
||||
B, k_max = 4, 5
|
||||
decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
|
||||
decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
|
||||
_force_stop_head_logit(decoder, 50.0)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
@@ -283,13 +286,13 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(sto
|
||||
assert not sec_valid.any()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stop_sampling", ["greedy", "sample"])
|
||||
def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(stop_sampling):
|
||||
@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
|
||||
def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_sampling):
|
||||
"""A stop_head pinned to a large negative logit never fires under either
|
||||
policy, so every row is capped at k_max (the safety cap, not a modeling
|
||||
ceiling)."""
|
||||
B, k_max = 4, 5
|
||||
decoder = _stage2_ar_stop_token("physical", "flow", stop_sampling=stop_sampling, k_max=k_max)
|
||||
decoder = _stage2_ar_stop_token("physical", "flow", n_sec_sampling=n_sec_sampling, k_max=k_max)
|
||||
_force_stop_head_logit(decoder, -50.0)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
@@ -336,3 +339,60 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
|
||||
stage1_out = torch.randn(3, X_DIM)
|
||||
with pytest.raises(AssertionError):
|
||||
sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
|
||||
|
||||
|
||||
# ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
|
||||
|
||||
|
||||
def _force_n_sec_head_bias(decoder: Stage2Autoregressive, bias: torch.Tensor) -> None:
|
||||
"""Zeroes n_sec_head's weights and pins its bias, so predict_n_sec
|
||||
returns `bias` (broadcast over the batch) as logits regardless of
|
||||
conditioning — mirrors `_force_stop_head_logit`."""
|
||||
assert decoder.n_sec_head is not None
|
||||
last_linear = decoder.n_sec_head[-1]
|
||||
with torch.no_grad():
|
||||
last_linear.weight.zero_()
|
||||
last_linear.bias.copy_(bias)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n_sec_sampling", ["greedy", "sample"])
|
||||
def test_resolve_n_sec_head_mode_sharply_peaked_logits_pick_dominant_class(n_sec_sampling):
|
||||
"""A logit vector overwhelmingly favoring one class gives the same
|
||||
answer under both policies — greedy because it's the argmax, sample
|
||||
because softmax puts ~all mass on it."""
|
||||
B, k_max = 8, 5
|
||||
decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling=n_sec_sampling)
|
||||
bias = torch.full((k_max + 1,), -50.0)
|
||||
bias[2] = 50.0
|
||||
_force_n_sec_head_bias(decoder, bias)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
|
||||
assert n_sec is not None
|
||||
assert torch.equal(n_sec, torch.full((B,), 2, dtype=torch.long))
|
||||
|
||||
|
||||
def test_resolve_n_sec_head_mode_greedy_is_deterministic_under_flat_logits():
|
||||
B, k_max = 32, 5
|
||||
decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="greedy")
|
||||
_force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
|
||||
assert n_sec is not None
|
||||
assert n_sec.unique().numel() == 1
|
||||
|
||||
|
||||
def test_resolve_n_sec_head_mode_sample_varies_under_flat_logits():
|
||||
"""Under a flat logit vector, a categorical draw across a large batch
|
||||
should hit more than one class — the whole point of gitea #86: greedy
|
||||
always collapses to one, sample should not."""
|
||||
torch.manual_seed(0)
|
||||
B, k_max = 256, 5
|
||||
decoder = _stage2_ar("physical", "flow", k_max=k_max, n_sec_sampling="sample")
|
||||
_force_n_sec_head_bias(decoder, torch.zeros(k_max + 1))
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec = resolve_n_sec(decoder, decoder, cond_cont, cond_cat, stage1_out, None)
|
||||
assert n_sec is not None
|
||||
assert n_sec.unique().numel() > 1
|
||||
|
||||
Reference in New Issue
Block a user