9 Commits

Author SHA1 Message Date
gitea-actions bb8d16caba chore: update changelog for v0.3.11 [skip ci] 2026-08-26 12:33:51 +00:00
gitea-actions c8a1b4f25d chore: bump version 0.3.10 -> 0.3.11 [skip ci] 2026-08-26 12:33:46 +00:00
lars 23efd6d9ff Merge pull request 'feat(analysis): per-step secondary multiplicity plots' (#85) from analysis/per-step-secondary-multiplicity into master
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (push) Successful in 28s
CI / Format (ruff format) (push) Successful in 38s
CI / Type check (ty) (push) Successful in 40s
CI / Tests (push) Successful in 3m37s
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 2m4s
Reviewed-on: #85
2026-08-26 14:28:16 +02:00
lars 9fa6420183 feat(analysis): per-step secondary multiplicity plots
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 2m3s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 1m19s
CI / Lint (ruff check) (push) Successful in 3m55s
CI / Format (ruff format) (push) Successful in 3m54s
CI / Format (ruff format) (pull_request) Successful in 5m37s
CI / Lint (ruff check) (pull_request) Successful in 5m42s
CI / Tests (push) Successful in 8m42s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
CI / Tests (pull_request) Successful in 5m3s
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
Replace the event-level n_sec confusion matrix with two step-resolved
secondary-multiplicity comparisons:

- sec_count_per_step: overlay histogram of how many secondaries a single
  step emits, rollout series vs reference.
- sec_count_per_step_by_species: heatmap of per-step multiplicity of one
  species (zero row included) against species, drawn as one panel per
  rollout plus a reference panel, raw counts on a log color scale.

Both are backed by a new sources.secondaries_by_step view, which tags each
secondary with its emitting step — (event_id, parent_id, birth position)
on the rollout side, the row index on the reference side — so neither plot
needs a join against the step frame. Steps that emitted nothing are
recovered by subtraction from the chunk's step count, keeping both specs
sum-mergeable across condor chunks.

The rollout multiplicity is derived from the actual secondary birth rows
rather than the n_sec_pred column, which records the predicted count
before the per-event max-tracks cap.

_render_heatmap gained reference-panel and log-color support;
marginal_distance_summary sets neither key and is unchanged.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-26 14:13:18 +02:00
gitea-actions b66574877b chore: update changelog for v0.3.10 [skip ci] 2026-08-26 08:15:52 +00:00
gitea-actions 9b77e04731 chore: bump version 0.3.9 -> 0.3.10 [skip ci] 2026-08-26 08:15:51 +00:00
lars 8dee2feab7 Merge pull request 'docs: bring README and CLAUDE.md in line with v0.3.9' (#82) from docs/sync-readme-claude-md into master
CI / Lint (ruff check) (push) Successful in 38s
CI / Format (ruff format) (push) Successful in 37s
CI / Type check (ty) (push) Successful in 38s
CI / Sync project version with tag (push) Has been skipped
CI / Tests (push) Successful in 2m49s
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 33s
Reviewed-on: #82
2026-08-26 10:05:54 +02:00
lars f2da0642b2 docs: bring README and CLAUDE.md in line with v0.3.9
CI / Format (ruff format) (push) Successful in 36s
CI / Lint (ruff check) (push) Successful in 41s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 32s
CI / Lint (ruff check) (pull_request) Successful in 49s
CI / Format (ruff format) (pull_request) Successful in 48s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 40s
CI / Tests (push) Successful in 5m50s
CI / Tests (pull_request) Successful in 4m36s
CI / Bump version, tag, and update changelog on merge to master (push) Has been skipped
CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped
CLAUDE.md still described the pre-v0.3.0 codebase: the Stage-2
autoregressive redesign as "designed, not implemented", a monolithic
network.py, a single global model.conditioning switch, and WGAN as
"implemented, not yet tested".

- Architecture rewritten around the actual giant/model split
  (layers/encoders/trunks/routers/history/objectives/models/builders/
  _legacy/summary; network.py is now a re-export shim), plus
  cond_layout.py, checkpoint_io.py, _migration.py, data/setup_cache.py
  and giant/training/.
- Conditioning documented per axis (conditioning.particle /
  conditioning.material, each physical|embedding|onehot, freely mixed).
- Stage 2 documented with both decoders, n_sec.mode, teacher forcing,
  stage1_context and the three particle_type.target options.
- Roadmap: v0.3.0 recorded as implemented/released; WGAN and MoE routing
  as implemented but unvalidated, with the router retrain as next step.
- Analysis: run dir is <cwd>/analysis_runs/analysis_<id>, plus
  variables/reduced/runtime_estimate and analyze list/merge-one/metrics.
- Added giant model summary, configs/, and the CI-automated version and
  changelog bump.

README drift fixes only: project tree for the model/analysis/training
splits, analyze run-dir default, missing subcommands, --precision and
--stage2-stage1-context, the extras list, and two accuracy fixes
(--router configures stage 1 only; --conditioning sets two independent
axes at once).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-08-26 10:04:06 +02:00
lars 1e92902c8d Backfill CHANGELOG.md for v0.2.0-v0.3.2
CI / Format (ruff format) (push) Successful in 31s
CI / Lint (ruff check) (push) Successful in 33s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 30s
CI / Tests (push) Successful in 2m53s
CI / Bump version, tag, and update changelog on merge to master (push) Successful in 13s
The automated changelog (gitea #50) deliberately started fresh with no
backfill; this reverses that call now that it's wanted. v0.2.0-v0.3.2 are
generated from tag history via git-cliff/cliff.toml, matching the format of
existing entries. v0.3.3 was bumped but never tagged, so its commits stay
folded into the existing v0.3.4 entry. The v0.2.0 range (198 uncurated
pre-automation commits) is hand-curated to drop duplicate commits and
dev-log noise (WIP markers, incomplete-validation runs, repeated
"Apply ruff format").
2026-08-24 15:30:14 +02:00
14 changed files with 837 additions and 208 deletions
+1 -1
View File
@@ -1,5 +1,5 @@
[tool.bumpversion] [tool.bumpversion]
current_version = "0.3.9" current_version = "0.3.11"
parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)" parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
serialize = ["{major}.{minor}.{patch}"] serialize = ["{major}.{minor}.{patch}"]
search = "{current_version}" search = "{current_version}"
+490 -1
View File
@@ -1,5 +1,19 @@
# Changelog # Changelog
## [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 ## [0.3.9] - 2026-08-24
### Added ### Added
@@ -65,4 +79,479 @@
- Document CI_TOKEN's write:repository scope requirement [gitea #50](https://git.larsbogner.de/lars/giant/issues/50) - 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
+57 -29
View File
@@ -7,18 +7,20 @@ This file provides guidance to Claude Code (claude.ai/code) when working with co
```bash ```bash
uv sync --extra cpu # install dependencies with CPU-only torch (standard/default) 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 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) uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle (giant rollout)
pytest # run tests pytest # run tests
giant new-run --hidden-dim 512 --lr 3e-4 # scaffold a config.toml + run dir ahead of training 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 # train (defaults: stage 1 flow, stage 2 wgan + autoregressive)
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline) giant train path/to/steps.parquet --mode flow # set both stages' generative objective at once
giant train path/to/steps.parquet --mode wgan # train (WGAN-GP, single-pass eval; implemented, not yet tested) 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 (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap) 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 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 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 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 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, dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
# bump-schema, status, update-manifest, create-manifest, # bump-schema, status, update-manifest, create-manifest,
# make-root, build-geometry-oracle, warm-cache, hparam-scan # 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. `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 ### Lint and type checking
```bash ```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. 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 ## Compute environment
Work on this repo happens across three kinds of machine: 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 ## 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. **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`): `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:
- **`"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 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.
- **`"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. - **`"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. `conditioning.share_stages` decides whether the two stages get one shared encoder instance or two identically-configured independent ones.
- **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.
`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 ## 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 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: **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.
- **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 ~6065% 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`.
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 N1 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.
+47 -12
View File
@@ -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) 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 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 train path/to/steps.parquet # train (flow + wgan by default)
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt 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). 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 ## Data
@@ -64,12 +65,24 @@ giant/
│ ├── data/ │ ├── data/
│ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads) │ │ ├── loader.py # parquet → numpy arrays (incl. streaming/chunked reads)
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode │ │ ├── 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/ │ ├── 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 │ │ ├── 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 │ ├── 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 │ ├── 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) │ ├── materials.py # material name → (Z_eff, A_eff, density, X0, λ_int)
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect │ ├── 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 │ │ ├── trainers.py # StageSpec + flow/ddpm and WGAN-GP per-stage trainers
│ │ ├── stage2_inputs.py# ground-truth stage-2 targets + autoregressive/teacher-forcing inputs │ │ ├── stage2_inputs.py# ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
│ │ ├── metrics.py # MetricsCollector: metrics.csv columns, W&B logging, progress/summary │ │ ├── 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) │ │ └── checkpoint.py # checkpoint assembly/restore (format unchanged since v0.2)
│ ├── sample.py # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling │ ├── 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 │ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
│ ├── rollout.py # autoregressive shower rollout driver │ ├── rollout.py # autoregressive shower rollout driver
│ ├── validate.py # step-level marginal + KL-divergence validation │ ├── 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) │ ├── analysis/ # rollout-vs-reference analysis pipeline (see `giant analyze` below)
│ │ ├── sources.py # canonical LazyFrames + secondary view │ │ ├── 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, ...) │ │ ├── reduce.py # streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
│ │ ├── grouping.py # fixed bin edges + energy/pdg/material group sets │ │ ├── grouping.py # fixed bin edges + energy/pdg/material group sets
│ │ ├── context.py # resolves grouping into `shared.json` once per run │ │ ├── context.py # resolves grouping into `shared.json` once per run
│ │ ├── catalog.py # declarative PlotSpec registry │ │ ├── reduced.py # Partial/Reduced — the compact JSON a compute job emits
│ │ ├── condor.py # prep / compute-one / submit-description plumbing │ │ ├── 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) │ │ └── 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`) ├── giant/tools/ # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status, │ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
│ │ # update-manifest, create-manifest, make-root, │ │ # 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 # 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 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 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. `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 ## Training, prediction, rollout
@@ -137,11 +163,13 @@ Useful flags on `giant train`:
- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture) - `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
- `--conditioning {physical,embedding,onehot}` — conditioning representation - `--conditioning {physical,embedding,onehot}` — conditioning representation
- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing - `--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 - `--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 - `--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 - `--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. `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): - `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
```bash ```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 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 ## Development
+157 -92
View File
@@ -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 ``finalize`` merges each rollout's chunks independently and assembles a
``Reduced.payload`` keyed the same way: ``"series": {name: ...}`` for the ``Reduced.payload`` keyed the same way: ``"series": {name: ...}`` for the
rollouts, ``"reference": ...`` as one distinguished entry (omitted on rollouts, ``"reference": ...`` as one distinguished entry (omitted on
rollout-only plots like ``leakage_fraction``). The two heatmap-shaped specs rollout-only plots like ``leakage_fraction``). The heatmap-shaped specs
(``marginal_distance_summary``, ``n_sec_confusion``) and the router (``marginal_distance_summary``, ``sec_count_per_step_by_species``) and the
diagnostics are inherently one-matrix/one-checkpoint per rollout, so their router diagnostics are inherently one-matrix/one-checkpoint per rollout, so their
``"series"`` entries are whole per-rollout artifacts (a matrix, a gating ``"series"`` entries are whole per-rollout artifacts (a matrix, a gating
dict) rather than a single number/array ``render.py`` draws those as one dict) rather than a single number/array ``render.py`` draws those as one
panel per rollout instead of one line/bar per rollout. panel per rollout instead of one line/bar per rollout.
@@ -60,7 +60,6 @@ from giant.analysis.reduce import (
leakage_fraction, leakage_fraction,
profile_finalize, profile_finalize,
profile_partial, profile_partial,
sec_count_by_event,
species_share, species_share,
sum_merge, sum_merge,
transverse_expr, transverse_expr,
@@ -72,7 +71,15 @@ from giant.analysis.router_gating import (
compute_router_share_by_process, compute_router_share_by_process,
compute_router_specialization, 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.type_embedding_distance import compute_type_embedding_l1_distance
from giant.analysis.variables import RANGED_VARS, cos_scatter_expr 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))) 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: def _containment_depths(mat: np.ndarray, edges: np.ndarray, quantile: float) -> np.ndarray:
"""Per-event depth containing ``quantile`` of that event's deposited energy. """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: def _sec_energy_partial(b: Bundle) -> dict:
edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1) edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1)
return { 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) # router diagnostics (not chunked — already bounded/subsampled)
# --------------------------------------------------------------------------- # ---------------------------------------------------------------------------
@@ -1077,6 +1136,18 @@ def build_catalog() -> list[PlotSpec]:
compute_partial=_sec_count_per_species_partial, compute_partial=_sec_count_per_species_partial,
finalize=_sec_count_per_species_finalize, 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( PlotSpec(
"sec_energy", "sec_energy",
"secondaries", "secondaries",
@@ -1089,12 +1160,6 @@ def build_catalog() -> list[PlotSpec]:
compute_partial=_sec_cos_angle_partial, compute_partial=_sec_cos_angle_partial,
finalize=_sec_cos_angle_finalize, finalize=_sec_cos_angle_finalize,
), ),
PlotSpec(
"n_sec_confusion",
"secondaries",
compute_partial=_n_sec_confusion_partial,
finalize=_n_sec_confusion_finalize,
),
PlotSpec( PlotSpec(
"router_gating", "router_gating",
"model", "model",
-17
View File
@@ -271,20 +271,3 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
escaped = per_event["escaped"].fill_null(0.0).to_numpy() escaped = per_event["escaped"].fill_null(0.0).to_numpy()
total = deposited + escaped total = deposited + escaped
return np.where(total > 0, escaped / total, 0.0) 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
+1 -1
View File
@@ -27,7 +27,7 @@ from pathlib import Path
# "router_specialization" max gate weight vs energy (one scalar trend line # "router_specialization" max gate weight vs energy (one scalar trend line
# summarizing "router_gating"), per rollout with an enabled router # summarizing "router_gating"), per rollout with an enabled router
# "heatmap" row x col matrix + colorbar, one panel per rollout (a # "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) # "unavailable" plot not applicable to this run (e.g. no MoE checkpoint)
+11 -3
View File
@@ -27,6 +27,7 @@ from pathlib import Path
import numpy as np import numpy as np
import plotstyle as ps import plotstyle as ps
from matplotlib.colors import LogNorm
import yaml import yaml
from giant.analysis.reduced import Reduced 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): def _render_heatmap(r: Reduced, params: dict):
series = r.payload["series"] series = dict(r.payload["series"])
row_labels = r.payload["row_labels"] row_labels = r.payload["row_labels"]
col_labels = r.payload["col_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) names = list(series)
norm = LogNorm(vmin=1) if r.payload.get("log_color") else None
fig, axes = ps.new_figure( fig, axes = ps.new_figure(
"slide-16x9" if len(names) > 1 else "thesis-single", "slide-16x9" if len(names) > 1 else "thesis-single",
title=r.title, title=r.title,
@@ -397,8 +404,9 @@ def _render_heatmap(r: Reduced, params: dict):
origin="upper", origin="upper",
aspect="auto", aspect="auto",
cmap=r.payload.get("cmap", "viridis"), cmap=r.payload.get("cmap", "viridis"),
vmin=r.payload.get("vmin"), norm=norm,
vmax=r.payload.get("vmax"), 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_xticks(range(len(col_labels)))
ax.set_xticklabels(col_labels, rotation=45, ha="right") ax.set_xticklabels(col_labels, rotation=45, ha="right")
+31
View File
@@ -236,3 +236,34 @@ def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
pl.col("sec_dz_list").alias("sdz"), 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"))
)
+1 -1
View File
@@ -1,6 +1,6 @@
[project] [project]
name = "giant" name = "giant"
version = "0.3.9" version = "0.3.11"
description = "Geant4 step-function surrogate via conditional flow matching" description = "Geant4 step-function surrogate via conditional flow matching"
readme = "README.md" readme = "README.md"
requires-python = ">=3.12" requires-python = ">=3.12"
+11 -11
View File
@@ -13,6 +13,7 @@ from giant.analysis.sources import (
open_side, open_side,
physical_steps, physical_steps,
secondaries, secondaries,
secondaries_by_step,
) )
from giant.data.loader import EVENT_ID_FILE_STRIDE 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] assert t["pdg"].to_list() == [22, 22]
def test_sec_count_by_event_zero_fills_events_with_no_secondaries(): def test_secondaries_by_step_keys_each_secondary_to_its_emitting_step():
r_phys = physical_steps(_rollout_frame(), Side.rollout) r = secondaries_by_step(_rollout_frame(), Side.rollout).collect()
r_sec = secondaries(_rollout_frame(), Side.rollout) assert r["pdg"].to_list() == [22]
ev, n = R.sec_count_by_event(r_phys, r_sec) # the rollout key is (event_id, parent_id, birth position) — the parent
# event 1 has one secondary track; event 2 has none and must still appear (as 0), # step's post_pos, copied verbatim onto the child's birth row.
# not silently drop out of a plain group_by on the secondaries frame alone. assert r["step_key"][0] == {"event_id": 1, "parent_id": 0, "pre_x": 0.0, "pre_y": 0.0, "pre_z": 1.0}
assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 0}
t_all = _reference_frame() t = secondaries_by_step(_reference_frame(), Side.reference).collect()
t_sec = secondaries(t_all, Side.reference) assert t["pdg"].to_list() == [22, 22]
ev, n = R.sec_count_by_event(t_all, t_sec) # one row per emitting step; the empty-list step drops out entirely
assert dict(zip(ev.tolist(), n.tolist())) == {1: 1, 2: 1} assert [k["_row"] for k in t["step_key"]] == [0, 2]
def test_leakage_fraction(): def test_leakage_fraction():
+26 -38
View File
@@ -10,10 +10,10 @@ from giant.analysis.catalog import (
Bundle, Bundle,
PlotSpec, PlotSpec,
_containment_depths, _containment_depths,
_integer_confusion,
_ks_statistic, _ks_statistic,
) )
from giant.analysis.context import Context, build_context from giant.analysis.context import Context, build_context
from giant.analysis.grouping import pdg_label
from giant.analysis.sources import RolloutSpec from giant.analysis.sources import RolloutSpec
from tests.test_analysis_reduce import _reference_frame, _rollout_frame 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_ # data-dependent edges (event_total_edep), concat-then-mean/std (shower_
# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a # longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
# ratio (species_edep_share), a chunkable=False passthrough (router_gating), # ratio (species_edep_share), a chunkable=False passthrough (router_gating),
# nested sum-merge into a scorecard (marginal_distance_summary), concat-then- # sum-mergeable-with-a-zero-fill-denominator (sec_count_per_step{,_by_species}),
# event-id-join (n_sec_confusion), and concat-then-per-event-derived-quantity # 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). # (shower_containment_depth_90, reusing the profile matrix's own merge shape).
_CHUNK_EQUIVALENCE_IDS = [ _CHUNK_EQUIVALENCE_IDS = [
"marginal_edep", "marginal_edep",
@@ -170,9 +171,10 @@ _CHUNK_EQUIVALENCE_IDS = [
"shower_longitudinal", "shower_longitudinal",
"leakage_fraction", "leakage_fraction",
"sec_count_per_species", "sec_count_per_species",
"sec_count_per_step",
"sec_count_per_step_by_species",
"router_gating", "router_gating",
"marginal_distance_summary", "marginal_distance_summary",
"n_sec_confusion",
"shower_containment_depth_90", "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 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(): def test_containment_depths_simple_ramp():
# one event, edep concentrated in the first bin -> 90%/95% containment # 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. # 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] assert depths.tolist() == [1.0]
def test_n_sec_confusion_spec(bundle): def test_sec_count_per_step_counts_empty_steps(bundle):
spec = get_spec("n_sec_confusion") spec = get_spec("sec_count_per_step")
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx) r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
assert r.payload["row_labels"] == r.payload["col_labels"] == ["0", "1+"] # reference: 3 steps, two of which emit exactly one secondary
assert r.payload["series"]["rollout"] == [[0, 0], [1, 1]] 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): def test_sec_count_per_step_by_species_zero_row_is_per_species(bundle):
spec = get_spec("n_sec_confusion") spec = get_spec("sec_count_per_step_by_species")
r = spec.finalize([spec.compute_partial(two_bundle)], two_bundle.ctx) r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
assert list(r.payload["series"]) == ["flow", "wgan"] cols = r.payload["col_labels"]
# both rollouts share the same fixture data here, so their matrices (and ref = r.payload["reference"]
# the shared label set) must be identical. g = cols.index(pdg_label(22))
assert r.payload["series"]["flow"] == r.payload["series"]["wgan"] # 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
+3 -1
View File
@@ -304,13 +304,15 @@ def test_render_one_of_each_kind(tmp_path: Path):
"hm1", "hm1",
"secondaries", "secondaries",
"heatmap", "heatmap",
"Confusion (single rollout)", "Heatmap (single rollout)",
"predicted", "predicted",
{ {
"series": {"flow": [[1, 0], [0, 1]]}, "series": {"flow": [[1, 0], [0, 1]]},
"reference": [[2, 0], [0, 1]],
"row_labels": ["0", "1+"], "row_labels": ["0", "1+"],
"col_labels": ["0", "1+"], "col_labels": ["0", "1+"],
"cbar_label": "count", "cbar_label": "count",
"log_color": True,
}, },
), ),
] ]
Generated
+1 -1
View File
@@ -675,7 +675,7 @@ wheels = [
[[package]] [[package]]
name = "giant" name = "giant"
version = "0.3.9" version = "0.3.11"
source = { editable = "." } source = { editable = "." }
dependencies = [ dependencies = [
{ name = "numpy" }, { name = "numpy" },