Compare commits
9 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| bb8d16caba | |||
| c8a1b4f25d | |||
| 23efd6d9ff | |||
| 9fa6420183 | |||
| b66574877b | |||
| 9b77e04731 | |||
| 8dee2feab7 | |||
| f2da0642b2 | |||
| 1e92902c8d |
+1
-1
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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 ~60–65% gate weight at the highest energies plotted — so eval-time top-1 (Voronoi) dispatch is choosing among near-ties rather than real specialists. Two contributors were identified: the missing load-balancing loss (`lambda_balance=0.0`), and `EnergyRouter`'s center init (`torch.linspace(-2, 2, n_experts)`) assuming a roughly uniform z-normalized energy distribution, which real energy spectra don't match. **Fixed (2026-07-27):** `EnergyRouter` now accepts an optional `centers_init` (backward compatible — omitting it keeps the old linspace), and `giant train` auto-populates it from real data quantiles via a reservoir sample collected during the existing normalizer-fitting pass in `giant/pipeline.py` (no extra file scan), for `--router-type energy` only. The routing *strategy* itself may still be sound, but the specific benchmarked config wasn't. **Next step before further evaluation: retrain with `lambda_balance > 0` and the new quantile-seeded centers (and consider more epochs / a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Full writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
|
|
||||||
|
|
||||||
A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
|
v0.2 configs and checkpoints are auto-migrated (`config.migrate_config`, `model._legacy._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`). **v0.2 checkpoint-loading support has no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
|
||||||
|
|
||||||
**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`) — v0.2 checkpoint-loading support has **no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
|
**Faster-eval architectures — both implemented, neither validated.** Target is a ~10× native-Geant4 eval budget; no eval-latency number exists for any configuration yet, so that budget is unverified across the board.
|
||||||
|
- **WGAN-GP** (`--stage2-generator wgan`, now the stage-2 default): first rollout benchmark 2026-08-03 failed with secondary-species mode collapse — the failure v0.3.0 was designed to address. **No post-v0.3.0 benchmark has been run.** Writeup: `/home/lars/knowledge-base/experiments/giant-wgan-physical-rollout-validation.md`.
|
||||||
|
- **MoE routing trunk** (`--router`): first rollout benchmark 2026-07-22 diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. Cause identified as a bad config, not a bad idea: `lambda_balance=0.0` (no load-balancing loss) plus `EnergyRouter`'s `torch.linspace(-2, 2, n_experts)` center init assuming a roughly uniform z-normalized energy distribution — so the ten experts overlapped across ~5 decades of energy instead of partitioning it, and eval-time top-1 dispatch chose among near-ties rather than real specialists. Both prerequisites are fixed in code (quantile-seeded `centers_init` from `pipeline.py`, `lambda_balance` exposed). **Next step: retrain with `lambda_balance > 0` and quantile-seeded centers (consider a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
|
||||||
|
|
||||||
**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.
|
A sampling-calorimeter (multi-material) dataset track is still open and unblocked, not yet started. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
|
||||||
|
|
||||||
|
**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2. Partway between "needs major features" and feature-complete — not ready to merge yet.
|
||||||
|
|||||||
@@ -10,6 +10,7 @@ A conditional generative model that replaces the Geant4 step function: given a p
|
|||||||
uv sync --extra cpu # install deps (CPU torch; use --extra cuda for GPU)
|
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
@@ -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",
|
||||||
|
|||||||
@@ -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
|
|
||||||
|
|||||||
@@ -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)
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
@@ -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")
|
||||||
|
|||||||
@@ -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
@@ -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"
|
||||||
|
|||||||
@@ -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
@@ -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
|
||||||
|
|||||||
@@ -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,
|
||||||
},
|
},
|
||||||
),
|
),
|
||||||
]
|
]
|
||||||
|
|||||||
Reference in New Issue
Block a user