diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml index c2c0765..e69fc65 100644 --- a/.gitea/workflows/ci.yml +++ b/.gitea/workflows/ci.yml @@ -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 diff --git a/.gitignore b/.gitignore index 377b0b7..95aeaf8 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,6 @@ checkpoints/ # Scratch working directory /scratchpad/ + +# giant analyze run directories (shared.json, reduced/, plots/, condor logs) +/analysis_runs/ diff --git a/CLAUDE.md b/CLAUDE.md index 0025238..e2efcde 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -13,6 +13,8 @@ pytest # run tests giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir for a new training run giant train path/to/steps.parquet --mode flow # train (flow matching) giant train path/to/steps.parquet --mode ddpm # train (DDPM baseline) +giant train path/to/steps.parquet --mode wgan # train (WGAN-GP, single-pass eval; implemented, not yet tested) +giant train path/to/steps.parquet --router --router-type energy # MoE routing trunk (implemented; first rollout benchmark failed with lambda_balance=0, retrain needed — see Roadmap) giant train-submit path/to/steps.parquet --config run/config.toml --accounting-group cms # train as a remote-GPU HTCondor job (TOpAS/NEMO2) giant predict path/to/steps.parquet --checkpoint ckpt/best.pt # per-step predictions giant rollout path/to/steps.parquet --checkpoint ckpt/best.pt --geometry oracle.pkl # full showers @@ -37,6 +39,14 @@ uv run ty check . # type check Part of the `dev` extra. Run these periodically (not just at commit time) to catch drift early. +## Compute environment + +Work on this repo happens across three kinds of machine: + +- **Local dev machines** (laptop + desktop, identical): repo at `~/Programming/giant`, no access to `/ceph` — datasets, training results, and models aren't reachable here. +- **Portal machines** (`portal1`, `deepthought`, `deepthought2`, `bms1`, `bms2`, `bms3`): repo lives under `/work`, and `/ceph` holds ROOT/parquet files and trained models. **These are shared with other users** — stay strictly within `/work/lbogner` and `/ceph/lbogner`, and keep resource usage to roughly a quarter of CPU/RAM and a single GPU so as not to disturb other users' jobs. +- **HTCondor worker nodes**: never run or SSH onto these directly — the only sanctioned path is submitting jobs through condor (`giant analyze submit`, and the in-progress remote-GPU train/rollout submission on `condor-gpu-train-rollout`). `/ceph` is available there; `/work` is only sometimes mounted, depending on the node. + ## Architecture GIANT is a conditional generative surrogate for the Geant4 step function. It replaces the stochastic physics engine: given a pre-step particle state (conditioning), it samples a post-step outcome — now including the variable-length list of secondary particles the step produces (Phase 2, see Roadmap). @@ -59,9 +69,13 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep **Samplers** (`giant/sample.py`): DDPM, DDIM, and flow matching (ODE integration, ~10 steps). Flow matching is the primary mode. +**WGAN-GP mode (`--mode wgan`, implemented, not yet tested):** a throwaway fast-eval alternative to the flow/DDPM samplers above — single forward pass instead of ~10 ODE steps. Dedicated noise-conditioned generators (`WGANGenerator`/`WGANSecondaryGenerator`, `giant/model/network.py`) stand in for `DenoisingMLP`/`SecondaryDecoder`, trained against `Critic`/`SecondaryCritic` discriminators with the gradient-penalty loss in `giant/model/wgan.py` (Gulrajani et al. 2017); `sample_wgan` (`giant/sample.py`) does the single-pass draw at inference. Not yet validated against the flow-matching baseline. + +**MoE routing trunk (`--router`, implemented; first rollout benchmark shows the experts don't specialize — see Roadmap):** an alternative to `DenoisingMLP`'s monolithic `ResBlock` trunk — a `Router` (`giant/model/network.py`, `ROUTER_REGISTRY`/`build_router`) gates between small per-expert `ResBlock` stacks (`Expert`), soft-mixed over all experts at train time but **top-1 dispatched at eval time** (each row runs exactly one small expert), which is the actual inference-speed win. Router types gate on different conditioning axes: `EnergyRouter`/`PdgRouter` read a quantity already known at inference time, `ProcessRouter` runs its own small classifier over pre-step conditioning (since process isn't known upfront); `ComposedRouter` gates jointly over multiple axes (outer-product expert cells) via repeated `--router-axis "type:key=val,..."` flags. Config lives under `model.router` (`giant/config.py`), deep-merged one level so `router.enabled` alone doesn't drop the rest of the defaults. + **Validation** (`giant/validate.py`): step-level marginal comparisons. -**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_/`) holding `shared.json`, `run_meta.json`, `reduced/`, `plots/`. **Compute/render split:** `giant analyze submit rollout.yaml` runs `prep` then submits one HTCondor job per plot (`compute-one --run-dir`, polars/numpy only — no LaTeX on workers), each writing a small `reduced/.json`; the local `giant analyze render ` turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`. +**Analysis** (`giant/analysis/`, `giant analyze` CLI): a lean, streaming rollout-vs-reference plotting pipeline that compares one autoregressive `giant rollout` (for a given checkpoint) against a held-out miniCaloSim reference steps file, and produces publication-styled PDFs assembled into an HTML gallery. It exploits the fact that rollout output and a raw reference file share a world-frame physical column subset under identical names (`pre_*`/`post_*`/`edep`/`step_length`/`pdg`/`material`/`event_id`), so no ALR/local-frame decode is needed — everything is world-frame mm/MeV. Structure: `sources.py` (canonical LazyFrames + synthetic-termination-row filtering + the secondary view, which is `generation>0 & step_no==0` rollout tracks vs exploded `sec_*_list` reference columns), `reduce.py` (the streaming primitives — a single `hist1d` `group_by([group,bin]).len()` pass, per-event scalars, edep-weighted depth/transverse profiles, species share, leakage), `grouping.py`/`context.py` (fixed bin edges + energy-quantile/pdg/material group sets resolved once by `prep` into `shared.json`, so every compute job is one pass with no range scan), `catalog.py` (the declarative `PlotSpec` registry — marginals × {overall,energy,pdg,material}, per-event totals, shower profiles, species/leakage, secondaries), and `render.py` (the only module importing ETPlot's `plotstyle`/LaTeX; dispatches on `Reduced.kind`, writes PDFs + `metadata.yaml`). **Input is a `giant rollout` YAML sidecar** (`condor.py:load_rollout_yaml`): its `output`/`dataset` keys name the rollout parquet and the seed file (= the reference truth), and the rest of the YAML (checkpoint, geometry oracle, cutoffs) flows into each plot's gallery metadata. `prep` derives its own **run directory** next to the rollout parquet (`<...>/analysis_/`) 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/__.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 ` first joins every plot's chunk partials into `reduced/.json` (`merge_all`, a no-op join when `N=1`), then turns those into the styled PDF/gallery tree. See `giant/analysis/__init__.py`. **Shower rollout** (`giant/rollout.py`, `giant rollout` CLI): autoregressively steps the two-stage model into a full shower — each primary post-step becomes the next pre-step, secondaries are pushed as new tracks, and per-step `material`/`layer_id` come from a `GeometryOracle` (`giant/geometry.py`, built via `dwarf build-geometry-oracle`) that learns position → (material, layer_id) from data and flags detector escape by nearest-neighbour distance. Tracks terminate on energy cutoff, per-track max steps, escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction. @@ -75,4 +89,10 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep **Physical-property conditioning (implemented):** `model.conditioning = "physical" | "embedding"` (see above) replaces the learned PDG/material embeddings with a small MLP over particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, and Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. `"embedding"` stays available as the generalization-comparison baseline. `giant/materials.py`'s table is already filled with real values for every material the geometry produces. **Not yet done:** the actual held-out-material/species generalization comparison against the `"embedding"` baseline is unrun — the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes including nuclear/ion codes) is the natural dataset for that experiment. -**Next directions** (parallel, not yet built): faster-eval architectures measured against a ~10× native-Geant4 budget — a Wasserstein-GAN throwaway (single-pass eval) and a mixture-of-experts / routing tree of small nets selected per call (pdg / energy / process), with soft/differentiable gating on continuous routing axes; a sampling-calorimeter (multi-material) dataset. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`). +**Faster-eval architectures (implemented, validation in progress):** both tracks below target a ~10× native-Geant4 eval budget and are now wired into `giant train`/`giant/model/network.py`, but neither has a validated result yet — treat both as unproven until the corresponding analysis run says otherwise: +- **WGAN-GP** (`--mode wgan`, see Architecture above): implemented, **not yet tested** — no rollout-vs-reference analysis run against it yet. +- **MoE routing trunk** (`--router`, see Architecture above): implemented, **first rollout benchmark done (2026-07-22), result: needs retraining with a different router config, not abandoned.** A 10-expert `EnergyRouter` run (`n_experts=10`, `temperature=0.5`, `learn_centers=true`, **`lambda_balance=0.0`**, only 20 fine-tuning epochs resumed from a non-routed checkpoint) diverged badly from Geant4 on step granularity, secondary species, and shower shape, despite roughly matching bulk total deposited energy. The `router_gating` diagnostic plot points at the likely cause: the ten experts overlap heavily across ~5 decades of pre-step energy instead of partitioning it — even the top-energy expert only reaches ~60–65% gate weight at the highest energies plotted — so eval-time top-1 (Voronoi) dispatch is choosing among near-ties rather than real specialists. Two contributors were identified: the missing load-balancing loss (`lambda_balance=0.0`), and `EnergyRouter`'s center init (`torch.linspace(-2, 2, n_experts)`) assuming a roughly uniform z-normalized energy distribution, which real energy spectra don't match. **Fixed (2026-07-27):** `EnergyRouter` now accepts an optional `centers_init` (backward compatible — omitting it keeps the old linspace), and `giant train` auto-populates it from real data quantiles via a reservoir sample collected during the existing normalizer-fitting pass in `giant/pipeline.py` (no extra file scan), for `--router-type energy` only. The routing *strategy* itself may still be sound, but the specific benchmarked config wasn't. **Next step before further evaluation: retrain with `lambda_balance > 0` and the new quantile-seeded centers (and consider more epochs / a from-scratch run rather than a short fine-tune), then re-check whether `router_gating` sharpens up.** Full writeup: `/home/lars/knowledge-base/experiments/giant-router-energy-rollout-validation.md`. + +A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`). + +**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet. diff --git a/docs/phase2_plan.md b/docs/phase2_plan.md deleted file mode 100644 index 596edec..0000000 --- a/docs/phase2_plan.md +++ /dev/null @@ -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 diff --git a/giant/analysis/__init__.py b/giant/analysis/__init__.py index a646f80..48c4343 100644 --- a/giant/analysis/__init__.py +++ b/giant/analysis/__init__.py @@ -2,7 +2,8 @@ Compares one autoregressive ``giant rollout`` against a held-out miniCaloSim reference file, producing publication-styled comparison plots generated in -parallel on HTCondor (one job per plot, compute/render split). +parallel on HTCondor (one job per plot x data chunk, compute/merge/render +split). Only ``render`` (and the ``render`` CLI path) imports plotstyle/LaTeX; everything re-exported here is plotstyle-free so it runs on a compute worker. Import @@ -17,11 +18,14 @@ from giant.analysis.condor import ( compute_reduced, derive_run_dir, load_rollout_yaml, + merge_all, + merge_one, prep, write_submit, ) from giant.analysis.context import Context, build_context -from giant.analysis.reduced import Reduced +from giant.analysis.reduced import Partial, Reduced +from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runtime_s from giant.analysis.sources import Side __all__ = [ @@ -34,10 +38,15 @@ __all__ = [ "compute_reduced", "derive_run_dir", "load_rollout_yaml", + "merge_all", + "merge_one", "prep", "write_submit", "Context", "build_context", + "Partial", "Reduced", "Side", + "RUNTIME_SAFETY_MARGIN", + "estimate_runtime_s", ] diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index a91398a..fc665e3 100644 --- a/giant/analysis/catalog.py +++ b/giant/analysis/catalog.py @@ -2,9 +2,22 @@ Each spec knows its stable ``id`` (used for the reduced-data filename, the PDF stem and the condor queue item), its gallery ``family`` (subdirectory), and a -``compute(bundle) -> Reduced`` that runs the streaming reduction. Rendering lives -in ``render.py`` and dispatches on ``Reduced.kind`` — the catalog itself never -imports plotstyle, so ``compute-one`` jobs stay LaTeX-free. +``compute_partial(bundle) -> dict`` / ``finalize(parts, ctx) -> Reduced`` pair +that together run the streaming reduction. ``compute_partial`` runs once per +``(plot, chunk)`` condor job against a ``Bundle`` whose four LazyFrames are +already filtered to that chunk (see ``Bundle.open``'s ``chunk`` argument); it +returns a small JSON-safe partial artifact — either a raw sum-mergeable count +dict (histograms/species sums against fixed edges) or a raw per-event/ +per-secondary array to be concatenated (anything that derives its own edges or +a mean/std from the full dataset). ``finalize`` merges the per-chunk partials +(in chunk order) and does the actual histogramming/edge-selection/mean-std +collapse, once, over the merged data — for ``n_chunks=1`` this reproduces +exactly what a single unchunked pass would produce. Specs marked +``chunkable=False`` (the router ones) always run as a single chunk regardless +of the configured chunk count. + +Rendering lives in ``render.py`` and dispatches on ``Reduced.kind`` — the +catalog itself never imports plotstyle, so ``compute-one`` jobs stay LaTeX-free. The registry is built by expanding parametric families (marginals over variable x grouping, secondaries, ...) into concrete specs. @@ -12,7 +25,7 @@ variable x grouping, secondaries, ...) into concrete specs. from __future__ import annotations -from dataclasses import dataclass +from dataclasses import asdict, dataclass from typing import Callable import numpy as np @@ -32,11 +45,18 @@ from giant.analysis.reduce import ( event_scalars, hist1d, leakage_fraction, + profile_finalize, + profile_partial, species_share, + sum_merge, transverse_expr, - weighted_profile, ) from giant.analysis.reduced import Reduced +from giant.analysis.router_gating import ( + compute_router_gating, + compute_router_share_by_pdg, + compute_router_share_by_process, +) from giant.analysis.sources import Side, open_side, physical_steps, secondaries from giant.analysis.variables import RANGED_VARS, cos_scatter_expr @@ -50,17 +70,39 @@ class Bundle: t_all: pl.LazyFrame # reference, all rows r_phys: pl.LazyFrame # rollout, physical steps only t_phys: pl.LazyFrame # reference, physical steps only + checkpoint: str | None = None # from the rollout YAML; router_gating only @classmethod - def open(cls, rollout, reference, ctx: Context) -> "Bundle": + def open( + cls, + rollout, + reference, + ctx: Context, + checkpoint=None, + chunk: tuple[int, int] | None = None, + ) -> "Bundle": + """Open both sides, optionally restricted to one event-disjoint chunk. + + ``chunk = (chunk_index, n_chunks)`` filters both sides to + ``event_id % n_chunks == chunk_index`` *before* deriving the physical/ + secondary views, so every downstream reduction (which is either + row-local or a ``group_by("event_id")``) sees a self-contained, + event-disjoint slice — no cross-chunk lookups are ever needed. + """ r_all = open_side(rollout, Side.rollout) t_all = open_side(reference, Side.reference) + if chunk is not None: + idx, n = chunk + pred = pl.col("event_id") % n == idx + r_all = r_all.filter(pred) + t_all = t_all.filter(pred) return cls( ctx=ctx, r_all=r_all, t_all=t_all, r_phys=physical_steps(r_all, Side.rollout), t_phys=physical_steps(t_all, Side.reference), + checkpoint=checkpoint, ) @@ -68,7 +110,26 @@ class Bundle: class PlotSpec: id: str family: str - compute: Callable[[Bundle], Reduced] + compute_partial: Callable[[Bundle], dict] + finalize: Callable[[list[dict], Context], Reduced] + chunkable: bool = True + + +def _unchunkable( + compute: Callable[[Bundle], Reduced], +) -> tuple[Callable[[Bundle], dict], Callable[[list[dict], Context], Reduced]]: + """Wrap a whole-dataset ``compute(bundle) -> Reduced`` as a trivial + ``(compute_partial, finalize)`` pair, for specs marked ``chunkable=False`` + (which always run as a single chunk, so ``parts`` is always one element). + """ + + def partial(b: Bundle) -> dict: + return {"reduced": asdict(compute(b))} + + def finalize(parts: list[dict], ctx: Context) -> Reduced: + return Reduced(**parts[0]["reduced"]) + + return partial, finalize # --------------------------------------------------------------------------- @@ -83,6 +144,20 @@ def _counts(h: dict, key, nbins: int) -> list[int]: return h.get(key, np.zeros(nbins, dtype=np.int64)).astype(np.int64).tolist() +def _partial_hist( + lf: pl.LazyFrame, value: pl.Expr, edges: np.ndarray, group: pl.Expr | None = None +) -> dict[str, list[int]]: + """One chunk's raw ``hist1d`` result as a JSON-safe, sum-mergeable dict.""" + nb = len(edges) - 1 + h = hist1d(lf, value, edges, group=group) + return {str(k): _counts(h, k, nb) for k in h} + + +def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]: + """One group's merged counts (zero-filled if the group never appeared).""" + return list(merged.get(str(key), [0] * nbins)) + + def _np_hist_pair( r: np.ndarray, t: np.ndarray, nbins: int ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: @@ -125,12 +200,21 @@ def _marginal_edges(ctx: Context, var: str) -> np.ndarray: # --------------------------------------------------------------------------- -def _marginal_overall(b: Bundle, var: str) -> Reduced: - label, expr = _var(var) +def _marginal_overall_partial(b: Bundle, var: str) -> dict: + _, expr = _var(var) edges = _marginal_edges(b.ctx, var) + return { + "r": _partial_hist(b.r_phys, expr, edges), + "t": _partial_hist(b.t_phys, expr, edges), + } + + +def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Reduced: + label, _ = _var(var) + edges = _marginal_edges(ctx, var) nb = len(edges) - 1 - r = hist1d(b.r_phys, expr, edges) - t = hist1d(b.t_phys, expr, edges) + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) return Reduced( id=f"marginal_{var}", family="marginals", @@ -139,8 +223,8 @@ def _marginal_overall(b: Bundle, var: str) -> Reduced: xlabel=label, payload={ "edges": edges.tolist(), - _ROLL: _counts(r, 0, nb), - _REF: _counts(t, 0, nb), + _ROLL: _finalize_counts(r, 0, nb), + _REF: _finalize_counts(t, 0, nb), "log_y": True, }, ) @@ -153,31 +237,55 @@ def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr: ) -def _marginal_grouped(b: Bundle, var: str, axis: str) -> Reduced: - label, expr = _var(var) +def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict: + _, expr = _var(var) edges = _marginal_edges(b.ctx, var) - nb = len(edges) - 1 - groups: dict[str, dict] = {} - if axis == "pdg": r = hist1d(b.r_phys, expr, edges, group=pl.col("pdg")) t = hist1d(b.t_phys, expr, edges, group=pl.col("pdg")) - for k in b.ctx.top_pdgs: - groups[pdg_label(k)] = {_ROLL: _counts(r, k, nb), _REF: _counts(t, k, nb)} elif axis == "material": r = hist1d(b.r_phys, expr, edges, group=pl.col("material")) t = hist1d(b.t_phys, expr, edges, group=pl.col("material")) - for m in b.ctx.materials: - groups[material_label(m)] = { - _ROLL: _counts(r, m, nb), - _REF: _counts(t, m, nb), - } else: # energy e_edges = np.asarray(b.ctx.energy_edges) r = hist1d(b.r_phys, expr, edges, group=_energy_group_expr(b.r_phys, e_edges)) t = hist1d(b.t_phys, expr, edges, group=_energy_group_expr(b.t_phys, e_edges)) + nb = len(edges) - 1 + return { + "r": {str(k): _counts(r, k, nb) for k in r}, + "t": {str(k): _counts(t, k, nb) for k in t}, + } + + +def _marginal_grouped_finalize( + parts: list[dict], ctx: Context, var: str, axis: str +) -> Reduced: + label, _ = _var(var) + edges = _marginal_edges(ctx, var) + nb = len(edges) - 1 + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) + groups: dict[str, dict] = {} + + if axis == "pdg": + for k in ctx.top_pdgs: + groups[pdg_label(k)] = { + _ROLL: _finalize_counts(r, k, nb), + _REF: _finalize_counts(t, k, nb), + } + elif axis == "material": + for m in ctx.materials: + groups[material_label(m)] = { + _ROLL: _finalize_counts(r, m, nb), + _REF: _finalize_counts(t, m, nb), + } + else: # energy + e_edges = np.asarray(ctx.energy_edges) for bi, lbl in enumerate(energy_bin_labels(e_edges)): - groups[lbl] = {_ROLL: _counts(r, bi, nb), _REF: _counts(t, bi, nb)} + groups[lbl] = { + _ROLL: _finalize_counts(r, bi, nb), + _REF: _finalize_counts(t, bi, nb), + } return Reduced( id=f"marginal_{var}_by_{axis}", @@ -194,13 +302,19 @@ def _marginal_grouped(b: Bundle, var: str, axis: str) -> Reduced: # --------------------------------------------------------------------------- -def _event_scalar( - b: Bundle, spec_id: str, title: str, xlabel: str, col: str, use_all: bool -) -> Reduced: +def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict: r_lf, t_lf = (b.r_all, b.t_all) if use_all else (b.r_phys, b.t_phys) r = event_scalars(r_lf)[col].to_numpy() t = event_scalars(t_lf)[col].to_numpy() - edges, rc, tc = _np_hist_pair(r, t, b.ctx.n_marginal_bins) + return {"r": r.tolist(), "t": t.tolist()} + + +def _event_scalar_finalize( + parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str +) -> Reduced: + r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts]) + t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts]) + edges, rc, tc = _np_hist_pair(r, t, ctx.n_marginal_bins) return Reduced( id=spec_id, family="event", @@ -216,18 +330,26 @@ def _event_scalar( ) -def _event_total_edep_by_energy(b: Bundle) -> Reduced: - e_edges = np.asarray(b.ctx.energy_edges) +def _event_total_edep_by_energy_partial(b: Bundle) -> dict: r = event_scalars(b.r_all) t = event_scalars(b.t_all) - r_bin = np.clip( - np.digitize(r["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2 - ) - t_bin = np.clip( - np.digitize(t["incident_E"].to_numpy(), e_edges[1:-1]), 0, len(e_edges) - 2 - ) - r_val, t_val = r["total_edep"].to_numpy(), t["total_edep"].to_numpy() - edges, _, _ = _np_hist_pair(r_val, t_val, b.ctx.n_marginal_bins) + return { + "r_incident": r["incident_E"].to_list(), + "r_edep": r["total_edep"].to_list(), + "t_incident": t["incident_E"].to_list(), + "t_edep": t["total_edep"].to_list(), + } + + +def _event_total_edep_by_energy_finalize(parts: list[dict], ctx: Context) -> Reduced: + e_edges = np.asarray(ctx.energy_edges) + r_inc = np.concatenate([np.asarray(p["r_incident"], dtype=float) for p in parts]) + r_val = np.concatenate([np.asarray(p["r_edep"], dtype=float) for p in parts]) + t_inc = np.concatenate([np.asarray(p["t_incident"], dtype=float) for p in parts]) + t_val = np.concatenate([np.asarray(p["t_edep"], dtype=float) for p in parts]) + r_bin = np.clip(np.digitize(r_inc, e_edges[1:-1]), 0, len(e_edges) - 2) + t_bin = np.clip(np.digitize(t_inc, e_edges[1:-1]), 0, len(e_edges) - 2) + edges, _, _ = _np_hist_pair(r_val, t_val, ctx.n_marginal_bins) groups: dict[str, dict] = {} for bi, lbl in enumerate(energy_bin_labels(e_edges)): rc = np.histogram(r_val[r_bin == bi], edges)[0] @@ -251,15 +373,54 @@ def _event_total_edep_by_energy(b: Bundle) -> Reduced: # --------------------------------------------------------------------------- -def _profile( - b: Bundle, spec_id: str, title: str, xlabel: str, coord_fn, edges_key: str -) -> Reduced: +def _profile_partial(b: Bundle, coord_fn, edges_key: str) -> dict: edges = np.asarray(getattr(b.ctx, edges_key)) - r_ea, t_ea = entry_axis(b.r_all), entry_axis(b.t_all) - r_lf = attach_entry_axis(b.r_all, r_ea) - t_lf = attach_entry_axis(b.t_all, t_ea) - r_mean, r_std = weighted_profile(r_lf, coord_fn(), edges, pl.col("edep")) - t_mean, t_std = weighted_profile(t_lf, coord_fn(), edges, pl.col("edep")) + r_lf = attach_entry_axis(b.r_all, entry_axis(b.r_all)) + t_lf = attach_entry_axis(b.t_all, entry_axis(b.t_all)) + r_ids, r_mat = profile_partial(r_lf, coord_fn(), edges, pl.col("edep")) + t_ids, t_mat = profile_partial(t_lf, coord_fn(), edges, pl.col("edep")) + return { + "r_ids": r_ids.tolist(), + "r_mat": r_mat.tolist(), + "t_ids": t_ids.tolist(), + "t_mat": t_mat.tolist(), + } + + +def _assert_event_disjoint(id_lists: list[list[int]], spec_id: str, side: str) -> None: + """Guard the chunking invariant profiles depend on: no event in two chunks. + + A violation would silently double-count that event in the merged mean/RMS + with no other symptom, so this is worth a loud failure rather than a + quietly-wrong plot. + """ + seen: set[int] = set() + for ids in id_lists: + overlap = seen & set(ids) + if overlap: + raise ValueError( + f"{spec_id} ({side}): event_id(s) {sorted(overlap)[:5]} appear " + "in more than one chunk — chunking must be event-disjoint" + ) + seen.update(ids) + + +def _profile_finalize( + parts: list[dict], + ctx: Context, + spec_id: str, + title: str, + xlabel: str, + edges_key: str, +) -> Reduced: + edges = np.asarray(getattr(ctx, edges_key)) + nb = len(edges) - 1 + _assert_event_disjoint([p["r_ids"] for p in parts], spec_id, "rollout") + _assert_event_disjoint([p["t_ids"] for p in parts], spec_id, "reference") + r_mats = [np.asarray(p["r_mat"], dtype=float).reshape(-1, nb) for p in parts] + t_mats = [np.asarray(p["t_mat"], dtype=float).reshape(-1, nb) for p in parts] + r_mean, r_std = profile_finalize(r_mats) + t_mean, t_std = profile_finalize(t_mats) return Reduced( id=spec_id, family="shower", @@ -282,14 +443,21 @@ def _profile( # --------------------------------------------------------------------------- -def _species_share(b: Bundle) -> Reduced: +def _species_share_partial(b: Bundle) -> dict: r = species_share(b.r_all) t = species_share(b.t_all) - r_map = dict(zip(r["pdg"].to_list(), r["total_edep"].to_list())) - t_map = dict(zip(t["pdg"].to_list(), t["total_edep"].to_list())) + return { + "r": {str(k): v for k, v in zip(r["pdg"].to_list(), r["total_edep"].to_list())}, + "t": {str(k): v for k, v in zip(t["pdg"].to_list(), t["total_edep"].to_list())}, + } + + +def _species_share_finalize(parts: list[dict], ctx: Context) -> Reduced: + r_map = sum_merge([p["r"] for p in parts]) + t_map = sum_merge([p["t"] for p in parts]) r_tot = sum(r_map.values()) or 1.0 t_tot = sum(t_map.values()) or 1.0 - labels = [pdg_label(k) for k in b.ctx.top_pdgs] + labels = [pdg_label(k) for k in ctx.top_pdgs] return Reduced( id="species_edep_share", family="species", @@ -298,19 +466,22 @@ def _species_share(b: Bundle) -> Reduced: xlabel="species", payload={ "labels": labels, - _ROLL: [r_map.get(k, 0.0) / r_tot for k in b.ctx.top_pdgs], - _REF: [t_map.get(k, 0.0) / t_tot for k in b.ctx.top_pdgs], + _ROLL: [r_map.get(str(k), 0.0) / r_tot for k in ctx.top_pdgs], + _REF: [t_map.get(str(k), 0.0) / t_tot for k in ctx.top_pdgs], "ylabel": "fraction of total deposited energy", }, ) -def _leakage(b: Bundle) -> Reduced: +def _leakage_partial(b: Bundle) -> dict: frac = leakage_fraction(b.r_all) + return {"frac": frac.tolist()} + + +def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced: + frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts]) edges = np.linspace( - 0.0, - max(float(frac.max()) if len(frac) else 1.0, 1e-3), - b.ctx.n_marginal_bins + 1, + 0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1 ) counts = np.histogram(frac, edges)[0] return Reduced( @@ -340,7 +511,7 @@ def _sec_frames(b: Bundle): ) -def _sec_count_per_event(b: Bundle) -> Reduced: +def _sec_count_per_event_partial(b: Bundle) -> dict: r_sec, t_sec = _sec_frames(b) r = ( r_sec.group_by("event_id") @@ -354,9 +525,13 @@ def _sec_count_per_event(b: Bundle) -> Reduced: .collect(engine="streaming")["n"] .to_numpy() ) - edges, rc, tc = _np_hist_pair( - r.astype(float), t.astype(float), min(b.ctx.n_marginal_bins, 40) - ) + return {"r": r.tolist(), "t": t.tolist()} + + +def _sec_count_per_event_finalize(parts: list[dict], ctx: Context) -> Reduced: + r = np.concatenate([np.asarray(p["r"], dtype=float) for p in parts]) + t = np.concatenate([np.asarray(p["t"], dtype=float) for p in parts]) + edges, rc, tc = _np_hist_pair(r, t, min(ctx.n_marginal_bins, 40)) return Reduced( id="sec_count_per_event", family="secondaries", @@ -372,32 +547,21 @@ def _sec_count_per_event(b: Bundle) -> Reduced: ) -def _sec_count_per_species(b: Bundle) -> Reduced: +def _counts_by_pdg(sec_lf: pl.LazyFrame) -> dict[str, int]: + df = sec_lf.group_by("pdg").agg(pl.len().alias("n")).collect(engine="streaming") + return {str(k): v for k, v in zip(df["pdg"].to_list(), df["n"].to_list())} + + +def _sec_count_per_species_partial(b: Bundle) -> dict: r_sec, t_sec = _sec_frames(b) - r = dict( - zip( - *[ - r_sec.group_by("pdg") - .agg(pl.len().alias("n")) - .collect(engine="streaming")[c] - .to_list() - for c in ("pdg", "n") - ] - ) - ) - t = dict( - zip( - *[ - t_sec.group_by("pdg") - .agg(pl.len().alias("n")) - .collect(engine="streaming")[c] - .to_list() - for c in ("pdg", "n") - ] - ) - ) + return {"r": _counts_by_pdg(r_sec), "t": _counts_by_pdg(t_sec)} + + +def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced: + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[ - : len(b.ctx.top_pdgs) + : len(ctx.top_pdgs) ] return Reduced( id="sec_count_per_species", @@ -406,7 +570,7 @@ def _sec_count_per_species(b: Bundle) -> Reduced: title="Secondary count by species", xlabel="species", payload={ - "labels": [pdg_label(k) for k in keys], + "labels": [pdg_label(int(k)) for k in keys], _ROLL: [float(r.get(k, 0)) for k in keys], _REF: [float(t.get(k, 0)) for k in keys], "ylabel": "secondary count", @@ -414,12 +578,20 @@ def _sec_count_per_species(b: Bundle) -> Reduced: ) -def _sec_energy(b: Bundle) -> Reduced: +def _sec_energy_partial(b: Bundle) -> dict: r_sec, t_sec = _sec_frames(b) edges = np.linspace(*b.ctx.sec_energy_range, b.ctx.n_sec_bins + 1) - r = hist1d(r_sec, pl.col("energy"), edges) - t = hist1d(t_sec, pl.col("energy"), edges) + return { + "r": _partial_hist(r_sec, pl.col("energy"), edges), + "t": _partial_hist(t_sec, pl.col("energy"), edges), + } + + +def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced: + edges = np.linspace(*ctx.sec_energy_range, ctx.n_sec_bins + 1) nb = len(edges) - 1 + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) return Reduced( id="sec_energy", family="secondaries", @@ -428,27 +600,34 @@ def _sec_energy(b: Bundle) -> Reduced: xlabel="secondary energy [MeV]", payload={ "edges": edges.tolist(), - _ROLL: _counts(r, 0, nb), - _REF: _counts(t, 0, nb), + _ROLL: _finalize_counts(r, 0, nb), + _REF: _finalize_counts(t, 0, nb), "log_y": True, }, ) -def _sec_cos_angle(b: Bundle) -> Reduced: +def _sec_cos_angle_partial(b: Bundle) -> dict: edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1) - nb = len(edges) - 1 cos = ( pl.col("sdx") * pl.col("axis_x") + pl.col("sdy") * pl.col("axis_y") + pl.col("sdz") * pl.col("axis_z") ).clip(-1.0, 1.0) - def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> list[int]: + def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]: ea = entry_axis(steps_lf) - return _counts(hist1d(attach_entry_axis(sec_lf, ea), cos, edges), 0, nb) + return _partial_hist(attach_entry_axis(sec_lf, ea), cos, edges) r_sec, t_sec = _sec_frames(b) + return {"r": _side(r_sec, b.r_phys), "t": _side(t_sec, b.t_all)} + + +def _sec_cos_angle_finalize(parts: list[dict], ctx: Context) -> Reduced: + edges = np.linspace(-1.0, 1.0, ctx.n_sec_bins + 1) + nb = len(edges) - 1 + r = sum_merge([p["r"] for p in parts]) + t = sum_merge([p["t"] for p in parts]) return Reduced( id="sec_cos_angle", family="secondaries", @@ -457,13 +636,30 @@ def _sec_cos_angle(b: Bundle) -> Reduced: xlabel="cos of emission angle", payload={ "edges": edges.tolist(), - _ROLL: _side(r_sec, b.r_phys), - _REF: _side(t_sec, b.t_all), + _ROLL: _finalize_counts(r, 0, nb), + _REF: _finalize_counts(t, 0, nb), "log_y": False, }, ) +# --------------------------------------------------------------------------- +# router diagnostics (not chunked — already bounded/subsampled) +# --------------------------------------------------------------------------- + +_router_gating_partial, _router_gating_finalize = _unchunkable( + lambda b: compute_router_gating(b.checkpoint, b.r_phys, b.t_phys) +) +_router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable( + lambda b: compute_router_share_by_pdg( + b.checkpoint, b.r_phys, b.t_phys, b.ctx.top_pdgs + ) +) +_router_share_process_partial, _router_share_process_finalize = _unchunkable( + lambda b: compute_router_share_by_process(b.checkpoint, b.t_phys) +) + + # --------------------------------------------------------------------------- # registry assembly # --------------------------------------------------------------------------- @@ -479,7 +675,12 @@ def build_catalog() -> list[PlotSpec]: for var in MARGINAL_VARS: specs.append( PlotSpec( - f"marginal_{var}", "marginals", lambda b, v=var: _marginal_overall(b, v) + f"marginal_{var}", + "marginals", + compute_partial=lambda b, v=var: _marginal_overall_partial(b, v), + finalize=lambda parts, ctx, v=var: _marginal_overall_finalize( + parts, ctx, v + ), ) ) for axis in GROUPING_AXES: @@ -487,7 +688,12 @@ def build_catalog() -> list[PlotSpec]: PlotSpec( f"marginal_{var}_by_{axis}", "marginals", - lambda b, v=var, a=axis: _marginal_grouped(b, v, a), + compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial( + b, v, a + ), + finalize=lambda parts, ctx, v=var, a=axis: ( + _marginal_grouped_finalize(parts, ctx, v, a) + ), ) ) @@ -495,70 +701,136 @@ def build_catalog() -> list[PlotSpec]: PlotSpec( "event_total_edep", "event", - lambda b: _event_scalar( - b, + compute_partial=lambda b: _event_scalar_partial( + b, "total_edep", use_all=True + ), + finalize=lambda parts, ctx: _event_scalar_finalize( + parts, + ctx, "event_total_edep", "Total deposited energy per event", "total deposited energy [MeV]", - "total_edep", - use_all=True, ), ), - PlotSpec("event_total_edep_by_energy", "event", _event_total_edep_by_energy), + PlotSpec( + "event_total_edep_by_energy", + "event", + compute_partial=_event_total_edep_by_energy_partial, + finalize=_event_total_edep_by_energy_finalize, + ), PlotSpec( "event_mean_length", "event", - lambda b: _event_scalar( - b, + compute_partial=lambda b: _event_scalar_partial( + b, "mean_length", use_all=False + ), + finalize=lambda parts, ctx: _event_scalar_finalize( + parts, + ctx, "event_mean_length", "Mean step length per event", "mean step length [mm]", - "mean_length", - use_all=False, ), ), PlotSpec( "event_n_steps", "event", - lambda b: _event_scalar( - b, + compute_partial=lambda b: _event_scalar_partial( + b, "n_steps", use_all=False + ), + finalize=lambda parts, ctx: _event_scalar_finalize( + parts, + ctx, "event_n_steps", "Number of steps per event", "steps per event", - "n_steps", - use_all=False, ), ), PlotSpec( "shower_longitudinal", "shower", - lambda b: _profile( - b, + compute_partial=lambda b: _profile_partial(b, depth_expr, "depth_edges"), + finalize=lambda parts, ctx: _profile_finalize( + parts, + ctx, "shower_longitudinal", "Longitudinal shower profile", "depth along shower axis [mm]", - depth_expr, "depth_edges", ), ), PlotSpec( "shower_transverse", "shower", - lambda b: _profile( - b, + compute_partial=lambda b: _profile_partial( + b, transverse_expr, "transverse_edges" + ), + finalize=lambda parts, ctx: _profile_finalize( + parts, + ctx, "shower_transverse", "Transverse shower profile", "radius from shower axis [mm]", - transverse_expr, "transverse_edges", ), ), - PlotSpec("species_edep_share", "species", _species_share), - PlotSpec("leakage_fraction", "species", _leakage), - PlotSpec("sec_count_per_event", "secondaries", _sec_count_per_event), - PlotSpec("sec_count_per_species", "secondaries", _sec_count_per_species), - PlotSpec("sec_energy", "secondaries", _sec_energy), - PlotSpec("sec_cos_angle", "secondaries", _sec_cos_angle), + PlotSpec( + "species_edep_share", + "species", + compute_partial=_species_share_partial, + finalize=_species_share_finalize, + ), + PlotSpec( + "leakage_fraction", + "species", + compute_partial=_leakage_partial, + finalize=_leakage_finalize, + ), + PlotSpec( + "sec_count_per_event", + "secondaries", + compute_partial=_sec_count_per_event_partial, + finalize=_sec_count_per_event_finalize, + ), + PlotSpec( + "sec_count_per_species", + "secondaries", + compute_partial=_sec_count_per_species_partial, + finalize=_sec_count_per_species_finalize, + ), + PlotSpec( + "sec_energy", + "secondaries", + compute_partial=_sec_energy_partial, + finalize=_sec_energy_finalize, + ), + PlotSpec( + "sec_cos_angle", + "secondaries", + compute_partial=_sec_cos_angle_partial, + finalize=_sec_cos_angle_finalize, + ), + PlotSpec( + "router_gating", + "model", + compute_partial=_router_gating_partial, + finalize=_router_gating_finalize, + chunkable=False, + ), + PlotSpec( + "router_share_by_pdg", + "model", + compute_partial=_router_share_pdg_partial, + finalize=_router_share_pdg_finalize, + chunkable=False, + ), + PlotSpec( + "router_share_by_process", + "model", + compute_partial=_router_share_process_partial, + finalize=_router_share_process_finalize, + chunkable=False, + ), ] return specs diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index e2ef763..60e1424 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -15,18 +15,25 @@ next to the rollout parquet, and lays everything out under it: /shared.json fixed bin edges / group sets (prep) /run_meta.json resolved rollout/reference paths + plot metadata - /reduced/.json one per compute job + /reduced_partial/__.json one per (plot, chunk) job + /reduced/.json merged, per plot /plots//.pdf rendered locally -Job model (one condor job per plot, compute/render split): +Job model (one condor job per (plot, chunk), compute/merge/render split): 1. ``prep`` runs once on the submit node — reads the YAML, resolves the shared - context from a subsample, writes ``shared.json`` + ``run_meta.json``. -2. one job per catalog id runs ``giant analyze compute-one --run-dir`` on a - worker — a single streaming pass writing ``reduced/.json`` (polars/numpy - only, no LaTeX). -3. a final *local* ``giant analyze render`` turns those into the styled PDF + - gallery tree (that step imports plotstyle/LaTeX). + context from a subsample, writes ``shared.json`` + ``run_meta.json`` + (including the run's configured ``n_chunks``). +2. one job per catalog id x chunk index runs ``giant analyze compute-one + --run-dir`` on a worker — a single streaming pass over that + ``event_id``-disjoint chunk, writing ``reduced_partial/__.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/.json``, then renders those into the styled PDF + gallery tree + (that step imports plotstyle/LaTeX). Files on ``/ceph`` or ``/work`` are reached via ``ProvidesETPResources``; no HTCondor file transfer of the multi-GB inputs. @@ -35,13 +42,17 @@ HTCondor file transfer of the multi-GB inputs. from __future__ import annotations import json -from dataclasses import dataclass +from dataclasses import dataclass, field from pathlib import Path +import polars as pl import yaml from giant.analysis.catalog import Bundle, catalog_ids, get_spec from giant.analysis.context import Context, build_context +from giant.analysis.reduced import Partial +from giant.analysis.runtime_estimate import estimate_runtime_s +from giant.analysis.sources import Side, open_side # Keys copied verbatim from a rollout YAML into each plot's gallery metadata. _PLOT_META_KEYS = ( @@ -54,12 +65,22 @@ _PLOT_META_KEYS = ( "max_steps", "steps", "max_tracks_per_event", + "escape_threshold", + "n_events", "n_seed_events", "timestamp", "comment", + "weights", + "batch_size", + "device", + "rollout_seed", + "n_rows", + "termination_reason_counts", "model_config", "training_epoch", "best_val_loss", + "training_config", + "training_meta", ) @@ -82,13 +103,26 @@ def load_rollout_yaml(path: str | Path) -> dict: return d -def derive_run_dir(rollout_yaml: dict, run_dir: str | Path | None = None) -> Path: - """Analysis output directory, next to the rollout parquet unless overridden.""" +def derive_run_dir( + rollout_yaml: dict, + run_dir: str | Path | None = None, + default_base: str | Path | None = None, +) -> Path: + """Analysis output directory. + + Precedence: an explicit ``run_dir`` always wins. Otherwise + ``default_base / analysis_`` if ``default_base`` is given (the CLI + passes the repo's gitignored ``analysis_runs/``, so run directories don't + pile up on ``/ceph`` next to the rollout parquet). Falls back to next to + the rollout parquet — the original convention — for callers that don't + care where the run directory lives. + """ if run_dir is not None: return Path(run_dir) rollout = Path(rollout_yaml["output"]) tag = str(rollout_yaml.get("prediction_id") or rollout.stem)[:8] - return rollout.parent / f"analysis_{tag}" + base = Path(default_base) if default_base is not None else rollout.parent + return base / f"analysis_{tag}" def _plot_meta(rollout_yaml: dict) -> dict: @@ -104,6 +138,12 @@ class RunMeta: run_dir: str title: str plot_meta: dict + n_chunks: int = 1 + # rollout+reference row count of each event_id-disjoint chunk, and the + # dataset total — inputs to `runtime_estimate.estimate_runtime_s`. Empty/0 + # on run directories written before this field existed. + rows_per_chunk: list[int] = field(default_factory=list) + total_rows: int = 0 def save(self, path: str | Path) -> None: Path(path).write_text(json.dumps(self.__dict__, indent=2)) @@ -113,23 +153,58 @@ class RunMeta: return cls(**json.loads(Path(path).read_text())) +def _rows_per_chunk( + rollout: str | Path, reference: str | Path, n_chunks: int +) -> list[int]: + """Rollout+reference row count of each ``event_id % n_chunks`` chunk. + + One cheap streaming ``group_by`` per side (just the ``event_id`` column) — + the sizing input every job's estimated walltime + (``runtime_estimate.estimate_runtime_s``) is computed from. + """ + + def counts(lf: pl.LazyFrame) -> pl.DataFrame: + return ( + lf.select((pl.col("event_id") % n_chunks).alias("_c")) + .group_by("_c") + .agg(pl.len().alias("n")) + .collect(engine="streaming") + ) + + out = [0] * n_chunks + for lf in (open_side(rollout, Side.rollout), open_side(reference, Side.reference)): + df = counts(lf) + for c, n in zip(df["_c"].to_list(), df["n"].to_list()): + out[c] += n + return out + + def prep( rollout_yaml: str | Path, run_dir: str | Path | None = None, + n_chunks: int = 1, + default_base: str | Path | None = None, **ctx_kwargs, ) -> Path: """Read the rollout YAML, build the shared context, and lay out the run dir. Writes ``shared.json`` + ``run_meta.json`` and returns the run directory. + ``n_chunks`` is the run-level chunk count every ``compute-one``/``merge-one`` + job reads back out of ``run_meta.json`` (via ``RunMeta.n_chunks``), so it is + resolved once here rather than re-passed (and risking disagreement) at every + later step. See ``derive_run_dir`` for how ``run_dir``/``default_base`` + resolve the actual directory. """ y = load_rollout_yaml(rollout_yaml) - run_path = derive_run_dir(y, run_dir) + run_path = derive_run_dir(y, run_dir, default_base=default_base) run_path.mkdir(parents=True, exist_ok=True) rollout, reference = y["output"], y["dataset"] ctx = build_context(rollout, reference, **ctx_kwargs) ctx.save(run_path / "shared.json") + rows_per_chunk = _rows_per_chunk(rollout, reference, n_chunks) + ckpt = Path(y.get("checkpoint", "")).name or "rollout" RunMeta( rollout=str(rollout), @@ -137,12 +212,15 @@ def prep( run_dir=str(run_path), title=f"GIANT rollout analysis — {ckpt}", plot_meta=_plot_meta(y), + n_chunks=n_chunks, + rows_per_chunk=rows_per_chunk, + total_rows=sum(rows_per_chunk), ).save(run_path / "run_meta.json") return run_path # --------------------------------------------------------------------------- -# per-plot compute (what each condor job runs) +# per-(plot, chunk) compute (what each condor job runs) # --------------------------------------------------------------------------- @@ -152,18 +230,42 @@ def compute_reduced( reference: str | Path, shared: str | Path, out: str | Path, + checkpoint: str | None = None, + chunk_index: int = 0, + n_chunks: int = 1, ) -> Path: - """Core: run one plot's reduction against explicit paths → ``Reduced`` JSON.""" + """Core: run one (plot, chunk)'s partial reduction against explicit paths. + + Writes a ``Partial`` JSON — the raw, not-yet-merged output of + ``PlotSpec.compute_partial`` — never a finished ``Reduced``; ``merge_one`` + is what combines every chunk's ``Partial`` for a plot into the final + ``Reduced``. Specs with ``chunkable=False`` always run as a single chunk + regardless of ``n_chunks``. + """ ctx = Context.load(shared) - bundle = Bundle.open(rollout, reference, ctx) - reduced = get_spec(spec_id).compute(bundle) + spec = get_spec(spec_id) + effective_n = n_chunks if spec.chunkable else 1 + if not (0 <= chunk_index < effective_n): + raise ValueError( + f"{spec_id}: chunk_index={chunk_index} out of range for " + f"n_chunks={effective_n} (chunkable={spec.chunkable})" + ) + bundle = Bundle.open( + rollout, reference, ctx, checkpoint=checkpoint, chunk=(chunk_index, effective_n) + ) + partial = Partial( + id=spec_id, + family=spec.family, + chunk=chunk_index, + data=spec.compute_partial(bundle), + ) out = Path(out) - reduced.save(out) + partial.save(out) return out -def compute_one(spec_id: str, run_dir: str | Path) -> Path: - """Run one plot's reduction from a prepped run directory.""" +def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path: + """Run one (plot, chunk)'s partial reduction from a prepped run directory.""" run_path = Path(run_dir) meta = RunMeta.load(run_path / "run_meta.json") return compute_reduced( @@ -171,10 +273,56 @@ def compute_one(spec_id: str, run_dir: str | Path) -> Path: meta.rollout, meta.reference, run_path / "shared.json", - run_path / "reduced" / f"{spec_id}.json", + run_path / "reduced_partial" / f"{spec_id}__{chunk_index}.json", + checkpoint=meta.plot_meta.get("checkpoint"), + chunk_index=chunk_index, + n_chunks=meta.n_chunks, ) +# --------------------------------------------------------------------------- +# per-plot merge (the join step ``render`` runs before rendering) +# --------------------------------------------------------------------------- + + +def merge_one(spec_id: str, run_dir: str | Path) -> Path: + """Merge every chunk's partial for one plot into the final ``Reduced`` JSON. + + Fails loudly if fewer partials exist than the run's configured chunk count + for this plot — that is what catches an incomplete/failed condor job + instead of silently rendering a plot from partial data. Idempotent: safe + to call again (e.g. from ``render_run``) once all chunks are in. + """ + run_path = Path(run_dir) + meta = RunMeta.load(run_path / "run_meta.json") + ctx = Context.load(run_path / "shared.json") + spec = get_spec(spec_id) + effective_n = meta.n_chunks if spec.chunkable else 1 + + partial_dir = run_path / "reduced_partial" + found = { + p.chunk: p + for p in (Partial.load(jf) for jf in partial_dir.glob(f"{spec_id}__*.json")) + } + missing = sorted(set(range(effective_n)) - set(found)) + if missing: + raise FileNotFoundError( + f"{spec_id}: missing chunk partial(s) {missing} of {effective_n} " + f"under {partial_dir} — did every compute-one job finish?" + ) + + parts = [found[k].data for k in range(effective_n)] + reduced = spec.finalize(parts, ctx) + out = run_path / "reduced" / f"{spec_id}.json" + reduced.save(out) + return out + + +def merge_all(run_dir: str | Path) -> list[Path]: + """Merge every catalog plot's chunk partials into ``reduced/.json``.""" + return [merge_one(spec_id, run_dir) for spec_id in catalog_ids()] + + # --------------------------------------------------------------------------- # submit description # --------------------------------------------------------------------------- @@ -185,21 +333,21 @@ class SubmitConfig: run_dir: Path accounting_group: str repo_dir: Path - docker_image: str = "mschnepf/slc7-condocker" - request_memory_mb: int = 4096 + docker_image: str = "cverstege/alma9-gridjob" + request_memory_mb: int = 8192 request_cpus: int = 1 - request_walltime_s: int = 3600 remote: bool = False # +RemoteJob (grid I/O) vs ProvidesETPResources (local files) + n_chunks: int = 1 # per-plot data chunks; ignored for chunkable=False specs _WRAPPER = """#!/bin/bash set -euo pipefail cd {repo_dir} -exec uv run giant analyze compute-one --id "$1" --run-dir {run_dir} +exec {repo_dir}/.venv/bin/giant analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir} """ -def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str: +def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str: reqs_attrs = ( "+RemoteJob = True\n" if cfg.remote @@ -209,39 +357,75 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, ids_file: Path) -> str "universe = docker\n" f"docker_image = {cfg.docker_image}\n" f"executable = {wrapper}\n" - "arguments = $(plotid)\n" + "arguments = $(plotid) $(chunk)\n" "should_transfer_files = YES\n" "when_to_transfer_output = ON_EXIT\n" f"request_memory = {cfg.request_memory_mb}\n" f"request_cpus = {cfg.request_cpus}\n" - f"+RequestWalltime = {cfg.request_walltime_s}\n" + "+RequestWalltime = $(walltime)\n" f"accounting_group = {cfg.accounting_group}\n" f"{reqs_attrs}" - f"output = {cfg.run_dir}/logs/$(plotid).out\n" - f"error = {cfg.run_dir}/logs/$(plotid).err\n" + f"output = {cfg.run_dir}/logs/$(plotid)__$(chunk).out\n" + f"error = {cfg.run_dir}/logs/$(plotid)__$(chunk).err\n" f"log = {cfg.run_dir}/logs/condor.log\n" - f"queue plotid from {ids_file}\n" + f"queue plotid,chunk,walltime from {jobs_file}\n" ) -def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path: - """Write the wrapper script, plot-id list, and HTCondor submit description. +def _job_walltimes( + run_dir: Path, ids: list[str], n_chunks: int +) -> list[tuple[str, int, int]]: + """``(spec_id, chunk, walltime_s)`` for every job, sized from ``run_meta.json``. + Row counts come from ``prep``'s ``RunMeta.rows_per_chunk``/``total_rows``; + ``chunkable=False`` specs (router diagnostics) always use the dataset + total since they run as a single job regardless of ``n_chunks``. + """ + meta = RunMeta.load(run_dir / "run_meta.json") + jobs: list[tuple[str, int, int]] = [] + for spec_id in ids: + chunkable = get_spec(spec_id).chunkable + chunks = range(n_chunks) if chunkable else [0] + for chunk in chunks: + n_rows = meta.rows_per_chunk[chunk] if chunkable else meta.total_rows + jobs.append((spec_id, chunk, estimate_runtime_s(spec_id, n_rows))) + return jobs + + +def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path: + """Write the wrapper script, (plot, chunk) job list, and HTCondor submit + description. + + Each catalog id gets ``cfg.n_chunks`` jobs, except ``chunkable=False`` + specs (the router diagnostics), which always get exactly one regardless of + ``cfg.n_chunks``. Every job's ``+RequestWalltime`` is estimated from its + chunk's row count (``runtime_estimate.estimate_runtime_s``, requires + ``run_meta.json`` from ``prep`` to already carry ``rows_per_chunk``). Returns the submit description path (``/analyze.sub``). Does not submit — call ``condor_submit`` on the returned file. """ + venv_giant = cfg.repo_dir / ".venv" / "bin" / "giant" + if not venv_giant.exists(): + raise FileNotFoundError( + f"{venv_giant} not found — condor jobs run it directly (no `uv` on " + f"the worker image), so run `uv sync --extra cpu` in {cfg.repo_dir} " + "before submitting." + ) + ids = ids or catalog_ids() run_dir = cfg.run_dir (run_dir / "logs").mkdir(parents=True, exist_ok=True) (run_dir / "reduced").mkdir(parents=True, exist_ok=True) + (run_dir / "reduced_partial").mkdir(parents=True, exist_ok=True) wrapper = run_dir / "run_compute.sh" wrapper.write_text(_WRAPPER.format(repo_dir=cfg.repo_dir, run_dir=run_dir)) wrapper.chmod(0o755) - ids_file = run_dir / "plotids.txt" - ids_file.write_text("\n".join(ids) + "\n") + jobs = _job_walltimes(run_dir, ids, cfg.n_chunks) + jobs_file = run_dir / "jobs.txt" + jobs_file.write_text("\n".join(f"{i},{k},{w}" for i, k, w in jobs) + "\n") sub = run_dir / "analyze.sub" - sub.write_text(_submit_description(cfg, wrapper, ids_file)) + sub.write_text(_submit_description(cfg, wrapper, jobs_file)) return sub diff --git a/giant/analysis/reduce.py b/giant/analysis/reduce.py index 51bf035..15089f6 100644 --- a/giant/analysis/reduce.py +++ b/giant/analysis/reduce.py @@ -16,6 +16,8 @@ efficiency" section): from __future__ import annotations +from typing import Any + import numpy as np import polars as pl @@ -58,6 +60,24 @@ def hist1d( return out +def sum_merge(dicts: list[dict[str, Any]]) -> dict[str, Any]: + """Elementwise-sum a list of sum-mergeable count/total dicts (JSON-safe keys). + + Used to merge chunked ``hist1d``/``species_share``-style partials, whose + values bin/group against edges or keys fixed by ``Context`` — a chunk's raw + count dict is exactly a partial sum, so merging is a plain elementwise sum + over the union of keys (a key absent from some chunk is all-zero there). + Values may be per-bin count lists or plain scalar totals; both round-trip + through ``np.asarray``/``.tolist()`` unchanged in shape. + """ + out: dict[str, np.ndarray] = {} + for d in dicts: + for k, v in d.items(): + arr = np.asarray(v) + out[k] = arr.copy() if k not in out else out[k] + arr + return {k: v.tolist() for k, v in out.items()} + + # --------------------------------------------------------------------------- # Per-event scalar observables (one bounded group_by pass) # --------------------------------------------------------------------------- @@ -149,18 +169,19 @@ def transverse_expr() -> pl.Expr: return (tx**2 + ty**2 + tz**2).sqrt() -def weighted_profile( +def profile_partial( lf: pl.LazyFrame, coord: pl.Expr, edges: np.ndarray, weight: pl.Expr, ) -> tuple[np.ndarray, np.ndarray]: - """Event-averaged, ``weight``-summed profile of ``coord``, with an event-RMS band. + """One chunk's per-event x bin ``weight``-sum matrix: ``(event_ids, matrix)``. - One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin); - collapsed in numpy to the per-bin mean over events and its event-to-event std - (the band). ``coord``/``weight`` require the entry/axis columns attached. - Returns ``(mean, std)``, each length ``len(edges)-1``. + One streaming ``group_by(event_id, bin)`` sums ``weight`` per (event, bin). + A chunk's matrix rows are only the events present in that chunk, so chunks' + matrices stack cleanly with no cross-chunk lookup — this requires chunking + to be event-disjoint (every row of an event lands in one chunk). + ``coord``/``weight`` require the entry/axis columns attached. """ lo, hi, nbins = float(edges[0]), float(edges[-1]), len(edges) - 1 grid = ( @@ -177,7 +198,35 @@ def weighted_profile( uniq, inv = np.unique(ev, return_inverse=True) mat = np.zeros((len(uniq), nbins), dtype=np.float64) np.add.at(mat, (inv, grid["_b"].to_numpy()), grid["_ws"].to_numpy()) - return mat.mean(axis=0), mat.std(axis=0) + return uniq, mat + + +def profile_finalize(mats: list[np.ndarray]) -> tuple[np.ndarray, np.ndarray]: + """Collapse per-chunk per-event x bin matrices into the final mean/std profile. + + Chunks are event-disjoint, so row-wise concatenation of their matrices + reconstructs the full per-event matrix; the mean/event-RMS collapse must + happen once over that full matrix — an average of per-chunk means/stds + would be wrong (chunks generally hold different numbers of events). + Returns ``(mean, std)``, each length ``nbins``. + """ + full = np.concatenate(mats, axis=0) + return full.mean(axis=0), full.std(axis=0) + + +def weighted_profile( + lf: pl.LazyFrame, + coord: pl.Expr, + edges: np.ndarray, + weight: pl.Expr, +) -> tuple[np.ndarray, np.ndarray]: + """Single-pass profile of ``coord`` (mean +/- event-RMS band over events). + + Convenience wrapper for the unchunked (whole-dataset) case; ``mean_std + + profile_partial`` is what a chunked compute/finalize split uses instead. + """ + _, mat = profile_partial(lf, coord, edges, weight) + return profile_finalize([mat]) # --------------------------------------------------------------------------- diff --git a/giant/analysis/reduced.py b/giant/analysis/reduced.py index 947641e..2188800 100644 --- a/giant/analysis/reduced.py +++ b/giant/analysis/reduced.py @@ -12,11 +12,14 @@ from dataclasses import asdict, dataclass, field from pathlib import Path # Reduced.kind values: -# "overlay_hist" rollout vs reference density histogram over shared edges -# "grouped_hist" one panel per group (energy/pdg/material), each an overlay -# "profile" edep-weighted mean +/- event-RMS vs depth/radius, two series -# "bar" per-category rollout vs reference bars (share / counts) -# "single_hist" one series only (e.g. rollout leakage; reference has none) +# "overlay_hist" rollout vs reference density histogram over shared edges +# "grouped_hist" one panel per group (energy/pdg/material), each an overlay +# "profile" edep-weighted mean +/- event-RMS vs depth/radius, two series +# "bar" per-category rollout vs reference bars (share / counts) +# "single_hist" one series only (e.g. rollout leakage; reference has none) +# "router_gating" stacked mean MoE gate weight vs energy, rollout + reference +# "router_share" stacked bar of MoE top-1 dispatch share by category +# "unavailable" plot not applicable to this run (e.g. non-MoE checkpoint) @dataclass @@ -36,3 +39,27 @@ class Reduced: @classmethod def load(cls, path: str | Path) -> "Reduced": return cls(**json.loads(Path(path).read_text())) + + +@dataclass +class Partial: + """The raw, not-yet-finalized output of one ``(plot, chunk)`` compute job. + + ``data`` holds whatever shape that plot's ``PlotSpec.compute_partial`` + returns — a raw sum-mergeable count dict, or a raw per-event/per-secondary + array to be concatenated across chunks — never a finished histogram/profile. + ``PlotSpec.finalize`` is the only thing that knows how to interpret it. + """ + + id: str + family: str + chunk: int + data: dict + + def save(self, path: str | Path) -> None: + Path(path).parent.mkdir(parents=True, exist_ok=True) + Path(path).write_text(json.dumps(asdict(self))) + + @classmethod + def load(cls, path: str | Path) -> "Partial": + return cls(**json.loads(Path(path).read_text())) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index d5066c8..3941981 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -41,15 +41,46 @@ def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> Non ax.set_yscale("log") -def _nn_params(run_meta: dict) -> dict: - """Flatten the rollout's model/training provenance for the figure subtitle.""" - params = { - k: v for k, v in (run_meta.get("model_config") or {}).items() if v is not None - } +def _router_summary(model_config: dict) -> str: + r = model_config.get("router") or {} + if not r.get("enabled"): + return "off" + return f"{r.get('type', '?')}×{r.get('n_experts', '?')}" + + +def _figure_params(run_meta: dict) -> dict: + """Curated run identity for the figure subtitle (``new_figure(params=...)``). + + ``run_meta``/each plot's own ``.yaml`` (see ``_plot_metadata``) already + carry every threaded model/training/rollout/dataset parameter for + after-the-fact lookup — this picks only the handful that matter for + telling figures apart at a glance while flipping through a gallery, since + the subtitle is one unwrapped line of text. The last slot is + architecture-conditional: flow/ddpm runs show the ODE ``steps`` used for + this rollout, wgan runs show ``noise_dim`` instead since wgan sampling is + single-pass and has no ODE step count. + """ + mc = run_meta.get("model_config") or {} + mode = mc.get("mode") + params: dict = {} + if mc.get("hidden_dim") is not None: + params["hidden_dim"] = mc["hidden_dim"] + if mc.get("n_blocks") is not None: + params["n_blocks"] = mc["n_blocks"] + if mode is not None: + params["mode"] = mode + if mc.get("conditioning") is not None: + params["conditioning"] = mc["conditioning"] + params["router"] = _router_summary(mc) if run_meta.get("training_epoch") is not None: params["epoch"] = run_meta["training_epoch"] if run_meta.get("best_val_loss") is not None: params["best_val_loss"] = round(run_meta["best_val_loss"], 4) + if mode == "wgan": + if mc.get("noise_dim") is not None: + params["noise_dim"] = mc["noise_dim"] + elif run_meta.get("steps") is not None: + params["steps"] = run_meta["steps"] return params @@ -137,18 +168,95 @@ def _render_bar(r: Reduced, params: dict): return fig +def _render_router_gating(r: Reduced, params: dict): + n_experts = r.payload["n_experts"] + log_x = r.payload.get("log_x", False) + fig, axes = ps.new_figure( + "slide-16x9", title=r.title, params=params, nrows=1, ncols=2, squeeze=False + ) + flat = axes.ravel() + for ax, key in zip(flat, ("rollout", "reference")): + side = r.payload.get(key, {}) + centers = np.asarray(side.get("centers", [])) + means = np.asarray(side.get("means", [])) + if len(centers) and means.size: + cum = np.zeros(len(centers)) + for i in range(n_experts): + ax.fill_between( + centers, cum, cum + means[:, i], alpha=0.7, label=f"expert {i}" + ) + cum = cum + means[:, i] + if log_x: + ax.set_xscale("log") + ax.set_ylim(0, 1) + ax.set_title(_SERIES_LABELS[key], fontsize=8) + ax.set_xlabel(r.xlabel) + flat[0].set_ylabel("mean gate weight") + ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router") + return fig + + +def _render_router_share(r: Reduced, params: dict): + categories = r.payload["categories"] + n_experts = r.payload["n_experts"] + x = np.arange(len(categories)) + present = [k for k in ("rollout", "reference") if k in r.payload] + fig, axes = ps.new_figure( + "slide-16x9", + title=r.title, + params=params, + nrows=1, + ncols=len(present), + squeeze=False, + ) + flat = axes.ravel() + for ax, key in zip(flat, present): + side = r.payload[key] + shares = np.array([side[c] for c in categories]) # (n_cat, n_experts) + bottom = np.zeros(len(categories)) + for i in range(n_experts): + ax.bar(x, shares[:, i], bottom=bottom, label=f"expert {i}") + bottom += shares[:, i] + ax.set_xticks(x) + ax.set_xticklabels(categories, rotation=45, ha="right") + ax.set_ylim(0, 1) + ax.set_title(_SERIES_LABELS[key], fontsize=8) + flat[0].set_ylabel("share of rows dispatched to expert") + ps.style_legend(flat[0], title=f"{r.payload.get('router_type', '')} router") + return fig + + +def _render_unavailable(r: Reduced, params: dict): + fig, ax = ps.new_figure("thesis-single", title=r.title, params=params) + ax.axis("off") + ax.text( + 0.5, + 0.5, + r.payload.get("note", "not available"), + ha="center", + va="center", + wrap=True, + fontsize=10, + transform=ax.transAxes, + ) + return fig + + _RENDERERS = { "overlay_hist": _render_overlay, "single_hist": _render_single, "grouped_hist": _render_grouped, "profile": _render_profile, "bar": _render_bar, + "router_gating": _render_router_gating, + "router_share": _render_router_share, + "unavailable": _render_unavailable, } def render(r: Reduced, run_meta: dict | None = None): """Build the matplotlib figure for one reduced artifact (dispatch on kind).""" - return _RENDERERS[r.kind](r, _nn_params(run_meta or {})) + return _RENDERERS[r.kind](r, _figure_params(run_meta or {})) def _plot_metadata(r: Reduced, run_meta: dict) -> dict: @@ -161,6 +269,11 @@ def _plot_metadata(r: Reduced, run_meta: dict) -> dict: meta.update(r.meta) if "note" in r.payload: meta["note"] = r.payload["note"] + if run_meta: + # Every threaded model/training/rollout/dataset parameter, so a + # single plot's metadata is self-contained for later comparison + # without cross-referencing the run's root metadata.yaml. + meta["parameters"] = {k: v for k, v in run_meta.items() if k != "title"} return meta @@ -226,12 +339,16 @@ def render_all( def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]: """Render a prepped run directory: ``/reduced`` → ``/plots``. - Pulls the rollout provenance (checkpoint, paths, cutoffs) from - ``run_meta.json`` into every plot's gallery metadata. + First joins every plot's chunk partials (``reduced_partial/__*.json``) + into ``reduced/.json`` via ``merge_all`` — a no-op merge when the run + wasn't chunked (``n_chunks=1``) — then pulls the rollout provenance + (checkpoint, paths, cutoffs) from ``run_meta.json`` into every plot's + gallery metadata and renders. """ - from giant.analysis.condor import RunMeta + from giant.analysis.condor import RunMeta, merge_all run_dir = Path(run_dir) + merge_all(run_dir) meta = RunMeta.load(run_dir / "run_meta.json") run_meta = { "title": meta.title, diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py new file mode 100644 index 0000000..6dbc3f8 --- /dev/null +++ b/giant/analysis/router_gating.py @@ -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}, + }, + ) diff --git a/giant/analysis/runtime_estimate.py b/giant/analysis/runtime_estimate.py new file mode 100644 index 0000000..f5279f7 --- /dev/null +++ b/giant/analysis/runtime_estimate.py @@ -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)) diff --git a/giant/analysis/sources.py b/giant/analysis/sources.py index d71b129..2248c90 100644 --- a/giant/analysis/sources.py +++ b/giant/analysis/sources.py @@ -106,18 +106,27 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame: can push their own narrow projection into the parquet read — the single biggest lever on a larger-than-RAM file. ``pl.LazyFrame`` inputs pass straight through (used by tests). + + ``pdg`` is cast to a canonical ``Int64`` here: the rollout writer and the + reference file's upstream ROOT→parquet conversion don't agree on integer + width, and an uncast mismatch only surfaces later as a ``pl.concat`` + ``SchemaError`` (e.g. in ``build_context``'s pdg-count merge). """ if isinstance(source, pl.LazyFrame): - return source + return source.with_columns(pl.col("pdg").cast(pl.Int64)) path = Path(source) if side is Side.rollout: _check_rollout_metadata(path) - return pl.scan_parquet(path) - # The reference (a rollout's seed `dataset`) may be a directory of parquet - # shards rather than a single file — scan them all. - if path.is_dir(): - return pl.scan_parquet(str(path / "**/*.parquet")) - return pl.scan_parquet(path) + lf = pl.scan_parquet(path) + else: + # The reference (a rollout's seed `dataset`) may be a directory of + # parquet shards rather than a single file — scan them all. + lf = ( + pl.scan_parquet(str(path / "**/*.parquet")) + if path.is_dir() + else pl.scan_parquet(path) + ) + return lf.with_columns(pl.col("pdg").cast(pl.Int64)) def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame: diff --git a/giant/cli.py b/giant/cli.py index 6d8c163..3ee1781 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -1188,6 +1188,9 @@ def rollout( ) raise typer.Exit(1) + gconfig.warn_if_checkpoint_config_mismatch(checkpoint) + training_cfg = gconfig.load_checkpoint_config(checkpoint) + model_cfg = ckpt["model_config"] conditioning = model_cfg.get("conditioning", "embedding") pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()} @@ -1270,17 +1273,26 @@ def rollout( "max_steps": max_steps, "steps": steps, "max_tracks_per_event": max_tracks_per_event, + "escape_threshold": escape_threshold, + "n_events": n_events, "n_seed_events": int(len(seeds["event_id"])), - "model_config": { - "mode": model_cfg.get("mode", "flow"), - "hidden_dim": model_cfg.get("hidden_dim"), - "n_blocks": model_cfg.get("n_blocks"), - "emb_dim": model_cfg.get("emb_dim"), - "dropout": model_cfg.get("dropout"), - "conditioning": conditioning, - }, + "weights": weights.value, + "batch_size": batch_size, + "device": str(_device), + "rollout_seed": seed, + "n_rows": summary["n_rows"], + "termination_reason_counts": summary["termination_reason_counts"], + # Full architecture spec baked into the checkpoint — includes the + # entire router sub-dict, not just a hand-picked subset, so any + # model knob (router type/n_experts, noise_dim, vocab sizes, ...) + # is available downstream without touching this command again. + "model_config": dict(model_cfg), "training_epoch": ckpt.get("epoch"), "best_val_loss": ckpt.get("best_val_loss"), + # [train]/[meta] from the sibling config.toml (giant.config.save_config) + # — empty dicts if the checkpoint has no config.toml next to it. + "training_config": dict(training_cfg.get("train", {})), + "training_meta": dict(training_cfg.get("meta", {})), } ) ref_path.write_text(yaml.dump(ref, default_flow_style=False, sort_keys=False)) @@ -1310,12 +1322,18 @@ def analyze_prep( typer.Option( "--run-dir", "-o", - help="Override the run directory (default: next to the rollout parquet)", + help="Override the run directory (default: /analysis_runs/analysis_)", ), ] = None, n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4, n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50, top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6, + chunks: Annotated[ + int, + typer.Option( + "--chunks", help="Split each plot's data into this many event_id chunks" + ), + ] = 1, ) -> None: """Read the rollout YAML → shared.json + run_meta.json in the run directory.""" from giant.analysis import prep @@ -1323,6 +1341,8 @@ def analyze_prep( path = prep( rollout_yaml, run_dir, + n_chunks=chunks, + default_base=Path.cwd() / "analysis_runs", n_energy_bins=n_energy_bins, n_marginal_bins=n_marginal_bins, top_k_pdg=top_k_pdg, @@ -1338,11 +1358,34 @@ def analyze_compute_one( run_dir: Annotated[ Path, typer.Option("--run-dir", help="Run directory from `analyze prep`") ], + chunk: Annotated[ + int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)") + ] = 0, ) -> None: - """Run one plot's streaming reduction (this is what each condor job runs).""" + """Run one (plot, chunk)'s streaming reduction (this is what each condor job runs).""" from giant.analysis import compute_one - path = compute_one(id, run_dir) + path = compute_one(id, run_dir, chunk_index=chunk) + typer.echo(f"wrote {path}") + + +@analyze_app.command("merge-one") +def analyze_merge_one( + id: Annotated[ + str, typer.Option("--id", help="Catalog plot id (see `analyze list`)") + ], + run_dir: Annotated[ + Path, typer.Option("--run-dir", help="Run directory from `analyze prep`") + ], +) -> None: + """Merge one plot's chunk partials into its final reduced JSON. + + Runs automatically as part of `analyze render`; useful standalone to + debug a specific plot without re-rendering everything. + """ + from giant.analysis import merge_one + + path = merge_one(id, run_dir) typer.echo(f"wrote {path}") @@ -1380,26 +1423,48 @@ def analyze_submit( accounting_group: Annotated[str, typer.Option("--accounting-group")], run_dir: Annotated[ Optional[Path], - typer.Option("--run-dir", "-o", help="Override the run directory"), + typer.Option( + "--run-dir", + "-o", + help="Override the run directory (default: /analysis_runs/analysis_)", + ), ] = None, docker_image: Annotated[ str, typer.Option("--docker-image") - ] = "mschnepf/slc7-condocker", - request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 4096, + ] = "cverstege/alma9-gridjob", + request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 8192, remote: Annotated[ bool, typer.Option("--remote/--local", help="+RemoteJob vs ProvidesETPResources"), ] = False, + chunks: Annotated[ + int, + typer.Option( + "--chunks", + help="Split each plot's data into this many event_id chunks/jobs", + ), + ] = 1, + n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4, + n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50, + top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6, dry_run: Annotated[ bool, typer.Option("--dry-run", help="Write files but don't condor_submit") ] = False, ) -> None: - """prep + write the HTCondor submit description (one job per plot), then submit.""" + """prep + write the HTCondor submit description (one job per plot x chunk), then submit.""" import subprocess from giant.analysis import SubmitConfig, prep, write_submit - path = prep(rollout_yaml, run_dir) + path = prep( + rollout_yaml, + run_dir, + n_chunks=chunks, + default_base=Path.cwd() / "analysis_runs", + n_energy_bins=n_energy_bins, + n_marginal_bins=n_marginal_bins, + top_k_pdg=top_k_pdg, + ) cfg = SubmitConfig( run_dir=path, accounting_group=accounting_group, @@ -1407,6 +1472,7 @@ def analyze_submit( docker_image=docker_image, request_memory_mb=request_memory, remote=remote, + n_chunks=chunks, ) sub = write_submit(cfg) typer.echo(f"run directory: {path}") diff --git a/giant/config.py b/giant/config.py index 47f92c7..98a078e 100644 --- a/giant/config.py +++ b/giant/config.py @@ -189,6 +189,21 @@ def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None: ) +def load_checkpoint_config(ckpt_path: str | Path) -> dict: + """Load the full ``[train]``/``[model]``/``[meta]`` config.toml written + alongside a checkpoint by ``save_config``. + + Returns ``{}`` if no config.toml sits next to the checkpoint (older runs, + or a checkpoint moved without its sidecar) — this is best-effort + provenance for threading into a rollout's YAML sidecar, not a hard + requirement for using the checkpoint itself. + """ + config_path = Path(ckpt_path).parent / "config.toml" + if not config_path.exists(): + return {} + return load_toml(config_path) + + def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None: """Look for a config.toml next to a checkpoint and warn on a git_hash mismatch. diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 9b74ea7..a4a5222 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -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], diff --git a/giant/model/network.py b/giant/model/network.py index 5596a17..5345747 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -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: diff --git a/giant/pipeline.py b/giant/pipeline.py index c06a0d0..6518059 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -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, diff --git a/scripts/profile_analysis_costs.py b/scripts/profile_analysis_costs.py new file mode 100644 index 0000000..de97a2b --- /dev/null +++ b/scripts/profile_analysis_costs.py @@ -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() diff --git a/test-cuda.py b/test-cuda.py deleted file mode 100644 index 2b95e64..0000000 --- a/test-cuda.py +++ /dev/null @@ -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.") diff --git a/tests/test_catalog.py b/tests/test_catalog.py index 02d449b..c8a47f0 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -2,21 +2,30 @@ from __future__ import annotations +import numpy as np import pytest from giant.analysis import build_catalog, catalog_ids, get_spec -from giant.analysis.catalog import Bundle -from giant.analysis.context import build_context +from giant.analysis.catalog import Bundle, PlotSpec +from giant.analysis.context import Context, build_context from tests.test_analysis_reduce import _reference_frame, _rollout_frame -@pytest.fixture(scope="module") -def bundle() -> Bundle: +def _build_ctx() -> Context: r, t = _rollout_frame(), _reference_frame() - ctx = build_context( + return build_context( r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000 ) - return Bundle.open(r, t, ctx) + + +@pytest.fixture(scope="module") +def ctx() -> Context: + return _build_ctx() + + +@pytest.fixture(scope="module") +def bundle(ctx: Context) -> Bundle: + return Bundle.open(_rollout_frame(), _reference_frame(), ctx) def test_catalog_ids_unique_and_nonempty(): @@ -36,7 +45,7 @@ def test_get_spec_roundtrip_and_unknown(): def test_every_spec_computes_valid_reduced(bundle: Bundle): for spec in build_catalog(): - r = spec.compute(bundle) + r = spec.finalize([spec.compute_partial(bundle)], bundle.ctx) assert r.id == spec.id assert r.kind in { "overlay_hist", @@ -44,6 +53,9 @@ def test_every_spec_computes_valid_reduced(bundle: Bundle): "profile", "bar", "single_hist", + "router_gating", + "router_share", + "unavailable", } assert r.title and r.xlabel _validate_payload(r) @@ -67,3 +79,75 @@ def _validate_payload(r) -> None: assert len(p[k]) == n elif r.kind == "bar": assert len(p["labels"]) == len(p["rollout"]) == len(p["reference"]) + elif r.kind == "unavailable": + assert p["note"] + elif r.kind == "router_gating": + for side in ("rollout", "reference"): + if side in p: + assert len(p[side]["centers"]) == len(p[side]["means"]) + elif r.kind == "router_share": + for cat in p["categories"]: + for side in ("rollout", "reference"): + if side in p: + assert cat in p[side] + + +# --------------------------------------------------------------------------- +# chunked (compute_partial x N -> finalize) must match the unchunked (N=1) result +# --------------------------------------------------------------------------- + +# One representative id per merge shape: sum-mergeable (marginal_edep, +# sec_count_per_species via pdg-keyed sums), concat-then-finalize with +# data-dependent edges (event_total_edep), concat-then-mean/std (shower_ +# longitudinal), concat-then-max-edge (leakage_fraction), pdg-keyed sum with a +# ratio (species_edep_share), and a chunkable=False passthrough (router_gating). +_CHUNK_EQUIVALENCE_IDS = [ + "marginal_edep", + "species_edep_share", + "event_total_edep", + "shower_longitudinal", + "leakage_fraction", + "sec_count_per_species", + "router_gating", +] + + +def _assert_payload_close(a, b, path: str = "payload") -> None: + """Recursively compare two JSON-shaped payloads (float-tolerant).""" + assert type(a) is type(b), f"{path}: {type(a)} != {type(b)}" + if isinstance(a, dict): + assert set(a) == set(b), f"{path}: key mismatch {set(a)} != {set(b)}" + for k in a: + _assert_payload_close(a[k], b[k], f"{path}.{k}") + elif isinstance(a, list): + assert len(a) == len(b), f"{path}: length mismatch" + for i, (x, y) in enumerate(zip(a, b)): + _assert_payload_close(x, y, f"{path}[{i}]") + elif isinstance(a, float): + assert np.isclose(a, b, atol=1e-9), f"{path}: {a} != {b}" + else: + assert a == b, f"{path}: {a} != {b}" + + +@pytest.mark.parametrize("spec_id", _CHUNK_EQUIVALENCE_IDS) +def test_chunked_matches_unchunked(ctx: Context, spec_id: str): + """A plot computed over N event-disjoint chunks then merged must equal the + same plot computed in one unchunked pass — the core chunking correctness + guarantee (see the analysis-rollout-plots chunking plan).""" + spec: PlotSpec = get_spec(spec_id) + r, t = _rollout_frame(), _reference_frame() + + unchunked_bundle = Bundle.open(r, t, ctx) + unchunked = spec.finalize([spec.compute_partial(unchunked_bundle)], ctx) + + # 4 chunks over only 2 distinct event_ids also exercises empty chunks. + n_chunks = 4 if spec.chunkable else 1 + parts = [ + spec.compute_partial(Bundle.open(r, t, ctx, chunk=(k, n_chunks))) + for k in range(n_chunks) + ] + chunked = spec.finalize(parts, ctx) + + assert chunked.id == unchunked.id + assert chunked.kind == unchunked.kind + _assert_payload_close(unchunked.payload, chunked.payload) diff --git a/tests/test_condor.py b/tests/test_condor.py index 1d6ff3e..c4259ac 100644 --- a/tests/test_condor.py +++ b/tests/test_condor.py @@ -16,11 +16,13 @@ from giant.analysis import ( compute_reduced, derive_run_dir, load_rollout_yaml, + merge_one, prep, write_submit, ) +from giant.analysis.catalog import get_spec from giant.analysis.condor import Context -from giant.analysis.reduced import Reduced +from giant.analysis.reduced import Partial, Reduced from giant.constants import PREDICT_COORD_METADATA_KEY, ROLLOUT_COORD_VALUE from tests.test_analysis_reduce import _reference_frame, _rollout_frame @@ -51,11 +53,22 @@ def _write_inputs(tmp_path: Path) -> Path: return yaml_path -def _prep(rollout_yaml: Path, run_dir: str | Path | None = None) -> Path: +def _fake_venv(repo_dir: Path) -> None: + """Stand in for a `uv sync`'d venv: write_submit checks `.venv/bin/giant` exists.""" + giant = repo_dir / ".venv" / "bin" / "giant" + giant.parent.mkdir(parents=True, exist_ok=True) + giant.write_text("#!/bin/bash\n") + giant.chmod(0o755) + + +def _prep( + rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1 +) -> Path: """``prep`` with small test-sized context bins/sampling.""" return prep( rollout_yaml, run_dir, + n_chunks=chunks, n_energy_bins=2, n_marginal_bins=8, top_k_pdg=3, @@ -76,6 +89,15 @@ def test_derive_run_dir_next_to_rollout(): assert derive_run_dir(y, "/somewhere") == Path("/somewhere") +def test_derive_run_dir_default_base(): + y = {"output": "/data/roll.parquet", "prediction_id": "abcd1234ef", "dataset": "d"} + assert derive_run_dir(y, default_base="/work/lbogner/giant2/analysis_runs") == Path( + "/work/lbogner/giant2/analysis_runs/analysis_abcd1234" + ) + # an explicit run_dir still wins over default_base + assert derive_run_dir(y, "/somewhere", default_base="/other") == Path("/somewhere") + + def test_prep_lays_out_run_dir(tmp_path: Path): yaml_path = _write_inputs(tmp_path) run_dir = _prep(yaml_path) @@ -87,15 +109,25 @@ def test_prep_lays_out_run_dir(tmp_path: Path): assert meta.reference.endswith("reference.parquet") assert meta.plot_meta["checkpoint"] == "/ckpt/best.pt" assert "best.pt" in meta.title + assert meta.n_chunks == 1 + assert meta.rows_per_chunk == [meta.total_rows] # single chunk holds everything + assert meta.total_rows == 8 # 5 rollout rows + 3 reference rows + + +def test_prep_splits_rows_per_chunk(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path), chunks=2) + meta = RunMeta.load(run_dir / "run_meta.json") + assert len(meta.rows_per_chunk) == 2 + assert sum(meta.rows_per_chunk) == meta.total_rows == 8 def test_compute_one_from_run_dir(tmp_path: Path): run_dir = _prep(_write_inputs(tmp_path)) out = compute_one("marginal_edep", run_dir) - assert out == run_dir / "reduced" / "marginal_edep.json" - reduced = Reduced.load(out) - assert reduced.id == "marginal_edep" - assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1 + assert out == run_dir / "reduced_partial" / "marginal_edep__0.json" + partial = Partial.load(out) + assert partial.id == "marginal_edep" and partial.chunk == 0 + assert "r" in partial.data and "t" in partial.data def test_compute_reduced_explicit_paths(tmp_path: Path): @@ -108,30 +140,131 @@ def test_compute_reduced_explicit_paths(tmp_path: Path): run_dir / "shared.json", tmp_path / "r.json", ) - assert Reduced.load(out).id == "marginal_step_length" + assert Partial.load(out).id == "marginal_step_length" + + +def test_merge_one_produces_reduced(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path)) + compute_one("marginal_edep", run_dir) + out = merge_one("marginal_edep", run_dir) + assert out == run_dir / "reduced" / "marginal_edep.json" + reduced = Reduced.load(out) + assert reduced.id == "marginal_edep" + assert len(reduced.payload["rollout"]) == len(reduced.payload["edges"]) - 1 + + +def test_merge_one_fails_loudly_on_missing_chunk(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path), chunks=2) + compute_one("marginal_edep", run_dir, chunk_index=0) # chunk 1 never computed + with pytest.raises(FileNotFoundError, match="missing chunk"): + merge_one("marginal_edep", run_dir) + + +def test_chunked_compute_and_merge_matches_unchunked(tmp_path: Path): + (tmp_path / "a").mkdir() + (tmp_path / "b").mkdir() + unchunked_dir = _prep(_write_inputs(tmp_path / "a")) + compute_one("marginal_step_length", unchunked_dir) + unchunked = Reduced.load(merge_one("marginal_step_length", unchunked_dir)) + + chunked_dir = _prep(_write_inputs(tmp_path / "b"), chunks=2) + for k in range(2): + compute_one("marginal_step_length", chunked_dir, chunk_index=k) + chunked = Reduced.load(merge_one("marginal_step_length", chunked_dir)) + + assert chunked.payload == unchunked.payload + + +def test_compute_reduced_rejects_out_of_range_chunk(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path)) # n_chunks=1 (default) + with pytest.raises(ValueError, match="out of range"): + compute_one("marginal_edep", run_dir, chunk_index=1) def test_write_submit_description(tmp_path: Path): run_dir = _prep(_write_inputs(tmp_path)) + _fake_venv(tmp_path) cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path) txt = write_submit(cfg).read_text() assert "universe = docker" in txt - assert "docker_image = mschnepf/slc7-condocker" in txt + assert "docker_image = cverstege/alma9-gridjob" in txt assert "requirements = TARGET.ProvidesETPResources" in txt assert "accounting_group = cms" in txt - assert "queue plotid from" in txt - assert (run_dir / "plotids.txt").read_text().split() == catalog_ids() + assert "+RequestWalltime = $(walltime)" in txt + assert "queue plotid,chunk,walltime from" in txt + jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()] + assert [i for i, _, _ in jobs] == catalog_ids() + assert all(k == "0" for _, k, _ in jobs) # n_chunks=1 default + assert all(int(w) > 0 for _, _, w in jobs) wrapper = run_dir / "run_compute.sh" assert wrapper.exists() and (wrapper.stat().st_mode & 0o111) body = wrapper.read_text() - assert "giant analyze compute-one --id" in body and "--run-dir" in body + assert "giant analyze compute-one --id" in body + assert "--chunk" in body and "--run-dir" in body + + +def test_write_submit_requires_synced_venv(tmp_path: Path): + run_dir = _prep(_write_inputs(tmp_path)) + cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path) + with pytest.raises(FileNotFoundError, match="uv sync"): + write_submit(cfg) def test_write_submit_remote_flag(tmp_path: Path): run_dir = _prep(_write_inputs(tmp_path)) + _fake_venv(tmp_path) cfg = SubmitConfig( run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, remote=True ) txt = write_submit(cfg).read_text() assert "+RemoteJob = True" in txt assert "ProvidesETPResources" not in txt + + +def test_write_submit_chunks_respect_chunkable(tmp_path: Path): + assert get_spec("router_gating").chunkable is False + run_dir = _prep(_write_inputs(tmp_path), chunks=4) + _fake_venv(tmp_path) + cfg = SubmitConfig( + run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4 + ) + write_submit(cfg) + jobs = [line.split(",") for line in (run_dir / "jobs.txt").read_text().split()] + counts: dict[str, int] = {} + for spec_id, _, _ in jobs: + counts[spec_id] = counts.get(spec_id, 0) + 1 + assert counts["marginal_edep"] == 4 + assert counts["router_gating"] == 1 # chunkable=False, ignores n_chunks + + +def test_estimate_runtime_s_scales_with_rows_and_margin(): + from giant.analysis import RUNTIME_SAFETY_MARGIN, estimate_runtime_s + from giant.analysis.runtime_estimate import _FIXED_OVERHEAD_S + + assert RUNTIME_SAFETY_MARGIN > 0 + small = estimate_runtime_s("marginal_edep", 1_000) + large = estimate_runtime_s("marginal_edep", 100_000_000) + assert small >= (1 + RUNTIME_SAFETY_MARGIN) * _FIXED_OVERHEAD_S + assert large > small # bigger chunk -> longer estimate + + +def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path): + """A chunked run's later job walltimes track that chunk's row count.""" + from giant.analysis.runtime_estimate import estimate_runtime_s + + run_dir = _prep(_write_inputs(tmp_path), chunks=2) + meta = RunMeta.load(run_dir / "run_meta.json") + _fake_venv(tmp_path) + cfg = SubmitConfig( + run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=2 + ) + write_submit(cfg) + jobs = { + (i, int(k)): int(w) + for i, k, w in ( + line.split(",") for line in (run_dir / "jobs.txt").read_text().split() + ) + } + for chunk in range(2): + expected = estimate_runtime_s("marginal_edep", meta.rows_per_chunk[chunk]) + assert jobs[("marginal_edep", chunk)] == expected diff --git a/tests/test_router.py b/tests/test_router.py index 8f3fa2a..62a8c86 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -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) diff --git a/tests/test_router_gating.py b/tests/test_router_gating.py new file mode 100644 index 0000000..51123d3 --- /dev/null +++ b/tests/test_router_gating.py @@ -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 diff --git a/tests/test_transforms.py b/tests/test_transforms.py index a7af975..c9bc7ac 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -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", + )