Merge remote-tracking branch 'origin/master' into condor-gpu-train-rollout
CI / Lint (ruff check) (push) Successful in 26s
CI / Format (ruff format) (push) Successful in 26s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 21s
CI / Tests (push) Successful in 53s
CI / Lint (ruff check) (pull_request) Successful in 25s
CI / Format (ruff format) (pull_request) Successful in 25s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 21s
CI / Tests (pull_request) Successful in 53s
CI / Lint (ruff check) (push) Successful in 26s
CI / Format (ruff format) (push) Successful in 26s
CI / Sync project version with tag (push) Has been skipped
CI / Type check (ty) (push) Successful in 21s
CI / Tests (push) Successful in 53s
CI / Lint (ruff check) (pull_request) Successful in 25s
CI / Format (ruff format) (pull_request) Successful in 25s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 21s
CI / Tests (pull_request) Successful in 53s
# Conflicts: # CLAUDE.md
This commit is contained in:
+57
-22
@@ -3,76 +3,111 @@ name: CI
|
||||
"on":
|
||||
push:
|
||||
branches: ["**"]
|
||||
tags: ["**"]
|
||||
pull_request:
|
||||
branches: [master]
|
||||
|
||||
env:
|
||||
UV_CACHE_DIR: /uv-cache
|
||||
|
||||
jobs:
|
||||
ruff-check:
|
||||
name: Lint (ruff check)
|
||||
runs-on: ubuntu-latest
|
||||
container:
|
||||
image: docker.gitea.com/runner-images:ubuntu-latest
|
||||
volumes:
|
||||
- /srv/act-runner-cache/uv:/uv-cache
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
enable-cache: false
|
||||
- run: |
|
||||
echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
|
||||
echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
|
||||
- run: uv sync --extra cpu --extra dev
|
||||
- run: uv run ruff check .
|
||||
|
||||
ruff-format:
|
||||
name: Format (ruff format)
|
||||
runs-on: ubuntu-latest
|
||||
container:
|
||||
image: docker.gitea.com/runner-images:ubuntu-latest
|
||||
volumes:
|
||||
- /srv/act-runner-cache/uv:/uv-cache
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
enable-cache: false
|
||||
- run: |
|
||||
echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
|
||||
echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
|
||||
- run: uv sync --extra cpu --extra dev
|
||||
- run: uv run ruff format --check .
|
||||
|
||||
type-check:
|
||||
name: Type check (ty)
|
||||
runs-on: ubuntu-latest
|
||||
container:
|
||||
image: docker.gitea.com/runner-images:ubuntu-latest
|
||||
volumes:
|
||||
- /srv/act-runner-cache/uv:/uv-cache
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
enable-cache: false
|
||||
- run: |
|
||||
echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
|
||||
echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
|
||||
- run: uv sync --extra cpu --extra dev
|
||||
- run: uv run ty check .
|
||||
|
||||
test:
|
||||
name: Tests
|
||||
needs: [ruff-check, type-check]
|
||||
runs-on: ubuntu-latest
|
||||
container:
|
||||
image: docker.gitea.com/runner-images:ubuntu-latest
|
||||
volumes:
|
||||
- /srv/act-runner-cache/uv:/uv-cache
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
with:
|
||||
enable-cache: true
|
||||
enable-cache: false
|
||||
- run: |
|
||||
echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
|
||||
echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
|
||||
- run: uv sync --extra cpu --extra dev
|
||||
- run: uv run pytest
|
||||
|
||||
build:
|
||||
name: Bump version, build & publish wheel
|
||||
needs: [ruff-check, ruff-format, type-check, test]
|
||||
if: github.event_name == 'push' && github.ref == 'refs/heads/master'
|
||||
sync-version-on-tag:
|
||||
name: Sync project version with tag
|
||||
if: startsWith(github.ref, 'refs/tags/')
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v4
|
||||
with:
|
||||
token: ${{ secrets.CI_TOKEN }}
|
||||
- uses: astral-sh/setup-uv@v5
|
||||
- name: Bump patch version
|
||||
- name: Check tag against project version, update if they differ
|
||||
run: |
|
||||
git config user.name "gitea-actions"
|
||||
git config user.email "actions@git.larsbogner.de"
|
||||
uv version --bump patch --no-sync
|
||||
NEW_VERSION=$(uv version --short)
|
||||
git add pyproject.toml uv.lock
|
||||
git commit -m "chore: bump version to ${NEW_VERSION} [skip ci]"
|
||||
git push
|
||||
- run: uv build
|
||||
- name: Publish to Gitea package registry
|
||||
env:
|
||||
TWINE_USERNAME: ${{ secrets.PACKAGE_USERNAME }}
|
||||
TWINE_PASSWORD: ${{ secrets.CI_TOKEN }}
|
||||
run: uvx twine upload --repository-url https://git.larsbogner.de/api/packages/lars/pypi dist/*
|
||||
TAG_VERSION="${GITHUB_REF_NAME#v}"
|
||||
CURRENT_VERSION=$(uv version --short)
|
||||
if [ "$TAG_VERSION" != "$CURRENT_VERSION" ]; then
|
||||
echo "Tag version ($TAG_VERSION) != project version ($CURRENT_VERSION); updating pyproject.toml"
|
||||
uv version "$TAG_VERSION" --no-sync
|
||||
git config user.name "gitea-actions"
|
||||
git config user.email "actions@git.larsbogner.de"
|
||||
git add pyproject.toml uv.lock
|
||||
git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME} [skip ci]"
|
||||
git push origin HEAD:master
|
||||
git push origin ":refs/tags/${GITHUB_REF_NAME}"
|
||||
git tag -f "${GITHUB_REF_NAME}" HEAD
|
||||
git push origin "refs/tags/${GITHUB_REF_NAME}"
|
||||
else
|
||||
echo "Tag version matches project version ($CURRENT_VERSION)"
|
||||
fi
|
||||
|
||||
@@ -17,3 +17,6 @@ checkpoints/
|
||||
|
||||
# Scratch working directory
|
||||
/scratchpad/
|
||||
|
||||
# giant analyze run directories (shared.json, reduced/, plots/, condor logs)
|
||||
/analysis_runs/
|
||||
|
||||
@@ -13,6 +13,8 @@ pytest # run tests
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir for a new training run
|
||||
giant train path/to/steps.parquet --mode flow # train (flow matching)
|
||||
giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline)
|
||||
giant train path/to/steps.parquet --mode wgan # train (WGAN-GP, single-pass eval; implemented, not yet tested)
|
||||
giant train path/to/steps.parquet --router --router-type energy # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap)
|
||||
giant train-submit path/to/steps.parquet --config run/config.toml --accounting-group cms # train as a remote-GPU HTCondor job (TOpAS/NEMO2)
|
||||
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
|
||||
@@ -37,6 +39,14 @@ uv run ty check . # type check
|
||||
|
||||
Part of the `dev` extra. Run these periodically (not just at commit time) to catch drift early.
|
||||
|
||||
## Compute environment
|
||||
|
||||
Work on this repo happens across three kinds of machine:
|
||||
|
||||
- **Local dev machines** (laptop + desktop, identical): repo at `~/Programming/giant`, no access to `/ceph` — datasets, training results, and models aren't reachable here.
|
||||
- **Portal machines** (`portal1`, `deepthought`, `deepthought2`, `bms1`, `bms2`, `bms3`): repo lives under `/work`, and `/ceph` holds ROOT/parquet files and trained models. **These are shared with other users** — stay strictly within `/work/lbogner` and `/ceph/lbogner`, and keep resource usage to roughly a quarter of CPU/RAM and a single GPU so as not to disturb other users' jobs.
|
||||
- **HTCondor worker nodes**: never run or SSH onto these directly — the only sanctioned path is submitting jobs through condor (`giant analyze submit`, and the in-progress remote-GPU train/rollout submission on `condor-gpu-train-rollout`). `/ceph` is available there; `/work` is only sometimes mounted, depending on the node.
|
||||
|
||||
## 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).
|
||||
@@ -59,9 +69,13 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
|
||||
|
||||
**Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode.
|
||||
|
||||
**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.
|
||||
|
||||
**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.
|
||||
|
||||
**Validation** (`giant/validate.py`): step-level marginal comparisons.
|
||||
|
||||
**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. 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 + 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`, 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), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_<id>/`) holding `shared.json`, `run_meta.json`, `reduced/`, `plots/`. **Compute/render split:** `giant analyze submit rollout.yaml` runs `prep` then submits one HTCondor job per plot (`compute-one --run-dir`, polars/numpy only — no LaTeX on workers), each writing a small `reduced/<id>.json`; the local `giant analyze render <run_dir>` turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`.
|
||||
**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. 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 + 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`, 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), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_<id>/`) holding `shared.json`, `run_meta.json`, `reduced_partial/`, `reduced/`, `plots/`. **Compute/merge/render split:** `giant analyze submit rollout.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`) 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 (`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`.
|
||||
|
||||
**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.
|
||||
|
||||
@@ -75,4 +89,10 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
|
||||
|
||||
**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.
|
||||
|
||||
**Next directions** (parallel, not yet built): faster-eval architectures measured against a ~10× native-Geant4 budget — a Wasserstein-GAN throwaway (single-pass eval) and a mixture-of-experts / routing tree of small nets selected per call (pdg / energy / process), with soft/differentiable gating on continuous routing axes; a sampling-calorimeter (multi-material) dataset. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
|
||||
**Faster-eval architectures (implemented, validation in progress):** both tracks below target a ~10× native-Geant4 eval budget and are now wired into `giant train`/`giant/model/network.py`, but neither has a validated result yet — treat both as unproven until the corresponding analysis run says otherwise:
|
||||
- **WGAN-GP** (`--mode wgan`, see Architecture above): implemented, **not yet tested** — no rollout-vs-reference analysis run against it yet.
|
||||
- **MoE routing trunk** (`--router`, see Architecture above): implemented, **first rollout benchmark done (2026-07-22), result: needs retraining with a different router config, not abandoned.** A 10-expert `EnergyRouter` run (`n_experts=10`, `temperature=0.5`, `learn_centers=true`, **`lambda_balance=0.0`**, only 20 fine-tuning epochs resumed from a non-routed checkpoint) diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. The `router_gating` diagnostic plot points at the likely cause: the ten experts overlap heavily across ~5 decades of pre-step energy instead of partitioning it — even the top-energy expert only reaches ~60–65% gate weight at the highest energies plotted — so eval-time top-1 (Voronoi) dispatch is choosing among near-ties rather than real specialists. Two contributors were identified: the missing load-balancing loss (`lambda_balance=0.0`), and `EnergyRouter`'s center init (`torch.linspace(-2, 2, n_experts)`) assuming a roughly uniform z-normalized energy distribution, which real energy spectra don't match. **Fixed (2026-07-27):** `EnergyRouter` now accepts an optional `centers_init` (backward compatible — omitting it keeps the old linspace), and `giant train` auto-populates it from real data quantiles via a reservoir sample collected during the existing normalizer-fitting pass in `giant/pipeline.py` (no extra file scan), for `--router-type energy` only. The routing *strategy* itself may still be sound, but the specific benchmarked config wasn't. **Next step before further evaluation: retrain with `lambda_balance > 0` and the new quantile-seeded centers (and consider more epochs / a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Full writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`.
|
||||
|
||||
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`).
|
||||
|
||||
**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.
|
||||
|
||||
@@ -1,172 +0,0 @@
|
||||
# Phase 2: Secondary Particle Prediction
|
||||
|
||||
## Context
|
||||
|
||||
Phase 1 takes `n_sec` (secondary count) and `e_sec` (total secondary energy) as **conditioning inputs**. Phase 2 must instead **predict** them, making the surrogate self-contained for shower rollout. Per Jan's 2026-06-29 decision: hard discrete `n_sec` integer head; escalation to Gumbel-Softmax only if empirically needed.
|
||||
|
||||
Two-stage factorization:
|
||||
- **Stage 1**: existing 9D flow model (reduced conditioning: drop `n_sec` + `log(e_sec)`) + a new discrete `n_sec` classification head
|
||||
- **Stage 2**: non-AR flow matching over `K_MAX` secondary slots simultaneously, each slot predicting `(stick_break_logit, dir_local_3D, type_emb)` — conditioned on pre-step state + Stage 1 output; padded slots masked from loss
|
||||
|
||||
Training: joint, combined loss `L = L_flow_s1 + λ_nsec * L_nsec + λ_s2 * L_flow_s2`.
|
||||
|
||||
---
|
||||
|
||||
## Prerequisite: Determine K_MAX
|
||||
|
||||
Before implementing, run a quick analysis over existing parquet files to find `max(n_sec)` and the 99th percentile. Expected to be 5–20 for EM shower steps. Set `K_MAX` as a constant in `giant/constants.py` (suggest 15 as a starting point, revise from data).
|
||||
|
||||
---
|
||||
|
||||
## New Branch
|
||||
|
||||
```bash
|
||||
git checkout -b phase2-secondary-prediction master
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Part A — Data Pipeline
|
||||
|
||||
### A1. `scripts/steps_to_parquet.py`
|
||||
|
||||
Extend `_add_secondary_energy` to also collect per-secondary attributes from the spawning tree join:
|
||||
- For each `child_track_id`, look up the child's first step → get `pdg`, `pre_E`, `pre_dx/dy/dz`
|
||||
- Emit list columns in the parquet: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`, `sec_dy_list`, `sec_dz_list`
|
||||
- Lists are sorted **descending by energy** at write time
|
||||
- Truncate to `K_MAX` entries if needed (flag if any row truncated)
|
||||
|
||||
Re-run ROOT→parquet conversion after this change.
|
||||
|
||||
### A2. `giant/data/loader.py`
|
||||
|
||||
In `_df_to_dict`: read the five new list columns. Pad each to length `K_MAX` with zeros (energy) / sentinel values (pdg → 0, dir → (0,0,1)). Return as fixed-shape arrays `(N, K_MAX)` / `(N, K_MAX, 3)`.
|
||||
|
||||
Also return a boolean validity mask `sec_valid` of shape `(N, K_MAX)`: `True` for slots `i < n_sec`.
|
||||
|
||||
### A3. `giant/data/transforms.py`
|
||||
|
||||
Add `encode_secondaries(sec_pdg_list, sec_E_list, sec_dir_list, sec_valid, e_sec, pdg_emb_weight, pre_dir, K_MAX)`:
|
||||
1. **Direction**: call existing `local_frame_rotation` per slot
|
||||
2. **Energy (stick-breaking)**:
|
||||
- Slot 0: `f_0 = E_0 / e_sec` → logit `log(f_0/(1-f_0))` (clamped)
|
||||
- Slot i: `f_i = E_i / (e_sec - sum(E_0..E_{i-1}))` → logit
|
||||
- Last valid slot: logit = large positive constant (takes all remaining budget)
|
||||
- Padding slots (beyond `n_sec`): set logit = 0, masked out of loss anyway
|
||||
3. **Type embedding**: index into `pdg_emb_weight` (the PDG embedding table weights) to get the target embedding vector for each secondary's `pdg`. Shape `(K_MAX, emb_dim)`.
|
||||
|
||||
Returns `sec_targets: (K_MAX, 1 + 3 + emb_dim)` and `sec_valid: (K_MAX,)`.
|
||||
|
||||
Inverse (`decode_secondaries`): sigmoid stick-breaking fractions → energies, inv local frame rotation → world dirs, nearest-neighbor lookup in PDG embedding table → pdg code.
|
||||
|
||||
### A4. `giant/data/dataset.py`
|
||||
|
||||
Update `build_features` and `StreamingStepsDataset.__iter__` to also yield `sec_targets` and `sec_valid` alongside the existing `(cond_cont, cond_cat, x1)` batch items.
|
||||
|
||||
---
|
||||
|
||||
## Part B — Constants (`giant/constants.py`)
|
||||
|
||||
- `COND_DIM`: 10 → **8** (remove `n_sec` and `log(e_sec)`)
|
||||
- Add `K_MAX: int` (set after data analysis, e.g. 15)
|
||||
- Add `SEC_SLOT_DIM: int` (= 4 + `emb_dim` = 20 for default emb_dim=16; 1 stick + 3 dir + 16 type)
|
||||
- Add `SEC_DIM: int = K_MAX * SEC_SLOT_DIM` (flattened Stage 2 target dimension)
|
||||
- Update `LOCAL_TARGET_NAMES` (Stage 1 only, still 9D)
|
||||
|
||||
---
|
||||
|
||||
## Part C — Model (`giant/model/network.py`)
|
||||
|
||||
### C1. `DenoisingMLP` — Stage 1 (minimal changes)
|
||||
|
||||
- `ConditionEncoder.cont_dim` drops from 10 to 8 (COND_DIM change propagates automatically)
|
||||
- Add `n_sec_head = nn.Sequential(Linear(cond_out_dim, hidden_dim//2), SiLU(), Linear(hidden_dim//2, K_MAX + 1))` applied to `c_emb` (the condition encoding, not the diffused latent)
|
||||
- Add method `predict_n_sec(cond_cont, cond_cat) -> Tensor[B, K_MAX+1]` — no diffusion, just encode conditioning and run the head
|
||||
|
||||
### C2. `SecondaryDecoder` — Stage 2 (new class)
|
||||
|
||||
Architecture mirrors `DenoisingMLP` but:
|
||||
- **Input**: `x_t` of shape `(B, SEC_DIM)` (flattened K_MAX secondary slots)
|
||||
- **Conditioning**: pre-step state (8D cont + 2 cat → same ConditionEncoder as Stage 1) concatenated with Stage 1 output (9D normalized target, detached from Stage 1 loss for stability initially). Total cond dim to the ResBlocks: `time_dim + cond_s1_out_dim + 9`
|
||||
- **Output**: vector field of shape `(B, SEC_DIM)`
|
||||
- Uses same `ResBlock` / `SinusoidalEmbedding` / `ConditionEncoder` building blocks
|
||||
|
||||
A `SecondaryConditionEncoder` wraps the base `ConditionEncoder` and concatenates the Stage 1 output:
|
||||
```python
|
||||
class SecondaryConditionEncoder(nn.Module):
|
||||
# base: ConditionEncoder(pdg_vocab, mat_vocab, 8, emb_dim, cond_out_dim)
|
||||
# stage1_proj: Linear(X_DIM, stage1_cond_dim)
|
||||
# mlp: fuses both
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Part D — Loss / Training
|
||||
|
||||
### `giant/model/schedule.py`
|
||||
|
||||
Add `flow_matching_loss_masked(model, x1, cond_cont, cond_cat, mask)`:
|
||||
- Same as `flow_matching_loss` but divides by `mask.sum()` instead of `B * SEC_DIM`, zeroing out padded slots before averaging. `mask` shape: `(B, K_MAX)`, broadcast over slot dims.
|
||||
|
||||
### `giant/train.py`
|
||||
|
||||
Batch now unpacks as `(cond_cont, cond_cat, x1_s1, n_sec_target, x1_s2, sec_mask)`.
|
||||
|
||||
Combined loss per batch:
|
||||
```
|
||||
L_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat)
|
||||
L_nsec = cross_entropy(stage1_model.predict_n_sec(cond_cont, cond_cat), n_sec_target)
|
||||
L_s2 = flow_matching_loss_masked(sec_decoder, x1_s2, cond_cont, cond_cat, stage1_detached, sec_mask)
|
||||
L = L_s1 + lambda_nsec * L_nsec + lambda_s2 * L_s2
|
||||
```
|
||||
|
||||
Config adds `lambda_nsec` (suggest 0.1) and `lambda_s2` (suggest 1.0) under `[train]`.
|
||||
|
||||
Both `stage1_model` and `sec_decoder` share a single `optimizer` (AdamW over all parameters).
|
||||
|
||||
Checkpoint saves both `stage1_model.state_dict()` and `sec_decoder.state_dict()`, plus `K_MAX` and `SEC_SLOT_DIM` in `model_config`.
|
||||
|
||||
### `giant/pipeline.py`
|
||||
|
||||
- Compute `K_MAX` from data (max `n_sec` over training events) before constructing models
|
||||
- Build both `DenoisingMLP` and `SecondaryDecoder`, pass both to `run_training`
|
||||
|
||||
---
|
||||
|
||||
## Part E — Sampling (`giant/sample.py`)
|
||||
|
||||
```python
|
||||
def sample_stage1(model, cond_cont, cond_cat, steps=10):
|
||||
# Euler ODE → primary sample (9D), + argmax n_sec head
|
||||
...
|
||||
|
||||
def sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec, steps=10):
|
||||
# Euler ODE on SEC_DIM → decode stick-breaking → energies
|
||||
# inv_local_frame_rotation → world-frame dirs
|
||||
# nearest-neighbor in pdg_emb_weight → pdg codes
|
||||
# mask slots >= n_sec
|
||||
...
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
## Part F — Wiring
|
||||
|
||||
- **`giant/validate.py`**: add secondary-specific marginals (n_sec distribution, species distribution, energy fraction per slot)
|
||||
- **`giant/cli.py`**: `predict` command loads both checkpoints, calls both samplers, appends secondary columns to output parquet
|
||||
|
||||
---
|
||||
|
||||
## Type embedding design note
|
||||
|
||||
The type embedding target at training is `pdg_emb.weight[sec_pdg_idx]` (the Stage 1 PDG embedding table rows). Gradients flow into the embedding table from both the conditioning path (input PDG) and the secondary type loss — this is intentional; the shared embedding space is the bridge. At inference, snap: `argmin_k ||pred_emb - pdg_emb.weight[k]||`.
|
||||
|
||||
---
|
||||
|
||||
## Verification
|
||||
|
||||
1. `uv run pytest` — existing tests pass (Stage 1 shape/interface unchanged beyond COND_DIM)
|
||||
2. Unit tests for `encode_secondaries` / `decode_secondaries` (round-trip: energies sum to `e_sec`, directions are unit vectors)
|
||||
3. Unit test for `flow_matching_loss_masked`: verify padded slots contribute zero gradient
|
||||
4. Short training run (1–2 epochs): confirm all three loss components decrease
|
||||
5. Sampling smoke test: verify `sum(sec_E) ≈ e_sec` per sample, all directions unit-normed
|
||||
@@ -2,7 +2,8 @@
|
||||
|
||||
Compares one autoregressive ``giant rollout`` against a held-out miniCaloSim
|
||||
reference file, producing publication-styled comparison plots generated in
|
||||
parallel on HTCondor (one job per plot, compute/render split).
|
||||
parallel on HTCondor (one job per plot x data chunk, compute/merge/render
|
||||
split).
|
||||
|
||||
Only ``render`` (and the ``render`` CLI path) imports plotstyle/LaTeX; everything
|
||||
re-exported here is plotstyle-free so it runs on a compute worker. Import
|
||||
@@ -17,11 +18,14 @@ from giant.analysis.condor import (
|
||||
compute_reduced,
|
||||
derive_run_dir,
|
||||
load_rollout_yaml,
|
||||
merge_all,
|
||||
merge_one,
|
||||
prep,
|
||||
write_submit,
|
||||
)
|
||||
from giant.analysis.context import Context, build_context
|
||||
from giant.analysis.reduced import Reduced
|
||||
from giant.analysis.reduced import Partial, Reduced
|
||||
from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
|
||||
from giant.analysis.sources import Side
|
||||
|
||||
__all__ = [
|
||||
@@ -34,10 +38,15 @@ __all__ = [
|
||||
"compute_reduced",
|
||||
"derive_run_dir",
|
||||
"load_rollout_yaml",
|
||||
"merge_all",
|
||||
"merge_one",
|
||||
"prep",
|
||||
"write_submit",
|
||||
"Context",
|
||||
"build_context",
|
||||
"Partial",
|
||||
"Reduced",
|
||||
"Side",
|
||||
"RUNTIME_SAFETY_MARGIN",
|
||||
"estimate_runtime_s",
|
||||
]
|
||||
|
||||
+397
-125
@@ -2,9 +2,22 @@
|
||||
|
||||
Each spec knows its stable ``id`` (used for the reduced-data filename, the PDF
|
||||
stem and the condor queue item), its gallery ``family`` (subdirectory), and a
|
||||
``compute(bundle) -> Reduced`` that runs the streaming reduction. Rendering lives
|
||||
in ``render.py`` and dispatches on ``Reduced.kind`` — the catalog itself never
|
||||
imports plotstyle, so ``compute-one`` jobs stay LaTeX-free.
|
||||
``compute_partial(bundle) -> dict`` / ``finalize(parts, ctx) -> Reduced`` pair
|
||||
that together run the streaming reduction. ``compute_partial`` runs once per
|
||||
``(plot, chunk)`` condor job against a ``Bundle`` whose four LazyFrames are
|
||||
already filtered to that chunk (see ``Bundle.open``'s ``chunk`` argument); it
|
||||
returns a small JSON-safe partial artifact — either a raw sum-mergeable count
|
||||
dict (histograms/species sums against fixed edges) or a raw per-event/
|
||||
per-secondary array to be concatenated (anything that derives its own edges or
|
||||
a mean/std from the full dataset). ``finalize`` merges the per-chunk partials
|
||||
(in chunk order) and does the actual histogramming/edge-selection/mean-std
|
||||
collapse, once, over the merged data — for ``n_chunks=1`` this reproduces
|
||||
exactly what a single unchunked pass would produce. Specs marked
|
||||
``chunkable=False`` (the router ones) always run as a single chunk regardless
|
||||
of the configured chunk count.
|
||||
|
||||
Rendering lives in ``render.py`` and dispatches on ``Reduced.kind`` — the
|
||||
catalog itself never imports plotstyle, so ``compute-one`` jobs stay LaTeX-free.
|
||||
|
||||
The registry is built by expanding parametric families (marginals over
|
||||
variable x grouping, secondaries, ...) into concrete specs.
|
||||
@@ -12,7 +25,7 @@ variable x grouping, secondaries, ...) into concrete specs.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Callable
|
||||
|
||||
import numpy as np
|
||||
@@ -32,11 +45,18 @@ from giant.analysis.reduce import (
|
||||
event_scalars,
|
||||
hist1d,
|
||||
leakage_fraction,
|
||||
profile_finalize,
|
||||
profile_partial,
|
||||
species_share,
|
||||
sum_merge,
|
||||
transverse_expr,
|
||||
weighted_profile,
|
||||
)
|
||||
from giant.analysis.reduced import Reduced
|
||||
from giant.analysis.router_gating import (
|
||||
compute_router_gating,
|
||||
compute_router_share_by_pdg,
|
||||
compute_router_share_by_process,
|
||||
)
|
||||
from giant.analysis.sources import Side, open_side, physical_steps, secondaries
|
||||
from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
|
||||
|
||||
@@ -50,17 +70,39 @@ class Bundle:
|
||||
t_all: pl.LazyFrame # reference, all rows
|
||||
r_phys: pl.LazyFrame # rollout, physical steps only
|
||||
t_phys: pl.LazyFrame # reference, physical steps only
|
||||
checkpoint: str | None = None # from the rollout YAML; router_gating only
|
||||
|
||||
@classmethod
|
||||
def open(cls, rollout, reference, ctx: Context) -> "Bundle":
|
||||
def open(
|
||||
cls,
|
||||
rollout,
|
||||
reference,
|
||||
ctx: Context,
|
||||
checkpoint=None,
|
||||
chunk: tuple[int, int] | None = None,
|
||||
) -> "Bundle":
|
||||
"""Open both sides, optionally restricted to one event-disjoint chunk.
|
||||
|
||||
``chunk = (chunk_index, n_chunks)`` filters both sides to
|
||||
``event_id % n_chunks == chunk_index`` *before* deriving the physical/
|
||||
secondary views, so every downstream reduction (which is either
|
||||
row-local or a ``group_by("event_id")``) sees a self-contained,
|
||||
event-disjoint slice — no cross-chunk lookups are ever needed.
|
||||
"""
|
||||
r_all = open_side(rollout, Side.rollout)
|
||||
t_all = open_side(reference, Side.reference)
|
||||
if chunk is not None:
|
||||
idx, n = chunk
|
||||
pred = pl.col("event_id") % n == idx
|
||||
r_all = r_all.filter(pred)
|
||||
t_all = t_all.filter(pred)
|
||||
return cls(
|
||||
ctx=ctx,
|
||||
r_all=r_all,
|
||||
t_all=t_all,
|
||||
r_phys=physical_steps(r_all, Side.rollout),
|
||||
t_phys=physical_steps(t_all, Side.reference),
|
||||
checkpoint=checkpoint,
|
||||
)
|
||||
|
||||
|
||||
@@ -68,7 +110,26 @@ class Bundle:
|
||||
class PlotSpec:
|
||||
id: str
|
||||
family: str
|
||||
compute: Callable[[Bundle], Reduced]
|
||||
compute_partial: Callable[[Bundle], dict]
|
||||
finalize: Callable[[list[dict], Context], Reduced]
|
||||
chunkable: bool = True
|
||||
|
||||
|
||||
def _unchunkable(
|
||||
compute: Callable[[Bundle], Reduced],
|
||||
) -> tuple[Callable[[Bundle], dict], Callable[[list[dict], Context], Reduced]]:
|
||||
"""Wrap a whole-dataset ``compute(bundle) -> Reduced`` as a trivial
|
||||
``(compute_partial, finalize)`` pair, for specs marked ``chunkable=False``
|
||||
(which always run as a single chunk, so ``parts`` is always one element).
|
||||
"""
|
||||
|
||||
def partial(b: Bundle) -> dict:
|
||||
return {"reduced": asdict(compute(b))}
|
||||
|
||||
def finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
return Reduced(**parts[0]["reduced"])
|
||||
|
||||
return partial, finalize
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -83,6 +144,20 @@ def _counts(h: dict, key, nbins: int) -> list[int]:
|
||||
return h.get(key, np.zeros(nbins, dtype=np.int64)).astype(np.int64).tolist()
|
||||
|
||||
|
||||
def _partial_hist(
|
||||
lf: pl.LazyFrame, value: pl.Expr, edges: np.ndarray, group: pl.Expr | None = None
|
||||
) -> dict[str, list[int]]:
|
||||
"""One chunk's raw ``hist1d`` result as a JSON-safe, sum-mergeable dict."""
|
||||
nb = len(edges) - 1
|
||||
h = hist1d(lf, value, edges, group=group)
|
||||
return {str(k): _counts(h, k, nb) for k in h}
|
||||
|
||||
|
||||
def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]:
|
||||
"""One group's merged counts (zero-filled if the group never appeared)."""
|
||||
return list(merged.get(str(key), [0] * nbins))
|
||||
|
||||
|
||||
def _np_hist_pair(
|
||||
r: np.ndarray, t: np.ndarray, nbins: int
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
@@ -125,12 +200,21 @@ def _marginal_edges(ctx: Context, var: str) -> np.ndarray:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _marginal_overall(b: Bundle, var: str) -> Reduced:
|
||||
label, expr = _var(var)
|
||||
def _marginal_overall_partial(b: Bundle, var: str) -> dict:
|
||||
_, expr = _var(var)
|
||||
edges = _marginal_edges(b.ctx, var)
|
||||
return {
|
||||
"r": _partial_hist(b.r_phys, expr, edges),
|
||||
"t": _partial_hist(b.t_phys, expr, edges),
|
||||
}
|
||||
|
||||
|
||||
def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced:
|
||||
label, _ = _var(var)
|
||||
edges = _marginal_edges(ctx, var)
|
||||
nb = len(edges) - 1
|
||||
r = hist1d(b.r_phys, expr, edges)
|
||||
t = hist1d(b.t_phys, expr, edges)
|
||||
r = sum_merge([p["r"] for p in parts])
|
||||
t = sum_merge([p["t"] for p in parts])
|
||||
return Reduced(
|
||||
id=f"marginal_{var}",
|
||||
family="marginals",
|
||||
@@ -139,8 +223,8 @@ def _marginal_overall(b: Bundle, var: str) -> Reduced:
|
||||
xlabel=label,
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: _counts(r, 0, nb),
|
||||
_REF: _counts(t, 0, nb),
|
||||
_ROLL: _finalize_counts(r, 0, nb),
|
||||
_REF: _finalize_counts(t, 0, nb),
|
||||
"log_y": True,
|
||||
},
|
||||
)
|
||||
@@ -153,31 +237,55 @@ def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
|
||||
)
|
||||
|
||||
|
||||
def _marginal_grouped(b: Bundle, var: str, axis: str) -> Reduced:
|
||||
label, expr = _var(var)
|
||||
def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
|
||||
_, expr = _var(var)
|
||||
edges = _marginal_edges(b.ctx, var)
|
||||
nb = len(edges) - 1
|
||||
groups: dict[str, dict] = {}
|
||||
|
||||
if axis == "pdg":
|
||||
r = hist1d(b.r_phys, expr, edges, group=pl.col("pdg"))
|
||||
t = hist1d(b.t_phys, expr, edges, group=pl.col("pdg"))
|
||||
for k in b.ctx.top_pdgs:
|
||||
groups[pdg_label(k)] = {_ROLL: _counts(r, k, nb), _REF: _counts(t, k, nb)}
|
||||
elif axis == "material":
|
||||
r = hist1d(b.r_phys, expr, edges, group=pl.col("material"))
|
||||
t = hist1d(b.t_phys, expr, edges, group=pl.col("material"))
|
||||
for m in b.ctx.materials:
|
||||
groups[material_label(m)] = {
|
||||
_ROLL: _counts(r, m, nb),
|
||||
_REF: _counts(t, m, nb),
|
||||
}
|
||||
else: # energy
|
||||
e_edges = np.asarray(b.ctx.energy_edges)
|
||||
r = hist1d(b.r_phys, expr, edges, group=_energy_group_expr(b.r_phys, e_edges))
|
||||
t = hist1d(b.t_phys, expr, edges, group=_energy_group_expr(b.t_phys, e_edges))
|
||||
nb = len(edges) - 1
|
||||
return {
|
||||
"r": {str(k): _counts(r, k, nb) for k in r},
|
||||
"t": {str(k): _counts(t, k, nb) for k in t},
|
||||
}
|
||||
|
||||
|
||||
def _marginal_grouped_finalize(
|
||||
parts: list[dict], ctx: Context, var: str, axis: str
|
||||
) -> Reduced:
|
||||
label, _ = _var(var)
|
||||
edges = _marginal_edges(ctx, var)
|
||||
nb = len(edges) - 1
|
||||
r = sum_merge([p["r"] for p in parts])
|
||||
t = sum_merge([p["t"] for p in parts])
|
||||
groups: dict[str, dict] = {}
|
||||
|
||||
if axis == "pdg":
|
||||
for k in ctx.top_pdgs:
|
||||
groups[pdg_label(k)] = {
|
||||
_ROLL: _finalize_counts(r, k, nb),
|
||||
_REF: _finalize_counts(t, k, nb),
|
||||
}
|
||||
elif axis == "material":
|
||||
for m in ctx.materials:
|
||||
groups[material_label(m)] = {
|
||||
_ROLL: _finalize_counts(r, m, nb),
|
||||
_REF: _finalize_counts(t, m, nb),
|
||||
}
|
||||
else: # energy
|
||||
e_edges = np.asarray(ctx.energy_edges)
|
||||
for bi, lbl in enumerate(energy_bin_labels(e_edges)):
|
||||
groups[lbl] = {_ROLL: _counts(r, bi, nb), _REF: _counts(t, bi, nb)}
|
||||
groups[lbl] = {
|
||||
_ROLL: _finalize_counts(r, bi, nb),
|
||||
_REF: _finalize_counts(t, bi, nb),
|
||||
}
|
||||
|
||||
return Reduced(
|
||||
id=f"marginal_{var}_by_{axis}",
|
||||
@@ -194,13 +302,19 @@ def _marginal_grouped(b: Bundle, var: str, axis: str) -> Reduced:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _event_scalar(
|
||||
b: Bundle, spec_id: str, title: str, xlabel: str, col: str, use_all: bool
|
||||
) -> Reduced:
|
||||
def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
|
||||
r_lf, t_lf = (b.r_all, b.t_all) if use_all else (b.r_phys, b.t_phys)
|
||||
r = event_scalars(r_lf)[col].to_numpy()
|
||||
t = event_scalars(t_lf)[col].to_numpy()
|
||||
edges, rc, tc = _np_hist_pair(r, t, b.ctx.n_marginal_bins)
|
||||
return {"r": r.tolist(), "t": t.tolist()}
|
||||
|
||||
|
||||
def _event_scalar_finalize(
|
||||
parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str
|
||||
) -> Reduced:
|
||||
r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
|
||||
t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
|
||||
edges, rc, tc = _np_hist_pair(r, t, ctx.n_marginal_bins)
|
||||
return Reduced(
|
||||
id=spec_id,
|
||||
family="event",
|
||||
@@ -216,18 +330,26 @@ def _event_scalar(
|
||||
)
|
||||
|
||||
|
||||
def _event_total_edep_by_energy(b: Bundle) -> Reduced:
|
||||
e_edges = np.asarray(b.ctx.energy_edges)
|
||||
def _event_total_edep_by_energy_partial(b: Bundle) -> dict:
|
||||
r = event_scalars(b.r_all)
|
||||
t = event_scalars(b.t_all)
|
||||
r_bin = np.clip(
|
||||
np.digitize(r["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2
|
||||
)
|
||||
t_bin = np.clip(
|
||||
np.digitize(t["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2
|
||||
)
|
||||
r_val, t_val = r["total_edep"].to_numpy(), t["total_edep"].to_numpy()
|
||||
edges, _, _ = _np_hist_pair(r_val, t_val, b.ctx.n_marginal_bins)
|
||||
return {
|
||||
"r_incident": r["incident_E"].to_list(),
|
||||
"r_edep": r["total_edep"].to_list(),
|
||||
"t_incident": t["incident_E"].to_list(),
|
||||
"t_edep": t["total_edep"].to_list(),
|
||||
}
|
||||
|
||||
|
||||
def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
e_edges = np.asarray(ctx.energy_edges)
|
||||
r_inc = np.concatenate([np.asarray(p["r_incident"], dtype=float) for p in parts])
|
||||
r_val = np.concatenate([np.asarray(p["r_edep"], dtype=float) for p in parts])
|
||||
t_inc = np.concatenate([np.asarray(p["t_incident"], dtype=float) for p in parts])
|
||||
t_val = np.concatenate([np.asarray(p["t_edep"], dtype=float) for p in parts])
|
||||
r_bin = np.clip(np.digitize(r_inc, e_edges[1:-1]), 0, len(e_edges) - 2)
|
||||
t_bin = np.clip(np.digitize(t_inc, e_edges[1:-1]), 0, len(e_edges) - 2)
|
||||
edges, _, _ = _np_hist_pair(r_val, t_val, ctx.n_marginal_bins)
|
||||
groups: dict[str, dict] = {}
|
||||
for bi, lbl in enumerate(energy_bin_labels(e_edges)):
|
||||
rc = np.histogram(r_val[r_bin == bi], edges)[0]
|
||||
@@ -251,15 +373,54 @@ def _event_total_edep_by_energy(b: Bundle) -> Reduced:
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _profile(
|
||||
b: Bundle, spec_id: str, title: str, xlabel: str, coord_fn, edges_key: str
|
||||
) -> Reduced:
|
||||
def _profile_partial(b: Bundle, coord_fn, edges_key: str) -> dict:
|
||||
edges = np.asarray(getattr(b.ctx, edges_key))
|
||||
r_ea, t_ea = entry_axis(b.r_all), entry_axis(b.t_all)
|
||||
r_lf = attach_entry_axis(b.r_all, r_ea)
|
||||
t_lf = attach_entry_axis(b.t_all, t_ea)
|
||||
r_mean, r_std = weighted_profile(r_lf, coord_fn(), edges, pl.col("edep"))
|
||||
t_mean, t_std = weighted_profile(t_lf, coord_fn(), edges, pl.col("edep"))
|
||||
r_lf = attach_entry_axis(b.r_all, entry_axis(b.r_all))
|
||||
t_lf = attach_entry_axis(b.t_all, entry_axis(b.t_all))
|
||||
r_ids, r_mat = profile_partial(r_lf, coord_fn(), edges, pl.col("edep"))
|
||||
t_ids, t_mat = profile_partial(t_lf, coord_fn(), edges, pl.col("edep"))
|
||||
return {
|
||||
"r_ids": r_ids.tolist(),
|
||||
"r_mat": r_mat.tolist(),
|
||||
"t_ids": t_ids.tolist(),
|
||||
"t_mat": t_mat.tolist(),
|
||||
}
|
||||
|
||||
|
||||
def _assert_event_disjoint(id_lists: list[list[int]], spec_id: str, side: str) -> None:
|
||||
"""Guard the chunking invariant profiles depend on: no event in two chunks.
|
||||
|
||||
A violation would silently double-count that event in the merged mean/RMS
|
||||
with no other symptom, so this is worth a loud failure rather than a
|
||||
quietly-wrong plot.
|
||||
"""
|
||||
seen: set[int] = set()
|
||||
for ids in id_lists:
|
||||
overlap = seen & set(ids)
|
||||
if overlap:
|
||||
raise ValueError(
|
||||
f"{spec_id} ({side}): event_id(s) {sorted(overlap)[:5]} appear "
|
||||
"in more than one chunk — chunking must be event-disjoint"
|
||||
)
|
||||
seen.update(ids)
|
||||
|
||||
|
||||
def _profile_finalize(
|
||||
parts: list[dict],
|
||||
ctx: Context,
|
||||
spec_id: str,
|
||||
title: str,
|
||||
xlabel: str,
|
||||
edges_key: str,
|
||||
) -> Reduced:
|
||||
edges = np.asarray(getattr(ctx, edges_key))
|
||||
nb = len(edges) - 1
|
||||
_assert_event_disjoint([p["r_ids"] for p in parts], spec_id, "rollout")
|
||||
_assert_event_disjoint([p["t_ids"] for p in parts], spec_id, "reference")
|
||||
r_mats = [np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts]
|
||||
t_mats = [np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts]
|
||||
r_mean, r_std = profile_finalize(r_mats)
|
||||
t_mean, t_std = profile_finalize(t_mats)
|
||||
return Reduced(
|
||||
id=spec_id,
|
||||
family="shower",
|
||||
@@ -282,14 +443,21 @@ def _profile(
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _species_share(b: Bundle) -> Reduced:
|
||||
def _species_share_partial(b: Bundle) -> dict:
|
||||
r = species_share(b.r_all)
|
||||
t = species_share(b.t_all)
|
||||
r_map = dict(zip(r["pdg"].to_list(), r["total_edep"].to_list()))
|
||||
t_map = dict(zip(t["pdg"].to_list(), t["total_edep"].to_list()))
|
||||
return {
|
||||
"r": {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())},
|
||||
"t": {str(k): v for k, v in zip(t["pdg"].to_list(), t["total_edep"].to_list())},
|
||||
}
|
||||
|
||||
|
||||
def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
r_map = sum_merge([p["r"] for p in parts])
|
||||
t_map = sum_merge([p["t"] for p in parts])
|
||||
r_tot = sum(r_map.values()) or 1.0
|
||||
t_tot = sum(t_map.values()) or 1.0
|
||||
labels = [pdg_label(k) for k in b.ctx.top_pdgs]
|
||||
labels = [pdg_label(k) for k in ctx.top_pdgs]
|
||||
return Reduced(
|
||||
id="species_edep_share",
|
||||
family="species",
|
||||
@@ -298,19 +466,22 @@ def _species_share(b: Bundle) -> Reduced:
|
||||
xlabel="species",
|
||||
payload={
|
||||
"labels": labels,
|
||||
_ROLL: [r_map.get(k, 0.0) / r_tot for k in b.ctx.top_pdgs],
|
||||
_REF: [t_map.get(k, 0.0) / t_tot for k in b.ctx.top_pdgs],
|
||||
_ROLL: [r_map.get(str(k), 0.0) / r_tot for k in ctx.top_pdgs],
|
||||
_REF: [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs],
|
||||
"ylabel": "fraction of total deposited energy",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _leakage(b: Bundle) -> Reduced:
|
||||
def _leakage_partial(b: Bundle) -> dict:
|
||||
frac = leakage_fraction(b.r_all)
|
||||
return {"frac": frac.tolist()}
|
||||
|
||||
|
||||
def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts])
|
||||
edges = np.linspace(
|
||||
0.0,
|
||||
max(float(frac.max()) if len(frac) else 1.0, 1e-3),
|
||||
b.ctx.n_marginal_bins + 1,
|
||||
0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1
|
||||
)
|
||||
counts = np.histogram(frac, edges)[0]
|
||||
return Reduced(
|
||||
@@ -340,7 +511,7 @@ def _sec_frames(b: Bundle):
|
||||
)
|
||||
|
||||
|
||||
def _sec_count_per_event(b: Bundle) -> Reduced:
|
||||
def _sec_count_per_event_partial(b: Bundle) -> dict:
|
||||
r_sec, t_sec = _sec_frames(b)
|
||||
r = (
|
||||
r_sec.group_by("event_id")
|
||||
@@ -354,9 +525,13 @@ def _sec_count_per_event(b: Bundle) -> Reduced:
|
||||
.collect(engine="streaming")["n"]
|
||||
.to_numpy()
|
||||
)
|
||||
edges, rc, tc = _np_hist_pair(
|
||||
r.astype(float), t.astype(float), min(b.ctx.n_marginal_bins, 40)
|
||||
)
|
||||
return {"r": r.tolist(), "t": t.tolist()}
|
||||
|
||||
|
||||
def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts])
|
||||
t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts])
|
||||
edges, rc, tc = _np_hist_pair(r, t, min(ctx.n_marginal_bins, 40))
|
||||
return Reduced(
|
||||
id="sec_count_per_event",
|
||||
family="secondaries",
|
||||
@@ -372,32 +547,21 @@ def _sec_count_per_event(b: Bundle) -> Reduced:
|
||||
)
|
||||
|
||||
|
||||
def _sec_count_per_species(b: Bundle) -> Reduced:
|
||||
def _counts_by_pdg(sec_lf: pl.LazyFrame) -> dict[str, int]:
|
||||
df = sec_lf.group_by("pdg").agg(pl.len().alias("n")).collect(engine="streaming")
|
||||
return {str(k): v for k, v in zip(df["pdg"].to_list(), df["n"].to_list())}
|
||||
|
||||
|
||||
def _sec_count_per_species_partial(b: Bundle) -> dict:
|
||||
r_sec, t_sec = _sec_frames(b)
|
||||
r = dict(
|
||||
zip(
|
||||
*[
|
||||
r_sec.group_by("pdg")
|
||||
.agg(pl.len().alias("n"))
|
||||
.collect(engine="streaming")[c]
|
||||
.to_list()
|
||||
for c in ("pdg", "n")
|
||||
]
|
||||
)
|
||||
)
|
||||
t = dict(
|
||||
zip(
|
||||
*[
|
||||
t_sec.group_by("pdg")
|
||||
.agg(pl.len().alias("n"))
|
||||
.collect(engine="streaming")[c]
|
||||
.to_list()
|
||||
for c in ("pdg", "n")
|
||||
]
|
||||
)
|
||||
)
|
||||
return {"r": _counts_by_pdg(r_sec), "t": _counts_by_pdg(t_sec)}
|
||||
|
||||
|
||||
def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
r = sum_merge([p["r"] for p in parts])
|
||||
t = sum_merge([p["t"] for p in parts])
|
||||
keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[
|
||||
: len(b.ctx.top_pdgs)
|
||||
: len(ctx.top_pdgs)
|
||||
]
|
||||
return Reduced(
|
||||
id="sec_count_per_species",
|
||||
@@ -406,7 +570,7 @@ def _sec_count_per_species(b: Bundle) -> Reduced:
|
||||
title="Secondary count by species",
|
||||
xlabel="species",
|
||||
payload={
|
||||
"labels": [pdg_label(k) for k in keys],
|
||||
"labels": [pdg_label(int(k)) for k in keys],
|
||||
_ROLL: [float(r.get(k, 0)) for k in keys],
|
||||
_REF: [float(t.get(k, 0)) for k in keys],
|
||||
"ylabel": "secondary count",
|
||||
@@ -414,12 +578,20 @@ def _sec_count_per_species(b: Bundle) -> Reduced:
|
||||
)
|
||||
|
||||
|
||||
def _sec_energy(b: Bundle) -> Reduced:
|
||||
def _sec_energy_partial(b: Bundle) -> dict:
|
||||
r_sec, t_sec = _sec_frames(b)
|
||||
edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1)
|
||||
r = hist1d(r_sec, pl.col("energy"), edges)
|
||||
t = hist1d(t_sec, pl.col("energy"), edges)
|
||||
return {
|
||||
"r": _partial_hist(r_sec, pl.col("energy"), edges),
|
||||
"t": _partial_hist(t_sec, pl.col("energy"), edges),
|
||||
}
|
||||
|
||||
|
||||
def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
edges = np.linspace(*ctx.sec_energy_range, ctx.n_sec_bins + 1)
|
||||
nb = len(edges) - 1
|
||||
r = sum_merge([p["r"] for p in parts])
|
||||
t = sum_merge([p["t"] for p in parts])
|
||||
return Reduced(
|
||||
id="sec_energy",
|
||||
family="secondaries",
|
||||
@@ -428,27 +600,34 @@ def _sec_energy(b: Bundle) -> Reduced:
|
||||
xlabel="secondary energy [MeV]",
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: _counts(r, 0, nb),
|
||||
_REF: _counts(t, 0, nb),
|
||||
_ROLL: _finalize_counts(r, 0, nb),
|
||||
_REF: _finalize_counts(t, 0, nb),
|
||||
"log_y": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _sec_cos_angle(b: Bundle) -> Reduced:
|
||||
def _sec_cos_angle_partial(b: Bundle) -> dict:
|
||||
edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1)
|
||||
nb = len(edges) - 1
|
||||
cos = (
|
||||
pl.col("sdx") * pl.col("axis_x")
|
||||
+ pl.col("sdy") * pl.col("axis_y")
|
||||
+ pl.col("sdz") * pl.col("axis_z")
|
||||
).clip(-1.0, 1.0)
|
||||
|
||||
def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> list[int]:
|
||||
def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]:
|
||||
ea = entry_axis(steps_lf)
|
||||
return _counts(hist1d(attach_entry_axis(sec_lf, ea), cos, edges), 0, nb)
|
||||
return _partial_hist(attach_entry_axis(sec_lf, ea), cos, edges)
|
||||
|
||||
r_sec, t_sec = _sec_frames(b)
|
||||
return {"r": _side(r_sec, b.r_phys), "t": _side(t_sec, b.t_all)}
|
||||
|
||||
|
||||
def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
edges = np.linspace(-1.0, 1.0, ctx.n_sec_bins + 1)
|
||||
nb = len(edges) - 1
|
||||
r = sum_merge([p["r"] for p in parts])
|
||||
t = sum_merge([p["t"] for p in parts])
|
||||
return Reduced(
|
||||
id="sec_cos_angle",
|
||||
family="secondaries",
|
||||
@@ -457,13 +636,30 @@ def _sec_cos_angle(b: Bundle) -> Reduced:
|
||||
xlabel="cos of emission angle",
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: _side(r_sec, b.r_phys),
|
||||
_REF: _side(t_sec, b.t_all),
|
||||
_ROLL: _finalize_counts(r, 0, nb),
|
||||
_REF: _finalize_counts(t, 0, nb),
|
||||
"log_y": False,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# router diagnostics (not chunked — already bounded/subsampled)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_router_gating_partial, _router_gating_finalize = _unchunkable(
|
||||
lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys)
|
||||
)
|
||||
_router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable(
|
||||
lambda b: compute_router_share_by_pdg(
|
||||
b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs
|
||||
)
|
||||
)
|
||||
_router_share_process_partial, _router_share_process_finalize = _unchunkable(
|
||||
lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# registry assembly
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -479,7 +675,12 @@ def build_catalog() -> list[PlotSpec]:
|
||||
for var in MARGINAL_VARS:
|
||||
specs.append(
|
||||
PlotSpec(
|
||||
f"marginal_{var}", "marginals", lambda b, v=var: _marginal_overall(b, v)
|
||||
f"marginal_{var}",
|
||||
"marginals",
|
||||
compute_partial=lambda b, v=var: _marginal_overall_partial(b, v),
|
||||
finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(
|
||||
parts, ctx, v
|
||||
),
|
||||
)
|
||||
)
|
||||
for axis in GROUPING_AXES:
|
||||
@@ -487,7 +688,12 @@ def build_catalog() -> list[PlotSpec]:
|
||||
PlotSpec(
|
||||
f"marginal_{var}_by_{axis}",
|
||||
"marginals",
|
||||
lambda b, v=var, a=axis: _marginal_grouped(b, v, a),
|
||||
compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(
|
||||
b, v, a
|
||||
),
|
||||
finalize=lambda parts, ctx, v=var, a=axis: (
|
||||
_marginal_grouped_finalize(parts, ctx, v, a)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -495,70 +701,136 @@ def build_catalog() -> list[PlotSpec]:
|
||||
PlotSpec(
|
||||
"event_total_edep",
|
||||
"event",
|
||||
lambda b: _event_scalar(
|
||||
b,
|
||||
compute_partial=lambda b: _event_scalar_partial(
|
||||
b, "total_edep", use_all=True
|
||||
),
|
||||
finalize=lambda parts, ctx: _event_scalar_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
"event_total_edep",
|
||||
"Total deposited energy per event",
|
||||
"total deposited energy [MeV]",
|
||||
"total_edep",
|
||||
use_all=True,
|
||||
),
|
||||
),
|
||||
PlotSpec("event_total_edep_by_energy", "event", _event_total_edep_by_energy),
|
||||
PlotSpec(
|
||||
"event_total_edep_by_energy",
|
||||
"event",
|
||||
compute_partial=_event_total_edep_by_energy_partial,
|
||||
finalize=_event_total_edep_by_energy_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"event_mean_length",
|
||||
"event",
|
||||
lambda b: _event_scalar(
|
||||
b,
|
||||
compute_partial=lambda b: _event_scalar_partial(
|
||||
b, "mean_length", use_all=False
|
||||
),
|
||||
finalize=lambda parts, ctx: _event_scalar_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
"event_mean_length",
|
||||
"Mean step length per event",
|
||||
"mean step length [mm]",
|
||||
"mean_length",
|
||||
use_all=False,
|
||||
),
|
||||
),
|
||||
PlotSpec(
|
||||
"event_n_steps",
|
||||
"event",
|
||||
lambda b: _event_scalar(
|
||||
b,
|
||||
compute_partial=lambda b: _event_scalar_partial(
|
||||
b, "n_steps", use_all=False
|
||||
),
|
||||
finalize=lambda parts, ctx: _event_scalar_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
"event_n_steps",
|
||||
"Number of steps per event",
|
||||
"steps per event",
|
||||
"n_steps",
|
||||
use_all=False,
|
||||
),
|
||||
),
|
||||
PlotSpec(
|
||||
"shower_longitudinal",
|
||||
"shower",
|
||||
lambda b: _profile(
|
||||
b,
|
||||
compute_partial=lambda b: _profile_partial(b, depth_expr, "depth_edges"),
|
||||
finalize=lambda parts, ctx: _profile_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
"shower_longitudinal",
|
||||
"Longitudinal shower profile",
|
||||
"depth along shower axis [mm]",
|
||||
depth_expr,
|
||||
"depth_edges",
|
||||
),
|
||||
),
|
||||
PlotSpec(
|
||||
"shower_transverse",
|
||||
"shower",
|
||||
lambda b: _profile(
|
||||
b,
|
||||
compute_partial=lambda b: _profile_partial(
|
||||
b, transverse_expr, "transverse_edges"
|
||||
),
|
||||
finalize=lambda parts, ctx: _profile_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
"shower_transverse",
|
||||
"Transverse shower profile",
|
||||
"radius from shower axis [mm]",
|
||||
transverse_expr,
|
||||
"transverse_edges",
|
||||
),
|
||||
),
|
||||
PlotSpec("species_edep_share", "species", _species_share),
|
||||
PlotSpec("leakage_fraction", "species", _leakage),
|
||||
PlotSpec("sec_count_per_event", "secondaries", _sec_count_per_event),
|
||||
PlotSpec("sec_count_per_species", "secondaries", _sec_count_per_species),
|
||||
PlotSpec("sec_energy", "secondaries", _sec_energy),
|
||||
PlotSpec("sec_cos_angle", "secondaries", _sec_cos_angle),
|
||||
PlotSpec(
|
||||
"species_edep_share",
|
||||
"species",
|
||||
compute_partial=_species_share_partial,
|
||||
finalize=_species_share_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"leakage_fraction",
|
||||
"species",
|
||||
compute_partial=_leakage_partial,
|
||||
finalize=_leakage_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"sec_count_per_event",
|
||||
"secondaries",
|
||||
compute_partial=_sec_count_per_event_partial,
|
||||
finalize=_sec_count_per_event_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"sec_count_per_species",
|
||||
"secondaries",
|
||||
compute_partial=_sec_count_per_species_partial,
|
||||
finalize=_sec_count_per_species_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"sec_energy",
|
||||
"secondaries",
|
||||
compute_partial=_sec_energy_partial,
|
||||
finalize=_sec_energy_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"sec_cos_angle",
|
||||
"secondaries",
|
||||
compute_partial=_sec_cos_angle_partial,
|
||||
finalize=_sec_cos_angle_finalize,
|
||||
),
|
||||
PlotSpec(
|
||||
"router_gating",
|
||||
"model",
|
||||
compute_partial=_router_gating_partial,
|
||||
finalize=_router_gating_finalize,
|
||||
chunkable=False,
|
||||
),
|
||||
PlotSpec(
|
||||
"router_share_by_pdg",
|
||||
"model",
|
||||
compute_partial=_router_share_pdg_partial,
|
||||
finalize=_router_share_pdg_finalize,
|
||||
chunkable=False,
|
||||
),
|
||||
PlotSpec(
|
||||
"router_share_by_process",
|
||||
"model",
|
||||
compute_partial=_router_share_process_partial,
|
||||
finalize=_router_share_process_finalize,
|
||||
chunkable=False,
|
||||
),
|
||||
]
|
||||
return specs
|
||||
|
||||
|
||||
+220
-36
@@ -15,18 +15,25 @@ next to the rollout parquet, and lays everything out under it:
|
||||
|
||||
<run_dir>/shared.json fixed bin edges / group sets (prep)
|
||||
<run_dir>/run_meta.json resolved rollout/reference paths + plot metadata
|
||||
<run_dir>/reduced/<id>.json one per compute job
|
||||
<run_dir>/reduced_partial/<id>__<chunk>.json one per (plot, chunk) job
|
||||
<run_dir>/reduced/<id>.json merged, per plot
|
||||
<run_dir>/plots/<family>/<id>.pdf rendered locally
|
||||
|
||||
Job model (one condor job per plot, compute/render split):
|
||||
Job model (one condor job per (plot, chunk), compute/merge/render split):
|
||||
|
||||
1. ``prep`` runs once on the submit node — reads the YAML, resolves the shared
|
||||
context from a subsample, writes ``shared.json`` + ``run_meta.json``.
|
||||
2. one job per catalog id runs ``giant analyze compute-one --run-dir`` on a
|
||||
worker — a single streaming pass writing ``reduced/<id>.json`` (polars/numpy
|
||||
only, no LaTeX).
|
||||
3. a final *local* ``giant analyze render`` turns those into the styled PDF +
|
||||
gallery tree (that step imports plotstyle/LaTeX).
|
||||
context from a subsample, writes ``shared.json`` + ``run_meta.json``
|
||||
(including the run's configured ``n_chunks``).
|
||||
2. one job per catalog id x chunk index runs ``giant analyze compute-one
|
||||
--run-dir`` on a worker — a single streaming pass over that
|
||||
``event_id``-disjoint chunk, writing ``reduced_partial/<id>__<chunk>.json``
|
||||
(polars/numpy only, no LaTeX). Specs marked ``chunkable=False``
|
||||
(``PlotSpec``, ``catalog.py``) always run as a single chunk.
|
||||
3. a *local* ``giant analyze render`` first merges every plot's chunk partials
|
||||
(``merge_all`` — sums/concatenates them and re-derives any data-dependent
|
||||
histogram edges or mean/std, per ``PlotSpec.finalize``) into
|
||||
``reduced/<id>.json``, then renders those into the styled PDF + gallery tree
|
||||
(that step imports plotstyle/LaTeX).
|
||||
|
||||
Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no
|
||||
HTCondor file transfer of the multi-GB inputs.
|
||||
@@ -35,13 +42,17 @@ HTCondor file transfer of the multi-GB inputs.
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
import yaml
|
||||
|
||||
from giant.analysis.catalog import Bundle, catalog_ids, get_spec
|
||||
from giant.analysis.context import Context, build_context
|
||||
from giant.analysis.reduced import Partial
|
||||
from giant.analysis.runtime_estimate import estimate_runtime_s
|
||||
from giant.analysis.sources import Side, open_side
|
||||
|
||||
# Keys copied verbatim from a rollout YAML into each plot's gallery metadata.
|
||||
_PLOT_META_KEYS = (
|
||||
@@ -54,12 +65,22 @@ _PLOT_META_KEYS = (
|
||||
"max_steps",
|
||||
"steps",
|
||||
"max_tracks_per_event",
|
||||
"escape_threshold",
|
||||
"n_events",
|
||||
"n_seed_events",
|
||||
"timestamp",
|
||||
"comment",
|
||||
"weights",
|
||||
"batch_size",
|
||||
"device",
|
||||
"rollout_seed",
|
||||
"n_rows",
|
||||
"termination_reason_counts",
|
||||
"model_config",
|
||||
"training_epoch",
|
||||
"best_val_loss",
|
||||
"training_config",
|
||||
"training_meta",
|
||||
)
|
||||
|
||||
|
||||
@@ -82,13 +103,26 @@ def load_rollout_yaml(path: str | Path) -> dict:
|
||||
return d
|
||||
|
||||
|
||||
def derive_run_dir(rollout_yaml: dict, run_dir: str | Path | None = None) -> Path:
|
||||
"""Analysis output directory, next to the rollout parquet unless overridden."""
|
||||
def derive_run_dir(
|
||||
rollout_yaml: dict,
|
||||
run_dir: str | Path | None = None,
|
||||
default_base: str | Path | None = None,
|
||||
) -> Path:
|
||||
"""Analysis output directory.
|
||||
|
||||
Precedence: an explicit ``run_dir`` always wins. Otherwise
|
||||
``default_base / analysis_<tag>`` if ``default_base`` is given (the CLI
|
||||
passes the repo's gitignored ``analysis_runs/``, so run directories don't
|
||||
pile up on ``/ceph`` next to the rollout parquet). Falls back to next to
|
||||
the rollout parquet — the original convention — for callers that don't
|
||||
care where the run directory lives.
|
||||
"""
|
||||
if run_dir is not None:
|
||||
return Path(run_dir)
|
||||
rollout = Path(rollout_yaml["output"])
|
||||
tag = str(rollout_yaml.get("prediction_id") or rollout.stem)[:8]
|
||||
return rollout.parent / f"analysis_{tag}"
|
||||
base = Path(default_base) if default_base is not None else rollout.parent
|
||||
return base / f"analysis_{tag}"
|
||||
|
||||
|
||||
def _plot_meta(rollout_yaml: dict) -> dict:
|
||||
@@ -104,6 +138,12 @@ class RunMeta:
|
||||
run_dir: str
|
||||
title: str
|
||||
plot_meta: dict
|
||||
n_chunks: int = 1
|
||||
# rollout+reference row count of each event_id-disjoint chunk, and the
|
||||
# dataset total — inputs to `runtime_estimate.estimate_runtime_s`. Empty/0
|
||||
# on run directories written before this field existed.
|
||||
rows_per_chunk: list[int] = field(default_factory=list)
|
||||
total_rows: int = 0
|
||||
|
||||
def save(self, path: str | Path) -> None:
|
||||
Path(path).write_text(json.dumps(self.__dict__, indent=2))
|
||||
@@ -113,23 +153,58 @@ class RunMeta:
|
||||
return cls(**json.loads(Path(path).read_text()))
|
||||
|
||||
|
||||
def _rows_per_chunk(
|
||||
rollout: str | Path, reference: str | Path, n_chunks: int
|
||||
) -> list[int]:
|
||||
"""Rollout+reference row count of each ``event_id % n_chunks`` chunk.
|
||||
|
||||
One cheap streaming ``group_by`` per side (just the ``event_id`` column) —
|
||||
the sizing input every job's estimated walltime
|
||||
(``runtime_estimate.estimate_runtime_s``) is computed from.
|
||||
"""
|
||||
|
||||
def counts(lf: pl.LazyFrame) -> pl.DataFrame:
|
||||
return (
|
||||
lf.select((pl.col("event_id") % n_chunks).alias("_c"))
|
||||
.group_by("_c")
|
||||
.agg(pl.len().alias("n"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
|
||||
out = [0] * n_chunks
|
||||
for lf in (open_side(rollout, Side.rollout), open_side(reference, Side.reference)):
|
||||
df = counts(lf)
|
||||
for c, n in zip(df["_c"].to_list(), df["n"].to_list()):
|
||||
out[c] += n
|
||||
return out
|
||||
|
||||
|
||||
def prep(
|
||||
rollout_yaml: str | Path,
|
||||
run_dir: str | Path | None = None,
|
||||
n_chunks: int = 1,
|
||||
default_base: str | Path | None = None,
|
||||
**ctx_kwargs,
|
||||
) -> Path:
|
||||
"""Read the rollout YAML, build the shared context, and lay out the run dir.
|
||||
|
||||
Writes ``shared.json`` + ``run_meta.json`` and returns the run directory.
|
||||
``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one``
|
||||
job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is
|
||||
resolved once here rather than re-passed (and risking disagreement) at every
|
||||
later step. See ``derive_run_dir`` for how ``run_dir``/``default_base``
|
||||
resolve the actual directory.
|
||||
"""
|
||||
y = load_rollout_yaml(rollout_yaml)
|
||||
run_path = derive_run_dir(y, run_dir)
|
||||
run_path = derive_run_dir(y, run_dir, default_base=default_base)
|
||||
run_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
rollout, reference = y["output"], y["dataset"]
|
||||
ctx = build_context(rollout, reference, **ctx_kwargs)
|
||||
ctx.save(run_path / "shared.json")
|
||||
|
||||
rows_per_chunk = _rows_per_chunk(rollout, reference, n_chunks)
|
||||
|
||||
ckpt = Path(y.get("checkpoint", "")).name or "rollout"
|
||||
RunMeta(
|
||||
rollout=str(rollout),
|
||||
@@ -137,12 +212,15 @@ def prep(
|
||||
run_dir=str(run_path),
|
||||
title=f"GIANT rollout analysis — {ckpt}",
|
||||
plot_meta=_plot_meta(y),
|
||||
n_chunks=n_chunks,
|
||||
rows_per_chunk=rows_per_chunk,
|
||||
total_rows=sum(rows_per_chunk),
|
||||
).save(run_path / "run_meta.json")
|
||||
return run_path
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# per-plot compute (what each condor job runs)
|
||||
# per-(plot, chunk) compute (what each condor job runs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@@ -152,18 +230,42 @@ def compute_reduced(
|
||||
reference: str | Path,
|
||||
shared: str | Path,
|
||||
out: str | Path,
|
||||
checkpoint: str | None = None,
|
||||
chunk_index: int = 0,
|
||||
n_chunks: int = 1,
|
||||
) -> Path:
|
||||
"""Core: run one plot's reduction against explicit paths → ``Reduced`` JSON."""
|
||||
"""Core: run one (plot, chunk)'s partial reduction against explicit paths.
|
||||
|
||||
Writes a ``Partial`` JSON — the raw, not-yet-merged output of
|
||||
``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one``
|
||||
is what combines every chunk's ``Partial`` for a plot into the final
|
||||
``Reduced``. Specs with ``chunkable=False`` always run as a single chunk
|
||||
regardless of ``n_chunks``.
|
||||
"""
|
||||
ctx = Context.load(shared)
|
||||
bundle = Bundle.open(rollout, reference, ctx)
|
||||
reduced = get_spec(spec_id).compute(bundle)
|
||||
spec = get_spec(spec_id)
|
||||
effective_n = n_chunks if spec.chunkable else 1
|
||||
if not (0 <= chunk_index < effective_n):
|
||||
raise ValueError(
|
||||
f"{spec_id}: chunk_index={chunk_index} out of range for "
|
||||
f"n_chunks={effective_n} (chunkable={spec.chunkable})"
|
||||
)
|
||||
bundle = Bundle.open(
|
||||
rollout, reference, ctx, checkpoint=checkpoint, chunk=(chunk_index, effective_n)
|
||||
)
|
||||
partial = Partial(
|
||||
id=spec_id,
|
||||
family=spec.family,
|
||||
chunk=chunk_index,
|
||||
data=spec.compute_partial(bundle),
|
||||
)
|
||||
out = Path(out)
|
||||
reduced.save(out)
|
||||
partial.save(out)
|
||||
return out
|
||||
|
||||
|
||||
def compute_one(spec_id: str, run_dir: str | Path) -> Path:
|
||||
"""Run one plot's reduction from a prepped run directory."""
|
||||
def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path:
|
||||
"""Run one (plot, chunk)'s partial reduction from a prepped run directory."""
|
||||
run_path = Path(run_dir)
|
||||
meta = RunMeta.load(run_path / "run_meta.json")
|
||||
return compute_reduced(
|
||||
@@ -171,10 +273,56 @@ def compute_one(spec_id: str, run_dir: str | Path) -> Path:
|
||||
meta.rollout,
|
||||
meta.reference,
|
||||
run_path / "shared.json",
|
||||
run_path / "reduced" / f"{spec_id}.json",
|
||||
run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json",
|
||||
checkpoint=meta.plot_meta.get("checkpoint"),
|
||||
chunk_index=chunk_index,
|
||||
n_chunks=meta.n_chunks,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# per-plot merge (the join step ``render`` runs before rendering)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def merge_one(spec_id: str, run_dir: str | Path) -> Path:
|
||||
"""Merge every chunk's partial for one plot into the final ``Reduced`` JSON.
|
||||
|
||||
Fails loudly if fewer partials exist than the run's configured chunk count
|
||||
for this plot — that is what catches an incomplete/failed condor job
|
||||
instead of silently rendering a plot from partial data. Idempotent: safe
|
||||
to call again (e.g. from ``render_run``) once all chunks are in.
|
||||
"""
|
||||
run_path = Path(run_dir)
|
||||
meta = RunMeta.load(run_path / "run_meta.json")
|
||||
ctx = Context.load(run_path / "shared.json")
|
||||
spec = get_spec(spec_id)
|
||||
effective_n = meta.n_chunks if spec.chunkable else 1
|
||||
|
||||
partial_dir = run_path / "reduced_partial"
|
||||
found = {
|
||||
p.chunk: p
|
||||
for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json"))
|
||||
}
|
||||
missing = sorted(set(range(effective_n)) - set(found))
|
||||
if missing:
|
||||
raise FileNotFoundError(
|
||||
f"{spec_id}: missing chunk partial(s) {missing} of {effective_n} "
|
||||
f"under {partial_dir} — did every compute-one job finish?"
|
||||
)
|
||||
|
||||
parts = [found[k].data for k in range(effective_n)]
|
||||
reduced = spec.finalize(parts, ctx)
|
||||
out = run_path / "reduced" / f"{spec_id}.json"
|
||||
reduced.save(out)
|
||||
return out
|
||||
|
||||
|
||||
def merge_all(run_dir: str | Path) -> list[Path]:
|
||||
"""Merge every catalog plot's chunk partials into ``reduced/<id>.json``."""
|
||||
return [merge_one(spec_id, run_dir) for spec_id in catalog_ids()]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# submit description
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -185,21 +333,21 @@ class SubmitConfig:
|
||||
run_dir: Path
|
||||
accounting_group: str
|
||||
repo_dir: Path
|
||||
docker_image: str = "mschnepf/slc7-condocker"
|
||||
request_memory_mb: int = 4096
|
||||
docker_image: str = "cverstege/alma9-gridjob"
|
||||
request_memory_mb: int = 8192
|
||||
request_cpus: int = 1
|
||||
request_walltime_s: int = 3600
|
||||
remote: bool = False # +RemoteJob (grid I/O) vs ProvidesETPResources (local files)
|
||||
n_chunks: int = 1 # per-plot data chunks; ignored for chunkable=False specs
|
||||
|
||||
|
||||
_WRAPPER = """#!/bin/bash
|
||||
set -euo pipefail
|
||||
cd {repo_dir}
|
||||
exec uv run giant analyze compute-one --id "$1" --run-dir {run_dir}
|
||||
exec {repo_dir}/.venv/bin/giant analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
|
||||
"""
|
||||
|
||||
|
||||
def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str:
|
||||
def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
|
||||
reqs_attrs = (
|
||||
"+RemoteJob = True\n"
|
||||
if cfg.remote
|
||||
@@ -209,39 +357,75 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str
|
||||
"universe = docker\n"
|
||||
f"docker_image = {cfg.docker_image}\n"
|
||||
f"executable = {wrapper}\n"
|
||||
"arguments = $(plotid)\n"
|
||||
"arguments = $(plotid) $(chunk)\n"
|
||||
"should_transfer_files = YES\n"
|
||||
"when_to_transfer_output = ON_EXIT\n"
|
||||
f"request_memory = {cfg.request_memory_mb}\n"
|
||||
f"request_cpus = {cfg.request_cpus}\n"
|
||||
f"+RequestWalltime = {cfg.request_walltime_s}\n"
|
||||
"+RequestWalltime = $(walltime)\n"
|
||||
f"accounting_group = {cfg.accounting_group}\n"
|
||||
f"{reqs_attrs}"
|
||||
f"output = {cfg.run_dir}/logs/$(plotid).out\n"
|
||||
f"error = {cfg.run_dir}/logs/$(plotid).err\n"
|
||||
f"output = {cfg.run_dir}/logs/$(plotid)__$(chunk).out\n"
|
||||
f"error = {cfg.run_dir}/logs/$(plotid)__$(chunk).err\n"
|
||||
f"log = {cfg.run_dir}/logs/condor.log\n"
|
||||
f"queue plotid from {ids_file}\n"
|
||||
f"queue plotid,chunk,walltime from {jobs_file}\n"
|
||||
)
|
||||
|
||||
|
||||
def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
|
||||
"""Write the wrapper script, plot-id list, and HTCondor submit description.
|
||||
def _job_walltimes(
|
||||
run_dir: Path, ids: list[str], n_chunks: int
|
||||
) -> list[tuple[str, int, int]]:
|
||||
"""``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``.
|
||||
|
||||
Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``;
|
||||
``chunkable=False`` specs (router diagnostics) always use the dataset
|
||||
total since they run as a single job regardless of ``n_chunks``.
|
||||
"""
|
||||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||||
jobs: list[tuple[str, int, int]] = []
|
||||
for spec_id in ids:
|
||||
chunkable = get_spec(spec_id).chunkable
|
||||
chunks = range(n_chunks) if chunkable else [0]
|
||||
for chunk in chunks:
|
||||
n_rows = meta.rows_per_chunk[chunk] if chunkable else meta.total_rows
|
||||
jobs.append((spec_id, chunk, estimate_runtime_s(spec_id, n_rows)))
|
||||
return jobs
|
||||
|
||||
|
||||
def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
|
||||
"""Write the wrapper script, (plot, chunk) job list, and HTCondor submit
|
||||
description.
|
||||
|
||||
Each catalog id gets ``cfg.n_chunks`` jobs, except ``chunkable=False``
|
||||
specs (the router diagnostics), which always get exactly one regardless of
|
||||
``cfg.n_chunks``. Every job's ``+RequestWalltime`` is estimated from its
|
||||
chunk's row count (``runtime_estimate.estimate_runtime_s``, requires
|
||||
``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``).
|
||||
Returns the submit description path (``<run_dir>/analyze.sub``). Does not
|
||||
submit — call ``condor_submit`` on the returned file.
|
||||
"""
|
||||
venv_giant = cfg.repo_dir / ".venv" / "bin" / "giant"
|
||||
if not venv_giant.exists():
|
||||
raise FileNotFoundError(
|
||||
f"{venv_giant} not found — condor jobs run it directly (no `uv` on "
|
||||
f"the worker image), so run `uv sync --extra cpu` in {cfg.repo_dir} "
|
||||
"before submitting."
|
||||
)
|
||||
|
||||
ids = ids or catalog_ids()
|
||||
run_dir = cfg.run_dir
|
||||
(run_dir / "logs").mkdir(parents=True, exist_ok=True)
|
||||
(run_dir / "reduced").mkdir(parents=True, exist_ok=True)
|
||||
(run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True)
|
||||
|
||||
wrapper = run_dir / "run_compute.sh"
|
||||
wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, run_dir=run_dir))
|
||||
wrapper.chmod(0o755)
|
||||
|
||||
ids_file = run_dir / "plotids.txt"
|
||||
ids_file.write_text("\n".join(ids) + "\n")
|
||||
jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
|
||||
jobs_file = run_dir / "jobs.txt"
|
||||
jobs_file.write_text("\n".join(f"{i},{k},{w}" for i, k, w in jobs) + "\n")
|
||||
|
||||
sub = run_dir / "analyze.sub"
|
||||
sub.write_text(_submit_description(cfg, wrapper, ids_file))
|
||||
sub.write_text(_submit_description(cfg, wrapper, jobs_file))
|
||||
return sub
|
||||
|
||||
@@ -16,6 +16,8 @@ efficiency" section):
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
@@ -58,6 +60,24 @@ def hist1d(
|
||||
return out
|
||||
|
||||
|
||||
def sum_merge(dicts: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
"""Elementwise-sum a list of sum-mergeable count/total dicts (JSON-safe keys).
|
||||
|
||||
Used to merge chunked ``hist1d``/``species_share``-style partials, whose
|
||||
values bin/group against edges or keys fixed by ``Context`` — a chunk's raw
|
||||
count dict is exactly a partial sum, so merging is a plain elementwise sum
|
||||
over the union of keys (a key absent from some chunk is all-zero there).
|
||||
Values may be per-bin count lists or plain scalar totals; both round-trip
|
||||
through ``np.asarray``/``.tolist()`` unchanged in shape.
|
||||
"""
|
||||
out: dict[str, np.ndarray] = {}
|
||||
for d in dicts:
|
||||
for k, v in d.items():
|
||||
arr = np.asarray(v)
|
||||
out[k] = arr.copy() if k not in out else out[k] + arr
|
||||
return {k: v.tolist() for k, v in out.items()}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-event scalar observables (one bounded group_by pass)
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -149,18 +169,19 @@ def transverse_expr() -> pl.Expr:
|
||||
return (tx**2 + ty**2 + tz**2).sqrt()
|
||||
|
||||
|
||||
def weighted_profile(
|
||||
def profile_partial(
|
||||
lf: pl.LazyFrame,
|
||||
coord: pl.Expr,
|
||||
edges: np.ndarray,
|
||||
weight: pl.Expr,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Event-averaged, ``weight``-summed profile of ``coord``, with an event-RMS band.
|
||||
"""One chunk's per-event x bin ``weight``-sum matrix: ``(event_ids, matrix)``.
|
||||
|
||||
One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin);
|
||||
collapsed in numpy to the per-bin mean over events and its event-to-event std
|
||||
(the band). ``coord``/``weight`` require the entry/axis columns attached.
|
||||
Returns ``(mean, std)``, each length ``len(edges)-1``.
|
||||
One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin).
|
||||
A chunk's matrix rows are only the events present in that chunk, so chunks'
|
||||
matrices stack cleanly with no cross-chunk lookup — this requires chunking
|
||||
to be event-disjoint (every row of an event lands in one chunk).
|
||||
``coord``/``weight`` require the entry/axis columns attached.
|
||||
"""
|
||||
lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1
|
||||
grid = (
|
||||
@@ -177,7 +198,35 @@ def weighted_profile(
|
||||
uniq, inv = np.unique(ev, return_inverse=True)
|
||||
mat = np.zeros((len(uniq), nbins), dtype=np.float64)
|
||||
np.add.at(mat, (inv, grid["_b"].to_numpy()), grid["_ws"].to_numpy())
|
||||
return mat.mean(axis=0), mat.std(axis=0)
|
||||
return uniq, mat
|
||||
|
||||
|
||||
def profile_finalize(mats: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Collapse per-chunk per-event x bin matrices into the final mean/std profile.
|
||||
|
||||
Chunks are event-disjoint, so row-wise concatenation of their matrices
|
||||
reconstructs the full per-event matrix; the mean/event-RMS collapse must
|
||||
happen once over that full matrix — an average of per-chunk means/stds
|
||||
would be wrong (chunks generally hold different numbers of events).
|
||||
Returns ``(mean, std)``, each length ``nbins``.
|
||||
"""
|
||||
full = np.concatenate(mats, axis=0)
|
||||
return full.mean(axis=0), full.std(axis=0)
|
||||
|
||||
|
||||
def weighted_profile(
|
||||
lf: pl.LazyFrame,
|
||||
coord: pl.Expr,
|
||||
edges: np.ndarray,
|
||||
weight: pl.Expr,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Single-pass profile of ``coord`` (mean +/- event-RMS band over events).
|
||||
|
||||
Convenience wrapper for the unchunked (whole-dataset) case; ``mean_std +
|
||||
profile_partial`` is what a chunked compute/finalize split uses instead.
|
||||
"""
|
||||
_, mat = profile_partial(lf, coord, edges, weight)
|
||||
return profile_finalize([mat])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
@@ -12,11 +12,14 @@ from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
# Reduced.kind values:
|
||||
# "overlay_hist" rollout vs reference density histogram over shared edges
|
||||
# "grouped_hist" one panel per group (energy/pdg/material), each an overlay
|
||||
# "profile" edep-weighted mean +/- event-RMS vs depth/radius, two series
|
||||
# "bar" per-category rollout vs reference bars (share / counts)
|
||||
# "single_hist" one series only (e.g. rollout leakage; reference has none)
|
||||
# "overlay_hist" rollout vs reference density histogram over shared edges
|
||||
# "grouped_hist" one panel per group (energy/pdg/material), each an overlay
|
||||
# "profile" edep-weighted mean +/- event-RMS vs depth/radius, two series
|
||||
# "bar" per-category rollout vs reference bars (share / counts)
|
||||
# "single_hist" one series only (e.g. rollout leakage; reference has none)
|
||||
# "router_gating" stacked mean MoE gate weight vs energy, rollout + reference
|
||||
# "router_share" stacked bar of MoE top-1 dispatch share by category
|
||||
# "unavailable" plot not applicable to this run (e.g. non-MoE checkpoint)
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -36,3 +39,27 @@ class Reduced:
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "Reduced":
|
||||
return cls(**json.loads(Path(path).read_text()))
|
||||
|
||||
|
||||
@dataclass
|
||||
class Partial:
|
||||
"""The raw, not-yet-finalized output of one ``(plot, chunk)`` compute job.
|
||||
|
||||
``data`` holds whatever shape that plot's ``PlotSpec.compute_partial``
|
||||
returns — a raw sum-mergeable count dict, or a raw per-event/per-secondary
|
||||
array to be concatenated across chunks — never a finished histogram/profile.
|
||||
``PlotSpec.finalize`` is the only thing that knows how to interpret it.
|
||||
"""
|
||||
|
||||
id: str
|
||||
family: str
|
||||
chunk: int
|
||||
data: dict
|
||||
|
||||
def save(self, path: str | Path) -> None:
|
||||
Path(path).parent.mkdir(parents=True, exist_ok=True)
|
||||
Path(path).write_text(json.dumps(asdict(self)))
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "Partial":
|
||||
return cls(**json.loads(Path(path).read_text()))
|
||||
|
||||
+126
-9
@@ -41,15 +41,46 @@ def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> Non
|
||||
ax.set_yscale("log")
|
||||
|
||||
|
||||
def _nn_params(run_meta: dict) -> dict:
|
||||
"""Flatten the rollout's model/training provenance for the figure subtitle."""
|
||||
params = {
|
||||
k: v for k, v in (run_meta.get("model_config") or {}).items() if v is not None
|
||||
}
|
||||
def _router_summary(model_config: dict) -> str:
|
||||
r = model_config.get("router") or {}
|
||||
if not r.get("enabled"):
|
||||
return "off"
|
||||
return f"{r.get('type', '?')}×{r.get('n_experts', '?')}"
|
||||
|
||||
|
||||
def _figure_params(run_meta: dict) -> dict:
|
||||
"""Curated run identity for the figure subtitle (``new_figure(params=...)``).
|
||||
|
||||
``run_meta``/each plot's own ``<id>.yaml`` (see ``_plot_metadata``) already
|
||||
carry every threaded model/training/rollout/dataset parameter for
|
||||
after-the-fact lookup — this picks only the handful that matter for
|
||||
telling figures apart at a glance while flipping through a gallery, since
|
||||
the subtitle is one unwrapped line of text. The last slot is
|
||||
architecture-conditional: flow/ddpm runs show the ODE ``steps`` used for
|
||||
this rollout, wgan runs show ``noise_dim`` instead since wgan sampling is
|
||||
single-pass and has no ODE step count.
|
||||
"""
|
||||
mc = run_meta.get("model_config") or {}
|
||||
mode = mc.get("mode")
|
||||
params: dict = {}
|
||||
if mc.get("hidden_dim") is not None:
|
||||
params["hidden_dim"] = mc["hidden_dim"]
|
||||
if mc.get("n_blocks") is not None:
|
||||
params["n_blocks"] = mc["n_blocks"]
|
||||
if mode is not None:
|
||||
params["mode"] = mode
|
||||
if mc.get("conditioning") is not None:
|
||||
params["conditioning"] = mc["conditioning"]
|
||||
params["router"] = _router_summary(mc)
|
||||
if run_meta.get("training_epoch") is not None:
|
||||
params["epoch"] = run_meta["training_epoch"]
|
||||
if run_meta.get("best_val_loss") is not None:
|
||||
params["best_val_loss"] = round(run_meta["best_val_loss"], 4)
|
||||
if mode == "wgan":
|
||||
if mc.get("noise_dim") is not None:
|
||||
params["noise_dim"] = mc["noise_dim"]
|
||||
elif run_meta.get("steps") is not None:
|
||||
params["steps"] = run_meta["steps"]
|
||||
return params
|
||||
|
||||
|
||||
@@ -137,18 +168,95 @@ def _render_bar(r: Reduced, params: dict):
|
||||
return fig
|
||||
|
||||
|
||||
def _render_router_gating(r: Reduced, params: dict):
|
||||
n_experts = r.payload["n_experts"]
|
||||
log_x = r.payload.get("log_x", False)
|
||||
fig, axes = ps.new_figure(
|
||||
"slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False
|
||||
)
|
||||
flat = axes.ravel()
|
||||
for ax, key in zip(flat, ("rollout", "reference")):
|
||||
side = r.payload.get(key, {})
|
||||
centers = np.asarray(side.get("centers", []))
|
||||
means = np.asarray(side.get("means", []))
|
||||
if len(centers) and means.size:
|
||||
cum = np.zeros(len(centers))
|
||||
for i in range(n_experts):
|
||||
ax.fill_between(
|
||||
centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}"
|
||||
)
|
||||
cum = cum + means[:, i]
|
||||
if log_x:
|
||||
ax.set_xscale("log")
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_title(_SERIES_LABELS[key], fontsize=8)
|
||||
ax.set_xlabel(r.xlabel)
|
||||
flat[0].set_ylabel("mean gate weight")
|
||||
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_router_share(r: Reduced, params: dict):
|
||||
categories = r.payload["categories"]
|
||||
n_experts = r.payload["n_experts"]
|
||||
x = np.arange(len(categories))
|
||||
present = [k for k in ("rollout", "reference") if k in r.payload]
|
||||
fig, axes = ps.new_figure(
|
||||
"slide-16x9",
|
||||
title=r.title,
|
||||
params=params,
|
||||
nrows=1,
|
||||
ncols=len(present),
|
||||
squeeze=False,
|
||||
)
|
||||
flat = axes.ravel()
|
||||
for ax, key in zip(flat, present):
|
||||
side = r.payload[key]
|
||||
shares = np.array([side[c] for c in categories]) # (n_cat, n_experts)
|
||||
bottom = np.zeros(len(categories))
|
||||
for i in range(n_experts):
|
||||
ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}")
|
||||
bottom += shares[:, i]
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(categories, rotation=45, ha="right")
|
||||
ax.set_ylim(0, 1)
|
||||
ax.set_title(_SERIES_LABELS[key], fontsize=8)
|
||||
flat[0].set_ylabel("share of rows dispatched to expert")
|
||||
ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_unavailable(r: Reduced, params: dict):
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
ax.axis("off")
|
||||
ax.text(
|
||||
0.5,
|
||||
0.5,
|
||||
r.payload.get("note", "not available"),
|
||||
ha="center",
|
||||
va="center",
|
||||
wrap=True,
|
||||
fontsize=10,
|
||||
transform=ax.transAxes,
|
||||
)
|
||||
return fig
|
||||
|
||||
|
||||
_RENDERERS = {
|
||||
"overlay_hist": _render_overlay,
|
||||
"single_hist": _render_single,
|
||||
"grouped_hist": _render_grouped,
|
||||
"profile": _render_profile,
|
||||
"bar": _render_bar,
|
||||
"router_gating": _render_router_gating,
|
||||
"router_share": _render_router_share,
|
||||
"unavailable": _render_unavailable,
|
||||
}
|
||||
|
||||
|
||||
def render(r: Reduced, run_meta: dict | None = None):
|
||||
"""Build the matplotlib figure for one reduced artifact (dispatch on kind)."""
|
||||
return _RENDERERS[r.kind](r, _nn_params(run_meta or {}))
|
||||
return _RENDERERS[r.kind](r, _figure_params(run_meta or {}))
|
||||
|
||||
|
||||
def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
|
||||
@@ -161,6 +269,11 @@ def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
|
||||
meta.update(r.meta)
|
||||
if "note" in r.payload:
|
||||
meta["note"] = r.payload["note"]
|
||||
if run_meta:
|
||||
# Every threaded model/training/rollout/dataset parameter, so a
|
||||
# single plot's metadata is self-contained for later comparison
|
||||
# without cross-referencing the run's root metadata.yaml.
|
||||
meta["parameters"] = {k: v for k, v in run_meta.items() if k != "title"}
|
||||
return meta
|
||||
|
||||
|
||||
@@ -226,12 +339,16 @@ def render_all(
|
||||
def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
|
||||
"""Render a prepped run directory: ``<run_dir>/reduced`` → ``<run_dir>/plots``.
|
||||
|
||||
Pulls the rollout provenance (checkpoint, paths, cutoffs) from
|
||||
``run_meta.json`` into every plot's gallery metadata.
|
||||
First joins every plot's chunk partials (``reduced_partial/<id>__*.json``)
|
||||
into ``reduced/<id>.json`` via ``merge_all`` — a no-op merge when the run
|
||||
wasn't chunked (``n_chunks=1``) — then pulls the rollout provenance
|
||||
(checkpoint, paths, cutoffs) from ``run_meta.json`` into every plot's
|
||||
gallery metadata and renders.
|
||||
"""
|
||||
from giant.analysis.condor import RunMeta
|
||||
from giant.analysis.condor import RunMeta, merge_all
|
||||
|
||||
run_dir = Path(run_dir)
|
||||
merge_all(run_dir)
|
||||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||||
run_meta = {
|
||||
"title": meta.title,
|
||||
|
||||
@@ -0,0 +1,339 @@
|
||||
"""Router gating diagnostic: where a MoE checkpoint's decision boundaries sit.
|
||||
|
||||
Unlike everything else in this package, this reduction needs a live PyTorch
|
||||
model — soft expert gate weights aren't columns in a rollout/predict parquet,
|
||||
they only exist by calling `Router.gate(cond_cont, cond_cat)` (see
|
||||
`giant.model.network.Router`) against the checkpoint that produced the
|
||||
rollout. That's a deliberate, narrow exception to the rest of the catalog's
|
||||
"polars/numpy only" contract; it still runs fine as a `compute-one` HTCondor
|
||||
job since torch is already installed there (the same env trains checkpoints).
|
||||
|
||||
The routing axis is fixed to pre-step energy: every router type at least
|
||||
indirectly depends on it (`EnergyRouter` reads it directly; `PdgRouter` and
|
||||
`ProcessRouter` correlate with it through the physics), and it's the one axis
|
||||
a reader can interpret without knowing the checkpoint's specific router
|
||||
config. `x` is binned into equal-population (quantile) bins rather than
|
||||
equal-width ones, since energy is heavy-tailed and equal-width bins would
|
||||
leave the upper end almost empty. Mean gate weight per bin is stacked as
|
||||
filled areas per expert — since `gate` rows are a partition of unity, the
|
||||
stack always fills exactly to 1, and the crossover bands are the router's
|
||||
soft decision boundaries (where two experts' means cross ~0.5).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis.grouping import pdg_label
|
||||
from giant.analysis.reduced import Reduced
|
||||
|
||||
if TYPE_CHECKING:
|
||||
import torch
|
||||
|
||||
from giant.data.transforms import Normalizer
|
||||
|
||||
_SAMPLE_ROWS = 200_000
|
||||
_N_BINS = 40
|
||||
_TOP_K_PROCESS = 8
|
||||
|
||||
_COLS = (
|
||||
"pre_x",
|
||||
"pre_y",
|
||||
"pre_z",
|
||||
"pre_E",
|
||||
"pre_dx",
|
||||
"pre_dy",
|
||||
"pre_dz",
|
||||
"layer_id",
|
||||
"pdg",
|
||||
"material",
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _RouterHandle:
|
||||
router: "torch.nn.Module"
|
||||
pdg_map: dict[int, int]
|
||||
mat_map: dict[str, int]
|
||||
cond_normalizer: "Normalizer"
|
||||
conditioning: str
|
||||
router_type: str
|
||||
|
||||
|
||||
def load_router(checkpoint: str | Path) -> _RouterHandle | None:
|
||||
"""Load a checkpoint's Stage-1 router, or None if it isn't a MoE checkpoint."""
|
||||
import torch
|
||||
|
||||
from giant.data.transforms import Normalizer
|
||||
from giant.model.network import build_models
|
||||
|
||||
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
||||
model_cfg = ckpt.get("model_config") or {}
|
||||
router_cfg = model_cfg.get("router")
|
||||
if not router_cfg or not router_cfg.get("enabled"):
|
||||
return None
|
||||
|
||||
stage1, _ = build_models(model_cfg)
|
||||
stage1.load_state_dict(ckpt["model"])
|
||||
stage1.eval()
|
||||
|
||||
return _RouterHandle(
|
||||
router=stage1.router,
|
||||
pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()},
|
||||
mat_map={str(k): v for k, v in ckpt["mat_map"].items()},
|
||||
cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]),
|
||||
conditioning=model_cfg.get("conditioning", "embedding"),
|
||||
router_type=router_cfg["type"],
|
||||
)
|
||||
|
||||
|
||||
def _subsample(
|
||||
lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()
|
||||
) -> pl.DataFrame:
|
||||
total = lf.select(pl.len()).collect(engine="streaming").item()
|
||||
if total > n:
|
||||
threshold = int(n / total * 2**32)
|
||||
lf = lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
|
||||
return lf.select(*_COLS, *extra_cols).collect(engine="streaming")
|
||||
|
||||
|
||||
def _gate_for_df(
|
||||
handle: _RouterHandle, df: pl.DataFrame
|
||||
) -> tuple[pl.DataFrame, np.ndarray]:
|
||||
"""(filtered df, gate_weights) for rows in ``df`` with a known pdg/material.
|
||||
|
||||
Rows whose species or material never appeared in the checkpoint's
|
||||
training vocab can't be embedded — dropped here the same way
|
||||
`giant.rollout`'s own known-pdg gate drops them at inference. The
|
||||
returned df keeps every original column (filtered to the same rows), so
|
||||
callers can key gate weights by any of them (energy, pdg, process, ...).
|
||||
"""
|
||||
import torch
|
||||
|
||||
from giant.data.transforms import build_cond_features
|
||||
|
||||
known = np.array(
|
||||
[
|
||||
int(p) in handle.pdg_map and str(m) in handle.mat_map
|
||||
for p, m in zip(df["pdg"].to_list(), df["material"].to_list())
|
||||
]
|
||||
)
|
||||
if not known.any():
|
||||
return df.clear(), np.zeros((0, handle.router.n_experts))
|
||||
df = df.filter(pl.Series(known, dtype=pl.Boolean))
|
||||
|
||||
data = {
|
||||
"pre_pos": np.column_stack(
|
||||
[df["pre_x"].to_numpy(), df["pre_y"].to_numpy(), df["pre_z"].to_numpy()]
|
||||
),
|
||||
"pre_E": df["pre_E"].to_numpy(),
|
||||
"pre_dir": np.column_stack(
|
||||
[df["pre_dx"].to_numpy(), df["pre_dy"].to_numpy(), df["pre_dz"].to_numpy()]
|
||||
),
|
||||
"layer_id": df["layer_id"].to_numpy(),
|
||||
"pdg": df["pdg"].to_numpy(),
|
||||
"material": df["material"].to_numpy(),
|
||||
}
|
||||
cond_cont, cond_cat = build_cond_features(
|
||||
data,
|
||||
handle.pdg_map,
|
||||
handle.mat_map,
|
||||
cond_normalizer=handle.cond_normalizer,
|
||||
conditioning=handle.conditioning,
|
||||
)
|
||||
with torch.no_grad():
|
||||
gate = handle.router.gate(
|
||||
torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()
|
||||
).numpy()
|
||||
return df, gate
|
||||
|
||||
|
||||
def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict:
|
||||
order = np.argsort(x)
|
||||
x_sorted, g_sorted = x[order], gate[order]
|
||||
edges = np.quantile(x_sorted, np.linspace(0, 1, n_bins + 1))
|
||||
edges[-1] = np.nextafter(edges[-1], np.inf) # include the max value
|
||||
bin_idx = np.clip(np.digitize(x_sorted, edges[1:-1]), 0, n_bins - 1)
|
||||
|
||||
n_experts = gate.shape[1]
|
||||
centers = np.full(n_bins, np.nan)
|
||||
means = np.full((n_bins, n_experts), np.nan)
|
||||
for b in range(n_bins):
|
||||
mask = bin_idx == b
|
||||
if mask.any():
|
||||
centers[b] = x_sorted[mask].mean()
|
||||
means[b] = g_sorted[mask].mean(axis=0)
|
||||
valid = ~np.isnan(centers)
|
||||
return {"centers": centers[valid].tolist(), "means": means[valid].tolist()}
|
||||
|
||||
|
||||
def _top1_shares(
|
||||
categories: np.ndarray, idx: np.ndarray, order: list, n_experts: int
|
||||
) -> dict[str, list[float]]:
|
||||
"""Fraction of each category's rows hard-dispatched to each expert.
|
||||
|
||||
Uses `Router.top1` (argmax), not the soft `gate` mean — grouped top-1
|
||||
dispatch is what `_route_forward` actually runs in eval mode (rollout,
|
||||
predict), so this answers "which expert does a photon/Compton step
|
||||
actually go through", not just its average soft weight.
|
||||
"""
|
||||
shares: dict[str, list[float]] = {}
|
||||
for key in order:
|
||||
mask = categories == key
|
||||
total = int(mask.sum())
|
||||
if total == 0:
|
||||
shares[str(key)] = [0.0] * n_experts
|
||||
continue
|
||||
counts = np.bincount(idx[mask], minlength=n_experts)
|
||||
shares[str(key)] = (counts / total).tolist()
|
||||
return shares
|
||||
|
||||
|
||||
_NOTE_NOT_MOE = (
|
||||
"checkpoint has no enabled MoE router (model.router.enabled is "
|
||||
"false/absent) — nothing to show"
|
||||
)
|
||||
|
||||
_TITLES = {
|
||||
"router_gating": "Router gating (mixture-of-experts decision boundaries)",
|
||||
"router_share_by_pdg": "Router expert share by particle species",
|
||||
"router_share_by_process": "Router expert share by physics process",
|
||||
}
|
||||
|
||||
|
||||
def _unavailable(spec_id: str) -> Reduced:
|
||||
return Reduced(
|
||||
id=spec_id,
|
||||
family="model",
|
||||
kind="unavailable",
|
||||
title=_TITLES[spec_id],
|
||||
xlabel="n/a",
|
||||
payload={"note": _NOTE_NOT_MOE},
|
||||
)
|
||||
|
||||
|
||||
def compute_router_gating(
|
||||
checkpoint: str | Path | None,
|
||||
r_phys: pl.LazyFrame,
|
||||
t_phys: pl.LazyFrame,
|
||||
seed: int = 0,
|
||||
) -> Reduced:
|
||||
"""`Reduced` for the router-gating figure, or an explanatory note if n/a."""
|
||||
handle = load_router(checkpoint) if checkpoint else None
|
||||
if handle is None:
|
||||
return _unavailable("router_gating")
|
||||
|
||||
sides: dict[str, dict] = {}
|
||||
for name, lf in (("rollout", r_phys), ("reference", t_phys)):
|
||||
df = _subsample(lf, _SAMPLE_ROWS, seed)
|
||||
df, gate = _gate_for_df(handle, df)
|
||||
x = df["pre_E"].to_numpy()
|
||||
sides[name] = (
|
||||
_quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
|
||||
)
|
||||
|
||||
return Reduced(
|
||||
id="router_gating",
|
||||
family="model",
|
||||
kind="router_gating",
|
||||
title=_TITLES["router_gating"],
|
||||
xlabel="pre-step energy [MeV]",
|
||||
payload={
|
||||
"router_type": handle.router_type,
|
||||
"n_experts": handle.router.n_experts,
|
||||
"log_x": True,
|
||||
**sides,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def compute_router_share_by_pdg(
|
||||
checkpoint: str | Path | None,
|
||||
r_phys: pl.LazyFrame,
|
||||
t_phys: pl.LazyFrame,
|
||||
top_pdgs: list[int],
|
||||
seed: int = 0,
|
||||
) -> Reduced:
|
||||
"""Stacked-bar share of each particle species dispatched to each expert."""
|
||||
handle = load_router(checkpoint) if checkpoint else None
|
||||
if handle is None:
|
||||
return _unavailable("router_share_by_pdg")
|
||||
|
||||
labels = [pdg_label(p) for p in top_pdgs]
|
||||
sides: dict[str, dict] = {}
|
||||
for name, lf in (("rollout", r_phys), ("reference", t_phys)):
|
||||
df = _subsample(lf, _SAMPLE_ROWS, seed)
|
||||
df, gate = _gate_for_df(handle, df)
|
||||
if len(df):
|
||||
idx = gate.argmax(axis=1)
|
||||
shares = _top1_shares(
|
||||
df["pdg"].to_numpy(), idx, top_pdgs, handle.router.n_experts
|
||||
)
|
||||
else:
|
||||
shares = {str(p): [0.0] * handle.router.n_experts for p in top_pdgs}
|
||||
sides[name] = {labels[i]: shares[str(p)] for i, p in enumerate(top_pdgs)}
|
||||
|
||||
return Reduced(
|
||||
id="router_share_by_pdg",
|
||||
family="model",
|
||||
kind="router_share",
|
||||
title=_TITLES["router_share_by_pdg"],
|
||||
xlabel="particle species",
|
||||
payload={
|
||||
"router_type": handle.router_type,
|
||||
"n_experts": handle.router.n_experts,
|
||||
"categories": labels,
|
||||
**sides,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def compute_router_share_by_process(
|
||||
checkpoint: str | Path | None,
|
||||
t_phys: pl.LazyFrame,
|
||||
seed: int = 0,
|
||||
top_k: int = _TOP_K_PROCESS,
|
||||
) -> Reduced:
|
||||
"""Stacked-bar share of each physics process dispatched to each expert.
|
||||
|
||||
Reference-only: ``process`` is the true post-step physics process — a
|
||||
label the rollout side has no equivalent of (see
|
||||
`giant.model.network.ProcessRouter`, which predicts it from pre-step
|
||||
conditioning alone, never observes it at eval time). This plot instead
|
||||
checks *after the fact*, on real data, how well the router's conditioning
|
||||
-based dispatch lines up with the true process.
|
||||
"""
|
||||
handle = load_router(checkpoint) if checkpoint else None
|
||||
if handle is None:
|
||||
return _unavailable("router_share_by_process")
|
||||
|
||||
df = _subsample(t_phys, _SAMPLE_ROWS, seed, extra_cols=("process",))
|
||||
df, gate = _gate_for_df(handle, df)
|
||||
if len(df):
|
||||
counts = df["process"].value_counts().sort("count", descending=True)
|
||||
order = counts["process"].to_list()[:top_k]
|
||||
idx = gate.argmax(axis=1)
|
||||
shares = _top1_shares(
|
||||
df["process"].to_numpy(), idx, order, handle.router.n_experts
|
||||
)
|
||||
else:
|
||||
order, shares = [], {}
|
||||
|
||||
return Reduced(
|
||||
id="router_share_by_process",
|
||||
family="model",
|
||||
kind="router_share",
|
||||
title=_TITLES["router_share_by_process"],
|
||||
xlabel="physics process",
|
||||
payload={
|
||||
"router_type": handle.router_type,
|
||||
"n_experts": handle.router.n_experts,
|
||||
"categories": order,
|
||||
"reference": {p: shares[p] for p in order},
|
||||
},
|
||||
)
|
||||
@@ -0,0 +1,117 @@
|
||||
"""Per-(plot, chunk) HTCondor walltime estimates for `giant analyze submit`.
|
||||
|
||||
Each catalog spec's compute cost is close to linear in the number of input
|
||||
rows a `compute-one` job streams over — every spec is one (or a couple of)
|
||||
streaming `group_by` pass(es) over the chunk (see `catalog.py`/`reduce.py`).
|
||||
`_COST_MODEL` below is ``spec_id -> (intercept_s, seconds_per_row)``.
|
||||
``n_rows`` is the combined rollout+reference row count of the job's input:
|
||||
the chunk's row count for `chunkable=True` specs, the whole dataset's for the
|
||||
three `chunkable=False` router specs (they always run as a single job
|
||||
regardless of chunk count).
|
||||
|
||||
Calibrated 2026-07-27 from real HTCondor timings (`condor_history`
|
||||
``RemoteWallClockTime``) of a production run: prediction ``563f5ee3``
|
||||
(PbWO4, 50 GeV) analyzed with ``--chunks 4`` against
|
||||
``giant/analysis/runtime_estimate.py``'s prior (local-synthetic-only) model —
|
||||
see the ``analysis-rollout-plots`` branch history for the raw data. That run's
|
||||
4 chunks came out at nearly identical row counts (~63-64M rows each, ~254M
|
||||
total), so this real data has no genuine row-count spread to fit a slope
|
||||
against — instead each spec's ``per_row`` here is a single line through the
|
||||
origin (``intercept=0``) hitting that spec's *median* wall-clock time across
|
||||
its 4 chunks at that run's row count. A handful of (spec, chunk) pairs showed
|
||||
3-8x spikes in one chunk only (e.g. ``marginal_edep_by_material``: 88, 88, 90,
|
||||
722s) — almost certainly shared ``/ceph`` contention from ~130 jobs landing on
|
||||
the filesystem at once right after submission, not a real per-row cost, so
|
||||
the median (not the max) was fit to avoid baking that noise into a rate that
|
||||
would then wrongly scale up with a bigger dataset. `RUNTIME_SAFETY_MARGIN` is
|
||||
deliberately generous (4x total) specifically to absorb that kind of
|
||||
contention spike instead. Rerun this calibration (pull fresh
|
||||
`condor_history`/`run_meta.json`, refit) if the catalog changes or timings
|
||||
drift — a synthetic local rebaseline via `scripts/profile_analysis_costs.py`
|
||||
is a reasonable fallback when no real cluster data is available yet, but
|
||||
undershoots real wall time badly (it can't see docker pull / `/ceph` I/O
|
||||
latency), which is exactly why this file moved off it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
|
||||
# Multiplicative pad applied to every job's estimated walltime. The one knob
|
||||
# this feature was asked to expose. Set generously (4x total, i.e. 3.0 here)
|
||||
# to absorb the shared-/ceph-contention spikes described above rather than
|
||||
# encoding them into individual specs' per-row rates.
|
||||
RUNTIME_SAFETY_MARGIN = 3.00
|
||||
|
||||
# Fixed per-job overhead (docker start, `.venv/bin/giant` startup, initial
|
||||
# `/ceph` read latency) — calibrated as the fastest observed real spec
|
||||
# (`leakage_fraction`, median 51s) rounded up, since even the cheapest spec
|
||||
# streams the whole chunk once.
|
||||
_FIXED_OVERHEAD_S = 60.0
|
||||
|
||||
# Router diagnostics run a live torch checkpoint (bounded inference over
|
||||
# <=200k subsampled rows, independent of chunk size) instead of a row-based
|
||||
# scan. Calibrated from the 3 real router jobs' observed wall times (119, 66,
|
||||
# 124s) — max minus _FIXED_OVERHEAD_S, on top of it.
|
||||
_ROUTER_FIXED_S = 64.0
|
||||
_ROUTER_IDS = frozenset(
|
||||
{"router_gating", "router_share_by_pdg", "router_share_by_process"}
|
||||
)
|
||||
|
||||
# Conservative fallback for any catalog id not in _COST_MODEL (e.g. a plot
|
||||
# added after the last calibration run) — the most expensive fitted per-row
|
||||
# rate observed, plus a small constant pad.
|
||||
_DEFAULT_COST = (5.0, 3.0e-6)
|
||||
|
||||
# spec_id -> (intercept_s, seconds_per_row), fit 2026-07-27 from real
|
||||
# HTCondor `RemoteWallClockTime` (see module docstring for methodology).
|
||||
_COST_MODEL: dict[str, tuple[float, float]] = {
|
||||
"marginal_step_length": (0.0, 5.199e-07),
|
||||
"marginal_step_length_by_energy": (0.0, 1.678e-06),
|
||||
"marginal_step_length_by_pdg": (0.0, 4.569e-07),
|
||||
"marginal_step_length_by_material": (0.0, 2.269e-06),
|
||||
"marginal_edep": (0.0, 2.804e-06),
|
||||
"marginal_edep_by_energy": (0.0, 1.386e-06),
|
||||
"marginal_edep_by_pdg": (0.0, 4.490e-07),
|
||||
"marginal_edep_by_material": (0.0, 4.333e-07),
|
||||
"marginal_delta_e": (0.0, 5.042e-07),
|
||||
"marginal_delta_e_by_energy": (0.0, 1.678e-06),
|
||||
"marginal_delta_e_by_pdg": (0.0, 4.727e-07),
|
||||
"marginal_delta_e_by_material": (0.0, 4.490e-07),
|
||||
"marginal_post_E": (0.0, 4.805e-07),
|
||||
"marginal_post_E_by_energy": (0.0, 1.284e-06),
|
||||
"marginal_post_E_by_pdg": (0.0, 4.569e-07),
|
||||
"marginal_post_E_by_material": (0.0, 4.490e-07),
|
||||
"marginal_cos_scatter": (0.0, 5.436e-07),
|
||||
"marginal_cos_scatter_by_energy": (0.0, 1.363e-06),
|
||||
"marginal_cos_scatter_by_pdg": (0.0, 4.727e-07),
|
||||
"marginal_cos_scatter_by_material": (0.0, 4.727e-07),
|
||||
"event_total_edep": (0.0, 4.490e-07),
|
||||
"event_total_edep_by_energy": (0.0, 4.411e-07),
|
||||
"event_mean_length": (0.0, 4.333e-07),
|
||||
"event_n_steps": (0.0, 4.569e-07),
|
||||
"shower_longitudinal": (0.0, 2.348e-06),
|
||||
"shower_transverse": (0.0, 2.899e-06),
|
||||
"species_edep_share": (0.0, 4.333e-07),
|
||||
"leakage_fraction": (0.0, 0.0),
|
||||
"sec_count_per_event": (0.0, 4.727e-07),
|
||||
"sec_count_per_species": (0.0, 4.963e-07),
|
||||
"sec_energy": (0.0, 4.727e-07),
|
||||
"sec_cos_angle": (0.0, 2.749e-06),
|
||||
}
|
||||
|
||||
|
||||
def estimate_runtime_s(spec_id: str, n_rows: int) -> int:
|
||||
"""Estimated `+RequestWalltime` (seconds) for one (plot, chunk) job.
|
||||
|
||||
``n_rows`` is the rollout+reference row count of that job's input slice.
|
||||
Includes `_FIXED_OVERHEAD_S`/`_ROUTER_FIXED_S` and `RUNTIME_SAFETY_MARGIN`
|
||||
— callers should pass this straight through to the submit description.
|
||||
"""
|
||||
if spec_id in _ROUTER_IDS:
|
||||
compute_s = _ROUTER_FIXED_S
|
||||
else:
|
||||
intercept, per_row = _COST_MODEL.get(spec_id, _DEFAULT_COST)
|
||||
compute_s = intercept + per_row * n_rows
|
||||
total = _FIXED_OVERHEAD_S + compute_s
|
||||
return math.ceil(total * (1 + RUNTIME_SAFETY_MARGIN))
|
||||
@@ -106,18 +106,27 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
can push their own narrow projection into the parquet read — the single
|
||||
biggest lever on a larger-than-RAM file. ``pl.LazyFrame`` inputs pass straight
|
||||
through (used by tests).
|
||||
|
||||
``pdg`` is cast to a canonical ``Int64`` here: the rollout writer and the
|
||||
reference file's upstream ROOT→parquet conversion don't agree on integer
|
||||
width, and an uncast mismatch only surfaces later as a ``pl.concat``
|
||||
``SchemaError`` (e.g. in ``build_context``'s pdg-count merge).
|
||||
"""
|
||||
if isinstance(source, pl.LazyFrame):
|
||||
return source
|
||||
return source.with_columns(pl.col("pdg").cast(pl.Int64))
|
||||
path = Path(source)
|
||||
if side is Side.rollout:
|
||||
_check_rollout_metadata(path)
|
||||
return pl.scan_parquet(path)
|
||||
# The reference (a rollout's seed `dataset`) may be a directory of parquet
|
||||
# shards rather than a single file — scan them all.
|
||||
if path.is_dir():
|
||||
return pl.scan_parquet(str(path / "**/*.parquet"))
|
||||
return pl.scan_parquet(path)
|
||||
lf = pl.scan_parquet(path)
|
||||
else:
|
||||
# The reference (a rollout's seed `dataset`) may be a directory of
|
||||
# parquet shards rather than a single file — scan them all.
|
||||
lf = (
|
||||
pl.scan_parquet(str(path / "**/*.parquet"))
|
||||
if path.is_dir()
|
||||
else pl.scan_parquet(path)
|
||||
)
|
||||
return lf.with_columns(pl.col("pdg").cast(pl.Int64))
|
||||
|
||||
|
||||
def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
|
||||
+82
-16
@@ -1188,6 +1188,9 @@ def rollout(
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
|
||||
gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
|
||||
training_cfg = gconfig.load_checkpoint_config(checkpoint)
|
||||
|
||||
model_cfg = ckpt["model_config"]
|
||||
conditioning = model_cfg.get("conditioning", "embedding")
|
||||
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
|
||||
@@ -1270,17 +1273,26 @@ def rollout(
|
||||
"max_steps": max_steps,
|
||||
"steps": steps,
|
||||
"max_tracks_per_event": max_tracks_per_event,
|
||||
"escape_threshold": escape_threshold,
|
||||
"n_events": n_events,
|
||||
"n_seed_events": int(len(seeds["event_id"])),
|
||||
"model_config": {
|
||||
"mode": model_cfg.get("mode", "flow"),
|
||||
"hidden_dim": model_cfg.get("hidden_dim"),
|
||||
"n_blocks": model_cfg.get("n_blocks"),
|
||||
"emb_dim": model_cfg.get("emb_dim"),
|
||||
"dropout": model_cfg.get("dropout"),
|
||||
"conditioning": conditioning,
|
||||
},
|
||||
"weights": weights.value,
|
||||
"batch_size": batch_size,
|
||||
"device": str(_device),
|
||||
"rollout_seed": seed,
|
||||
"n_rows": summary["n_rows"],
|
||||
"termination_reason_counts": summary["termination_reason_counts"],
|
||||
# Full architecture spec baked into the checkpoint — includes the
|
||||
# entire router sub-dict, not just a hand-picked subset, so any
|
||||
# model knob (router type/n_experts, noise_dim, vocab sizes, ...)
|
||||
# is available downstream without touching this command again.
|
||||
"model_config": dict(model_cfg),
|
||||
"training_epoch": ckpt.get("epoch"),
|
||||
"best_val_loss": ckpt.get("best_val_loss"),
|
||||
# [train]/[meta] from the sibling config.toml (giant.config.save_config)
|
||||
# — empty dicts if the checkpoint has no config.toml next to it.
|
||||
"training_config": dict(training_cfg.get("train", {})),
|
||||
"training_meta": dict(training_cfg.get("meta", {})),
|
||||
}
|
||||
)
|
||||
ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False))
|
||||
@@ -1310,12 +1322,18 @@ def analyze_prep(
|
||||
typer.Option(
|
||||
"--run-dir",
|
||||
"-o",
|
||||
help="Override the run directory (default: next to the rollout parquet)",
|
||||
help="Override the run directory (default: <cwd>/analysis_runs/analysis_<id>)",
|
||||
),
|
||||
] = None,
|
||||
n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
|
||||
n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
|
||||
top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
|
||||
chunks: Annotated[
|
||||
int,
|
||||
typer.Option(
|
||||
"--chunks", help="Split each plot's data into this many event_id chunks"
|
||||
),
|
||||
] = 1,
|
||||
) -> None:
|
||||
"""Read the rollout YAML → shared.json + run_meta.json in the run directory."""
|
||||
from giant.analysis import prep
|
||||
@@ -1323,6 +1341,8 @@ def analyze_prep(
|
||||
path = prep(
|
||||
rollout_yaml,
|
||||
run_dir,
|
||||
n_chunks=chunks,
|
||||
default_base=Path.cwd() / "analysis_runs",
|
||||
n_energy_bins=n_energy_bins,
|
||||
n_marginal_bins=n_marginal_bins,
|
||||
top_k_pdg=top_k_pdg,
|
||||
@@ -1338,11 +1358,34 @@ def analyze_compute_one(
|
||||
run_dir: Annotated[
|
||||
Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
|
||||
],
|
||||
chunk: Annotated[
|
||||
int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")
|
||||
] = 0,
|
||||
) -> None:
|
||||
"""Run one plot's streaming reduction (this is what each condor job runs)."""
|
||||
"""Run one (plot, chunk)'s streaming reduction (this is what each condor job runs)."""
|
||||
from giant.analysis import compute_one
|
||||
|
||||
path = compute_one(id, run_dir)
|
||||
path = compute_one(id, run_dir, chunk_index=chunk)
|
||||
typer.echo(f"wrote {path}")
|
||||
|
||||
|
||||
@analyze_app.command("merge-one")
|
||||
def analyze_merge_one(
|
||||
id: Annotated[
|
||||
str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")
|
||||
],
|
||||
run_dir: Annotated[
|
||||
Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")
|
||||
],
|
||||
) -> None:
|
||||
"""Merge one plot's chunk partials into its final reduced JSON.
|
||||
|
||||
Runs automatically as part of `analyze render`; useful standalone to
|
||||
debug a specific plot without re-rendering everything.
|
||||
"""
|
||||
from giant.analysis import merge_one
|
||||
|
||||
path = merge_one(id, run_dir)
|
||||
typer.echo(f"wrote {path}")
|
||||
|
||||
|
||||
@@ -1380,26 +1423,48 @@ def analyze_submit(
|
||||
accounting_group: Annotated[str, typer.Option("--accounting-group")],
|
||||
run_dir: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option("--run-dir", "-o", help="Override the run directory"),
|
||||
typer.Option(
|
||||
"--run-dir",
|
||||
"-o",
|
||||
help="Override the run directory (default: <cwd>/analysis_runs/analysis_<id>)",
|
||||
),
|
||||
] = None,
|
||||
docker_image: Annotated[
|
||||
str, typer.Option("--docker-image")
|
||||
] = "mschnepf/slc7-condocker",
|
||||
request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 4096,
|
||||
] = "cverstege/alma9-gridjob",
|
||||
request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 8192,
|
||||
remote: Annotated[
|
||||
bool,
|
||||
typer.Option("--remote/--local", help="+RemoteJob vs ProvidesETPResources"),
|
||||
] = False,
|
||||
chunks: Annotated[
|
||||
int,
|
||||
typer.Option(
|
||||
"--chunks",
|
||||
help="Split each plot's data into this many event_id chunks/jobs",
|
||||
),
|
||||
] = 1,
|
||||
n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4,
|
||||
n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50,
|
||||
top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6,
|
||||
dry_run: Annotated[
|
||||
bool, typer.Option("--dry-run", help="Write files but don't condor_submit")
|
||||
] = False,
|
||||
) -> None:
|
||||
"""prep + write the HTCondor submit description (one job per plot), then submit."""
|
||||
"""prep + write the HTCondor submit description (one job per plot x chunk), then submit."""
|
||||
import subprocess
|
||||
|
||||
from giant.analysis import SubmitConfig, prep, write_submit
|
||||
|
||||
path = prep(rollout_yaml, run_dir)
|
||||
path = prep(
|
||||
rollout_yaml,
|
||||
run_dir,
|
||||
n_chunks=chunks,
|
||||
default_base=Path.cwd() / "analysis_runs",
|
||||
n_energy_bins=n_energy_bins,
|
||||
n_marginal_bins=n_marginal_bins,
|
||||
top_k_pdg=top_k_pdg,
|
||||
)
|
||||
cfg = SubmitConfig(
|
||||
run_dir=path,
|
||||
accounting_group=accounting_group,
|
||||
@@ -1407,6 +1472,7 @@ def analyze_submit(
|
||||
docker_image=docker_image,
|
||||
request_memory_mb=request_memory,
|
||||
remote=remote,
|
||||
n_chunks=chunks,
|
||||
)
|
||||
sub = write_submit(cfg)
|
||||
typer.echo(f"run directory: {path}")
|
||||
|
||||
@@ -189,6 +189,21 @@ def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None:
|
||||
)
|
||||
|
||||
|
||||
def load_checkpoint_config(ckpt_path: str | Path) -> dict:
|
||||
"""Load the full ``[train]``/``[model]``/``[meta]`` config.toml written
|
||||
alongside a checkpoint by ``save_config``.
|
||||
|
||||
Returns ``{}`` if no config.toml sits next to the checkpoint (older runs,
|
||||
or a checkpoint moved without its sidecar) — this is best-effort
|
||||
provenance for threading into a rollout's YAML sidecar, not a hard
|
||||
requirement for using the checkpoint itself.
|
||||
"""
|
||||
config_path = Path(ckpt_path).parent / "config.toml"
|
||||
if not config_path.exists():
|
||||
return {}
|
||||
return load_toml(config_path)
|
||||
|
||||
|
||||
def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None:
|
||||
"""Look for a config.toml next to a checkpoint and warn on a git_hash mismatch.
|
||||
|
||||
|
||||
@@ -228,6 +228,54 @@ class _WelfordAccumulator:
|
||||
return norm
|
||||
|
||||
|
||||
class _ReservoirSampler:
|
||||
"""Uniform random sample of a fixed capacity drawn from a data stream.
|
||||
|
||||
Algorithm R (Vitter 1985), vectorized per chunk so it stays cheap over
|
||||
hundreds of millions of rows: use to get a representative subsample of
|
||||
a column for a distribution estimate (e.g. quantiles) without
|
||||
materializing the full column.
|
||||
|
||||
sampler = _ReservoirSampler(capacity=100_000)
|
||||
for chunk in data:
|
||||
sampler.update(chunk)
|
||||
sample = sampler.sample
|
||||
"""
|
||||
|
||||
def __init__(self, capacity: int, seed: int = 0) -> None:
|
||||
self.capacity = capacity
|
||||
self.n_seen = 0
|
||||
self._rng = np.random.default_rng(seed)
|
||||
self._reservoir = np.empty(0, dtype=np.float64)
|
||||
|
||||
def update(self, values: np.ndarray) -> None:
|
||||
values = np.asarray(values, dtype=np.float64).reshape(-1)
|
||||
if values.size == 0:
|
||||
return
|
||||
n_before = self.n_seen
|
||||
if n_before < self.capacity:
|
||||
take = min(values.size, self.capacity - n_before)
|
||||
self._reservoir = np.concatenate([self._reservoir, values[:take]])
|
||||
values = values[take:]
|
||||
n_before += take
|
||||
self.n_seen = n_before + values.size
|
||||
if values.size == 0 or self.capacity == 0:
|
||||
return
|
||||
# remaining elements are past the fill phase: element at 1-based
|
||||
# stream position j replaces a uniformly random reservoir slot with
|
||||
# probability capacity/j, which yields a uniform sample overall.
|
||||
positions = n_before + np.arange(1, values.size + 1)
|
||||
accept = self._rng.random(values.size) < (self.capacity / positions)
|
||||
accept_idx = np.nonzero(accept)[0]
|
||||
if accept_idx.size > 0:
|
||||
slots = self._rng.integers(0, self.capacity, size=accept_idx.size)
|
||||
self._reservoir[slots] = values[accept_idx]
|
||||
|
||||
@property
|
||||
def sample(self) -> np.ndarray:
|
||||
return self._reservoir.astype(np.float32)
|
||||
|
||||
|
||||
def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
|
||||
"""World-frame unit vector pointing from pre_pos to post_pos.
|
||||
|
||||
@@ -530,11 +578,44 @@ def build_cond_features(
|
||||
cond_cat = np.column_stack([pdg_idx, mat_idx])
|
||||
|
||||
if cond_normalizer is not None:
|
||||
cond_cont = cond_normalizer.transform(cond_cont)
|
||||
cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning)
|
||||
|
||||
return cond_cont, cond_cat
|
||||
|
||||
|
||||
def _cond_normalizer_transform(
|
||||
cond_cont: np.ndarray, cond_normalizer: "Normalizer", conditioning: str
|
||||
) -> np.ndarray:
|
||||
"""Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed.
|
||||
|
||||
Checkpoints trained before physical-property conditioning (``COND_DIM``
|
||||
8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond
|
||||
normalizer, fit before ``build_cond_features`` grew the extra physical
|
||||
columns. In "embedding" mode those columns are never read by
|
||||
``ConditionEncoder`` (``giant/model/network.py``), so padding the missing
|
||||
entries with mean=0/std=1 is a safe no-op that keeps such checkpoints
|
||||
usable under the current, always-``COND_DIM``-wide contract. In
|
||||
"physical" mode the physical columns are load-bearing, so a mismatch
|
||||
there is a real incompatibility, not something to paper over.
|
||||
"""
|
||||
mean, std = cond_normalizer.mean, cond_normalizer.std
|
||||
assert mean is not None and std is not None, "Normalizer not fitted"
|
||||
width = cond_cont.shape[-1]
|
||||
if mean.shape[-1] < width:
|
||||
if conditioning != "embedding":
|
||||
raise ValueError(
|
||||
f"cond normalizer has {mean.shape[-1]} columns, expected "
|
||||
f"{width}, and conditioning={conditioning!r} reads the "
|
||||
"physical columns directly — this checkpoint predates "
|
||||
"physical-property conditioning and can't be safely padded; "
|
||||
"retrain it under the current code."
|
||||
)
|
||||
pad = width - mean.shape[-1]
|
||||
mean = np.concatenate([mean, np.zeros(pad, dtype=mean.dtype)])
|
||||
std = np.concatenate([std, np.ones(pad, dtype=std.dtype)])
|
||||
return ((cond_cont - mean) / std).astype(np.float32)
|
||||
|
||||
|
||||
def build_features(
|
||||
data: dict[str, np.ndarray],
|
||||
pdg_map: dict[int, int],
|
||||
|
||||
+19
-5
@@ -1,6 +1,7 @@
|
||||
import inspect
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
@@ -600,10 +601,14 @@ class EnergyRouter(Router):
|
||||
"""Soft turn-on gate over normalized pre-step log-energy.
|
||||
|
||||
Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or
|
||||
fixed) 1-D centers, initialized spread across [-2, 2] — roughly the
|
||||
z-normalized energy range. `gate(e) = softmax_i(-(e - c_i)^2 / tau)`,
|
||||
differentiable in e; as tau -> 0 this hardens to nearest-center
|
||||
(Voronoi) selection, which is exactly what `top1` uses at eval.
|
||||
fixed) 1-D centers. By default initialized spread evenly across
|
||||
[-2, 2] — an assumed-uniform z-normalized energy range that may not
|
||||
match the true (often skewed) distribution and can leave experts
|
||||
overlapping instead of partitioning the range; pass `centers_init` to
|
||||
seed them from data (e.g. energy quantiles) instead.
|
||||
`gate(e) = softmax_i(-(e - c_i)^2 / tau)`, differentiable in e; as
|
||||
tau -> 0 this hardens to nearest-center (Voronoi) selection, which is
|
||||
exactly what `top1` uses at eval.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -612,11 +617,20 @@ class EnergyRouter(Router):
|
||||
temperature: float = 0.5,
|
||||
learn_centers: bool = True,
|
||||
energy_idx: int = 3,
|
||||
centers_init: Sequence[float] | None = None,
|
||||
) -> None:
|
||||
super().__init__(n_experts)
|
||||
self.temperature = temperature
|
||||
self.energy_idx = energy_idx
|
||||
centers = torch.linspace(-2.0, 2.0, n_experts)
|
||||
if centers_init is None:
|
||||
centers = torch.linspace(-2.0, 2.0, n_experts)
|
||||
else:
|
||||
if len(centers_init) != n_experts:
|
||||
raise ValueError(
|
||||
f"centers_init has {len(centers_init)} values, "
|
||||
f"expected n_experts={n_experts}"
|
||||
)
|
||||
centers = torch.tensor(list(centers_init), dtype=torch.float32)
|
||||
if learn_centers:
|
||||
self.centers = nn.Parameter(centers)
|
||||
else:
|
||||
|
||||
+27
-1
@@ -20,7 +20,7 @@ from giant.data.loader import (
|
||||
build_index_maps_from_files,
|
||||
build_process_map_from_files,
|
||||
)
|
||||
from giant.data.transforms import build_features, _WelfordAccumulator
|
||||
from giant.data.transforms import build_features, _WelfordAccumulator, _ReservoirSampler
|
||||
from giant.data.dataset import make_event_split, StreamingStepsDataset
|
||||
from giant.model.network import build_models, build_critics
|
||||
from giant.train import train as run_training
|
||||
@@ -83,6 +83,18 @@ def run_train_job(
|
||||
cond_acc = _WelfordAccumulator(COND_DIM)
|
||||
tgt_acc = _WelfordAccumulator(X_DIM)
|
||||
sec_phys_acc = _WelfordAccumulator(PARTICLE_PHYS_DIM)
|
||||
# EnergyRouter's default center spread (linspace over [-2, 2]) assumes
|
||||
# the z-normalized energy column is roughly uniform, which real energy
|
||||
# spectra rarely are — collect a reservoir sample here (reusing this
|
||||
# same pass, not a second scan) so centers can instead be seeded from
|
||||
# actual data quantiles below.
|
||||
energy_router_active = (
|
||||
router_cfg.get("enabled") and router_cfg.get("type") == "energy"
|
||||
)
|
||||
energy_idx = router_cfg.get("energy_idx", 3)
|
||||
energy_sampler = (
|
||||
_ReservoirSampler(capacity=100_000) if energy_router_active else None
|
||||
)
|
||||
for path in files:
|
||||
for chunk in iter_file_chunks(path):
|
||||
mask = np.isin(chunk["event_id"], events_arr)
|
||||
@@ -99,6 +111,8 @@ def run_train_job(
|
||||
)
|
||||
cond_acc.update(cond_cont)
|
||||
tgt_acc.update(target_s1)
|
||||
if energy_sampler is not None:
|
||||
energy_sampler.update(cond_cont[:, energy_idx])
|
||||
sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
|
||||
sec_phys = sec_cont[:, :, 4:6][sec_valid]
|
||||
if len(sec_phys) > 0:
|
||||
@@ -107,6 +121,18 @@ def run_train_job(
|
||||
tgt_norm = tgt_acc.to_normalizer()
|
||||
sec_phys_norm = sec_phys_acc.to_normalizer()
|
||||
|
||||
if energy_sampler is not None and energy_sampler.n_seen > 0:
|
||||
assert cond_norm.mean is not None and cond_norm.std is not None
|
||||
normalized_sample = (
|
||||
energy_sampler.sample - cond_norm.mean[energy_idx]
|
||||
) / cond_norm.std[energy_idx]
|
||||
quantiles = np.linspace(0.0, 1.0, router_cfg["n_experts"])
|
||||
centers_init = np.quantile(normalized_sample, quantiles).astype(np.float32)
|
||||
router_cfg["centers_init"] = centers_init.tolist()
|
||||
echo(
|
||||
f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
|
||||
)
|
||||
|
||||
train_ds = StreamingStepsDataset(
|
||||
files=files,
|
||||
split_events=train_events,
|
||||
|
||||
@@ -0,0 +1,242 @@
|
||||
"""Benchmark `giant analyze compute-one`'s per-job cost against synthetic data.
|
||||
|
||||
Generates mock rollout+reference parquet files at a few row counts, times
|
||||
`compute_reduced` for every chunkable catalog spec at each size (a single
|
||||
chunk covering the whole mock file), fits a straight line (intercept, seconds
|
||||
per row) through the timings, and prints the result as a Python dict literal
|
||||
ready to paste into `giant/analysis/runtime_estimate.py::_COST_MODEL`.
|
||||
|
||||
The three `chunkable=False` router specs (`router_gating`,
|
||||
`router_share_by_pdg`, `router_share_by_process`) need a live MoE checkpoint
|
||||
to do any real work; without one (this machine has no `/ceph` access, so no
|
||||
real checkpoint) they short-circuit almost instantly and are excluded here —
|
||||
see `runtime_estimate.py`'s `_ROUTER_FIXED_S` for how those are handled
|
||||
instead.
|
||||
|
||||
Usage: ``uv run python scripts/profile_analysis_costs.py``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
from pathlib import Path
|
||||
from tempfile import TemporaryDirectory
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis.catalog import catalog_ids, get_spec
|
||||
from giant.analysis.condor import compute_reduced
|
||||
from giant.analysis.context import build_context
|
||||
|
||||
# Row counts (per side) to benchmark at. Kept in local memory/CPU range so the
|
||||
# whole sweep finishes in about a minute; the fit is linear so it extrapolates
|
||||
# fine to real multi-GB rollouts.
|
||||
SIDE_ROW_COUNTS = [20_000, 100_000, 500_000, 2_000_000]
|
||||
|
||||
_MATERIALS = ["G4_PbWO4", "G4_Pb", "G4_lAr", "G4_Si"]
|
||||
_PDGS = [11, -11, 22, 2112, 2212, 211, -211, 13]
|
||||
_ROUTER_IDS = {"router_gating", "router_share_by_pdg", "router_share_by_process"}
|
||||
|
||||
|
||||
def _unit_vectors(n: int, rng: np.random.Generator) -> np.ndarray:
|
||||
v = rng.normal(size=(n, 3))
|
||||
return v / np.linalg.norm(v, axis=1, keepdims=True)
|
||||
|
||||
|
||||
def _ragged_lists(k: np.ndarray, rng: np.random.Generator, lo: float, hi: float):
|
||||
total = int(k.sum())
|
||||
flat = rng.uniform(lo, hi, size=total)
|
||||
idx = np.cumsum(k)[:-1]
|
||||
return [arr.tolist() for arr in np.split(flat, idx)]
|
||||
|
||||
|
||||
def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
|
||||
rng = np.random.default_rng(seed)
|
||||
event_id = rng.integers(0, n_events, size=n)
|
||||
is_secondary = rng.random(n) < 0.15 # generation>0, step_no==0 birth rows
|
||||
is_synthetic = rng.random(n) < 0.05 # bookkeeping termination rows
|
||||
|
||||
pre_E = rng.lognormal(mean=3.0, sigma=1.5, size=n)
|
||||
edep = rng.uniform(0, 1, size=n) * pre_E * 0.3
|
||||
post_E = np.clip(pre_E - edep, 0.0, None)
|
||||
pre_dir = _unit_vectors(n, rng)
|
||||
post_dir = _unit_vectors(n, rng)
|
||||
pos = rng.uniform(-50, 300, size=(n, 3))
|
||||
step_length = rng.uniform(0.1, 10.0, size=n)
|
||||
post_pos = pos + pre_dir * step_length[:, None]
|
||||
|
||||
reasons = np.where(
|
||||
is_synthetic,
|
||||
rng.choice(["escaped", "energy_cutoff", "max_steps", "unknown_pdg"], size=n),
|
||||
"natural_end",
|
||||
)
|
||||
|
||||
return pl.DataFrame(
|
||||
{
|
||||
"event_id": event_id,
|
||||
"track_id": rng.integers(0, 5, size=n),
|
||||
"parent_id": np.where(is_secondary, 0, -1),
|
||||
"generation": is_secondary.astype(np.int64),
|
||||
"step_no": np.where(is_secondary, 0, rng.integers(0, 20, size=n)),
|
||||
"pdg": rng.choice(_PDGS, size=n),
|
||||
"pre_x": pos[:, 0],
|
||||
"pre_y": pos[:, 1],
|
||||
"pre_z": pos[:, 2],
|
||||
"pre_E": pre_E,
|
||||
"pre_dx": pre_dir[:, 0],
|
||||
"pre_dy": pre_dir[:, 1],
|
||||
"pre_dz": pre_dir[:, 2],
|
||||
"post_x": post_pos[:, 0],
|
||||
"post_y": post_pos[:, 1],
|
||||
"post_z": post_pos[:, 2],
|
||||
"post_E": post_E,
|
||||
"post_dx": post_dir[:, 0],
|
||||
"post_dy": post_dir[:, 1],
|
||||
"post_dz": post_dir[:, 2],
|
||||
"edep": np.where(
|
||||
is_synthetic, np.where(reasons == "escaped", 0.0, pre_E), edep
|
||||
),
|
||||
"step_length": np.where(is_synthetic, 0.0, step_length),
|
||||
"material": rng.choice(_MATERIALS, size=n),
|
||||
"layer_id": rng.integers(0, 30, size=n),
|
||||
"n_sec_pred": rng.integers(0, 4, size=n),
|
||||
"termination_reason": reasons,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
|
||||
rng = np.random.default_rng(seed + 1)
|
||||
event_id = rng.integers(0, n_events, size=n)
|
||||
pre_E = rng.lognormal(mean=3.0, sigma=1.5, size=n)
|
||||
edep = rng.uniform(0, 1, size=n) * pre_E * 0.3
|
||||
post_E = np.clip(pre_E - edep, 0.0, None)
|
||||
pre_dir = _unit_vectors(n, rng)
|
||||
post_dir = _unit_vectors(n, rng)
|
||||
pos = rng.uniform(-50, 300, size=(n, 3))
|
||||
step_length = rng.uniform(0.1, 10.0, size=n)
|
||||
post_pos = pos + pre_dir * step_length[:, None]
|
||||
|
||||
k = rng.poisson(0.3, size=n).clip(max=5).astype(np.int64)
|
||||
sec_pdg = _ragged_lists(k, rng, 0, 1) # placeholder, overwritten below
|
||||
sec_E = _ragged_lists(k, rng, 0.1, 50.0)
|
||||
sec_dx = _ragged_lists(k, rng, -1.0, 1.0)
|
||||
sec_dy = _ragged_lists(k, rng, -1.0, 1.0)
|
||||
sec_dz = _ragged_lists(k, rng, -1.0, 1.0)
|
||||
total = int(k.sum())
|
||||
flat_pdg = rng.choice(_PDGS, size=total).tolist()
|
||||
idx = np.cumsum(k)[:-1]
|
||||
sec_pdg = [list(x) for x in np.split(np.array(flat_pdg), idx)]
|
||||
|
||||
return pl.DataFrame(
|
||||
{
|
||||
"event_id": event_id,
|
||||
"track_id": rng.integers(0, 5, size=n),
|
||||
"step_no": rng.integers(0, 20, size=n),
|
||||
"pdg": rng.choice(_PDGS, size=n),
|
||||
"pre_x": pos[:, 0],
|
||||
"pre_y": pos[:, 1],
|
||||
"pre_z": pos[:, 2],
|
||||
"pre_E": pre_E,
|
||||
"pre_dx": pre_dir[:, 0],
|
||||
"pre_dy": pre_dir[:, 1],
|
||||
"pre_dz": pre_dir[:, 2],
|
||||
"post_x": post_pos[:, 0],
|
||||
"post_y": post_pos[:, 1],
|
||||
"post_z": post_pos[:, 2],
|
||||
"post_E": post_E,
|
||||
"post_dx": post_dir[:, 0],
|
||||
"post_dy": post_dir[:, 1],
|
||||
"post_dz": post_dir[:, 2],
|
||||
"edep": edep,
|
||||
"step_length": step_length,
|
||||
"material": rng.choice(_MATERIALS, size=n),
|
||||
"layer_id": rng.integers(0, 30, size=n),
|
||||
"process": rng.choice(["compt", "phot", "eBrem", "eIoni", "conv"], size=n),
|
||||
"sec_E_list": sec_E,
|
||||
"sec_pdg_list": sec_pdg,
|
||||
"sec_dx_list": sec_dx,
|
||||
"sec_dy_list": sec_dy,
|
||||
"sec_dz_list": sec_dz,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _time(
|
||||
spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path
|
||||
) -> float:
|
||||
t0 = time.perf_counter()
|
||||
compute_reduced(
|
||||
spec_id,
|
||||
rollout,
|
||||
reference,
|
||||
shared,
|
||||
out,
|
||||
checkpoint=None,
|
||||
chunk_index=0,
|
||||
n_chunks=1,
|
||||
)
|
||||
return time.perf_counter() - t0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
ids = [i for i in catalog_ids() if get_spec(i).chunkable]
|
||||
timings: dict[str, list[tuple[int, float]]] = {i: [] for i in ids}
|
||||
|
||||
with TemporaryDirectory(prefix="giant-profile-") as tmp:
|
||||
tmp_path = Path(tmp)
|
||||
for n_side in SIDE_ROW_COUNTS:
|
||||
n_events = max(n_side // 20, 10)
|
||||
rollout = tmp_path / f"rollout_{n_side}.parquet"
|
||||
reference = tmp_path / f"reference_{n_side}.parquet"
|
||||
_make_rollout(n_side, n_events, seed=0).write_parquet(rollout)
|
||||
_make_reference(n_side, n_events, seed=0).write_parquet(reference)
|
||||
|
||||
shared = tmp_path / f"shared_{n_side}.json"
|
||||
ctx = build_context(
|
||||
rollout,
|
||||
reference,
|
||||
n_energy_bins=4,
|
||||
n_marginal_bins=50,
|
||||
top_k_pdg=6,
|
||||
sample_rows=min(n_side, 200_000),
|
||||
)
|
||||
ctx.save(shared)
|
||||
|
||||
# warm the OS page cache so the timed pass measures compute, not
|
||||
# the one-time cold read of a freshly-written file.
|
||||
pl.scan_parquet(rollout).select(pl.len()).collect()
|
||||
pl.scan_parquet(reference).select(pl.len()).collect()
|
||||
|
||||
n_rows = 2 * n_side # rollout + reference rows in this "chunk"
|
||||
for spec_id in ids:
|
||||
out = tmp_path / f"{spec_id}_{n_side}.json"
|
||||
dt = _time(spec_id, rollout, reference, shared, out)
|
||||
timings[spec_id].append((n_rows, dt))
|
||||
print(f"{spec_id:35s} n_rows={n_rows:>9d} time={dt:7.3f}s")
|
||||
|
||||
rollout.unlink()
|
||||
reference.unlink()
|
||||
shared.unlink()
|
||||
|
||||
print("\n# spec_id -> (intercept_s, seconds_per_row), fit by least squares")
|
||||
print("_COST_MODEL: dict[str, tuple[float, float]] = {")
|
||||
for spec_id in ids:
|
||||
xs = np.array([n for n, _ in timings[spec_id]], dtype=float)
|
||||
ys = np.array([t for _, t in timings[spec_id]], dtype=float)
|
||||
slope, intercept = np.polyfit(xs, ys, 1)
|
||||
intercept = max(intercept, 0.0)
|
||||
slope = max(slope, 0.0)
|
||||
print(f' "{spec_id}": ({intercept:.6f}, {slope:.9f}),')
|
||||
print("}")
|
||||
|
||||
if _ROUTER_IDS:
|
||||
print(
|
||||
"\n# router_* specs excluded: need a live MoE checkpoint to do real\n"
|
||||
"# work, none available on this machine — see _ROUTER_FIXED_S instead."
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,19 +0,0 @@
|
||||
import torch
|
||||
import torch.version
|
||||
|
||||
print(f"PyTorch version: {torch.__version__}")
|
||||
print(f"CUDA available: {torch.cuda.is_available()}")
|
||||
|
||||
if torch.cuda.is_available():
|
||||
print(f"CUDA version: {torch.version.cuda}")
|
||||
print(f"Device count: {torch.cuda.device_count()}")
|
||||
print(f"Device name: {torch.cuda.get_device_name(0)}")
|
||||
|
||||
# Run a small tensor op on the GPU
|
||||
a = torch.randn(1000, 1000, device="cuda")
|
||||
b = torch.randn(1000, 1000, device="cuda")
|
||||
c = a @ b
|
||||
torch.cuda.synchronize()
|
||||
print(f"Matrix multiply: OK (result shape {c.shape}, device {c.device})")
|
||||
else:
|
||||
print("No CUDA device found — check driver/CUDA installation.")
|
||||
+91
-7
@@ -2,21 +2,30 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from giant.analysis import build_catalog, catalog_ids, get_spec
|
||||
from giant.analysis.catalog import Bundle
|
||||
from giant.analysis.context import build_context
|
||||
from giant.analysis.catalog import Bundle, PlotSpec
|
||||
from giant.analysis.context import Context, build_context
|
||||
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bundle() -> Bundle:
|
||||
def _build_ctx() -> Context:
|
||||
r, t = _rollout_frame(), _reference_frame()
|
||||
ctx = build_context(
|
||||
return build_context(
|
||||
r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
|
||||
)
|
||||
return Bundle.open(r, t, ctx)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ctx() -> Context:
|
||||
return _build_ctx()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def bundle(ctx: Context) -> Bundle:
|
||||
return Bundle.open(_rollout_frame(), _reference_frame(), ctx)
|
||||
|
||||
|
||||
def test_catalog_ids_unique_and_nonempty():
|
||||
@@ -36,7 +45,7 @@ def test_get_spec_roundtrip_and_unknown():
|
||||
|
||||
def test_every_spec_computes_valid_reduced(bundle: Bundle):
|
||||
for spec in build_catalog():
|
||||
r = spec.compute(bundle)
|
||||
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
||||
assert r.id == spec.id
|
||||
assert r.kind in {
|
||||
"overlay_hist",
|
||||
@@ -44,6 +53,9 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle):
|
||||
"profile",
|
||||
"bar",
|
||||
"single_hist",
|
||||
"router_gating",
|
||||
"router_share",
|
||||
"unavailable",
|
||||
}
|
||||
assert r.title and r.xlabel
|
||||
_validate_payload(r)
|
||||
@@ -67,3 +79,75 @@ def _validate_payload(r) -> None:
|
||||
assert len(p[k]) == n
|
||||
elif r.kind == "bar":
|
||||
assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"])
|
||||
elif r.kind == "unavailable":
|
||||
assert p["note"]
|
||||
elif r.kind == "router_gating":
|
||||
for side in ("rollout", "reference"):
|
||||
if side in p:
|
||||
assert len(p[side]["centers"]) == len(p[side]["means"])
|
||||
elif r.kind == "router_share":
|
||||
for cat in p["categories"]:
|
||||
for side in ("rollout", "reference"):
|
||||
if side in p:
|
||||
assert cat in p[side]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# chunked (compute_partial x N -> finalize) must match the unchunked (N=1) result
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
# One representative id per merge shape: sum-mergeable (marginal_edep,
|
||||
# sec_count_per_species via pdg-keyed sums), concat-then-finalize with
|
||||
# data-dependent edges (event_total_edep), concat-then-mean/std (shower_
|
||||
# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a
|
||||
# ratio (species_edep_share), and a chunkable=False passthrough (router_gating).
|
||||
_CHUNK_EQUIVALENCE_IDS = [
|
||||
"marginal_edep",
|
||||
"species_edep_share",
|
||||
"event_total_edep",
|
||||
"shower_longitudinal",
|
||||
"leakage_fraction",
|
||||
"sec_count_per_species",
|
||||
"router_gating",
|
||||
]
|
||||
|
||||
|
||||
def _assert_payload_close(a, b, path: str = "payload") -> None:
|
||||
"""Recursively compare two JSON-shaped payloads (float-tolerant)."""
|
||||
assert type(a) is type(b), f"{path}: {type(a)} != {type(b)}"
|
||||
if isinstance(a, dict):
|
||||
assert set(a) == set(b), f"{path}: key mismatch {set(a)} != {set(b)}"
|
||||
for k in a:
|
||||
_assert_payload_close(a[k], b[k], f"{path}.{k}")
|
||||
elif isinstance(a, list):
|
||||
assert len(a) == len(b), f"{path}: length mismatch"
|
||||
for i, (x, y) in enumerate(zip(a, b)):
|
||||
_assert_payload_close(x, y, f"{path}[{i}]")
|
||||
elif isinstance(a, float):
|
||||
assert np.isclose(a, b, atol=1e-9), f"{path}: {a} != {b}"
|
||||
else:
|
||||
assert a == b, f"{path}: {a} != {b}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS)
|
||||
def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
|
||||
"""A plot computed over N event-disjoint chunks then merged must equal the
|
||||
same plot computed in one unchunked pass — the core chunking correctness
|
||||
guarantee (see the analysis-rollout-plots chunking plan)."""
|
||||
spec: PlotSpec = get_spec(spec_id)
|
||||
r, t = _rollout_frame(), _reference_frame()
|
||||
|
||||
unchunked_bundle = Bundle.open(r, t, ctx)
|
||||
unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx)
|
||||
|
||||
# 4 chunks over only 2 distinct event_ids also exercises empty chunks.
|
||||
n_chunks = 4 if spec.chunkable else 1
|
||||
parts = [
|
||||
spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks)))
|
||||
for k in range(n_chunks)
|
||||
]
|
||||
chunked = spec.finalize(parts, ctx)
|
||||
|
||||
assert chunked.id == unchunked.id
|
||||
assert chunked.kind == unchunked.kind
|
||||
_assert_payload_close(unchunked.payload, chunked.payload)
|
||||
|
||||
+144
-11
@@ -16,11 +16,13 @@ from giant.analysis import (
|
||||
compute_reduced,
|
||||
derive_run_dir,
|
||||
load_rollout_yaml,
|
||||
merge_one,
|
||||
prep,
|
||||
write_submit,
|
||||
)
|
||||
from giant.analysis.catalog import get_spec
|
||||
from giant.analysis.condor import Context
|
||||
from giant.analysis.reduced import Reduced
|
||||
from giant.analysis.reduced import Partial, Reduced
|
||||
from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
|
||||
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
||||
|
||||
@@ -51,11 +53,22 @@ def _write_inputs(tmp_path: Path) -> Path:
|
||||
return yaml_path
|
||||
|
||||
|
||||
def _prep(rollout_yaml: Path, run_dir: str | Path | None = None) -> Path:
|
||||
def _fake_venv(repo_dir: Path) -> None:
|
||||
"""Stand in for a `uv sync`'d venv: write_submit checks `.venv/bin/giant` exists."""
|
||||
giant = repo_dir / ".venv" / "bin" / "giant"
|
||||
giant.parent.mkdir(parents=True, exist_ok=True)
|
||||
giant.write_text("#!/bin/bash\n")
|
||||
giant.chmod(0o755)
|
||||
|
||||
|
||||
def _prep(
|
||||
rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1
|
||||
) -> Path:
|
||||
"""``prep`` with small test-sized context bins/sampling."""
|
||||
return prep(
|
||||
rollout_yaml,
|
||||
run_dir,
|
||||
n_chunks=chunks,
|
||||
n_energy_bins=2,
|
||||
n_marginal_bins=8,
|
||||
top_k_pdg=3,
|
||||
@@ -76,6 +89,15 @@ def test_derive_run_dir_next_to_rollout():
|
||||
assert derive_run_dir(y, "/somewhere") == Path("/somewhere")
|
||||
|
||||
|
||||
def test_derive_run_dir_default_base():
|
||||
y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
|
||||
assert derive_run_dir(y, default_base="/work/lbogner/giant2/analysis_runs") == Path(
|
||||
"/work/lbogner/giant2/analysis_runs/analysis_abcd1234"
|
||||
)
|
||||
# an explicit run_dir still wins over default_base
|
||||
assert derive_run_dir(y, "/somewhere", default_base="/other") == Path("/somewhere")
|
||||
|
||||
|
||||
def test_prep_lays_out_run_dir(tmp_path: Path):
|
||||
yaml_path = _write_inputs(tmp_path)
|
||||
run_dir = _prep(yaml_path)
|
||||
@@ -87,15 +109,25 @@ def test_prep_lays_out_run_dir(tmp_path: Path):
|
||||
assert meta.reference.endswith("reference.parquet")
|
||||
assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt"
|
||||
assert "best.pt" in meta.title
|
||||
assert meta.n_chunks == 1
|
||||
assert meta.rows_per_chunk == [meta.total_rows] # single chunk holds everything
|
||||
assert meta.total_rows == 8 # 5 rollout rows + 3 reference rows
|
||||
|
||||
|
||||
def test_prep_splits_rows_per_chunk(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path), chunks=2)
|
||||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||||
assert len(meta.rows_per_chunk) == 2
|
||||
assert sum(meta.rows_per_chunk) == meta.total_rows == 8
|
||||
|
||||
|
||||
def test_compute_one_from_run_dir(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path))
|
||||
out = compute_one("marginal_edep", run_dir)
|
||||
assert out == run_dir / "reduced" / "marginal_edep.json"
|
||||
reduced = Reduced.load(out)
|
||||
assert reduced.id == "marginal_edep"
|
||||
assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
|
||||
assert out == run_dir / "reduced_partial" / "marginal_edep__0.json"
|
||||
partial = Partial.load(out)
|
||||
assert partial.id == "marginal_edep" and partial.chunk == 0
|
||||
assert "r" in partial.data and "t" in partial.data
|
||||
|
||||
|
||||
def test_compute_reduced_explicit_paths(tmp_path: Path):
|
||||
@@ -108,30 +140,131 @@ def test_compute_reduced_explicit_paths(tmp_path: Path):
|
||||
run_dir / "shared.json",
|
||||
tmp_path / "r.json",
|
||||
)
|
||||
assert Reduced.load(out).id == "marginal_step_length"
|
||||
assert Partial.load(out).id == "marginal_step_length"
|
||||
|
||||
|
||||
def test_merge_one_produces_reduced(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path))
|
||||
compute_one("marginal_edep", run_dir)
|
||||
out = merge_one("marginal_edep", run_dir)
|
||||
assert out == run_dir / "reduced" / "marginal_edep.json"
|
||||
reduced = Reduced.load(out)
|
||||
assert reduced.id == "marginal_edep"
|
||||
assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1
|
||||
|
||||
|
||||
def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path), chunks=2)
|
||||
compute_one("marginal_edep", run_dir, chunk_index=0) # chunk 1 never computed
|
||||
with pytest.raises(FileNotFoundError, match="missing chunk"):
|
||||
merge_one("marginal_edep", run_dir)
|
||||
|
||||
|
||||
def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path):
|
||||
(tmp_path / "a").mkdir()
|
||||
(tmp_path / "b").mkdir()
|
||||
unchunked_dir = _prep(_write_inputs(tmp_path / "a"))
|
||||
compute_one("marginal_step_length", unchunked_dir)
|
||||
unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir))
|
||||
|
||||
chunked_dir = _prep(_write_inputs(tmp_path / "b"), chunks=2)
|
||||
for k in range(2):
|
||||
compute_one("marginal_step_length", chunked_dir, chunk_index=k)
|
||||
chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir))
|
||||
|
||||
assert chunked.payload == unchunked.payload
|
||||
|
||||
|
||||
def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path)) # n_chunks=1 (default)
|
||||
with pytest.raises(ValueError, match="out of range"):
|
||||
compute_one("marginal_edep", run_dir, chunk_index=1)
|
||||
|
||||
|
||||
def test_write_submit_description(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path))
|
||||
_fake_venv(tmp_path)
|
||||
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
|
||||
txt = write_submit(cfg).read_text()
|
||||
assert "universe = docker" in txt
|
||||
assert "docker_image = mschnepf/slc7-condocker" in txt
|
||||
assert "docker_image = cverstege/alma9-gridjob" in txt
|
||||
assert "requirements = TARGET.ProvidesETPResources" in txt
|
||||
assert "accounting_group = cms" in txt
|
||||
assert "queue plotid from" in txt
|
||||
assert (run_dir / "plotids.txt").read_text().split() == catalog_ids()
|
||||
assert "+RequestWalltime = $(walltime)" in txt
|
||||
assert "queue plotid,chunk,walltime from" in txt
|
||||
jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
|
||||
assert [i for i, _, _ in jobs] == catalog_ids()
|
||||
assert all(k == "0" for _, k, _ in jobs) # n_chunks=1 default
|
||||
assert all(int(w) > 0 for _, _, w in jobs)
|
||||
wrapper = run_dir / "run_compute.sh"
|
||||
assert wrapper.exists() and (wrapper.stat().st_mode & 0o111)
|
||||
body = wrapper.read_text()
|
||||
assert "giant analyze compute-one --id" in body and "--run-dir" in body
|
||||
assert "giant analyze compute-one --id" in body
|
||||
assert "--chunk" in body and "--run-dir" in body
|
||||
|
||||
|
||||
def test_write_submit_requires_synced_venv(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path))
|
||||
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path)
|
||||
with pytest.raises(FileNotFoundError, match="uv sync"):
|
||||
write_submit(cfg)
|
||||
|
||||
|
||||
def test_write_submit_remote_flag(tmp_path: Path):
|
||||
run_dir = _prep(_write_inputs(tmp_path))
|
||||
_fake_venv(tmp_path)
|
||||
cfg = SubmitConfig(
|
||||
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True
|
||||
)
|
||||
txt = write_submit(cfg).read_text()
|
||||
assert "+RemoteJob = True" in txt
|
||||
assert "ProvidesETPResources" not in txt
|
||||
|
||||
|
||||
def test_write_submit_chunks_respect_chunkable(tmp_path: Path):
|
||||
assert get_spec("router_gating").chunkable is False
|
||||
run_dir = _prep(_write_inputs(tmp_path), chunks=4)
|
||||
_fake_venv(tmp_path)
|
||||
cfg = SubmitConfig(
|
||||
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
|
||||
)
|
||||
write_submit(cfg)
|
||||
jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()]
|
||||
counts: dict[str, int] = {}
|
||||
for spec_id, _, _ in jobs:
|
||||
counts[spec_id] = counts.get(spec_id, 0) + 1
|
||||
assert counts["marginal_edep"] == 4
|
||||
assert counts["router_gating"] == 1 # chunkable=False, ignores n_chunks
|
||||
|
||||
|
||||
def test_estimate_runtime_s_scales_with_rows_and_margin():
|
||||
from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
|
||||
from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S
|
||||
|
||||
assert RUNTIME_SAFETY_MARGIN > 0
|
||||
small = estimate_runtime_s("marginal_edep", 1_000)
|
||||
large = estimate_runtime_s("marginal_edep", 100_000_000)
|
||||
assert small >= (1 + RUNTIME_SAFETY_MARGIN) * _FIXED_OVERHEAD_S
|
||||
assert large > small # bigger chunk -> longer estimate
|
||||
|
||||
|
||||
def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
|
||||
"""A chunked run's later job walltimes track that chunk's row count."""
|
||||
from giant.analysis.runtime_estimate import estimate_runtime_s
|
||||
|
||||
run_dir = _prep(_write_inputs(tmp_path), chunks=2)
|
||||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||||
_fake_venv(tmp_path)
|
||||
cfg = SubmitConfig(
|
||||
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2
|
||||
)
|
||||
write_submit(cfg)
|
||||
jobs = {
|
||||
(i, int(k)): int(w)
|
||||
for i, k, w in (
|
||||
line.split(",") for line in (run_dir / "jobs.txt").read_text().split()
|
||||
)
|
||||
}
|
||||
for chunk in range(2):
|
||||
expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk])
|
||||
assert jobs[("marginal_edep", chunk)] == expected
|
||||
|
||||
@@ -97,6 +97,42 @@ def test_build_router_ignores_unrecognized_kwargs():
|
||||
assert router.temperature == 0.3
|
||||
|
||||
|
||||
def test_energy_router_default_centers_are_linspace():
|
||||
router = EnergyRouter(n_experts=4)
|
||||
torch.testing.assert_close(router.centers, torch.linspace(-2.0, 2.0, 4))
|
||||
|
||||
|
||||
def test_energy_router_centers_init_overrides_default():
|
||||
centers_init = [-1.0, 0.0, 0.5, 3.0]
|
||||
router = EnergyRouter(n_experts=4, centers_init=centers_init)
|
||||
torch.testing.assert_close(router.centers, torch.tensor(centers_init))
|
||||
|
||||
|
||||
def test_energy_router_centers_init_wrong_length_raises():
|
||||
try:
|
||||
EnergyRouter(n_experts=4, centers_init=[0.0, 1.0])
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for centers_init length mismatch")
|
||||
|
||||
|
||||
def test_energy_router_centers_init_respects_learn_centers_flag():
|
||||
learned = EnergyRouter(
|
||||
n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=True
|
||||
)
|
||||
fixed = EnergyRouter(
|
||||
n_experts=3, centers_init=[-1.0, 0.0, 1.0], learn_centers=False
|
||||
)
|
||||
assert isinstance(learned.centers, torch.nn.Parameter)
|
||||
assert not isinstance(fixed.centers, torch.nn.Parameter)
|
||||
|
||||
|
||||
def test_build_router_threads_centers_init_through_energy_router():
|
||||
centers_init = [-1.5, -0.5, 0.5, 1.5]
|
||||
router = build_router("energy", 4, centers_init=centers_init)
|
||||
torch.testing.assert_close(router.centers, torch.tensor(centers_init))
|
||||
|
||||
|
||||
def test_build_router_unknown_type_raises():
|
||||
try:
|
||||
build_router("nonexistent", 4)
|
||||
|
||||
@@ -0,0 +1,126 @@
|
||||
"""Tests for the MoE router-gating diagnostic (giant.analysis.router_gating)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
import torch
|
||||
|
||||
from giant.analysis.router_gating import (
|
||||
compute_router_gating,
|
||||
compute_router_share_by_pdg,
|
||||
compute_router_share_by_process,
|
||||
)
|
||||
from giant.data.transforms import Normalizer
|
||||
from giant.model.network import build_models
|
||||
|
||||
_PDG_MAP = {11: 0, 22: 1}
|
||||
_MAT_MAP = {"G4_PbWO4": 0, "G4_Pb": 1}
|
||||
|
||||
|
||||
def _model_cfg() -> dict:
|
||||
return {
|
||||
"router": {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 2,
|
||||
"temperature": 0.5,
|
||||
"learn_centers": True,
|
||||
"energy_idx": 3,
|
||||
},
|
||||
"pdg_vocab": len(_PDG_MAP),
|
||||
"mat_vocab": len(_MAT_MAP),
|
||||
"conditioning": "embedding",
|
||||
}
|
||||
|
||||
|
||||
def _write_checkpoint(tmp_path) -> str:
|
||||
cfg = _model_cfg()
|
||||
stage1, _ = build_models(cfg)
|
||||
norm = Normalizer()
|
||||
norm.mean = np.zeros(15, dtype=np.float32)
|
||||
norm.std = np.ones(15, dtype=np.float32)
|
||||
ckpt = {
|
||||
"model_config": cfg,
|
||||
"model": stage1.state_dict(),
|
||||
"pdg_map": _PDG_MAP,
|
||||
"mat_map": _MAT_MAP,
|
||||
"normalizer": {"cond": norm.to_dict()},
|
||||
}
|
||||
path = tmp_path / "ckpt.pt"
|
||||
torch.save(ckpt, path)
|
||||
return str(path)
|
||||
|
||||
|
||||
def _steps_frame(process: bool = False) -> pl.LazyFrame:
|
||||
n = 40
|
||||
rng = np.random.default_rng(0)
|
||||
pre_e = np.concatenate([rng.uniform(1, 10, n // 2), rng.uniform(100, 1000, n // 2)])
|
||||
pdg = np.where(np.arange(n) % 2 == 0, 11, 22)
|
||||
material = np.where(np.arange(n) % 3 == 0, "G4_Pb", "G4_PbWO4")
|
||||
data = {
|
||||
"event_id": np.arange(n),
|
||||
"pdg": pdg,
|
||||
"pre_x": np.zeros(n),
|
||||
"pre_y": np.zeros(n),
|
||||
"pre_z": np.zeros(n),
|
||||
"pre_E": pre_e,
|
||||
"pre_dx": np.zeros(n),
|
||||
"pre_dy": np.zeros(n),
|
||||
"pre_dz": np.ones(n),
|
||||
"post_x": np.zeros(n),
|
||||
"post_y": np.zeros(n),
|
||||
"post_z": np.ones(n),
|
||||
"post_E": pre_e * 0.5,
|
||||
"post_dx": np.zeros(n),
|
||||
"post_dy": np.zeros(n),
|
||||
"post_dz": np.ones(n),
|
||||
"edep": pre_e * 0.5,
|
||||
"step_length": np.ones(n),
|
||||
"material": material,
|
||||
"layer_id": np.zeros(n, dtype=np.int64),
|
||||
}
|
||||
if process:
|
||||
data["process"] = np.where(pdg == 11, "eIoni", "compt")
|
||||
return pl.DataFrame(data).lazy()
|
||||
|
||||
|
||||
def test_compute_router_gating_shapes(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
lf = _steps_frame()
|
||||
r = compute_router_gating(checkpoint, lf, lf)
|
||||
assert r.kind == "router_gating"
|
||||
assert r.payload["n_experts"] == 2
|
||||
for side in ("rollout", "reference"):
|
||||
means = r.payload[side]["means"]
|
||||
assert means, f"{side} produced no bins"
|
||||
assert all(abs(sum(row) - 1.0) < 1e-5 for row in means)
|
||||
|
||||
|
||||
def test_compute_router_gating_missing_checkpoint_is_unavailable():
|
||||
lf = _steps_frame()
|
||||
r = compute_router_gating(None, lf, lf)
|
||||
assert r.kind == "unavailable"
|
||||
assert "note" in r.payload
|
||||
assert r.title
|
||||
|
||||
|
||||
def test_compute_router_share_by_pdg(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
lf = _steps_frame()
|
||||
r = compute_router_share_by_pdg(checkpoint, lf, lf, top_pdgs=[11, 22])
|
||||
assert r.kind == "router_share"
|
||||
for side in ("rollout", "reference"):
|
||||
assert set(r.payload[side]) == {"e-", "gamma"}
|
||||
for shares in r.payload[side].values():
|
||||
assert abs(sum(shares) - 1.0) < 1e-5
|
||||
|
||||
|
||||
def test_compute_router_share_by_process(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
lf = _steps_frame(process=True)
|
||||
r = compute_router_share_by_process(checkpoint, lf)
|
||||
assert r.kind == "router_share"
|
||||
assert set(r.payload["categories"]) <= {"eIoni", "compt"}
|
||||
for shares in r.payload["reference"].values():
|
||||
assert abs(sum(shares) - 1.0) < 1e-5
|
||||
@@ -397,3 +397,46 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
|
||||
cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0]))
|
||||
)
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0])
|
||||
|
||||
|
||||
def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode():
|
||||
"""A pre-physical-conditioning checkpoint's cond normalizer is COND_DIM_BASE
|
||||
(8) wide, fit before build_cond_features grew the extra physical columns.
|
||||
In "embedding" mode those columns are never read downstream, so a legacy
|
||||
normalizer should be usable as-is (padded, not rejected)."""
|
||||
data = _minimal_step_data(3)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
legacy_norm = Normalizer()
|
||||
legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
|
||||
legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
|
||||
|
||||
cond_cont, _ = build_cond_features(
|
||||
data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding"
|
||||
)
|
||||
|
||||
assert cond_cont.shape[-1] == COND_DIM
|
||||
# padded physical columns are zero-filled pre-normalization and
|
||||
# mean=0/std=1 post-normalization, so they should come out as zero
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
|
||||
|
||||
|
||||
def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
|
||||
fake_material_props,
|
||||
):
|
||||
"""Unlike "embedding" mode, "physical" mode actually reads the physical
|
||||
columns, so a legacy 8-wide normalizer can't be silently padded — that
|
||||
would silently feed the network un-normalized physical properties."""
|
||||
data = _minimal_step_data(3)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
legacy_norm = Normalizer()
|
||||
legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32)
|
||||
legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32)
|
||||
|
||||
with pytest.raises(ValueError, match="predates physical-property conditioning"):
|
||||
build_cond_features(
|
||||
data,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
cond_normalizer=legacy_norm,
|
||||
conditioning="physical",
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user