Merge pull request 'Refactor the analysis plot creation with focus on rollout' (#16) from analysis-rollout-plots into master
Reviewed-on: #16
This commit was merged in pull request #16.
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/
|
||||
|
||||
@@ -12,8 +12,12 @@ uv sync --extra cpu --extra geometry # add scikit-learn for the geometry oracle
|
||||
pytest # run tests
|
||||
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 predict path/to/steps.parquet --checkpoint ckpt/best.pt # per-step predictions
|
||||
giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl # full showers
|
||||
giant analyze submit rollout.yaml --accounting-group cms # parallel rollout-vs-reference analysis on HTCondor
|
||||
giant analyze render <run_dir> --gallery # render PDFs + HTML gallery (run_dir from prep/submit)
|
||||
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
|
||||
# bump-schema, status, update-manifest, create-manifest,
|
||||
# make-root, build-geometry-oracle, hparam-scan
|
||||
@@ -32,6 +36,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).
|
||||
@@ -54,7 +66,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.
|
||||
|
||||
**Validation** (`giant/validate.py`): step-level marginal comparisons. `giant/analysis.py` is a fully-streaming (lazy polars) diagnostics module, sized for predict/rollout files larger than RAM, with no in-memory `SampleCollection` and no full-array materialization. It covers one-step-ahead `giant predict --coord local` output (`compute_event_observables_pl` + `plot_total_energy`/`plot_longitudinal_profile`/etc. for shower-level observables, plus the marginal/correlation/constraint tiers) and, via the `RolloutVsTruth` source type, a full autoregressive `giant rollout` shower compared against held-out truth data (`compute_rollout_vs_truth_observables_pl` for shower-level observables, reusing the same plot functions) — see the module docstring.
|
||||
**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_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.
|
||||
|
||||
@@ -66,4 +84,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.
|
||||
|
||||
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -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
|
||||
-2359
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,52 @@
|
||||
"""Rollout-vs-reference analysis: streaming compute + plotstyle rendering.
|
||||
|
||||
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 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
|
||||
``giant.analysis.render`` explicitly for the local render step.
|
||||
"""
|
||||
|
||||
from giant.analysis.catalog import build_catalog, catalog_ids, get_spec
|
||||
from giant.analysis.condor import (
|
||||
RunMeta,
|
||||
SubmitConfig,
|
||||
compute_one,
|
||||
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 Partial, Reduced
|
||||
from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s
|
||||
from giant.analysis.sources import Side
|
||||
|
||||
__all__ = [
|
||||
"build_catalog",
|
||||
"catalog_ids",
|
||||
"get_spec",
|
||||
"RunMeta",
|
||||
"SubmitConfig",
|
||||
"compute_one",
|
||||
"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",
|
||||
]
|
||||
@@ -0,0 +1,846 @@
|
||||
"""The declarative plot catalog: one ``PlotSpec`` per figure.
|
||||
|
||||
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_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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import Callable
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis.context import Context
|
||||
from giant.analysis.grouping import (
|
||||
energy_bin_labels,
|
||||
event_energy_bins,
|
||||
material_label,
|
||||
pdg_label,
|
||||
)
|
||||
from giant.analysis.reduce import (
|
||||
attach_entry_axis,
|
||||
depth_expr,
|
||||
entry_axis,
|
||||
event_scalars,
|
||||
hist1d,
|
||||
leakage_fraction,
|
||||
profile_finalize,
|
||||
profile_partial,
|
||||
species_share,
|
||||
sum_merge,
|
||||
transverse_expr,
|
||||
)
|
||||
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
|
||||
|
||||
|
||||
@dataclass
|
||||
class Bundle:
|
||||
"""Everything a compute runs against — built once per ``compute-one`` job."""
|
||||
|
||||
ctx: Context
|
||||
r_all: pl.LazyFrame # rollout, all rows (incl. synthetic termination rows)
|
||||
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,
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PlotSpec:
|
||||
id: str
|
||||
family: str
|
||||
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
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# small numpy/hist helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_ROLL = "rollout"
|
||||
_REF = "reference"
|
||||
|
||||
|
||||
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]:
|
||||
"""Shared-edge histogram of two small per-event arrays (robust range)."""
|
||||
both = np.concatenate([r, t]) if (len(r) or len(t)) else np.array([0.0, 1.0])
|
||||
lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
|
||||
if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
|
||||
lo, hi = lo - 0.5, hi + 0.5
|
||||
edges = np.linspace(lo, hi, nbins + 1)
|
||||
return edges, np.histogram(r, edges)[0], np.histogram(t, edges)[0]
|
||||
|
||||
|
||||
# Human-readable figure titles per marginal variable (the axis labels carry units;
|
||||
# these read cleanly as a title without them).
|
||||
_TITLE_NAMES = {
|
||||
"step_length": "Step length",
|
||||
"edep": "Deposited energy per step",
|
||||
"delta_e": "Energy loss per step",
|
||||
"post_E": "Post-step energy",
|
||||
"cos_scatter": "Scattering cosine",
|
||||
}
|
||||
|
||||
|
||||
def _var(var: str):
|
||||
"""(axis label, value expr) for a marginal variable name."""
|
||||
if var == "cos_scatter":
|
||||
return ("cos of scattering angle", cos_scatter_expr())
|
||||
label, expr = RANGED_VARS[var]
|
||||
return (label, expr)
|
||||
|
||||
|
||||
def _marginal_edges(ctx: Context, var: str) -> np.ndarray:
|
||||
if var == "cos_scatter":
|
||||
return np.linspace(-1.0, 1.0, ctx.n_marginal_bins + 1)
|
||||
return ctx.marginal_edges(var)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# marginals: variable x {overall, energy, pdg, material}
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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 = 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",
|
||||
kind="overlay_hist",
|
||||
title=_TITLE_NAMES[var],
|
||||
xlabel=label,
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: _finalize_counts(r, 0, nb),
|
||||
_REF: _finalize_counts(t, 0, nb),
|
||||
"log_y": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
|
||||
ids, bins = event_energy_bins(lf, edges)
|
||||
return pl.col("event_id").replace_strict(
|
||||
ids, bins, default=-1, return_dtype=pl.Int64
|
||||
)
|
||||
|
||||
|
||||
def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
|
||||
_, expr = _var(var)
|
||||
edges = _marginal_edges(b.ctx, var)
|
||||
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"))
|
||||
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"))
|
||||
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: _finalize_counts(r, bi, nb),
|
||||
_REF: _finalize_counts(t, bi, nb),
|
||||
}
|
||||
|
||||
return Reduced(
|
||||
id=f"marginal_{var}_by_{axis}",
|
||||
family="marginals",
|
||||
kind="grouped_hist",
|
||||
title=f"{_TITLE_NAMES[var]} by {axis}",
|
||||
xlabel=label,
|
||||
payload={"edges": edges.tolist(), "groups": groups, "log_y": True},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# per-event scalar observables
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
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()
|
||||
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",
|
||||
kind="overlay_hist",
|
||||
title=title,
|
||||
xlabel=xlabel,
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: rc.astype(np.int64).tolist(),
|
||||
_REF: tc.astype(np.int64).tolist(),
|
||||
"log_y": False,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _event_total_edep_by_energy_partial(b: Bundle) -> dict:
|
||||
r = event_scalars(b.r_all)
|
||||
t = event_scalars(b.t_all)
|
||||
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]
|
||||
tc = np.histogram(t_val[t_bin == bi], edges)[0]
|
||||
groups[lbl] = {
|
||||
_ROLL: rc.astype(np.int64).tolist(),
|
||||
_REF: tc.astype(np.int64).tolist(),
|
||||
}
|
||||
return Reduced(
|
||||
id="event_total_edep_by_energy",
|
||||
family="event",
|
||||
kind="grouped_hist",
|
||||
title="Total deposited energy per event by incident energy",
|
||||
xlabel="total deposited energy [MeV]",
|
||||
payload={"edges": edges.tolist(), "groups": groups, "log_y": False},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# shower shape profiles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _profile_partial(b: Bundle, coord_fn, edges_key: str) -> dict:
|
||||
edges = np.asarray(getattr(b.ctx, edges_key))
|
||||
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",
|
||||
kind="profile",
|
||||
title=title,
|
||||
xlabel=xlabel,
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
"rollout_mean": r_mean.tolist(),
|
||||
"rollout_std": r_std.tolist(),
|
||||
"reference_mean": t_mean.tolist(),
|
||||
"reference_std": t_std.tolist(),
|
||||
"ylabel": "mean deposited energy per event [MeV]",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# species share + leakage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _species_share_partial(b: Bundle) -> dict:
|
||||
r = species_share(b.r_all)
|
||||
t = species_share(b.t_all)
|
||||
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 ctx.top_pdgs]
|
||||
return Reduced(
|
||||
id="species_edep_share",
|
||||
family="species",
|
||||
kind="bar",
|
||||
title="Deposited-energy share by species",
|
||||
xlabel="species",
|
||||
payload={
|
||||
"labels": labels,
|
||||
_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_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), ctx.n_marginal_bins + 1
|
||||
)
|
||||
counts = np.histogram(frac, edges)[0]
|
||||
return Reduced(
|
||||
id="leakage_fraction",
|
||||
family="species",
|
||||
kind="single_hist",
|
||||
title="Escaped (leakage) energy fraction per shower",
|
||||
xlabel="escaped energy fraction",
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: counts.astype(np.int64).tolist(),
|
||||
"log_y": True,
|
||||
"note": "rollout only; the reference has no detector-escape concept",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# secondaries
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _sec_frames(b: Bundle):
|
||||
return (
|
||||
secondaries(b.r_phys, Side.rollout),
|
||||
secondaries(b.t_all, Side.reference),
|
||||
)
|
||||
|
||||
|
||||
def _sec_count_per_event_partial(b: Bundle) -> dict:
|
||||
r_sec, t_sec = _sec_frames(b)
|
||||
r = (
|
||||
r_sec.group_by("event_id")
|
||||
.agg(pl.len().alias("n"))
|
||||
.collect(engine="streaming")["n"]
|
||||
.to_numpy()
|
||||
)
|
||||
t = (
|
||||
t_sec.group_by("event_id")
|
||||
.agg(pl.len().alias("n"))
|
||||
.collect(engine="streaming")["n"]
|
||||
.to_numpy()
|
||||
)
|
||||
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",
|
||||
kind="overlay_hist",
|
||||
title="Number of secondaries per event",
|
||||
xlabel="secondaries per event",
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: rc.astype(np.int64).tolist(),
|
||||
_REF: tc.astype(np.int64).tolist(),
|
||||
"log_y": False,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
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(ctx.top_pdgs)
|
||||
]
|
||||
return Reduced(
|
||||
id="sec_count_per_species",
|
||||
family="secondaries",
|
||||
kind="bar",
|
||||
title="Secondary count by species",
|
||||
xlabel="species",
|
||||
payload={
|
||||
"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",
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
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",
|
||||
kind="overlay_hist",
|
||||
title="Secondary birth energy",
|
||||
xlabel="secondary energy [MeV]",
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_ROLL: _finalize_counts(r, 0, nb),
|
||||
_REF: _finalize_counts(t, 0, nb),
|
||||
"log_y": True,
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _sec_cos_angle_partial(b: Bundle) -> dict:
|
||||
edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 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) -> dict[str, list[int]]:
|
||||
ea = entry_axis(steps_lf)
|
||||
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",
|
||||
kind="overlay_hist",
|
||||
title="Secondary emission angle relative to the shower axis",
|
||||
xlabel="cos of emission angle",
|
||||
payload={
|
||||
"edges": edges.tolist(),
|
||||
_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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
MARGINAL_VARS = ["step_length", "edep", "delta_e", "post_E", "cos_scatter"]
|
||||
GROUPING_AXES = ["energy", "pdg", "material"]
|
||||
|
||||
|
||||
def build_catalog() -> list[PlotSpec]:
|
||||
"""All concrete plot specs, each with a unique id."""
|
||||
specs: list[PlotSpec] = []
|
||||
|
||||
for var in MARGINAL_VARS:
|
||||
specs.append(
|
||||
PlotSpec(
|
||||
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:
|
||||
specs.append(
|
||||
PlotSpec(
|
||||
f"marginal_{var}_by_{axis}",
|
||||
"marginals",
|
||||
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)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
specs += [
|
||||
PlotSpec(
|
||||
"event_total_edep",
|
||||
"event",
|
||||
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]",
|
||||
),
|
||||
),
|
||||
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",
|
||||
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]",
|
||||
),
|
||||
),
|
||||
PlotSpec(
|
||||
"event_n_steps",
|
||||
"event",
|
||||
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",
|
||||
),
|
||||
),
|
||||
PlotSpec(
|
||||
"shower_longitudinal",
|
||||
"shower",
|
||||
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_edges",
|
||||
),
|
||||
),
|
||||
PlotSpec(
|
||||
"shower_transverse",
|
||||
"shower",
|
||||
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_edges",
|
||||
),
|
||||
),
|
||||
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
|
||||
|
||||
|
||||
def catalog_ids() -> list[str]:
|
||||
return [s.id for s in build_catalog()]
|
||||
|
||||
|
||||
def get_spec(spec_id: str) -> PlotSpec:
|
||||
for s in build_catalog():
|
||||
if s.id == spec_id:
|
||||
return s
|
||||
raise KeyError(f"unknown plot id: {spec_id!r}")
|
||||
@@ -0,0 +1,431 @@
|
||||
"""HTCondor orchestration driven by a ``giant rollout`` YAML sidecar.
|
||||
|
||||
A rollout writes a YAML sidecar (``giant/cli.py:_write_prediction_ref`` +
|
||||
rollout extras) that already names both files we need and carries the run's
|
||||
provenance:
|
||||
|
||||
* ``output`` — the rollout steps parquet (the *generated* side)
|
||||
* ``dataset`` — the file the rollout was seeded from, i.e. the held-out real
|
||||
steps (the *reference* side)
|
||||
* ``checkpoint``, ``geometry_oracle``, ``energy_cutoff``, ``steps``, ... —
|
||||
metadata that flows straight into every plot's gallery ``metadata.yaml``.
|
||||
|
||||
So the analysis takes that one YAML as input, derives its own **run directory**
|
||||
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_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, 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``
|
||||
(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.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
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 = (
|
||||
"prediction_id",
|
||||
"checkpoint",
|
||||
"output",
|
||||
"dataset",
|
||||
"geometry_oracle",
|
||||
"energy_cutoff",
|
||||
"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",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# rollout-YAML → run directory
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def load_rollout_yaml(path: str | Path) -> dict:
|
||||
"""Load a ``giant rollout`` YAML sidecar, requiring the two file paths."""
|
||||
d = yaml.safe_load(Path(path).read_text())
|
||||
for key in ("output", "dataset"):
|
||||
if key not in d:
|
||||
raise ValueError(
|
||||
f"{path} is not a rollout YAML (missing {key!r}); expected the "
|
||||
"sidecar `giant rollout` writes next to the checkpoint"
|
||||
)
|
||||
if d.get("kind") not in (None, "rollout"):
|
||||
raise ValueError(f"{path} has kind={d.get('kind')!r}, not a rollout YAML")
|
||||
return d
|
||||
|
||||
|
||||
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]
|
||||
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:
|
||||
return {k: rollout_yaml[k] for k in _PLOT_META_KEYS if k in rollout_yaml}
|
||||
|
||||
|
||||
@dataclass
|
||||
class RunMeta:
|
||||
"""Resolved paths + plot metadata for one analysis run (``run_meta.json``)."""
|
||||
|
||||
rollout: str
|
||||
reference: str
|
||||
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))
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "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, 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),
|
||||
reference=str(reference),
|
||||
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, chunk) compute (what each condor job runs)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def compute_reduced(
|
||||
spec_id: str,
|
||||
rollout: str | Path,
|
||||
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, 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)
|
||||
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)
|
||||
partial.save(out)
|
||||
return out
|
||||
|
||||
|
||||
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(
|
||||
spec_id,
|
||||
meta.rollout,
|
||||
meta.reference,
|
||||
run_path / "shared.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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@dataclass
|
||||
class SubmitConfig:
|
||||
run_dir: Path
|
||||
accounting_group: str
|
||||
repo_dir: Path
|
||||
docker_image: str = "cverstege/alma9-gridjob"
|
||||
request_memory_mb: int = 8192
|
||||
request_cpus: int = 1
|
||||
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 {repo_dir}/.venv/bin/giant analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
|
||||
"""
|
||||
|
||||
|
||||
def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
|
||||
reqs_attrs = (
|
||||
"+RemoteJob = True\n"
|
||||
if cfg.remote
|
||||
else "requirements = TARGET.ProvidesETPResources\n"
|
||||
)
|
||||
return (
|
||||
"universe = docker\n"
|
||||
f"docker_image = {cfg.docker_image}\n"
|
||||
f"executable = {wrapper}\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"
|
||||
"+RequestWalltime = $(walltime)\n"
|
||||
f"accounting_group = {cfg.accounting_group}\n"
|
||||
f"{reqs_attrs}"
|
||||
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,chunk,walltime from {jobs_file}\n"
|
||||
)
|
||||
|
||||
|
||||
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)
|
||||
|
||||
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, jobs_file))
|
||||
return sub
|
||||
@@ -0,0 +1,199 @@
|
||||
"""Shared analysis context (the ``prep`` step): fixed bin edges + group sets.
|
||||
|
||||
Every histogram in the catalog bins against **fixed** edges so each compute job
|
||||
is a single streaming pass with no min/max range scan. Those edges — plus the
|
||||
energy-bin quantiles, the top PDG species and the material list to stratify by,
|
||||
and the shower depth/transverse ranges — are resolved *once* here, on the submit
|
||||
node, from a hash-subsample plus a few cheap exact ``group_by`` passes, and shipped
|
||||
in ``shared.json``. Tiny and self-describing; no per-event arrays.
|
||||
|
||||
plotstyle-free (runs on the submit node, but also importable by workers).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis.grouping import energy_bin_edges
|
||||
from giant.analysis.reduce import (
|
||||
attach_entry_axis,
|
||||
depth_expr,
|
||||
entry_axis,
|
||||
transverse_expr,
|
||||
)
|
||||
from giant.analysis.sources import Side, open_side, physical_steps, secondaries
|
||||
from giant.analysis.variables import RANGED_VARS
|
||||
|
||||
|
||||
@dataclass
|
||||
class Context:
|
||||
"""Resolved bin edges and grouping sets shared by every compute job."""
|
||||
|
||||
n_marginal_bins: int
|
||||
var_ranges: dict[str, tuple[float, float]] # ranged var -> (lo, hi)
|
||||
energy_edges: list[float]
|
||||
top_pdgs: list[int]
|
||||
materials: list[str]
|
||||
depth_edges: list[float]
|
||||
transverse_edges: list[float]
|
||||
sec_energy_range: tuple[float, float]
|
||||
n_sec_bins: int
|
||||
n_events: dict[str, int] = field(default_factory=dict)
|
||||
|
||||
# -- (de)serialization -------------------------------------------------
|
||||
def save(self, path: str | Path) -> None:
|
||||
Path(path).write_text(json.dumps(asdict(self), indent=2))
|
||||
|
||||
@classmethod
|
||||
def load(cls, path: str | Path) -> "Context":
|
||||
d = json.loads(Path(path).read_text())
|
||||
d["var_ranges"] = {k: tuple(v) for k, v in d["var_ranges"].items()}
|
||||
d["sec_energy_range"] = tuple(d["sec_energy_range"])
|
||||
return cls(**d)
|
||||
|
||||
# -- convenience -------------------------------------------------------
|
||||
def marginal_edges(self, var: str) -> np.ndarray:
|
||||
lo, hi = self.var_ranges[var]
|
||||
return np.linspace(lo, hi, self.n_marginal_bins + 1)
|
||||
|
||||
|
||||
_LO_Q, _HI_Q = 0.001, 0.999
|
||||
|
||||
|
||||
def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFrame:
|
||||
"""Hash-subsample ~``sample_rows`` rows (for range estimation only)."""
|
||||
n_total = lf.select(pl.len()).collect(engine="streaming").item()
|
||||
if n_total <= sample_rows:
|
||||
return lf
|
||||
threshold = int(sample_rows / n_total * 2**32)
|
||||
return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
|
||||
|
||||
|
||||
def _combined_quantiles(
|
||||
r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
|
||||
) -> tuple[float, float]:
|
||||
"""Robust (lo_q, hi_q) range over the union of two value samples."""
|
||||
both = np.concatenate([r_vals, t_vals])
|
||||
lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
|
||||
if not (hi - lo > 1e-6 * max(abs(hi), 1.0)):
|
||||
lo, hi = lo - 0.5, hi + 0.5
|
||||
return lo, hi
|
||||
|
||||
|
||||
def build_context(
|
||||
rollout: str | Path | pl.LazyFrame,
|
||||
reference: str | Path | pl.LazyFrame,
|
||||
*,
|
||||
n_energy_bins: int = 4,
|
||||
n_marginal_bins: int = 50,
|
||||
n_sec_bins: int = 40,
|
||||
top_k_pdg: int = 6,
|
||||
sample_rows: int = 1_000_000,
|
||||
seed: int = 0,
|
||||
) -> Context:
|
||||
"""Resolve the shared context from the two files (the ``prep`` step)."""
|
||||
r_all = open_side(rollout, Side.rollout)
|
||||
t_all = open_side(reference, Side.reference)
|
||||
r_lf = physical_steps(r_all, Side.rollout)
|
||||
t_lf = physical_steps(t_all, Side.reference)
|
||||
|
||||
# Ranged marginal variables: robust ranges over a shared row subsample.
|
||||
exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
|
||||
r_s = (
|
||||
_row_subsample(r_lf, sample_rows, seed)
|
||||
.select(exprs)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
t_s = (
|
||||
_row_subsample(t_lf, sample_rows, seed)
|
||||
.select(exprs)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
var_ranges = {
|
||||
name: _combined_quantiles(
|
||||
r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
|
||||
)
|
||||
for name in RANGED_VARS
|
||||
}
|
||||
|
||||
# Energy-bin edges from exact per-event incident energies (cheap group_by).
|
||||
def _incident(lf: pl.LazyFrame) -> np.ndarray:
|
||||
return (
|
||||
lf.group_by("event_id")
|
||||
.agg(pl.col("pre_E").max())
|
||||
.collect(engine="streaming")["pre_E"]
|
||||
.to_numpy()
|
||||
)
|
||||
|
||||
r_inc, t_inc = _incident(r_lf), _incident(t_lf)
|
||||
energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
|
||||
|
||||
# Top PDG species and material list (cheap single-column group_bys).
|
||||
def _counts(lf: pl.LazyFrame, col: str) -> pl.DataFrame:
|
||||
return lf.group_by(col).agg(pl.len().alias("n")).collect(engine="streaming")
|
||||
|
||||
pdg_counts = (
|
||||
pl.concat([_counts(r_lf, "pdg"), _counts(t_lf, "pdg")])
|
||||
.group_by("pdg")
|
||||
.agg(pl.col("n").sum())
|
||||
.sort("n", descending=True)
|
||||
)
|
||||
top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
|
||||
materials = sorted(
|
||||
set(_counts(r_lf, "material")["material"].to_list())
|
||||
| set(_counts(t_lf, "material")["material"].to_list())
|
||||
)
|
||||
|
||||
# Shower depth / transverse ranges from a subsampled proxy.
|
||||
def _proxy(lf: pl.LazyFrame) -> tuple[np.ndarray, np.ndarray]:
|
||||
ea = entry_axis(lf)
|
||||
sub = (
|
||||
attach_entry_axis(_row_subsample(lf, sample_rows, seed), ea)
|
||||
.select(depth_expr().alias("d"), transverse_expr().alias("t"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
return sub["d"].to_numpy(), sub["t"].to_numpy()
|
||||
|
||||
r_d, r_t = _proxy(r_lf)
|
||||
t_d, t_t = _proxy(t_lf)
|
||||
d_lo, d_hi = _combined_quantiles(r_d, t_d, _LO_Q, _HI_Q)
|
||||
depth_edges = np.linspace(d_lo, d_hi, n_marginal_bins + 1)
|
||||
t_hi = max(float(np.quantile(np.concatenate([r_t, t_t]), _HI_Q)), 1e-6)
|
||||
transverse_edges = np.linspace(0.0, t_hi, n_marginal_bins + 1)
|
||||
|
||||
# Secondary energy range.
|
||||
r_se = secondaries(r_lf, Side.rollout).select("energy")
|
||||
t_se = secondaries(t_all, Side.reference).select("energy")
|
||||
r_se = _row_sample_col(r_se, sample_rows, seed)
|
||||
t_se = _row_sample_col(t_se, sample_rows, seed)
|
||||
sec_energy_range = _combined_quantiles(r_se, t_se, _LO_Q, _HI_Q)
|
||||
|
||||
return Context(
|
||||
n_marginal_bins=n_marginal_bins,
|
||||
var_ranges=var_ranges,
|
||||
energy_edges=[float(x) for x in energy_edges],
|
||||
top_pdgs=top_pdgs,
|
||||
materials=materials,
|
||||
depth_edges=[float(x) for x in depth_edges],
|
||||
transverse_edges=[float(x) for x in transverse_edges],
|
||||
sec_energy_range=sec_energy_range,
|
||||
n_sec_bins=n_sec_bins,
|
||||
n_events={
|
||||
"rollout": len(r_inc),
|
||||
"reference": len(t_inc),
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _row_sample_col(lf: pl.LazyFrame, sample_rows: int, seed: int) -> np.ndarray:
|
||||
"""Collect a subsample of a single-column ``energy`` LazyFrame to numpy."""
|
||||
vals = lf.collect(engine="streaming")["energy"].to_numpy()
|
||||
if len(vals) > sample_rows:
|
||||
rng = np.random.default_rng(seed)
|
||||
vals = vals[rng.choice(len(vals), size=sample_rows, replace=False)]
|
||||
return vals
|
||||
@@ -0,0 +1,108 @@
|
||||
"""Grouping axes (overall / energy / pdg / material) and their labels.
|
||||
|
||||
Energy grouping is by the event's **incident (primary) energy** — the largest
|
||||
``pre_E`` in the event — so every step of a shower lands in one bin, the physically
|
||||
meaningful stratification for a calorimeter surrogate. The quantile bin *edges* are
|
||||
sized once in ``prep`` (from a subsample) and shipped in ``shared.json``; a compute
|
||||
job that needs them re-derives the small per-event ``event_id -> bin`` map itself
|
||||
(one bounded streaming ``group_by`` over ``pre_E``), so ``shared.json`` stays tiny.
|
||||
|
||||
Pure/plotstyle-free so it can run on the compute workers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
# Common electromagnetic/hadronic species; anything else falls back to its code.
|
||||
PDG_NAMES: dict[int, str] = {
|
||||
11: "e-",
|
||||
-11: "e+",
|
||||
22: "gamma",
|
||||
2112: "n",
|
||||
2212: "p",
|
||||
-2212: "pbar",
|
||||
111: "pi0",
|
||||
211: "pi+",
|
||||
-211: "pi-",
|
||||
13: "mu-",
|
||||
-13: "mu+",
|
||||
321: "K+",
|
||||
-321: "K-",
|
||||
130: "K0L",
|
||||
}
|
||||
|
||||
|
||||
def pdg_label(code: int) -> str:
|
||||
"""Human-readable species label for a PDG code (falls back to the code)."""
|
||||
code = int(code)
|
||||
if code in PDG_NAMES:
|
||||
return PDG_NAMES[code]
|
||||
if abs(code) > 1_000_000_000:
|
||||
return f"ion {code}"
|
||||
return str(code)
|
||||
|
||||
|
||||
def material_label(name: str) -> str:
|
||||
"""Display label for a Geant4 material, dropping the ``G4_`` prefix."""
|
||||
return name[3:] if name.startswith("G4_") else name
|
||||
|
||||
|
||||
def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
|
||||
"""Equal-population (quantile) bin edges over per-event incident energies.
|
||||
|
||||
Returns ``n_bins + 1`` monotonically non-decreasing edges. The top edge is
|
||||
nudged up so the largest value falls inside the last bin under a
|
||||
right-open convention. Degenerate (single-value) input widens by +/-0.5.
|
||||
"""
|
||||
incident_E = np.asarray(incident_E, dtype=np.float64)
|
||||
edges = np.quantile(incident_E, np.linspace(0.0, 1.0, n_bins + 1))
|
||||
edges = np.unique(edges)
|
||||
if edges.size < 2:
|
||||
v = edges[0] if edges.size else 0.0
|
||||
edges = np.array([v - 0.5, v + 0.5])
|
||||
edges[-1] = np.nextafter(edges[-1], np.inf)
|
||||
return edges
|
||||
|
||||
|
||||
def energy_bin_labels(edges: np.ndarray) -> list[str]:
|
||||
"""``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
|
||||
return [
|
||||
f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
|
||||
]
|
||||
|
||||
|
||||
def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
|
||||
"""Bin index of ``value`` under arbitrary (possibly non-uniform) ``edges``.
|
||||
|
||||
``bin = (#interior edges <= value)``, clipped to ``[0, n_bins-1]`` — matches
|
||||
``np.digitize(value, edges[1:-1])`` and works for the quantile energy edges.
|
||||
Vectorized as a sum of boolean comparisons; no per-row Python.
|
||||
"""
|
||||
interior = [float(e) for e in edges[1:-1]]
|
||||
n_bins = len(edges) - 1
|
||||
idx = pl.lit(0, dtype=pl.Int32)
|
||||
for e in interior:
|
||||
idx = idx + (value >= e).cast(pl.Int32)
|
||||
return idx.clip(0, n_bins - 1)
|
||||
|
||||
|
||||
def event_energy_bins(
|
||||
lf: pl.LazyFrame, edges: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.
|
||||
|
||||
Incident energy is ``max(pre_E)`` per event (the primary). One bounded
|
||||
streaming ``group_by``; the tiny per-event result is digitized in numpy.
|
||||
"""
|
||||
per_event = (
|
||||
lf.group_by("event_id")
|
||||
.agg(pl.col("pre_E").max().alias("incident_E"))
|
||||
.collect(engine="streaming")
|
||||
.sort("event_id")
|
||||
)
|
||||
event_ids = per_event["event_id"].to_numpy()
|
||||
incident = per_event["incident_E"].to_numpy()
|
||||
bin_idx = np.clip(np.digitize(incident, edges[1:-1]), 0, len(edges) - 2)
|
||||
return event_ids, bin_idx.astype(np.int64)
|
||||
@@ -0,0 +1,267 @@
|
||||
"""Streaming compute primitives — the reduce half of the analysis.
|
||||
|
||||
Everything here turns a (possibly larger-than-RAM) LazyFrame into a *compact*
|
||||
numpy/DataFrame artifact in bounded memory, and never imports plotstyle so it can
|
||||
run on an HTCondor worker. Efficiency rules (see the plan's "Histogram
|
||||
efficiency" section):
|
||||
|
||||
* ``hist1d`` is a single streaming ``group_by([group, bin]).len()`` pass against
|
||||
**fixed** edges (no min/max range pass) with a strict column projection — only
|
||||
the columns the value/group expressions reference are read from the parquet.
|
||||
* the multi-quantity reductions (``event_scalars``, profiles, ``species_share``,
|
||||
``leakage_fraction``) each emit *all* their outputs from one ``group_by``.
|
||||
* per-event -> per-row lookups (shower entry/axis) use ``replace_strict`` (a hash
|
||||
map applied as an expression, bounded memory), never a streaming join.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.constants import TERM_ESCAPED
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1-D histogram primitive
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
|
||||
"""Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
|
||||
return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
|
||||
|
||||
|
||||
def hist1d(
|
||||
lf: pl.LazyFrame,
|
||||
value: pl.Expr,
|
||||
edges: np.ndarray,
|
||||
group: pl.Expr | None = None,
|
||||
) -> dict[object, np.ndarray]:
|
||||
"""Streaming histogram of ``value`` over fixed uniform ``edges``, by ``group``.
|
||||
|
||||
Returns ``{group_key: counts}`` (counts is an ``int64`` array of length
|
||||
``len(edges)-1``). One hash pass; runtime is independent of group cardinality,
|
||||
so every pdg/material/energy stratum falls out together. Only the tiny
|
||||
``(n_groups x nbins)`` result is materialized.
|
||||
"""
|
||||
lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1
|
||||
group = pl.lit(0, dtype=pl.Int64) if group is None else group
|
||||
res = (
|
||||
lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
|
||||
.group_by("_g", "_b")
|
||||
.agg(pl.len().alias("_n"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
out: dict[object, np.ndarray] = {}
|
||||
for g, b, n in res.iter_rows():
|
||||
out.setdefault(g, np.zeros(nbins, dtype=np.int64))[b] = n
|
||||
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)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def event_scalars(lf: pl.LazyFrame) -> pl.DataFrame:
|
||||
"""One row per event: total/mean deposited energy, path length, step count.
|
||||
|
||||
Columns: ``event_id, total_edep, total_length, n_steps, mean_length,
|
||||
incident_E`` (incident = ``max(pre_E)``, the primary). The caller chooses
|
||||
whether ``lf`` includes the rollout's synthetic termination rows — pass the
|
||||
full scan for energy totals (they carry the deposited remainder), physical
|
||||
steps only for step-count / mean-length.
|
||||
"""
|
||||
return (
|
||||
lf.group_by("event_id")
|
||||
.agg(
|
||||
pl.col("edep").sum().alias("total_edep"),
|
||||
pl.col("step_length").sum().alias("total_length"),
|
||||
pl.len().alias("n_steps"),
|
||||
pl.col("pre_E").max().alias("incident_E"),
|
||||
)
|
||||
.with_columns((pl.col("total_length") / pl.col("n_steps")).alias("mean_length"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Shower shape: entry/axis + edep-weighted longitudinal / transverse profiles
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def entry_axis(lf: pl.LazyFrame) -> pl.DataFrame:
|
||||
"""Per-event shower entry point and axis (from the highest-``pre_E`` step).
|
||||
|
||||
One bounded ``group_by``: the primary's ``pre_pos`` becomes the entry point
|
||||
and its ``pre_dir`` the shower axis.
|
||||
"""
|
||||
return (
|
||||
lf.group_by("event_id")
|
||||
.agg(
|
||||
pl.col("pre_x").get(pl.col("pre_E").arg_max()).alias("entry_x"),
|
||||
pl.col("pre_y").get(pl.col("pre_E").arg_max()).alias("entry_y"),
|
||||
pl.col("pre_z").get(pl.col("pre_E").arg_max()).alias("entry_z"),
|
||||
pl.col("pre_dx").get(pl.col("pre_E").arg_max()).alias("axis_x"),
|
||||
pl.col("pre_dy").get(pl.col("pre_E").arg_max()).alias("axis_y"),
|
||||
pl.col("pre_dz").get(pl.col("pre_E").arg_max()).alias("axis_z"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
.sort("event_id")
|
||||
)
|
||||
|
||||
|
||||
_ENTRY_AXIS_COLS = ("entry_x", "entry_y", "entry_z", "axis_x", "axis_y", "axis_z")
|
||||
|
||||
|
||||
def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
|
||||
"""Broadcast each event's entry/axis onto its rows via ``replace_strict``.
|
||||
|
||||
A hash map applied as an expression — streams in bounded memory, unlike a
|
||||
join which would buffer the whole file-sized left side.
|
||||
"""
|
||||
ids = entry["event_id"].to_numpy()
|
||||
return lf.with_columns(
|
||||
pl.col("event_id")
|
||||
.replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64)
|
||||
.alias(col)
|
||||
for col in _ENTRY_AXIS_COLS
|
||||
)
|
||||
|
||||
|
||||
def depth_expr() -> pl.Expr:
|
||||
"""Signed distance of ``post_pos`` from the entry point along the shower axis."""
|
||||
dx = pl.col("post_x") - pl.col("entry_x")
|
||||
dy = pl.col("post_y") - pl.col("entry_y")
|
||||
dz = pl.col("post_z") - pl.col("entry_z")
|
||||
return dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
|
||||
|
||||
|
||||
def transverse_expr() -> pl.Expr:
|
||||
"""Perpendicular distance of ``post_pos`` from the shower axis."""
|
||||
dx = pl.col("post_x") - pl.col("entry_x")
|
||||
dy = pl.col("post_y") - pl.col("entry_y")
|
||||
dz = pl.col("post_z") - pl.col("entry_z")
|
||||
depth = dx * pl.col("axis_x") + dy * pl.col("axis_y") + dz * pl.col("axis_z")
|
||||
tx = dx - depth * pl.col("axis_x")
|
||||
ty = dy - depth * pl.col("axis_y")
|
||||
tz = dz - depth * pl.col("axis_z")
|
||||
return (tx**2 + ty**2 + tz**2).sqrt()
|
||||
|
||||
|
||||
def profile_partial(
|
||||
lf: pl.LazyFrame,
|
||||
coord: pl.Expr,
|
||||
edges: np.ndarray,
|
||||
weight: pl.Expr,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""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).
|
||||
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 = (
|
||||
lf.select(
|
||||
"event_id",
|
||||
_bin_expr(coord, lo, hi, nbins).alias("_b"),
|
||||
weight.alias("_w"),
|
||||
)
|
||||
.group_by("event_id", "_b")
|
||||
.agg(pl.col("_w").sum().alias("_ws"))
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
ev = grid["event_id"].to_numpy()
|
||||
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 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])
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Species contribution and leakage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def species_share(lf: pl.LazyFrame) -> pl.DataFrame:
|
||||
"""Total deposited energy per PDG species (``pdg, total_edep``), one pass."""
|
||||
return (
|
||||
lf.group_by("pdg")
|
||||
.agg(pl.col("edep").sum().alias("total_edep"))
|
||||
.collect(engine="streaming")
|
||||
.sort("total_edep", descending=True)
|
||||
)
|
||||
|
||||
|
||||
def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
|
||||
"""Per-event escaped-energy fraction (rollout only), one bounded pass.
|
||||
|
||||
Escaped rows carry the leaked energy in ``pre_E`` (``edep`` is 0 there); the
|
||||
fraction is ``escaped / (deposited + escaped)`` per event.
|
||||
"""
|
||||
per_event = (
|
||||
lf.group_by("event_id")
|
||||
.agg(
|
||||
pl.col("edep").sum().alias("deposited"),
|
||||
pl.col("pre_E")
|
||||
.filter(pl.col("termination_reason") == TERM_ESCAPED)
|
||||
.sum()
|
||||
.alias("escaped"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
deposited = per_event["deposited"].to_numpy()
|
||||
escaped = per_event["escaped"].fill_null(0.0).to_numpy()
|
||||
total = deposited + escaped
|
||||
return np.where(total > 0, escaped / total, 0.0)
|
||||
@@ -0,0 +1,65 @@
|
||||
"""The compact, self-describing artifact a compute job produces per plot.
|
||||
|
||||
Serialized as small JSON (no pickle, no per-event arrays) so it is trivially
|
||||
transferable off the batch worker and human-inspectable. ``render.py`` dispatches
|
||||
on ``kind`` and needs nothing but this file.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
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)
|
||||
# "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
|
||||
class Reduced:
|
||||
id: str
|
||||
family: str
|
||||
kind: str
|
||||
title: str
|
||||
xlabel: str
|
||||
payload: dict
|
||||
meta: dict = field(default_factory=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) -> "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()))
|
||||
@@ -0,0 +1,361 @@
|
||||
"""Render reduced artifacts to styled PDFs + gallery metadata (the local step).
|
||||
|
||||
This is the *only* module that imports ``plotstyle`` (ETPlot's KIT matplotlib
|
||||
theme), which renders through a real LaTeX toolchain — so it runs on the
|
||||
submit/login node, never on a compute worker. It reads nothing but the small
|
||||
``Reduced`` JSON files a run produced, so it is fully decoupled from the heavy
|
||||
streaming compute.
|
||||
|
||||
For each reduced artifact it writes ``<out>/<family>/<id>.pdf`` plus a sibling
|
||||
``<id>.yaml`` (per-plot gallery metadata) and a per-family ``metadata.yaml``.
|
||||
Optionally runs ``gallery generate`` to build the static HTML site.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import subprocess
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import plotstyle as ps
|
||||
import yaml
|
||||
|
||||
from giant.analysis.reduced import Reduced
|
||||
|
||||
_SERIES_LABELS = {"rollout": "rollout", "reference": "reference (Geant4)"}
|
||||
|
||||
|
||||
def _density(counts: list[int] | np.ndarray, edges: np.ndarray) -> np.ndarray:
|
||||
counts = np.asarray(counts, dtype=np.float64)
|
||||
total = counts.sum()
|
||||
if total == 0:
|
||||
return counts
|
||||
return counts / (total * (edges[1] - edges[0]))
|
||||
|
||||
|
||||
def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> None:
|
||||
for key in ("reference", "rollout"):
|
||||
if key in series:
|
||||
ax.stairs(_density(series[key], edges), edges, label=_SERIES_LABELS[key])
|
||||
if log_y:
|
||||
ax.set_yscale("log")
|
||||
|
||||
|
||||
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
|
||||
|
||||
|
||||
def _render_overlay(r: Reduced, params: dict):
|
||||
edges = np.asarray(r.payload["edges"])
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
_overlay(ax, edges, r.payload, r.payload.get("log_y", False))
|
||||
ax.set_xlabel(r.xlabel)
|
||||
ax.set_ylabel("density")
|
||||
ps.style_legend(ax, title="source")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_single(r: Reduced, params: dict):
|
||||
edges = np.asarray(r.payload["edges"])
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
ax.stairs(
|
||||
_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
|
||||
)
|
||||
if r.payload.get("log_y"):
|
||||
ax.set_yscale("log")
|
||||
ax.set_xlabel(r.xlabel)
|
||||
ax.set_ylabel("density")
|
||||
ps.style_legend(ax, title="source")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_grouped(r: Reduced, params: dict):
|
||||
edges = np.asarray(r.payload["edges"])
|
||||
groups = r.payload["groups"]
|
||||
labels = list(groups)
|
||||
n = len(labels)
|
||||
ncols = min(3, n) or 1
|
||||
nrows = (n + ncols - 1) // ncols
|
||||
fig, axes = ps.new_figure(
|
||||
"slide-16x9",
|
||||
title=r.title,
|
||||
params=params,
|
||||
nrows=nrows,
|
||||
ncols=ncols,
|
||||
squeeze=False,
|
||||
)
|
||||
flat = axes.ravel()
|
||||
for i, lbl in enumerate(labels):
|
||||
ax = flat[i]
|
||||
_overlay(ax, edges, groups[lbl], r.payload.get("log_y", False))
|
||||
ax.set_title(lbl, fontsize=8)
|
||||
ax.set_xlabel(r.xlabel)
|
||||
for j in range(n, len(flat)):
|
||||
flat[j].set_visible(False)
|
||||
ps.style_legend(flat[0], title="source")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_profile(r: Reduced, params: dict):
|
||||
edges = np.asarray(r.payload["edges"])
|
||||
centers = 0.5 * (edges[:-1] + edges[1:])
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
for key in ("reference", "rollout"):
|
||||
mean = np.asarray(r.payload[f"{key}_mean"])
|
||||
std = np.asarray(r.payload[f"{key}_std"])
|
||||
(line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
|
||||
ax.fill_between(
|
||||
centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
|
||||
)
|
||||
ax.set_xlabel(r.xlabel)
|
||||
ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
|
||||
ps.style_legend(ax, title="source")
|
||||
return fig
|
||||
|
||||
|
||||
def _render_bar(r: Reduced, params: dict):
|
||||
labels = r.payload["labels"]
|
||||
x = np.arange(len(labels))
|
||||
width = 0.4
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
ax.bar(
|
||||
x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
|
||||
)
|
||||
ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(labels, rotation=45, ha="right")
|
||||
ax.set_ylabel(r.payload.get("ylabel", "value"))
|
||||
ps.style_legend(ax, title="source")
|
||||
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, _figure_params(run_meta or {}))
|
||||
|
||||
|
||||
def _plot_metadata(r: Reduced, run_meta: dict) -> dict:
|
||||
meta = {
|
||||
"title": r.title,
|
||||
"description": f"Rollout vs reference: {r.title}.",
|
||||
"plot_type": r.kind,
|
||||
"family": r.family,
|
||||
}
|
||||
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
|
||||
|
||||
|
||||
def render_all(
|
||||
reduced_dir: str | Path,
|
||||
out_dir: str | Path,
|
||||
run_meta: dict | None = None,
|
||||
*,
|
||||
run_gallery: bool = False,
|
||||
) -> list[Path]:
|
||||
"""Render every reduced artifact under ``reduced_dir`` to a PDF tree.
|
||||
|
||||
Writes ``<out>/<family>/<id>.pdf`` + ``<id>.yaml`` and a per-family
|
||||
``metadata.yaml`` (carrying the run's checkpoint/paths as gallery params).
|
||||
Returns the list of PDF paths written.
|
||||
"""
|
||||
ps.use()
|
||||
run_meta = run_meta or {}
|
||||
reduced_dir, out_dir = Path(reduced_dir), Path(out_dir)
|
||||
pdfs: list[Path] = []
|
||||
families: set[str] = set()
|
||||
|
||||
for jf in sorted(reduced_dir.glob("*.json")):
|
||||
r = Reduced.load(jf)
|
||||
family_dir = out_dir / r.family
|
||||
family_dir.mkdir(parents=True, exist_ok=True)
|
||||
families.add(r.family)
|
||||
fig = render(r, run_meta)
|
||||
ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
|
||||
(family_dir / f"{r.id}.yaml").write_text(
|
||||
yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
|
||||
)
|
||||
pdfs.append(family_dir / f"{r.id}.pdf")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
plt.close(fig)
|
||||
|
||||
# Root + per-family gallery metadata.
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
(out_dir / "metadata.yaml").write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"title": run_meta.get("title", "GIANT rollout analysis"),
|
||||
"description": "Autoregressive rollout compared against held-out Geant4 reference steps.",
|
||||
"experiment": "GIANT",
|
||||
"parameters": {k: v for k, v in run_meta.items() if k != "title"},
|
||||
},
|
||||
sort_keys=False,
|
||||
)
|
||||
)
|
||||
for fam in families:
|
||||
(out_dir / fam / "metadata.yaml").write_text(
|
||||
yaml.safe_dump(
|
||||
{"title": fam, "description": f"{fam} plots."}, sort_keys=False
|
||||
)
|
||||
)
|
||||
|
||||
if run_gallery:
|
||||
subprocess.run(["gallery", "generate", "--source", str(out_dir)], check=True)
|
||||
return pdfs
|
||||
|
||||
|
||||
def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
|
||||
"""Render a prepped run directory: ``<run_dir>/reduced`` → ``<run_dir>/plots``.
|
||||
|
||||
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, merge_all
|
||||
|
||||
run_dir = Path(run_dir)
|
||||
merge_all(run_dir)
|
||||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||||
run_meta = {
|
||||
"title": meta.title,
|
||||
"rollout": meta.rollout,
|
||||
"reference": meta.reference,
|
||||
**meta.plot_meta,
|
||||
}
|
||||
return render_all(
|
||||
run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
|
||||
)
|
||||
@@ -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))
|
||||
@@ -0,0 +1,179 @@
|
||||
"""Canonical world-frame LazyFrame builders for the two sides of a comparison.
|
||||
|
||||
The analysis compares one autoregressive ``giant rollout`` (the *generated* side)
|
||||
against a raw miniCaloSim steps file (the *reference* / real side). Both carry a
|
||||
**shared world-frame physical column subset** under identical names, so no
|
||||
renaming or coordinate decode is needed — everything is already in world-frame
|
||||
mm / MeV:
|
||||
|
||||
event_id, track_id, step_no, pdg,
|
||||
pre_x, pre_y, pre_z, pre_E, pre_dx, pre_dy, pre_dz,
|
||||
post_x, post_y, post_z, post_E, post_dx, post_dy, post_dz,
|
||||
edep, step_length, material, layer_id
|
||||
|
||||
(rollout: ``rollout.py:_RECORD_KEYS``; reference: minicalosim ``RunAction.cc``
|
||||
Steps ntuple passed through by ``dwarf convert``.)
|
||||
|
||||
The two files differ in their *extra* columns — the rollout adds ``parent_id``,
|
||||
``generation``, ``n_sec_pred``, ``termination_reason``; the reference adds
|
||||
``process``, field columns, ``child_track_ids`` and the ``sec_*_list`` secondary
|
||||
birth-state lists. Those are only touched by the side-specific helpers here
|
||||
(synthetic-row filtering, the secondary view).
|
||||
|
||||
Nothing in this module (or ``reduce.py``) imports plotstyle — compute runs on
|
||||
HTCondor workers that have no LaTeX toolchain.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from enum import Enum
|
||||
from pathlib import Path
|
||||
|
||||
import polars as pl
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
from giant.constants import (
|
||||
PREDICT_COORD_METADATA_KEY,
|
||||
ROLLOUT_COORD_VALUE,
|
||||
TERM_ENERGY_CUTOFF,
|
||||
TERM_ESCAPED,
|
||||
TERM_MAX_STEPS,
|
||||
TERM_UNKNOWN_PDG,
|
||||
)
|
||||
|
||||
# The world-frame physical columns both sides share under identical names.
|
||||
PHYS_COLS: tuple[str, ...] = (
|
||||
"event_id",
|
||||
"pdg",
|
||||
"pre_x",
|
||||
"pre_y",
|
||||
"pre_z",
|
||||
"pre_E",
|
||||
"pre_dx",
|
||||
"pre_dy",
|
||||
"pre_dz",
|
||||
"post_x",
|
||||
"post_y",
|
||||
"post_z",
|
||||
"post_E",
|
||||
"post_dx",
|
||||
"post_dy",
|
||||
"post_dz",
|
||||
"edep",
|
||||
"step_length",
|
||||
"material",
|
||||
"layer_id",
|
||||
)
|
||||
|
||||
# Rollout rows written purely for bookkeeping (a track's forced stop): they carry
|
||||
# step_length=0, post_pos=pre_pos, and — for every reason but escaped — the
|
||||
# track's whole remaining pre_E dumped into edep so the shower still conserves
|
||||
# energy. They are not physical steps (the reference has no equivalent), so a
|
||||
# per-step marginal comparison must drop them; a per-event energy total must keep
|
||||
# them. See rollout.py's terminal-row handling.
|
||||
SYNTHETIC_TERMINATION_REASONS: frozenset[str] = frozenset(
|
||||
{TERM_ESCAPED, TERM_UNKNOWN_PDG, TERM_ENERGY_CUTOFF, TERM_MAX_STEPS}
|
||||
)
|
||||
|
||||
|
||||
class Side(str, Enum):
|
||||
"""Which of the two comparison inputs a file is."""
|
||||
|
||||
rollout = "rollout"
|
||||
reference = "reference"
|
||||
|
||||
|
||||
def _check_rollout_metadata(path: Path) -> None:
|
||||
"""Raise if ``path`` carries coord metadata that isn't the rollout tag.
|
||||
|
||||
A missing tag (older rollout output, predating tagging) is allowed through,
|
||||
matching ``giant rollout``'s own leniency; a tag that is present but wrong is
|
||||
a real mismatch and worth failing on before the column layout is trusted.
|
||||
"""
|
||||
metadata = pq.read_schema(path).metadata or {}
|
||||
coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
|
||||
if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
|
||||
raise ValueError(
|
||||
f"{path} is not a rollout file (coord={coord.decode()!r}); "
|
||||
"expected `giant rollout` output"
|
||||
)
|
||||
|
||||
|
||||
def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
"""Lazily scan one side's file, verifying the rollout tag when applicable.
|
||||
|
||||
Returns the *full* lazy scan (no column projection) so downstream reductions
|
||||
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.with_columns(pl.col("pdg").cast(pl.Int64))
|
||||
path = Path(source)
|
||||
if side is Side.rollout:
|
||||
_check_rollout_metadata(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:
|
||||
"""Real, physical steps only — drops the rollout's synthetic termination rows.
|
||||
|
||||
The predicate is pushed down so the dropped rows are never decoded. The
|
||||
reference has no such rows, so it is returned unchanged.
|
||||
"""
|
||||
if side is Side.reference:
|
||||
return lf
|
||||
return lf.filter(
|
||||
~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS))
|
||||
)
|
||||
|
||||
|
||||
def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
"""Per-secondary birth state, one row per produced secondary.
|
||||
|
||||
Canonical columns: ``event_id, energy, pdg, sdx, sdy, sdz`` (birth energy in
|
||||
MeV, PDG code, birth unit direction in the world frame). The two sides encode
|
||||
secondaries differently:
|
||||
|
||||
- rollout: each secondary is its own track, so its birth state is the row with
|
||||
``generation > 0`` and ``step_no == 0`` (``pre_E`` / ``pre_dir`` there).
|
||||
- reference: secondaries live in per-parent-step ``sec_*_list`` columns; the
|
||||
lists are exploded together and empty (no-secondary) steps drop out.
|
||||
"""
|
||||
if side is Side.rollout:
|
||||
return lf.filter((pl.col("generation") > 0) & (pl.col("step_no") == 0)).select(
|
||||
"event_id",
|
||||
pl.col("pre_E").alias("energy"),
|
||||
"pdg",
|
||||
pl.col("pre_dx").alias("sdx"),
|
||||
pl.col("pre_dy").alias("sdy"),
|
||||
pl.col("pre_dz").alias("sdz"),
|
||||
)
|
||||
lists = ["sec_E_list", "sec_pdg_list", "sec_dx_list", "sec_dy_list", "sec_dz_list"]
|
||||
return (
|
||||
lf.select("event_id", *lists)
|
||||
.explode(lists)
|
||||
.drop_nulls("sec_E_list")
|
||||
.select(
|
||||
"event_id",
|
||||
pl.col("sec_E_list").alias("energy"),
|
||||
pl.col("sec_pdg_list").alias("pdg"),
|
||||
pl.col("sec_dx_list").alias("sdx"),
|
||||
pl.col("sec_dy_list").alias("sdy"),
|
||||
pl.col("sec_dz_list").alias("sdz"),
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,28 @@
|
||||
"""Per-step value expressions shared by ``context`` (range sizing) and ``catalog``.
|
||||
|
||||
Kept separate from both so the range-sizing prep and the plot registry agree on
|
||||
exactly what each variable *is*, with no import cycle. plotstyle-free.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import polars as pl
|
||||
|
||||
# Ranged marginal variables: name -> (axis label, value expression). Their
|
||||
# histogram ranges are sized from data in `context.build_context`.
|
||||
RANGED_VARS: dict[str, tuple[str, pl.Expr]] = {
|
||||
"step_length": ("step length [mm]", pl.col("step_length")),
|
||||
"edep": ("deposited energy [MeV]", pl.col("edep")),
|
||||
"delta_e": ("energy loss [MeV]", pl.col("pre_E") - pl.col("post_E")),
|
||||
"post_E": ("post-step energy [MeV]", pl.col("post_E")),
|
||||
}
|
||||
|
||||
|
||||
def cos_scatter_expr() -> pl.Expr:
|
||||
"""cos of the scattering angle: ``pre_dir . post_dir`` (both unit), in [-1, 1]."""
|
||||
dot = (
|
||||
pl.col("pre_dx") * pl.col("post_dx")
|
||||
+ pl.col("pre_dy") * pl.col("post_dy")
|
||||
+ pl.col("pre_dz") * pl.col("post_dz")
|
||||
)
|
||||
return dot.clip(-1.0, 1.0)
|
||||
+204
-1
@@ -891,7 +891,7 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda
|
||||
|
||||
Streams conditioning columns and keeps the highest-pre_E step per event_id —
|
||||
the codebase's convention for the primary (a secondary always carries less
|
||||
energy than its parent). See giant/analysis.py:_entry_axis_and_bin_edges.
|
||||
energy than its parent). See giant/analysis/reduce.py:entry_axis.
|
||||
"""
|
||||
best_E: dict[int, float] = {}
|
||||
best: dict[int, tuple] = {}
|
||||
@@ -1024,6 +1024,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()}
|
||||
@@ -1104,7 +1107,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"])),
|
||||
"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))
|
||||
@@ -1114,5 +1136,186 @@ def rollout(
|
||||
typer.echo(f"reference: {ref_path}")
|
||||
|
||||
|
||||
analyze_app = typer.Typer(
|
||||
no_args_is_help=True,
|
||||
help="Rollout-vs-reference analysis: parallel compute on HTCondor + local render.",
|
||||
)
|
||||
app.add_typer(analyze_app, name="analyze")
|
||||
|
||||
|
||||
@analyze_app.command("prep")
|
||||
def analyze_prep(
|
||||
rollout_yaml: Annotated[
|
||||
Path,
|
||||
typer.Argument(
|
||||
help="giant rollout YAML sidecar (names the rollout + reference files)"
|
||||
),
|
||||
],
|
||||
run_dir: Annotated[
|
||||
Optional[Path],
|
||||
typer.Option(
|
||||
"--run-dir",
|
||||
"-o",
|
||||
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
|
||||
|
||||
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,
|
||||
)
|
||||
typer.echo(f"run directory: {path}")
|
||||
|
||||
|
||||
@analyze_app.command("compute-one")
|
||||
def analyze_compute_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`")
|
||||
],
|
||||
chunk: Annotated[
|
||||
int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")
|
||||
] = 0,
|
||||
) -> None:
|
||||
"""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, 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}")
|
||||
|
||||
|
||||
@analyze_app.command("list")
|
||||
def analyze_list() -> None:
|
||||
"""Print every catalog plot id."""
|
||||
from giant.analysis import catalog_ids
|
||||
|
||||
for pid in catalog_ids():
|
||||
typer.echo(pid)
|
||||
|
||||
|
||||
@analyze_app.command("render")
|
||||
def analyze_render(
|
||||
run_dir: Annotated[
|
||||
Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")
|
||||
],
|
||||
gallery: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
"--gallery/--no-gallery", help="Run `gallery generate` after rendering"
|
||||
),
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Render reduced artifacts to styled PDFs + gallery metadata (local; needs LaTeX)."""
|
||||
from giant.analysis.render import render_run
|
||||
|
||||
pdfs = render_run(run_dir, run_gallery=gallery)
|
||||
typer.echo(f"rendered {len(pdfs)} plots → {Path(run_dir) / 'plots'}")
|
||||
|
||||
|
||||
@analyze_app.command("submit")
|
||||
def analyze_submit(
|
||||
rollout_yaml: Annotated[Path, typer.Argument(help="giant rollout YAML sidecar")],
|
||||
accounting_group: Annotated[str, typer.Option("--accounting-group")],
|
||||
run_dir: Annotated[
|
||||
Optional[Path],
|
||||
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")
|
||||
] = "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 x chunk), then submit."""
|
||||
import subprocess
|
||||
|
||||
from giant.analysis import SubmitConfig, prep, write_submit
|
||||
|
||||
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,
|
||||
repo_dir=Path.cwd(),
|
||||
docker_image=docker_image,
|
||||
request_memory_mb=request_memory,
|
||||
remote=remote,
|
||||
n_chunks=chunks,
|
||||
)
|
||||
sub = write_submit(cfg)
|
||||
typer.echo(f"run directory: {path}")
|
||||
typer.echo(f"wrote submit description: {sub}")
|
||||
if dry_run:
|
||||
typer.echo("dry-run: not submitting")
|
||||
return
|
||||
subprocess.run(["condor_submit", str(sub)], check=True)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
app()
|
||||
|
||||
@@ -188,6 +188,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,
|
||||
|
||||
@@ -39,6 +39,9 @@ analysis = [
|
||||
"matplotlib>=3.8,<4",
|
||||
"polars>=1.0,<2",
|
||||
"ipykernel>=7.3.0",
|
||||
# KIT matplotlib theme, published from git.larsbogner.de. Only the local
|
||||
# `giant analyze render` step imports it; compute workers never do.
|
||||
"plotstyle>=1.0.0",
|
||||
]
|
||||
|
||||
[project.scripts]
|
||||
@@ -65,6 +68,7 @@ torch = [
|
||||
{ index = "pytorch-cpu", extra = "cpu" },
|
||||
{ index = "pytorch-cu118", extra = "cuda" },
|
||||
]
|
||||
plotstyle = { index = "larsbogner" }
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "pytorch-cpu"
|
||||
@@ -75,3 +79,8 @@ explicit = true
|
||||
name = "pytorch-cu118"
|
||||
url = "https://download.pytorch.org/whl/cu118"
|
||||
explicit = true
|
||||
|
||||
[[tool.uv.index]]
|
||||
name = "larsbogner"
|
||||
url = "https://git.larsbogner.de/api/packages/lars/pypi/simple/"
|
||||
explicit = true
|
||||
|
||||
@@ -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.")
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,171 @@
|
||||
"""Tests for the streaming compute primitives (giant.analysis.reduce/sources/grouping)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import numpy as np
|
||||
import polars as pl
|
||||
|
||||
from giant.analysis import grouping as G
|
||||
from giant.analysis import reduce as R
|
||||
from giant.analysis.sources import (
|
||||
SYNTHETIC_TERMINATION_REASONS,
|
||||
Side,
|
||||
physical_steps,
|
||||
secondaries,
|
||||
)
|
||||
|
||||
|
||||
def _rollout_frame() -> pl.LazyFrame:
|
||||
# event 1: primary (2 steps) + 1 secondary track + 1 escaped bookkeeping row
|
||||
# event 2: primary (1 step)
|
||||
return pl.DataFrame(
|
||||
{
|
||||
"event_id": [1, 1, 1, 1, 2],
|
||||
"track_id": [0, 0, 1, 0, 0],
|
||||
"parent_id": [-1, -1, 0, -1, -1],
|
||||
"generation": [0, 0, 1, 0, 0],
|
||||
"step_no": [0, 1, 0, 99, 0],
|
||||
"pdg": [11, 11, 22, 11, 11],
|
||||
"pre_x": [0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
"pre_y": [0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
"pre_z": [0.0, 1.0, 1.0, 2.0, 0.0],
|
||||
"pre_E": [100.0, 60.0, 20.0, 30.0, 50.0],
|
||||
"pre_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
|
||||
"pre_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
"pre_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
|
||||
"post_x": [0.0, 0.0, 1.0, 0.0, 0.0],
|
||||
"post_y": [0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
"post_z": [1.0, 2.0, 1.0, 2.0, 1.0],
|
||||
"post_E": [60.0, 30.0, 0.0, 0.0, 20.0],
|
||||
"post_dx": [0.0, 0.0, 1.0, 0.0, 0.0],
|
||||
"post_dy": [0.0, 0.0, 0.0, 0.0, 0.0],
|
||||
"post_dz": [1.0, 1.0, 0.0, 1.0, 1.0],
|
||||
"edep": [40.0, 30.0, 20.0, 0.0, 30.0],
|
||||
"step_length": [1.0, 1.0, 1.0, 0.0, 1.0],
|
||||
"material": ["G4_PbWO4"] * 5,
|
||||
"layer_id": [0, 1, 1, -1, 0],
|
||||
"n_sec_pred": [1, 0, 0, 0, 0],
|
||||
"termination_reason": [
|
||||
"",
|
||||
"natural_end",
|
||||
"natural_end",
|
||||
"escaped",
|
||||
"natural_end",
|
||||
],
|
||||
}
|
||||
).lazy()
|
||||
|
||||
|
||||
def _reference_frame() -> pl.LazyFrame:
|
||||
return pl.DataFrame(
|
||||
{
|
||||
"event_id": [1, 1, 2],
|
||||
"track_id": [0, 0, 0],
|
||||
"step_no": [0, 1, 0],
|
||||
"pdg": [11, 11, 11],
|
||||
"pre_x": [0.0, 0.0, 0.0],
|
||||
"pre_y": [0.0, 0.0, 0.0],
|
||||
"pre_z": [0.0, 1.0, 0.0],
|
||||
"pre_E": [100.0, 60.0, 50.0],
|
||||
"pre_dx": [0.0, 0.0, 0.0],
|
||||
"pre_dy": [0.0, 0.0, 0.0],
|
||||
"pre_dz": [1.0, 1.0, 1.0],
|
||||
"post_x": [0.0, 0.0, 0.0],
|
||||
"post_y": [0.0, 0.0, 0.0],
|
||||
"post_z": [1.0, 2.0, 1.0],
|
||||
"post_E": [60.0, 30.0, 20.0],
|
||||
"post_dx": [0.0, 0.0, 0.0],
|
||||
"post_dy": [0.0, 0.0, 0.0],
|
||||
"post_dz": [1.0, 1.0, 1.0],
|
||||
"edep": [40.0, 30.0, 30.0],
|
||||
"step_length": [1.0, 1.0, 1.0],
|
||||
"material": ["G4_PbWO4", "G4_PbWO4", "G4_Pb"],
|
||||
"layer_id": [0, 1, 0],
|
||||
"sec_E_list": [[20.0], [], [10.0]],
|
||||
"sec_pdg_list": [[22], [], [22]],
|
||||
"sec_dx_list": [[1.0], [], [0.0]],
|
||||
"sec_dy_list": [[0.0], [], [0.0]],
|
||||
"sec_dz_list": [[0.0], [], [1.0]],
|
||||
}
|
||||
).lazy()
|
||||
|
||||
|
||||
def test_hist1d_overall_and_grouped():
|
||||
lf = _rollout_frame()
|
||||
edges = np.linspace(0.0, 50.0, 6) # width 10
|
||||
h = R.hist1d(lf, pl.col("edep"), edges)
|
||||
# edep values: 40,30,20,0,30 -> bins [0),[10),[20),[30),[40)
|
||||
assert h[0].tolist() == [1, 0, 1, 2, 1]
|
||||
# grouped by pdg: pdg 22 has a single edep=20
|
||||
hg = R.hist1d(lf, pl.col("edep"), edges, group=pl.col("pdg"))
|
||||
assert hg[22].tolist() == [0, 0, 1, 0, 0]
|
||||
assert hg[11].sum() == 4
|
||||
|
||||
|
||||
def test_physical_steps_drops_synthetic_rollout_rows_only():
|
||||
lf = _rollout_frame()
|
||||
phys = physical_steps(lf, Side.rollout).collect()
|
||||
assert phys.height == 4 # dropped the escaped bookkeeping row
|
||||
assert "escaped" not in phys["termination_reason"].to_list()
|
||||
assert SYNTHETIC_TERMINATION_REASONS # non-empty guard
|
||||
# reference passes through unchanged
|
||||
ref = _reference_frame()
|
||||
assert physical_steps(ref, Side.reference).collect().height == ref.collect().height
|
||||
|
||||
|
||||
def test_event_scalars_totals_include_all_rows():
|
||||
lf = _rollout_frame()
|
||||
es = R.event_scalars(lf).sort("event_id")
|
||||
row1 = es.filter(pl.col("event_id") == 1).to_dicts()[0]
|
||||
assert row1["total_edep"] == 90.0 # 40+30+20+0
|
||||
assert row1["incident_E"] == 100.0
|
||||
assert row1["n_steps"] == 4
|
||||
|
||||
|
||||
def test_secondaries_rollout_vs_reference_align():
|
||||
r = secondaries(_rollout_frame(), Side.rollout).collect().sort("event_id")
|
||||
assert r["energy"].to_list() == [20.0] # only the generation>0, step_no==0 row
|
||||
assert r["pdg"].to_list() == [22]
|
||||
t = secondaries(_reference_frame(), Side.reference).collect().sort("event_id")
|
||||
# two secondaries (event 1 and event 2); empty list dropped
|
||||
assert sorted(t["energy"].to_list()) == [10.0, 20.0]
|
||||
assert t["pdg"].to_list() == [22, 22]
|
||||
|
||||
|
||||
def test_leakage_fraction():
|
||||
frac = R.leakage_fraction(_rollout_frame())
|
||||
# event 1: escaped pre_E=30, deposited=90 -> 30/120 = 0.25; event 2: 0
|
||||
assert sorted(round(f, 6) for f in frac) == [0.0, 0.25]
|
||||
|
||||
|
||||
def test_weighted_profile_matches_manual_bincount():
|
||||
lf = _rollout_frame()
|
||||
ea = R.entry_axis(lf)
|
||||
lf2 = R.attach_entry_axis(lf, ea)
|
||||
edges = np.linspace(0.0, 3.0, 4) # depth bins along +z
|
||||
mean, std = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
|
||||
assert mean.shape == (3,)
|
||||
# totals conserved: sum over bins == mean total edep per event
|
||||
assert np.isclose(mean.sum() * 1, (90.0 + 30.0) / 2) # 2 events
|
||||
|
||||
|
||||
def test_energy_bins_edges_and_event_map():
|
||||
incident = np.array([100.0, 100.0, 1000.0, 1000.0])
|
||||
edges = G.energy_bin_edges(incident, n_bins=2)
|
||||
assert len(edges) == 3 and edges[0] <= 100.0 < edges[-1]
|
||||
ids, bins = G.event_energy_bins(_rollout_frame(), edges)
|
||||
assert set(bins.tolist()) <= {0, 1}
|
||||
assert len(ids) == 2
|
||||
|
||||
|
||||
def test_digitize_expr_matches_numpy():
|
||||
edges = np.array([0.0, 10.0, 100.0, 1000.0])
|
||||
df = pl.DataFrame({"v": [5.0, 50.0, 500.0, 2000.0]})
|
||||
got = df.select(G.digitize_expr(pl.col("v"), edges).alias("b"))["b"].to_list()
|
||||
assert got == np.digitize([5.0, 50.0, 500.0, 2000.0], edges[1:-1]).tolist()
|
||||
|
||||
|
||||
def test_pdg_and_material_labels():
|
||||
assert G.pdg_label(22) == "gamma"
|
||||
assert G.pdg_label(999999) == "999999"
|
||||
assert G.material_label("G4_PbWO4") == "PbWO4"
|
||||
@@ -0,0 +1,153 @@
|
||||
"""Tests for the plot catalog: id uniqueness + every spec computes a valid Reduced."""
|
||||
|
||||
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, PlotSpec
|
||||
from giant.analysis.context import Context, build_context
|
||||
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
||||
|
||||
|
||||
def _build_ctx() -> Context:
|
||||
r, t = _rollout_frame(), _reference_frame()
|
||||
return build_context(
|
||||
r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
|
||||
)
|
||||
|
||||
|
||||
@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():
|
||||
ids = catalog_ids()
|
||||
assert ids and len(ids) == len(set(ids))
|
||||
# the required families are all present
|
||||
fams = {s.family for s in build_catalog()}
|
||||
assert {"marginals", "event", "shower", "species", "secondaries"} <= fams
|
||||
|
||||
|
||||
def test_get_spec_roundtrip_and_unknown():
|
||||
spec = get_spec("marginal_edep")
|
||||
assert spec.id == "marginal_edep" and spec.family == "marginals"
|
||||
with pytest.raises(KeyError):
|
||||
get_spec("does_not_exist")
|
||||
|
||||
|
||||
def test_every_spec_computes_valid_reduced(bundle: Bundle):
|
||||
for spec in build_catalog():
|
||||
r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx)
|
||||
assert r.id == spec.id
|
||||
assert r.kind in {
|
||||
"overlay_hist",
|
||||
"grouped_hist",
|
||||
"profile",
|
||||
"bar",
|
||||
"single_hist",
|
||||
"router_gating",
|
||||
"router_share",
|
||||
"unavailable",
|
||||
}
|
||||
assert r.title and r.xlabel
|
||||
_validate_payload(r)
|
||||
|
||||
|
||||
def _validate_payload(r) -> None:
|
||||
p = r.payload
|
||||
if r.kind == "overlay_hist":
|
||||
n = len(p["edges"]) - 1
|
||||
assert len(p["rollout"]) == n and len(p["reference"]) == n
|
||||
elif r.kind == "single_hist":
|
||||
assert len(p["rollout"]) == len(p["edges"]) - 1
|
||||
elif r.kind == "grouped_hist":
|
||||
n = len(p["edges"]) - 1
|
||||
assert p["groups"], "grouped hist must have at least one group"
|
||||
for g in p["groups"].values():
|
||||
assert len(g["rollout"]) == n and len(g["reference"]) == n
|
||||
elif r.kind == "profile":
|
||||
n = len(p["edges"]) - 1
|
||||
for k in ("rollout_mean", "rollout_std", "reference_mean", "reference_std"):
|
||||
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)
|
||||
@@ -0,0 +1,270 @@
|
||||
"""Tests for the rollout-YAML → run-directory flow, compute, and submit."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pyarrow.parquet as pq
|
||||
import pytest
|
||||
import yaml
|
||||
|
||||
from giant.analysis import (
|
||||
RunMeta,
|
||||
SubmitConfig,
|
||||
catalog_ids,
|
||||
compute_one,
|
||||
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 Partial, Reduced
|
||||
from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE
|
||||
from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
||||
|
||||
|
||||
def _write_inputs(tmp_path: Path) -> Path:
|
||||
"""Materialize rollout+reference parquet and a rollout YAML; return the YAML path."""
|
||||
rollout = tmp_path / "rollout.parquet"
|
||||
reference = tmp_path / "reference.parquet"
|
||||
tbl = _rollout_frame().collect().to_arrow()
|
||||
tbl = tbl.replace_schema_metadata({PREDICT_COORD_METADATA_KEY: ROLLOUT_COORD_VALUE})
|
||||
pq.write_table(tbl, rollout)
|
||||
_reference_frame().collect().write_parquet(reference)
|
||||
|
||||
yaml_path = tmp_path / "run.yaml"
|
||||
yaml_path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"prediction_id": "abcd1234ef",
|
||||
"output": str(rollout),
|
||||
"dataset": str(reference),
|
||||
"checkpoint": "/ckpt/best.pt",
|
||||
"kind": "rollout",
|
||||
"energy_cutoff": 0.1,
|
||||
"steps": 10,
|
||||
}
|
||||
)
|
||||
)
|
||||
return yaml_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,
|
||||
sample_rows=1000,
|
||||
)
|
||||
|
||||
|
||||
def test_load_rollout_yaml_requires_paths(tmp_path: Path):
|
||||
bad = tmp_path / "bad.yaml"
|
||||
bad.write_text(yaml.safe_dump({"output": "x.parquet"})) # no dataset
|
||||
with pytest.raises(ValueError):
|
||||
load_rollout_yaml(bad)
|
||||
|
||||
|
||||
def test_derive_run_dir_next_to_rollout():
|
||||
y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"}
|
||||
assert derive_run_dir(y) == Path("/data/analysis_abcd1234")
|
||||
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)
|
||||
assert run_dir == tmp_path / "analysis_abcd1234"
|
||||
assert (run_dir / "shared.json").exists()
|
||||
ctx = Context.load(run_dir / "shared.json")
|
||||
assert set(ctx.var_ranges) == {"step_length", "edep", "delta_e", "post_E"}
|
||||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||||
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_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):
|
||||
run_dir = _prep(_write_inputs(tmp_path))
|
||||
meta = RunMeta.load(run_dir / "run_meta.json")
|
||||
out = compute_reduced(
|
||||
"marginal_step_length",
|
||||
meta.rollout,
|
||||
meta.reference,
|
||||
run_dir / "shared.json",
|
||||
tmp_path / "r.json",
|
||||
)
|
||||
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 = cverstege/alma9-gridjob" in txt
|
||||
assert "requirements = TARGET.ProvidesETPResources" in txt
|
||||
assert "accounting_group = cms" in txt
|
||||
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
|
||||
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
|
||||
@@ -0,0 +1,92 @@
|
||||
"""Render smoke test — skipped where plotstyle / LaTeX is unavailable."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("plotstyle")
|
||||
|
||||
from giant.analysis.reduced import Reduced # noqa: E402
|
||||
|
||||
|
||||
def _try_render(reduced: list[Reduced], out: Path):
|
||||
from giant.analysis.render import render_all
|
||||
|
||||
for r in reduced:
|
||||
r.save(out / "reduced" / f"{r.id}.json")
|
||||
return render_all(out / "reduced", out / "plots")
|
||||
|
||||
|
||||
def test_render_one_of_each_kind(tmp_path: Path):
|
||||
reduced = [
|
||||
Reduced(
|
||||
"m",
|
||||
"marginals",
|
||||
"overlay_hist",
|
||||
"Overlay",
|
||||
"x",
|
||||
{
|
||||
"edges": [0, 1, 2, 3],
|
||||
"rollout": [1, 2, 3],
|
||||
"reference": [3, 2, 1],
|
||||
"log_y": False,
|
||||
},
|
||||
),
|
||||
Reduced(
|
||||
"g",
|
||||
"marginals",
|
||||
"grouped_hist",
|
||||
"Grouped",
|
||||
"x",
|
||||
{
|
||||
"edges": [0, 1, 2],
|
||||
"groups": {"a": {"rollout": [1, 2], "reference": [2, 1]}},
|
||||
"log_y": False,
|
||||
},
|
||||
),
|
||||
Reduced(
|
||||
"p",
|
||||
"shower",
|
||||
"profile",
|
||||
"Profile",
|
||||
"depth",
|
||||
{
|
||||
"edges": [0, 1, 2],
|
||||
"rollout_mean": [1, 2],
|
||||
"rollout_std": [0.1, 0.2],
|
||||
"reference_mean": [1.1, 1.9],
|
||||
"reference_std": [0.1, 0.1],
|
||||
"ylabel": "e",
|
||||
},
|
||||
),
|
||||
Reduced(
|
||||
"b",
|
||||
"species",
|
||||
"bar",
|
||||
"Bar",
|
||||
"species",
|
||||
{
|
||||
"labels": ["e-", "gamma"],
|
||||
"rollout": [0.6, 0.4],
|
||||
"reference": [0.5, 0.5],
|
||||
"ylabel": "frac",
|
||||
},
|
||||
),
|
||||
Reduced(
|
||||
"s",
|
||||
"species",
|
||||
"single_hist",
|
||||
"Single",
|
||||
"x",
|
||||
{"edges": [0, 1, 2], "rollout": [5, 1], "log_y": True},
|
||||
),
|
||||
]
|
||||
try:
|
||||
pdfs = _try_render(reduced, tmp_path)
|
||||
except RuntimeError as e: # LaTeX missing at render time
|
||||
pytest.skip(f"LaTeX rendering unavailable: {e}")
|
||||
assert len(pdfs) == len(reduced)
|
||||
assert all(p.exists() for p in pdfs)
|
||||
assert (tmp_path / "plots" / "metadata.yaml").exists()
|
||||
@@ -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",
|
||||
)
|
||||
|
||||
@@ -459,6 +459,7 @@ dependencies = [
|
||||
analysis = [
|
||||
{ name = "ipykernel" },
|
||||
{ name = "matplotlib" },
|
||||
{ name = "plotstyle" },
|
||||
{ name = "polars" },
|
||||
]
|
||||
convert = [
|
||||
@@ -477,6 +478,7 @@ dev = [
|
||||
{ name = "awkward" },
|
||||
{ name = "ipykernel" },
|
||||
{ name = "matplotlib" },
|
||||
{ name = "plotstyle" },
|
||||
{ name = "polars" },
|
||||
{ name = "pytest" },
|
||||
{ name = "ruff" },
|
||||
@@ -497,6 +499,7 @@ requires-dist = [
|
||||
{ name = "numpy", specifier = ">=1.26,<3" },
|
||||
{ name = "pandas", specifier = ">=2.2,<4" },
|
||||
{ name = "particle", specifier = ">=1.0,<2" },
|
||||
{ name = "plotstyle", marker = "extra == 'analysis'", specifier = ">=1.0.0", index = "https://git.larsbogner.de/api/packages/lars/pypi/simple/" },
|
||||
{ name = "polars", marker = "extra == 'analysis'", specifier = ">=1.0,<2" },
|
||||
{ name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" },
|
||||
{ name = "pyarrow", specifier = ">=16,<25" },
|
||||
@@ -1307,6 +1310,18 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/81/e6/cd9575ac904136b3cbf7aa7ee819ef86eedb7274e46f230e94ea4342e729/platformdirs-4.10.0-py3-none-any.whl", hash = "sha256:fb516cdb12eb0d857d0cd85a7c57cea4d060bee4578d6cf5a14dfdf8cbf8784a", size = 22743, upload-time = "2026-05-28T03:32:52.175Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "plotstyle"
|
||||
version = "1.0.0"
|
||||
source = { registry = "https://git.larsbogner.de/api/packages/lars/pypi/simple/" }
|
||||
dependencies = [
|
||||
{ name = "matplotlib" },
|
||||
]
|
||||
sdist = { url = "https://git.larsbogner.de/api/packages/lars/pypi/files/plotstyle/1.0.0/plotstyle-1.0.0.tar.gz", hash = "sha256:30f0c429077e5909ab1c89e0982c4351ffb71dc0688b177f78c35ed31b9bfc92" }
|
||||
wheels = [
|
||||
{ url = "https://git.larsbogner.de/api/packages/lars/pypi/files/plotstyle/1.0.0/plotstyle-1.0.0-py3-none-any.whl", hash = "sha256:f222a866fc81e7166cab78a8a30f2fc8e1436012c98812d39deb1103bd7e99ef" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pluggy"
|
||||
version = "1.6.0"
|
||||
|
||||
Reference in New Issue
Block a user