Compare commits
3 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 lars
|
313373cc10 | ||
|
lars
|
98b09d2b4b | ||
|
lars
|
d0381728ee |
@@ -82,11 +82,7 @@ jobs:
|
||||
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 --cov --cov-report=term-missing --cov-report=xml
|
||||
- uses: actions/upload-artifact@v3
|
||||
with:
|
||||
name: coverage-report
|
||||
path: coverage.xml
|
||||
- run: uv run pytest
|
||||
|
||||
sync-version-on-tag:
|
||||
name: Sync project version with tag
|
||||
|
||||
@@ -20,8 +20,3 @@ checkpoints/
|
||||
|
||||
# giant analyze run directories (shared.json, reduced/, plots/, condor logs)
|
||||
/analysis_runs/
|
||||
|
||||
# Coverage artifacts
|
||||
.coverage
|
||||
coverage.xml
|
||||
htmlcov/
|
||||
|
||||
@@ -22,7 +22,7 @@ giant analyze render <run_dir> --gallery # render PDFs + HTML
|
||||
dwarf --help # dataset/tooling CLI: convert, migrate, bump-gen,
|
||||
# bump-schema, status, update-manifest, create-manifest,
|
||||
# make-root, build-geometry-oracle, warm-cache, hparam-scan
|
||||
# (see giant/tools/dwarf.py)
|
||||
# (see scripts/dwarf.py)
|
||||
```
|
||||
|
||||
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x; newer torch requires newer NVIDIA drivers). Plain `uv sync` with no extra will not install torch at all; uv has no concept of a "default extra", so `--extra cpu` should always be included unless you need GPU support.
|
||||
@@ -91,6 +91,4 @@ GIANT is a conditional generative surrogate for the Geant4 step function. It rep
|
||||
|
||||
A sampling-calorimeter (multi-material) dataset is still a planned future direction, not yet built. See the knowledge base (`/home/lars/knowledge-base/meta/roadmap.md`).
|
||||
|
||||
**v0.3.0 — Stage-2 autoregressive redesign (designed, not implemented; branch `v0.3.0-stage2-autoregressive`):** the 2026-08-03 WGAN rollout benchmark failed specifically at the secondary-species level (zero photon secondaries, ~4M hallucinated `-14` muon antineutrinos). The agreed response pivots Stage 2 to **autoregressive generation** in descending-energy order with teacher forcing, and switches the particle-type representation back to **categorical** (top N−1 by training-set count + an "other" bucket), reversing the 2026-07-17 continuous `(log-mass, charge)` target. This requires a config break: `[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]` blocks replace the single global `train.mode` + `[model]`, so per-stage generators (`stage1 = flow` + `stage2 = wgan`), stage-2-only training, and one-shot-vs-autoregressive comparison are all expressible. `network.py` is refactored from ten permutation classes into composable parts (encoder × trunk × objective), which also makes routed WGAN work for the first time. This config break is why v0.2-shaped configs/checkpoints need migrating at all (`config.migrate_config`, `model.network._migrate_legacy_model_config`, both drawing on shared facts in `giant/_migration.py`) — v0.2 checkpoint-loading support has **no expiry decided yet**: `/ceph` still holds pre-v0.3.0 checkpoints and analysis runs referencing them, so don't delete or substantially alter either migration function or `tests/legacy/network_v02_snapshot.py` (the frozen v0.2 snapshot they're tested against) without an explicit decision to do so first.
|
||||
|
||||
**Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet.
|
||||
|
||||
@@ -1,29 +1,20 @@
|
||||
# giant
|
||||
|
||||
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
|
||||
**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate — a play on *Geant4* and the step function being the computationally heaviest part of the simulation.
|
||||
|
||||
A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.
|
||||
Conditional generative surrogate for the Geant4 step function. Given a pre-step particle state, the model samples a physically plausible post-step outcome — the primary's continuation plus the variable-length list of secondary particles it produces — replacing the stochastic Geant4 physics engine with a trained generative model. A trained checkpoint autoregressively rolls out full showers, stepping each primary and pushing secondaries as new tracks.
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
uv sync --extra cpu # install deps (CPU torch; use --extra cuda for GPU)
|
||||
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 # scaffold config.toml + run dir
|
||||
giant train path/to/steps.parquet # train (flow + wgan by default)
|
||||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||||
|
||||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl # needed for rollout
|
||||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||||
```
|
||||
|
||||
Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.
|
||||
Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency.
|
||||
|
||||
## Architecture
|
||||
|
||||
A **two-stage model**, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (`--stage1-generator`/`--stage2-generator`, or `--mode` to set both at once): `flow` (conditional flow matching, ODE-sampled in ~10 steps), `ddpm` (denoising diffusion), or `wgan` (single-pass WGAN-GP generator/critic).
|
||||
A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`):
|
||||
|
||||
**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
|
||||
- **`flow`** (default) — conditional flow matching (Lipman et al. 2022): an MLP learns a vector field mapping noise → step outcomes, sampled via ODE integration in ~10 steps.
|
||||
- **`ddpm`** — a standard denoising diffusion baseline for comparison (`giant/model/schedule.py:CosineSchedule`).
|
||||
- **`wgan`** — a single-pass Wasserstein-GAN-GP generator/critic (`giant/model/wgan.py`), trading iterative sampling for one forward pass; implemented, not yet validated against the flow-matching baseline.
|
||||
|
||||
**Stage 1 — primary (9D, `giant/constants.py:LOCAL_TARGET_NAMES`):**
|
||||
|
||||
| Index | Variable | Encoding |
|
||||
|-------|----------|----------|
|
||||
@@ -32,29 +23,36 @@ A **two-stage model**, checkpointed together. Either stage's outcome can be prod
|
||||
| 3–5 | `post_dir` in local frame | unit vector |
|
||||
| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
|
||||
|
||||
- Energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — conservation is architectural, not learned.
|
||||
- `post_dir`/`travel_dir` live in the frame where `pre_dir = ẑ`. `post_pos` isn't a target — it's reconstructed as `pre_pos + step_length * world_frame(travel_dir)`.
|
||||
The two energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — **energy conservation is built into the parametrization**, not left to the loss (`energy_simplex_decode`). Stage 1 also has a classifier head (`predict_n_sec`) predicting the number of secondaries `n_sec ∈ {0..K_MAX}` (`K_MAX = 15`) from the conditioning alone, no diffusion noise involved.
|
||||
|
||||
**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (`--stage2-decoder`):
|
||||
Both `post_dir` and `travel_dir` are expressed in the coordinate frame where `pre_dir = ẑ`, making the scattering distribution nearly azimuthally symmetric. `post_pos` itself is not a raw target — it's reconstructed at inference as `pre_pos + step_length * world_frame(travel_dir)`, so the two stay consistent by construction instead of being learned (and potentially diverging) independently.
|
||||
|
||||
- `autoregressive` — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (`markov`: previous token only, or `attention`: causal self-attention, KV-cached at inference)
|
||||
- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
|
||||
**Stage 2 — secondaries (`SecondaryDecoder`):** conditioned on the pre-step state *and* the Stage-1 outcome, a second net generates all `K_MAX` secondary slots at once — `(stick-breaking energy logit, local-frame direction, log-mass, charge)` per slot, ordered by descending energy; slots beyond the predicted `n_sec` are masked. Secondary energies are a stick-breaking partition of the `e_sec` budget from Stage 1, so the whole chain conserves energy. A secondary's mass/charge are regressed directly against its ground-truth PDG code's physical values (`giant.particles.particle_mass_charge`) and used as-is at inference — including for its own conditioning if it takes further steps in a rollout. No snapping to a known PDG code happens in the model path; `giant.particles.nearest_known_pdg` is a reporting-only lookup used to populate a nominal `pdg` label on output rows.
|
||||
|
||||
Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).
|
||||
**Conditioning (`--conditioning`, per-checkpoint):** pre-step position, log(pre-energy), pre-step direction, layer ID, plus particle/material physical properties — mass/charge (`giant/particles.py`) and Z_eff/A_eff/density/X0/λ_int (`giant/materials.py`). Two mutually exclusive modes:
|
||||
|
||||
**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above (`--conditioning`) — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
|
||||
- **`physical`** (default) — the physical-property columns are routed through small MLPs, computable for any PDG code / material, letting the surrogate generalize to species/materials outside the training menu.
|
||||
- **`embedding`** — the original design: a learned `nn.Embedding` per PDG code / material, kept as a generalization-comparison baseline (memorizes the training menu).
|
||||
|
||||
**MoE routing** (`--router`, either stage): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk.
|
||||
`n_sec` and `e_sec` are model outputs, not conditioning inputs — a rollout is self-contained and never injects ground truth.
|
||||
|
||||
**Mixture-of-experts routing (`--router`, opt-in):** `giant/model/network.py` also implements a pluggable `Router` contract (`ROUTER_REGISTRY`: `energy`, `pdg`, `process`, plus a `composed` router combining several axes) that splits `DenoisingMLP`/`SecondaryDecoder` into per-expert trunks, soft-gated in training and top-1 dispatched at eval. Implemented; first rollout benchmark needs a retrain with a load-balancing loss and better-seeded router centers (see Roadmap). See `--router-type`/`--n-experts`/`--router-axis` on `giant train`/`giant new-run`.
|
||||
|
||||
## Roadmap
|
||||
|
||||
**Phase 1 (done):** `n_sec` and total secondary energy `e_sec` were conditioning inputs; the model predicted only the 9D primary post-step (energy-conservation PoC).
|
||||
|
||||
**Phase 2 (implemented — baseline):** the two-stage model above predicts `n_sec` and each secondary's energy, direction, and species jointly with the primary, so a shower rollout is fully self-contained.
|
||||
|
||||
**Physical-property conditioning (implemented):** replaces learned PDG/material embeddings with physical-property MLPs (see above); Stage 2 predicts a secondary's mass/charge directly instead of a snapped species embedding. Not yet done: the held-out-material/species generalization comparison against the `embedding` baseline — the natural dataset for that is the 34GB multi-material dataset at the repo root (6 materials, 237 PDG codes).
|
||||
|
||||
**Faster-eval architectures (implemented, validation in progress):** both target a ~10× native-Geant4 eval budget. WGAN-GP (`--mode wgan`) has no rollout-vs-reference analysis run against it yet. The MoE router (`--router`) had its first rollout benchmark diverge from Geant4 despite matching bulk deposited energy — the experts weren't specializing (near-uniform gating), traced to a missing load-balance loss and a center-init that didn't match the real energy distribution; both are now fixable via `lambda_balance > 0` and quantile-seeded router centers, but a re-run to confirm hasn't happened yet.
|
||||
|
||||
A multi-material sampling-calorimeter dataset is a planned future direction, not yet built.
|
||||
|
||||
## Data
|
||||
|
||||
- Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from ROOT via `dwarf convert`. One row = one Geant4 step.
|
||||
- **Conditioning (pre-step) columns:** `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz` (direction), `material`, `layer_id`.
|
||||
- **Primary outcome (post-step) columns:** `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep` (energy deposited in this step), `e_sec` (total energy carried off by secondaries), `child_track_ids` (its length gives `n_sec`).
|
||||
- **Secondary columns**, one variable-length list per step: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX` (15) slots on load, ordered by descending energy.
|
||||
- **Optional:** `process` — the physics process that produced the step (e.g. `compt`, `phot`, `eBrem`); a post-step label used only as classifier supervision (`ProcessRouter`), never as conditioning.
|
||||
- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
|
||||
- Loading a directory or `.manifest` of multiple parquet files (each one Geant4 job, `event_id` restarting from 0) offsets each file's `event_id`s by a fixed per-file stride so ids stay globally unique across files.
|
||||
Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from a ROOT file via `dwarf convert`. Each row is one Geant4 step. Train/val split is by `event_id` (not row shuffle, and `--seed`-controlled) to avoid leaking correlated steps from the same shower.
|
||||
|
||||
## Project structure
|
||||
|
||||
@@ -66,20 +64,15 @@ giant/
|
||||
│ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode
|
||||
│ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch)
|
||||
│ ├── model/
|
||||
│ │ ├── network.py # ConditionEncoder, Stage1Model, Stage2OneShot/Stage2Autoregressive, Router/MoE, CriticModel
|
||||
│ │ ├── network.py # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic
|
||||
│ │ ├── schedule.py # CosineSchedule (DDPM) and flow matching utilities
|
||||
│ │ └── wgan.py # WGAN-GP gradient penalty / critic / generator losses
|
||||
│ ├── constants.py # output/conditioning dims, K_MAX, secondary slot layout, schema keys
|
||||
│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
|
||||
│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; nearest-known-PDG lookup
|
||||
│ ├── materials.py # material name → (Z_eff, A_eff, density, X0, λ_int)
|
||||
│ ├── config.py # default hyperparameters, TOML config merging, device autodetect
|
||||
│ ├── pipeline.py # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
|
||||
│ ├── training/ # two-stage training: loop, per-stage trainers, metrics, checkpointing
|
||||
│ │ ├── loop.py # epoch loop, graceful shutdown, best-checkpoint selection
|
||||
│ │ ├── trainers.py # StageSpec + flow/ddpm and WGAN-GP per-stage trainers
|
||||
│ │ ├── stage2_inputs.py# ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
|
||||
│ │ ├── metrics.py # MetricsCollector: metrics.csv columns, W&B logging, progress/summary
|
||||
│ │ └── checkpoint.py # checkpoint assembly/restore (format unchanged since v0.2)
|
||||
│ ├── train.py # two-stage training loop, checkpointing, graceful shutdown, W&B logging
|
||||
│ ├── sample.py # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
|
||||
│ ├── geometry.py # GeometryOracle: position → (material, layer_id, escaped) for rollout
|
||||
│ ├── rollout.py # autoregressive shower rollout driver
|
||||
@@ -93,7 +86,7 @@ giant/
|
||||
│ │ ├── condor.py # prep / compute-one / submit-description plumbing
|
||||
│ │ └── render.py # PDFs + HTML gallery (only module importing plotstyle/LaTeX)
|
||||
│ └── cli.py # `giant train` / `new-run` / `predict` / `rollout` / `analyze` Typer app
|
||||
├── giant/tools/ # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
|
||||
├── scripts/ # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
|
||||
│ ├── dwarf.py # Typer app: convert, migrate, bump-gen, bump-schema, status,
|
||||
│ │ # update-manifest, create-manifest, make-root,
|
||||
│ │ # build-geometry-oracle, warm-cache, hparam-scan
|
||||
@@ -113,48 +106,39 @@ giant/
|
||||
## Setup
|
||||
|
||||
```bash
|
||||
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
|
||||
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
|
||||
uv sync --extra cpu # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
|
||||
uv sync --extra cpu --extra dev # add dev tools (pytest, ruff, ty)
|
||||
uv sync --extra cpu --extra geometry # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
|
||||
```
|
||||
|
||||
`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||||
`cpu` and `cuda` are mutually exclusive extras selecting the torch build (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
|
||||
|
||||
## Training, prediction, rollout
|
||||
## Training, prediction, and rollout
|
||||
|
||||
```bash
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir
|
||||
giant train path/to/steps.parquet # train (flow stage 1 + wgan stage 2, default)
|
||||
giant new-run --hidden-dim 512 --lr 3e-4 --comment "..." # scaffold a config.toml + run dir for a new run
|
||||
giant train path/to/steps.parquet --mode flow # train (flow matching; also --mode ddpm / wgan)
|
||||
giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
|
||||
|
||||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl # position → material/layer_id
|
||||
# Full-shower rollout needs a geometry oracle (position → material/layer_id):
|
||||
dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl
|
||||
giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
|
||||
```
|
||||
|
||||
Useful flags on `giant train`:
|
||||
|
||||
- `--mode {flow,ddpm,wgan}` sets both stages' objective at once; `--stage1-generator`/`--stage2-generator` override per stage
|
||||
- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
|
||||
- `--conditioning {physical,embedding,onehot}` — conditioning representation
|
||||
- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
|
||||
- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
|
||||
- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
|
||||
|
||||
Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
|
||||
|
||||
`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.
|
||||
`train`/`predict` accept a TOML config file (`--config`) and CLI overrides for hyperparameters; see `--help` on any command for the full option list. `giant train --wandb` logs per-epoch metrics (the same ones written to `metrics.csv`) to Weights & Biases; requires `uv sync --extra wandb`. A repeat `giant train` against the same dataset (e.g. a hyperparameter sweep) reuses a cached setup-stage sidecar (vocab maps, event split, normalizer stats) unless `--no-cache-setup`/`--rebuild-setup-cache`; `dwarf warm-cache` precomputes it ahead of time. `giant rollout` seeds showers from the highest-energy entry step per event, then autoregressively steps the two-stage model to completion — pushing secondaries as new tracks and looking up `material`/`layer_id` from the oracle at each step. Tracks terminate on energy cutoff, per-track max steps, detector escape, or natural end; energy is deposited locally on every stop except escape (leakage), so showers conserve energy by construction.
|
||||
|
||||
## Validation and analysis
|
||||
|
||||
- `giant.validate.validate_marginals` — step-level marginal + KL-divergence checks during training (`--validate-every`)
|
||||
- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
|
||||
`giant.validate.validate_marginals` runs step-level marginal and KL-divergence checks during training (`--validate-every`).
|
||||
|
||||
For deeper rollout-vs-reference diagnostics — marginals stratified by energy/pdg/material, per-event totals, shower profiles, species share, leakage, and secondaries — `giant analyze` runs a streaming compute/render pipeline against a `giant rollout` YAML sidecar:
|
||||
|
||||
```bash
|
||||
giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot (compute only)
|
||||
giant analyze render <run_dir> --gallery # local: styled PDFs + HTML gallery (needs LaTeX)
|
||||
```
|
||||
|
||||
`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
|
||||
`<run_dir>` is derived next to the rollout parquet (`analyze prep`/`submit` print it). Compute jobs are polars/numpy only; only `render` imports plotstyle/LaTeX, so it always runs locally.
|
||||
|
||||
## Development
|
||||
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
"""Shared v0.2 -> v0.3 migration knowledge.
|
||||
|
||||
v0.3.0 broke the config format (single `[train]` + `[model]` -> `[conditioning]`/
|
||||
`[stage1_model]`/`[stage2_model]`/`[train]`), and that break has to be absorbed by two
|
||||
independent migration surfaces: `giant.config.migrate_config` (a v0.2 `config.toml`) and
|
||||
`giant.model.network._migrate_legacy_model_config` (a v0.2 checkpoint's flat
|
||||
`model_config` dict). Both translate the same v0.2 facts into the same v0.3 shape, so
|
||||
the facts live here once rather than as two hand-maintained copies — see issues.md
|
||||
Issue 6.
|
||||
|
||||
A dependency-free leaf module so neither `config.py` nor `network.py` has to import the
|
||||
other to share this.
|
||||
"""
|
||||
|
||||
# v0.2 model-shaped keys (config.toml's [model] table, or a checkpoint's flat
|
||||
# model_config dict — same key names in both) applied identically to both v0.3 stage
|
||||
# blocks, because v0.2 had only one trunk shape shared by both stages.
|
||||
V02_MODEL_KEY_TO_STAGES: tuple[tuple[str, str], ...] = (
|
||||
("hidden_dim", "hidden_dim"),
|
||||
("n_blocks", "n_res_blocks"),
|
||||
("dropout", "dropout"),
|
||||
)
|
||||
|
||||
# v0.2 architectural facts that had no corresponding config key at all — always true of
|
||||
# a v0.2 model, so both migration surfaces inject them unconditionally. Keyed by dotted
|
||||
# path relative to the migrated dict's root. NOTE: conditioning.*.n_layers (2) differs
|
||||
# from the v0.3 *default* (1) — not a typo, v0.2's conditioning MLP was always 2 layers
|
||||
# deep.
|
||||
V02_FIXED_FACTS: dict[str, object] = {
|
||||
"conditioning.out_dim": 128,
|
||||
"conditioning.particle.n_layers": 2,
|
||||
"conditioning.material.n_layers": 2,
|
||||
"stage1_model.active": True,
|
||||
"stage1_model.flow.time_dim": 64,
|
||||
"stage1_model.ddpm.time_dim": 64,
|
||||
"stage2_model.active": True,
|
||||
"stage2_model.flow.time_dim": 64,
|
||||
"stage2_model.ddpm.time_dim": 64,
|
||||
"stage2_model.context_dim": 64,
|
||||
"stage2_model.decoder": "one_shot",
|
||||
"stage2_model.particle_type.target": "physical",
|
||||
}
|
||||
|
||||
|
||||
def reject_legacy_router_expert_sizing(router_cfg: dict, *, source: str) -> None:
|
||||
"""Pop and validate v0.2's per-expert width/depth override, in place.
|
||||
|
||||
v0.3.0 removed per-expert sizing — experts always inherit the stage's
|
||||
hidden_dim/n_res_blocks — so a v0.2 router config/checkpoint that set a non-default
|
||||
`expert_hidden_dim`/`expert_n_blocks` describes experts with a different width/depth
|
||||
than the monolith, and can only be reproduced by v0.2 code. Silently dropping these
|
||||
keys (a router builder's kwarg filtering would do this for free) would resize the
|
||||
experts instead of refusing, so this raises loudly.
|
||||
|
||||
Always pops both keys, whether or not they were non-default, so callers can go on
|
||||
to use the (now-cleaned) `router_cfg` unconditionally. `source` names what's being
|
||||
migrated (e.g. "v0.2 config's model.router" or "this checkpoint's
|
||||
model_config.router") for the error message.
|
||||
"""
|
||||
expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0)
|
||||
expert_n_blocks = router_cfg.pop("expert_n_blocks", 0)
|
||||
if not (expert_hidden_dim or expert_n_blocks):
|
||||
return
|
||||
raise ValueError(
|
||||
f"{source} sets expert_hidden_dim/expert_n_blocks to a non-default value "
|
||||
f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed per-expert "
|
||||
"sizing (experts always inherit the stage's hidden_dim/n_res_blocks), so "
|
||||
"this router's experts have a different width/depth than the monolith. "
|
||||
"This checkpoint/config can only be loaded by v0.2 code."
|
||||
)
|
||||
+59
-37
@@ -58,7 +58,6 @@ from giant.analysis.router_gating import (
|
||||
compute_router_share_by_process,
|
||||
)
|
||||
from giant.analysis.sources import Side, open_side, physical_steps, secondaries
|
||||
from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
|
||||
from giant.analysis.variables import RANGED_VARS, cos_scatter_expr
|
||||
|
||||
|
||||
@@ -72,11 +71,6 @@ class Bundle:
|
||||
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
|
||||
# Diagnostic pre-aggregated at rollout time (giant.rollout.
|
||||
# L1DistCollector.summary()) — from the rollout YAML, type_embedding_l1_distance
|
||||
# only. Unlike checkpoint/router_gating, this needs no live model: it's
|
||||
# already a finished histogram, just passed through.
|
||||
type_embedding_l1_dist: dict | None = None
|
||||
|
||||
@classmethod
|
||||
def open(
|
||||
@@ -86,7 +80,6 @@ class Bundle:
|
||||
ctx: Context,
|
||||
checkpoint=None,
|
||||
chunk: tuple[int, int] | None = None,
|
||||
type_embedding_l1_dist: dict | None = None,
|
||||
) -> "Bundle":
|
||||
"""Open both sides, optionally restricted to one event-disjoint chunk.
|
||||
|
||||
@@ -110,7 +103,6 @@ class Bundle:
|
||||
r_phys=physical_steps(r_all, Side.rollout),
|
||||
t_phys=physical_steps(t_all, Side.reference),
|
||||
checkpoint=checkpoint,
|
||||
type_embedding_l1_dist=type_embedding_l1_dist,
|
||||
)
|
||||
|
||||
|
||||
@@ -166,7 +158,9 @@ def _finalize_counts(merged: dict[str, list], key, nbins: int) -> list[int]:
|
||||
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]:
|
||||
def _np_hist_pair(
|
||||
r: np.ndarray, t: np.ndarray, nbins: int
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Shared-edge histogram of two small per-event arrays (robust range)."""
|
||||
both = np.concatenate([r, t]) if (len(r) or len(t)) else np.array([0.0, 1.0])
|
||||
lo, hi = float(np.quantile(both, 0.001)), float(np.quantile(both, 0.999))
|
||||
@@ -238,7 +232,9 @@ def _marginal_overall_finalize(parts: list[dict], ctx: Context, var: str) -> Red
|
||||
|
||||
def _energy_group_expr(lf: pl.LazyFrame, edges: np.ndarray) -> pl.Expr:
|
||||
ids, bins = event_energy_bins(lf, edges)
|
||||
return pl.col("event_id").replace_strict(ids, bins, default=-1, return_dtype=pl.Int64)
|
||||
return pl.col("event_id").replace_strict(
|
||||
ids, bins, default=-1, return_dtype=pl.Int64
|
||||
)
|
||||
|
||||
|
||||
def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
|
||||
@@ -261,7 +257,9 @@ def _marginal_grouped_partial(b: Bundle, var: str, axis: str) -> dict:
|
||||
}
|
||||
|
||||
|
||||
def _marginal_grouped_finalize(parts: list[dict], ctx: Context, var: str, axis: str) -> Reduced:
|
||||
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
|
||||
@@ -311,7 +309,9 @@ def _event_scalar_partial(b: Bundle, col: str, use_all: bool) -> dict:
|
||||
return {"r": r.tolist(), "t": t.tolist()}
|
||||
|
||||
|
||||
def _event_scalar_finalize(parts: list[dict], ctx: Context, spec_id: str, title: str, xlabel: str) -> Reduced:
|
||||
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)
|
||||
@@ -480,7 +480,9 @@ def _leakage_partial(b: Bundle) -> dict:
|
||||
|
||||
def _leakage_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
frac = np.concatenate([np.asarray(p["frac"], dtype=float) for p in parts])
|
||||
edges = np.linspace(0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1)
|
||||
edges = np.linspace(
|
||||
0.0, max(float(frac.max()) if len(frac) else 1.0, 1e-3), ctx.n_marginal_bins + 1
|
||||
)
|
||||
counts = np.histogram(frac, edges)[0]
|
||||
return Reduced(
|
||||
id="leakage_fraction",
|
||||
@@ -511,8 +513,18 @@ def _sec_frames(b: Bundle):
|
||||
|
||||
def _sec_count_per_event_partial(b: Bundle) -> dict:
|
||||
r_sec, t_sec = _sec_frames(b)
|
||||
r = r_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
|
||||
t = t_sec.group_by("event_id").agg(pl.len().alias("n")).collect(engine="streaming")["n"].to_numpy()
|
||||
r = (
|
||||
r_sec.group_by("event_id")
|
||||
.agg(pl.len().alias("n"))
|
||||
.collect(engine="streaming")["n"]
|
||||
.to_numpy()
|
||||
)
|
||||
t = (
|
||||
t_sec.group_by("event_id")
|
||||
.agg(pl.len().alias("n"))
|
||||
.collect(engine="streaming")["n"]
|
||||
.to_numpy()
|
||||
)
|
||||
return {"r": r.tolist(), "t": t.tolist()}
|
||||
|
||||
|
||||
@@ -548,7 +560,9 @@ def _sec_count_per_species_partial(b: Bundle) -> dict:
|
||||
def _sec_count_per_species_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
r = sum_merge([p["r"] for p in parts])
|
||||
t = sum_merge([p["t"] for p in parts])
|
||||
keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[: len(ctx.top_pdgs)]
|
||||
keys = sorted(set(r) | set(t), key=lambda k: -(r.get(k, 0) + t.get(k, 0)))[
|
||||
: len(ctx.top_pdgs)
|
||||
]
|
||||
return Reduced(
|
||||
id="sec_count_per_species",
|
||||
family="secondaries",
|
||||
@@ -595,9 +609,11 @@ def _sec_energy_finalize(parts: list[dict], ctx: Context) -> Reduced:
|
||||
|
||||
def _sec_cos_angle_partial(b: Bundle) -> dict:
|
||||
edges = np.linspace(-1.0, 1.0, b.ctx.n_sec_bins + 1)
|
||||
cos = (pl.col("sdx") * pl.col("axis_x") + pl.col("sdy") * pl.col("axis_y") + pl.col("sdz") * pl.col("axis_z")).clip(
|
||||
-1.0, 1.0
|
||||
)
|
||||
cos = (
|
||||
pl.col("sdx") * pl.col("axis_x")
|
||||
+ pl.col("sdy") * pl.col("axis_y")
|
||||
+ pl.col("sdz") * pl.col("axis_z")
|
||||
).clip(-1.0, 1.0)
|
||||
|
||||
def _side(sec_lf: pl.LazyFrame, steps_lf: pl.LazyFrame) -> dict[str, list[int]]:
|
||||
ea = entry_axis(steps_lf)
|
||||
@@ -635,14 +651,13 @@ _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)
|
||||
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)
|
||||
)
|
||||
_type_embedding_l1_distance_partial, _type_embedding_l1_distance_finalize = _unchunkable(
|
||||
lambda b: compute_type_embedding_l1_distance(b.type_embedding_l1_dist)
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
@@ -663,7 +678,9 @@ def build_catalog() -> list[PlotSpec]:
|
||||
f"marginal_{var}",
|
||||
"marginals",
|
||||
compute_partial=lambda b, v=var: _marginal_overall_partial(b, v),
|
||||
finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(parts, ctx, v),
|
||||
finalize=lambda parts, ctx, v=var: _marginal_overall_finalize(
|
||||
parts, ctx, v
|
||||
),
|
||||
)
|
||||
)
|
||||
for axis in GROUPING_AXES:
|
||||
@@ -671,8 +688,12 @@ def build_catalog() -> list[PlotSpec]:
|
||||
PlotSpec(
|
||||
f"marginal_{var}_by_{axis}",
|
||||
"marginals",
|
||||
compute_partial=lambda b, v=var, a=axis: _marginal_grouped_partial(b, v, a),
|
||||
finalize=lambda parts, ctx, v=var, a=axis: _marginal_grouped_finalize(parts, ctx, v, a),
|
||||
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)
|
||||
),
|
||||
)
|
||||
)
|
||||
|
||||
@@ -680,7 +701,9 @@ def build_catalog() -> list[PlotSpec]:
|
||||
PlotSpec(
|
||||
"event_total_edep",
|
||||
"event",
|
||||
compute_partial=lambda b: _event_scalar_partial(b, "total_edep", use_all=True),
|
||||
compute_partial=lambda b: _event_scalar_partial(
|
||||
b, "total_edep", use_all=True
|
||||
),
|
||||
finalize=lambda parts, ctx: _event_scalar_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
@@ -698,7 +721,9 @@ def build_catalog() -> list[PlotSpec]:
|
||||
PlotSpec(
|
||||
"event_mean_length",
|
||||
"event",
|
||||
compute_partial=lambda b: _event_scalar_partial(b, "mean_length", use_all=False),
|
||||
compute_partial=lambda b: _event_scalar_partial(
|
||||
b, "mean_length", use_all=False
|
||||
),
|
||||
finalize=lambda parts, ctx: _event_scalar_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
@@ -710,7 +735,9 @@ def build_catalog() -> list[PlotSpec]:
|
||||
PlotSpec(
|
||||
"event_n_steps",
|
||||
"event",
|
||||
compute_partial=lambda b: _event_scalar_partial(b, "n_steps", use_all=False),
|
||||
compute_partial=lambda b: _event_scalar_partial(
|
||||
b, "n_steps", use_all=False
|
||||
),
|
||||
finalize=lambda parts, ctx: _event_scalar_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
@@ -735,7 +762,9 @@ def build_catalog() -> list[PlotSpec]:
|
||||
PlotSpec(
|
||||
"shower_transverse",
|
||||
"shower",
|
||||
compute_partial=lambda b: _profile_partial(b, transverse_expr, "transverse_edges"),
|
||||
compute_partial=lambda b: _profile_partial(
|
||||
b, transverse_expr, "transverse_edges"
|
||||
),
|
||||
finalize=lambda parts, ctx: _profile_finalize(
|
||||
parts,
|
||||
ctx,
|
||||
@@ -802,13 +831,6 @@ def build_catalog() -> list[PlotSpec]:
|
||||
finalize=_router_share_process_finalize,
|
||||
chunkable=False,
|
||||
),
|
||||
PlotSpec(
|
||||
"type_embedding_l1_distance",
|
||||
"model",
|
||||
compute_partial=_type_embedding_l1_distance_partial,
|
||||
finalize=_type_embedding_l1_distance_finalize,
|
||||
chunkable=False,
|
||||
),
|
||||
]
|
||||
return specs
|
||||
|
||||
|
||||
+21
-20
@@ -83,12 +83,6 @@ _PLOT_META_KEYS = (
|
||||
"best_val_loss",
|
||||
"training_config",
|
||||
"training_meta",
|
||||
# Diagnostic — only present when giant rollout ran under
|
||||
# stage2_model.particle_type.target="embedding" (see giant/cli.py's
|
||||
# rollout command and giant.rollout.L1DistCollector); absent otherwise,
|
||||
# which the type_embedding_l1_distance PlotSpec (catalog.py) reads as
|
||||
# "not applicable to this checkpoint".
|
||||
"type_embedding_l1_dist",
|
||||
)
|
||||
|
||||
|
||||
@@ -161,7 +155,9 @@ 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]:
|
||||
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) —
|
||||
@@ -251,7 +247,6 @@ def compute_reduced(
|
||||
checkpoint: str | None = None,
|
||||
chunk_index: int = 0,
|
||||
n_chunks: int = 1,
|
||||
type_embedding_l1_dist: dict | None = None,
|
||||
) -> Path:
|
||||
"""Core: run one (plot, chunk)'s partial reduction against explicit paths.
|
||||
|
||||
@@ -266,15 +261,11 @@ def compute_reduced(
|
||||
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 n_chunks={effective_n} (chunkable={spec.chunkable})"
|
||||
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),
|
||||
type_embedding_l1_dist=type_embedding_l1_dist,
|
||||
rollout, reference, ctx, checkpoint=checkpoint, chunk=(chunk_index, effective_n)
|
||||
)
|
||||
partial = Partial(
|
||||
id=spec_id,
|
||||
@@ -300,7 +291,6 @@ def compute_one(spec_id: str, run_dir: str | Path, chunk_index: int = 0) -> Path
|
||||
checkpoint=meta.plot_meta.get("checkpoint"),
|
||||
chunk_index=chunk_index,
|
||||
n_chunks=meta.n_chunks,
|
||||
type_embedding_l1_dist=meta.plot_meta.get("type_embedding_l1_dist"),
|
||||
)
|
||||
|
||||
|
||||
@@ -324,7 +314,10 @@ def merge_one(spec_id: str, run_dir: str | Path) -> Path:
|
||||
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"))}
|
||||
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(
|
||||
@@ -369,7 +362,11 @@ exec {giant_exe} analyze compute-one --id "$1" --chunk "$2" --run-dir {run_dir}
|
||||
|
||||
|
||||
def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> str:
|
||||
reqs_attrs = "+RemoteJob = True\n" if cfg.remote else "requirements = TARGET.ProvidesETPResources\n"
|
||||
reqs_attrs = (
|
||||
"+RemoteJob = True\n"
|
||||
if cfg.remote
|
||||
else "requirements = TARGET.ProvidesETPResources\n"
|
||||
)
|
||||
return (
|
||||
"universe = docker\n"
|
||||
f"docker_image = {cfg.docker_image}\n"
|
||||
@@ -389,7 +386,9 @@ def _submit_description(cfg: SubmitConfig, wrapper: Path, jobs_file: Path) -> st
|
||||
)
|
||||
|
||||
|
||||
def _job_walltimes(run_dir: Path, ids: list[str], n_chunks: int) -> list[tuple[str, int, int]]:
|
||||
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``;
|
||||
@@ -465,7 +464,9 @@ def write_submit(cfg: SubmitConfig, ids: list[str] | None = None) -> Path:
|
||||
(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, giant_exe=giant_exe, run_dir=run_dir))
|
||||
wrapper.write_text(
|
||||
_WRAPPER.format(repo_dir=cfg.repo_dir, giant_exe=giant_exe, run_dir=run_dir)
|
||||
)
|
||||
wrapper.chmod(0o755)
|
||||
|
||||
jobs = _job_walltimes(run_dir, ids, cfg.n_chunks)
|
||||
|
||||
@@ -74,7 +74,9 @@ def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFram
|
||||
return lf.filter((pl.col("pre_E").hash(seed=seed) % 2**32) < threshold)
|
||||
|
||||
|
||||
def _combined_quantiles(r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float) -> tuple[float, float]:
|
||||
def _combined_quantiles(
|
||||
r_vals: np.ndarray, t_vals: np.ndarray, lo_q: float, hi_q: float
|
||||
) -> tuple[float, float]:
|
||||
"""Robust (lo_q, hi_q) range over the union of two value samples."""
|
||||
both = np.concatenate([r_vals, t_vals])
|
||||
lo, hi = float(np.quantile(both, lo_q)), float(np.quantile(both, hi_q))
|
||||
@@ -102,15 +104,31 @@ def build_context(
|
||||
|
||||
# Ranged marginal variables: robust ranges over a shared row subsample.
|
||||
exprs = [e.alias(n) for n, (_, e) in RANGED_VARS.items()]
|
||||
r_s = _row_subsample(r_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
|
||||
t_s = _row_subsample(t_lf, sample_rows, seed).select(exprs).collect(engine="streaming")
|
||||
r_s = (
|
||||
_row_subsample(r_lf, sample_rows, seed)
|
||||
.select(exprs)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
t_s = (
|
||||
_row_subsample(t_lf, sample_rows, seed)
|
||||
.select(exprs)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
var_ranges = {
|
||||
name: _combined_quantiles(r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q) for name in RANGED_VARS
|
||||
name: _combined_quantiles(
|
||||
r_s[name].to_numpy(), t_s[name].to_numpy(), _LO_Q, _HI_Q
|
||||
)
|
||||
for name in RANGED_VARS
|
||||
}
|
||||
|
||||
# Energy-bin edges from exact per-event incident energies (cheap group_by).
|
||||
def _incident(lf: pl.LazyFrame) -> np.ndarray:
|
||||
return lf.group_by("event_id").agg(pl.col("pre_E").max()).collect(engine="streaming")["pre_E"].to_numpy()
|
||||
return (
|
||||
lf.group_by("event_id")
|
||||
.agg(pl.col("pre_E").max())
|
||||
.collect(engine="streaming")["pre_E"]
|
||||
.to_numpy()
|
||||
)
|
||||
|
||||
r_inc, t_inc = _incident(r_lf), _incident(t_lf)
|
||||
energy_edges = energy_bin_edges(np.concatenate([r_inc, t_inc]), n_energy_bins)
|
||||
@@ -127,7 +145,8 @@ def build_context(
|
||||
)
|
||||
top_pdgs = [int(x) for x in pdg_counts["pdg"].to_list()[:top_k_pdg]]
|
||||
materials = sorted(
|
||||
set(_counts(r_lf, "material")["material"].to_list()) | set(_counts(t_lf, "material")["material"].to_list())
|
||||
set(_counts(r_lf, "material")["material"].to_list())
|
||||
| set(_counts(t_lf, "material")["material"].to_list())
|
||||
)
|
||||
|
||||
# Shower depth / transverse ranges from a subsampled proxy.
|
||||
|
||||
@@ -68,7 +68,9 @@ def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
|
||||
|
||||
def energy_bin_labels(edges: np.ndarray) -> list[str]:
|
||||
"""``E in [lo, hi)`` labels for each bin defined by ``edges`` (MeV)."""
|
||||
return [f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)]
|
||||
return [
|
||||
f"E in [{edges[i]:.3g}, {edges[i + 1]:.3g}) MeV" for i in range(len(edges) - 1)
|
||||
]
|
||||
|
||||
|
||||
def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
|
||||
@@ -86,7 +88,9 @@ def digitize_expr(value: pl.Expr, edges: np.ndarray) -> pl.Expr:
|
||||
return idx.clip(0, n_bins - 1)
|
||||
|
||||
|
||||
def event_energy_bins(lf: pl.LazyFrame, edges: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
|
||||
def event_energy_bins(
|
||||
lf: pl.LazyFrame, edges: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Per-event incident-energy bin: ``(event_ids, bin_idx)`` numpy arrays.
|
||||
|
||||
Incident energy is ``max(pre_E)`` per event (the primary). One bounded
|
||||
|
||||
@@ -29,8 +29,17 @@ from giant.constants import TERM_ESCAPED
|
||||
|
||||
|
||||
def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr:
|
||||
"""Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins."""
|
||||
return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1)
|
||||
"""Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins.
|
||||
|
||||
Out-of-range values clamp into the edge bins, and the clamp deliberately
|
||||
happens in f64 *before* the integer cast: a rollout is free to emit a wildly
|
||||
out-of-range outlier (a step_length of 1e10 mm, say) or an inf, whose
|
||||
unclamped bin index overflows i32 and makes the cast fail outright. NaN has
|
||||
no edge to clamp to, so it becomes null and is dropped by the callers below
|
||||
— the same thing ``np.histogram`` does with it.
|
||||
"""
|
||||
idx = ((value - lo) / (hi - lo) * nbins).floor().clip(0, nbins - 1)
|
||||
return pl.when(idx.is_nan()).then(None).otherwise(idx).cast(pl.Int32)
|
||||
|
||||
|
||||
def hist1d(
|
||||
@@ -50,6 +59,7 @@ def hist1d(
|
||||
group = pl.lit(0, dtype=pl.Int64) if group is None else group
|
||||
res = (
|
||||
lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b"))
|
||||
.drop_nulls("_b")
|
||||
.group_by("_g", "_b")
|
||||
.agg(pl.len().alias("_n"))
|
||||
.collect(engine="streaming")
|
||||
@@ -142,7 +152,9 @@ def attach_entry_axis(lf: pl.LazyFrame, entry: pl.DataFrame) -> pl.LazyFrame:
|
||||
"""
|
||||
ids = entry["event_id"].to_numpy()
|
||||
return lf.with_columns(
|
||||
pl.col("event_id").replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64).alias(col)
|
||||
pl.col("event_id")
|
||||
.replace_strict(ids, entry[col].to_numpy(), return_dtype=pl.Float64)
|
||||
.alias(col)
|
||||
for col in _ENTRY_AXIS_COLS
|
||||
)
|
||||
|
||||
@@ -188,6 +200,7 @@ def profile_partial(
|
||||
_bin_expr(coord, lo, hi, nbins).alias("_b"),
|
||||
weight.alias("_w"),
|
||||
)
|
||||
.drop_nulls("_b")
|
||||
.group_by("event_id", "_b")
|
||||
.agg(pl.col("_w").sum().alias("_ws"))
|
||||
.collect(engine="streaming")
|
||||
@@ -252,7 +265,10 @@ def leakage_fraction(lf: pl.LazyFrame) -> np.ndarray:
|
||||
lf.group_by("event_id")
|
||||
.agg(
|
||||
pl.col("edep").sum().alias("deposited"),
|
||||
pl.col("pre_E").filter(pl.col("termination_reason") == TERM_ESCAPED).sum().alias("escaped"),
|
||||
pl.col("pre_E")
|
||||
.filter(pl.col("termination_reason") == TERM_ESCAPED)
|
||||
.sum()
|
||||
.alias("escaped"),
|
||||
)
|
||||
.collect(engine="streaming")
|
||||
)
|
||||
|
||||
+29
-57
@@ -41,47 +41,11 @@ def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> Non
|
||||
ax.set_yscale("log")
|
||||
|
||||
|
||||
def _router_summary(router_cfg: dict) -> str:
|
||||
if not router_cfg.get("enabled"):
|
||||
def _router_summary(model_config: dict) -> str:
|
||||
r = model_config.get("router") or {}
|
||||
if not r.get("enabled"):
|
||||
return "off"
|
||||
return f"{router_cfg.get('type', '?')}×{router_cfg.get('n_experts', '?')}"
|
||||
|
||||
|
||||
def _figure_params_v2(mc: dict, run_meta: dict) -> dict:
|
||||
"""`_figure_params` for a new-shape (nested) `model_config` — has a
|
||||
`stage1_model` key. Reports stage 1's architecture (the headline
|
||||
generator); stage 2's generator is only added (`mode_s2`) when it
|
||||
differs from stage 1's, since a mixed run (the `stage1=flow` +
|
||||
`stage2=wgan` case) is the interesting exception, not
|
||||
the common case."""
|
||||
s1 = mc["stage1_model"]
|
||||
s2 = mc.get("stage2_model") or {}
|
||||
mode = s1.get("generator")
|
||||
params: dict = {}
|
||||
if s1.get("hidden_dim") is not None:
|
||||
params["hidden_dim"] = s1["hidden_dim"]
|
||||
if s1.get("n_res_blocks") is not None:
|
||||
params["n_res_blocks"] = s1["n_res_blocks"]
|
||||
if mode is not None:
|
||||
params["mode"] = mode
|
||||
s2_mode = s2.get("generator")
|
||||
if s2_mode is not None and s2_mode != mode:
|
||||
params["mode_s2"] = s2_mode
|
||||
particle_type = ((mc.get("conditioning") or {}).get("particle") or {}).get("type")
|
||||
if particle_type is not None:
|
||||
params["conditioning"] = particle_type
|
||||
params["router"] = _router_summary(s1.get("router") or {})
|
||||
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":
|
||||
noise_dim = (s1.get("wgan") or {}).get("noise_dim")
|
||||
if noise_dim is not None:
|
||||
params["noise_dim"] = noise_dim
|
||||
elif run_meta.get("steps") is not None:
|
||||
params["steps"] = run_meta["steps"]
|
||||
return params
|
||||
return f"{r.get('type', '?')}×{r.get('n_experts', '?')}"
|
||||
|
||||
|
||||
def _figure_params(run_meta: dict) -> dict:
|
||||
@@ -95,14 +59,8 @@ def _figure_params(run_meta: dict) -> dict:
|
||||
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.
|
||||
|
||||
Handles both a v0.2 checkpoint's flat ``model_config`` and a v0.3.0
|
||||
nested one (has a ``stage1_model`` key — see ``_figure_params_v2``).
|
||||
"""
|
||||
mc = run_meta.get("model_config") or {}
|
||||
if "stage1_model" in mc:
|
||||
return _figure_params_v2(mc, run_meta)
|
||||
|
||||
mode = mc.get("mode")
|
||||
params: dict = {}
|
||||
if mc.get("hidden_dim") is not None:
|
||||
@@ -113,7 +71,7 @@ def _figure_params(run_meta: dict) -> dict:
|
||||
params["mode"] = mode
|
||||
if mc.get("conditioning") is not None:
|
||||
params["conditioning"] = mc["conditioning"]
|
||||
params["router"] = _router_summary(mc.get("router") or {})
|
||||
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:
|
||||
@@ -139,11 +97,11 @@ def _render_overlay(r: Reduced, params: dict):
|
||||
def _render_single(r: Reduced, params: dict):
|
||||
edges = np.asarray(r.payload["edges"])
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
ax.stairs(_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"])
|
||||
ax.stairs(
|
||||
_density(r.payload["rollout"], edges), edges, label=_SERIES_LABELS["rollout"]
|
||||
)
|
||||
if r.payload.get("log_y"):
|
||||
ax.set_yscale("log")
|
||||
if r.payload.get("log_x"):
|
||||
ax.set_xscale("log")
|
||||
ax.set_xlabel(r.xlabel)
|
||||
ax.set_ylabel("density")
|
||||
ps.style_legend(ax, title="source")
|
||||
@@ -185,7 +143,9 @@ def _render_profile(r: Reduced, params: dict):
|
||||
mean = np.asarray(r.payload[f"{key}_mean"])
|
||||
std = np.asarray(r.payload[f"{key}_std"])
|
||||
(line,) = ax.plot(centers, mean, label=_SERIES_LABELS[key])
|
||||
ax.fill_between(centers, mean - std, mean + std, alpha=0.2, color=line.get_color())
|
||||
ax.fill_between(
|
||||
centers, mean - std, mean + std, alpha=0.2, color=line.get_color()
|
||||
)
|
||||
ax.set_xlabel(r.xlabel)
|
||||
ax.set_ylabel(r.payload.get("ylabel", "mean deposited energy [MeV]"))
|
||||
ps.style_legend(ax, title="source")
|
||||
@@ -197,7 +157,9 @@ def _render_bar(r: Reduced, params: dict):
|
||||
x = np.arange(len(labels))
|
||||
width = 0.4
|
||||
fig, ax = ps.new_figure("thesis-single", title=r.title, params=params)
|
||||
ax.bar(x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"])
|
||||
ax.bar(
|
||||
x - width / 2, r.payload["reference"], width, label=_SERIES_LABELS["reference"]
|
||||
)
|
||||
ax.bar(x + width / 2, r.payload["rollout"], width, label=_SERIES_LABELS["rollout"])
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels(labels, rotation=45, ha="right")
|
||||
@@ -209,7 +171,9 @@ def _render_bar(r: Reduced, params: dict):
|
||||
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)
|
||||
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, {})
|
||||
@@ -218,7 +182,9 @@ def _render_router_gating(r: Reduced, params: dict):
|
||||
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}")
|
||||
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")
|
||||
@@ -337,7 +303,9 @@ def render_all(
|
||||
families.add(r.family)
|
||||
fig = render(r, run_meta)
|
||||
ps.savefig(fig, str(family_dir / r.id), formats=("pdf",))
|
||||
(family_dir / f"{r.id}.yaml").write_text(yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False))
|
||||
(family_dir / f"{r.id}.yaml").write_text(
|
||||
yaml.safe_dump(_plot_metadata(r, run_meta), sort_keys=False)
|
||||
)
|
||||
pdfs.append(family_dir / f"{r.id}.pdf")
|
||||
import matplotlib.pyplot as plt
|
||||
|
||||
@@ -358,7 +326,9 @@ def render_all(
|
||||
)
|
||||
for fam in families:
|
||||
(out_dir / fam / "metadata.yaml").write_text(
|
||||
yaml.safe_dump({"title": fam, "description": f"{fam} plots."}, sort_keys=False)
|
||||
yaml.safe_dump(
|
||||
{"title": fam, "description": f"{fam} plots."}, sort_keys=False
|
||||
)
|
||||
)
|
||||
|
||||
if run_gallery:
|
||||
@@ -386,4 +356,6 @@ def render_run(run_dir: str | Path, *, run_gallery: bool = False) -> list[Path]:
|
||||
"reference": meta.reference,
|
||||
**meta.plot_meta,
|
||||
}
|
||||
return render_all(run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery)
|
||||
return render_all(
|
||||
run_dir / "reduced", run_dir / "plots", run_meta, run_gallery=run_gallery
|
||||
)
|
||||
|
||||
@@ -61,8 +61,7 @@ class _RouterHandle:
|
||||
pdg_map: dict[int, int]
|
||||
mat_map: dict[str, int]
|
||||
cond_normalizer: "Normalizer"
|
||||
particle_conditioning: str
|
||||
material_conditioning: str
|
||||
conditioning: str
|
||||
router_type: str
|
||||
|
||||
|
||||
@@ -70,43 +69,32 @@ 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.checkpoint_io import conditioning_axes
|
||||
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 {}
|
||||
# New nested shape (has a "stage1_model" key) vs. a v0.2 checkpoint's
|
||||
# flat model_config.
|
||||
router_cfg = (
|
||||
(model_cfg.get("stage1_model") or {}).get("router") if "stage1_model" in model_cfg else model_cfg.get("router")
|
||||
)
|
||||
router_cfg = model_cfg.get("router")
|
||||
if not router_cfg or not router_cfg.get("enabled"):
|
||||
return None
|
||||
|
||||
built = build_models(model_cfg)
|
||||
stage1 = built["stage1"]
|
||||
if stage1 is None:
|
||||
return None
|
||||
stage1, _ = build_models(model_cfg)
|
||||
stage1.load_state_dict(ckpt["model"])
|
||||
stage1.eval()
|
||||
router = stage1.trunk.router
|
||||
if router is None:
|
||||
return None
|
||||
|
||||
particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
|
||||
return _RouterHandle(
|
||||
router=router,
|
||||
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"]),
|
||||
particle_conditioning=particle_conditioning,
|
||||
material_conditioning=material_conditioning,
|
||||
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:
|
||||
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)
|
||||
@@ -114,7 +102,9 @@ def _subsample(lf: pl.LazyFrame, n: int, seed: int, extra_cols: tuple = ()) -> p
|
||||
return lf.select(*_COLS, *extra_cols).collect(engine="streaming")
|
||||
|
||||
|
||||
def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame, np.ndarray]:
|
||||
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
|
||||
@@ -138,9 +128,13 @@ def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame,
|
||||
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_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()]),
|
||||
"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(),
|
||||
@@ -150,11 +144,12 @@ def _gate_for_df(handle: _RouterHandle, df: pl.DataFrame) -> tuple[pl.DataFrame,
|
||||
handle.pdg_map,
|
||||
handle.mat_map,
|
||||
cond_normalizer=handle.cond_normalizer,
|
||||
particle_conditioning=handle.particle_conditioning,
|
||||
material_conditioning=handle.material_conditioning,
|
||||
conditioning=handle.conditioning,
|
||||
)
|
||||
with torch.no_grad():
|
||||
gate = handle.router.gate(torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()).numpy()
|
||||
gate = handle.router.gate(
|
||||
torch.from_numpy(cond_cont).float(), torch.from_numpy(cond_cat).long()
|
||||
).numpy()
|
||||
return df, gate
|
||||
|
||||
|
||||
@@ -177,7 +172,9 @@ def _quantile_bins(x: np.ndarray, gate: np.ndarray, n_bins: int) -> dict:
|
||||
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]]:
|
||||
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
|
||||
@@ -197,7 +194,10 @@ def _top1_shares(categories: np.ndarray, idx: np.ndarray, order: list, n_experts
|
||||
return shares
|
||||
|
||||
|
||||
_NOTE_NOT_MOE = "checkpoint has no enabled MoE router (model.router.enabled is false/absent) — nothing to show"
|
||||
_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)",
|
||||
@@ -233,7 +233,9 @@ def compute_router_gating(
|
||||
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": []}
|
||||
sides[name] = (
|
||||
_quantile_bins(x, gate, _N_BINS) if len(x) else {"centers": [], "means": []}
|
||||
)
|
||||
|
||||
return Reduced(
|
||||
id="router_gating",
|
||||
@@ -269,7 +271,9 @@ def compute_router_share_by_pdg(
|
||||
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)
|
||||
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)}
|
||||
@@ -314,7 +318,9 @@ def compute_router_share_by_process(
|
||||
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)
|
||||
shares = _top1_shares(
|
||||
df["process"].to_numpy(), idx, order, handle.router.n_experts
|
||||
)
|
||||
else:
|
||||
order, shares = [], {}
|
||||
|
||||
|
||||
@@ -27,7 +27,7 @@ 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 `giant/tools/profile_analysis_costs.py`
|
||||
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.
|
||||
@@ -54,7 +54,9 @@ _FIXED_OVERHEAD_S = 60.0
|
||||
# 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"})
|
||||
_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
|
||||
|
||||
@@ -93,7 +93,10 @@ def _check_rollout_metadata(path: Path) -> None:
|
||||
metadata = pq.read_schema(path).metadata or {}
|
||||
coord = metadata.get(PREDICT_COORD_METADATA_KEY.encode())
|
||||
if coord is not None and coord.decode() != ROLLOUT_COORD_VALUE:
|
||||
raise ValueError(f"{path} is not a rollout file (coord={coord.decode()!r}); expected `giant rollout` output")
|
||||
raise ValueError(
|
||||
f"{path} is not a rollout file (coord={coord.decode()!r}); "
|
||||
"expected `giant rollout` output"
|
||||
)
|
||||
|
||||
|
||||
def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
@@ -118,7 +121,11 @@ def open_side(source: str | Path | pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
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)
|
||||
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))
|
||||
|
||||
|
||||
@@ -130,7 +137,9 @@ def physical_steps(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
"""
|
||||
if side is Side.reference:
|
||||
return lf
|
||||
return lf.filter(~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS)))
|
||||
return lf.filter(
|
||||
~pl.col("termination_reason").is_in(list(SYNTHETIC_TERMINATION_REASONS))
|
||||
)
|
||||
|
||||
|
||||
def secondaries(lf: pl.LazyFrame, side: Side) -> pl.LazyFrame:
|
||||
|
||||
@@ -1,70 +0,0 @@
|
||||
"""Secondary-type embedding-distance diagnostic.
|
||||
|
||||
Unlike every other diagnostic in this package, the data isn't derivable from
|
||||
a rollout/reference parquet at all — it's the L1 distance between each
|
||||
emitted secondary's *raw* predicted embedding vector (under
|
||||
`stage2_model.particle_type.target = "embedding"`) and the nearest row of the
|
||||
conditioning's embedding table it snapped to, which only exists transiently
|
||||
inside `giant rollout` (`giant.rollout.decode_secondary_identity`), never
|
||||
written to a column. So it's accumulated once, at rollout time
|
||||
(`giant.rollout.L1DistCollector`), and stashed as a pre-finished histogram
|
||||
summary in the rollout YAML sidecar (`type_embedding_l1_dist`) — this module
|
||||
just turns that summary into a `Reduced`, no parquet scan involved (a
|
||||
`chunkable=False` spec, like `router_gating`, but even cheaper: no live model
|
||||
call either).
|
||||
|
||||
A heavy right tail means the decoder is emitting vectors off the embedding
|
||||
manifold — the direct analogue of the species-collapse symptom the v0.3.0
|
||||
redesign exists to fix.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from giant.analysis.reduced import Reduced
|
||||
|
||||
_NOTE_NOT_APPLICABLE = (
|
||||
"not applicable: this rollout's checkpoint doesn't use "
|
||||
"stage2_model.particle_type.target='embedding' (or generated no "
|
||||
"secondaries), so giant rollout recorded no type_embedding_l1_dist "
|
||||
"diagnostic in its YAML sidecar"
|
||||
)
|
||||
|
||||
|
||||
def compute_type_embedding_l1_distance(l1_dist: dict | None) -> Reduced:
|
||||
"""`Reduced` for the type-embedding-distance figure, or an explanatory
|
||||
note if this checkpoint never populated the diagnostic.
|
||||
|
||||
`l1_dist`: `giant.rollout.L1DistCollector.summary()`'s dict, as recorded
|
||||
in the rollout YAML's `type_embedding_l1_dist` key (`Bundle.
|
||||
type_embedding_l1_dist`) — `{"n", "mean", "std", "min", "max",
|
||||
"hist_edges", "hist_counts"}`.
|
||||
"""
|
||||
if l1_dist is None:
|
||||
return Reduced(
|
||||
id="type_embedding_l1_distance",
|
||||
family="model",
|
||||
kind="unavailable",
|
||||
title="Secondary-type embedding L1 distance",
|
||||
xlabel="n/a",
|
||||
payload={"note": _NOTE_NOT_APPLICABLE},
|
||||
)
|
||||
|
||||
return Reduced(
|
||||
id="type_embedding_l1_distance",
|
||||
family="model",
|
||||
kind="single_hist",
|
||||
title="Secondary-type embedding L1 distance (predicted vector -> nearest PDG row)",
|
||||
xlabel="L1 distance",
|
||||
payload={
|
||||
"edges": l1_dist["hist_edges"],
|
||||
"rollout": l1_dist["hist_counts"],
|
||||
"log_y": True,
|
||||
"log_x": True,
|
||||
"note": (
|
||||
f"n={l1_dist['n']:,} mean={l1_dist['mean']:.4g} "
|
||||
f"std={l1_dist['std']:.4g} min={l1_dist['min']:.4g} "
|
||||
f"max={l1_dist['max']:.4g}; rollout only, no reference "
|
||||
"concept for a raw pre-decode vector"
|
||||
),
|
||||
},
|
||||
)
|
||||
@@ -1,234 +0,0 @@
|
||||
"""Load a trained checkpoint into ready-to-run models (giant.cli's `predict`/`rollout`).
|
||||
|
||||
Both commands need the same ~15 steps to go from a checkpoint path to two
|
||||
`eval()`-mode models plus their normalizers/vocab maps: load the pickle,
|
||||
validate it carries what current code expects, resolve which conditioning
|
||||
mode each axis was trained with, restore the top-N vocab maps (if the
|
||||
checkpoint used one-hot conditioning), rebuild the normalizers, construct the
|
||||
model from `model_config`, and load the requested (raw or EMA) weights. This
|
||||
used to be duplicated near-verbatim in both commands (issues.md Issue 5) —
|
||||
`load_for_inference` is the single implementation.
|
||||
|
||||
This module intentionally has no Typer dependency, so it can be unit-tested
|
||||
directly and imported from non-CLI code (`giant.analysis.router_gating`,
|
||||
lazily — see that module's docstring for why). Failures raise
|
||||
`CheckpointCompatibilityError` with the same wording the CLI has always
|
||||
shown; the CLI layer catches it and does the `typer.echo`/`Exit(1)`.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from giant import config as gconfig
|
||||
from giant.constants import K_MAX
|
||||
from giant.data.loader import TopNMap
|
||||
from giant.data.setup_cache import topnmap_from_json
|
||||
from giant.data.transforms import Normalizer
|
||||
from giant.model.network import build_models
|
||||
|
||||
|
||||
class CheckpointCompatibilityError(Exception):
|
||||
"""Checkpoint is missing something `load_for_inference` needs."""
|
||||
|
||||
|
||||
def conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]:
|
||||
"""(particle_conditioning, material_conditioning) for
|
||||
`giant.data.transforms.build_cond_features`/`build_features` — from
|
||||
either a v0.2 checkpoint's flat `model_config["conditioning"]` (one
|
||||
shared string, same for both axes) or a new-format one (independent
|
||||
`model_config["conditioning"]["particle"/"material"]["type"]` — the two
|
||||
axes are configured independently and may differ)."""
|
||||
raw = model_cfg.get("conditioning", default)
|
||||
if isinstance(raw, dict):
|
||||
return (
|
||||
raw.get("particle", {}).get("type", default),
|
||||
raw.get("material", {}).get("type", default),
|
||||
)
|
||||
return raw, raw
|
||||
|
||||
|
||||
def stage_cfg(model_cfg: dict, stage: str) -> dict:
|
||||
"""`model_cfg[f"{stage}_model"]` for a new-format model_config, `{}` for
|
||||
a v0.2 flat one (whose ddpm schedule always used `CosineSchedule`'s own
|
||||
default `T=1000` — never a config key — and which never had
|
||||
`particle_type` at all, so `{}` is the correct fallback for both
|
||||
`ddpm_steps`/`particle_type_other_policy` below)."""
|
||||
val = model_cfg.get(f"{stage}_model")
|
||||
return val if isinstance(val, dict) else {}
|
||||
|
||||
|
||||
def ddpm_steps(model_cfg: dict, stage: str) -> int:
|
||||
return stage_cfg(model_cfg, stage).get("ddpm", {}).get("n_steps", 1000)
|
||||
|
||||
|
||||
def particle_type_other_policy(model_cfg: dict) -> str:
|
||||
return stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("other_policy", "sample")
|
||||
|
||||
|
||||
def load_pdg_topn_map(ckpt: dict) -> TopNMap | None:
|
||||
"""`ckpt["pdg_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
|
||||
this checkpoint's conditioning/particle_type never needed one (see
|
||||
`giant.pipeline.run_setup_stage`, which only populates it when
|
||||
`conditioning.particle.type` or `stage2_model.particle_type.target` is
|
||||
`"onehot"`)."""
|
||||
raw = ckpt.get("pdg_topn_map")
|
||||
return topnmap_from_json(raw, axis="pdg") if raw is not None else None
|
||||
|
||||
|
||||
def load_mat_topn_map(ckpt: dict) -> TopNMap | None:
|
||||
"""`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if
|
||||
this checkpoint's `conditioning.material.type` was never `"onehot"` (see
|
||||
`giant.pipeline.run_setup_stage`)."""
|
||||
raw = ckpt.get("mat_topn_map")
|
||||
return topnmap_from_json(raw, axis="material") if raw is not None else None
|
||||
|
||||
|
||||
def load_sec_type_topn_map(ckpt: dict) -> TopNMap | None:
|
||||
"""`ckpt["sec_type_topn_map"]` as a `giant.data.loader.TopNMap`, or
|
||||
`None` if this checkpoint's `stage2_model.particle_type.target` was never
|
||||
`"onehot"` (see `giant.pipeline.run_setup_stage`).
|
||||
|
||||
Pre-gitea-#29 checkpoints have no `sec_type_topn_map` key at all — before
|
||||
#29, the secondary-species decode map and the conditioning PDG onehot map
|
||||
were always numerically the same map, saved once under `pdg_topn_map`.
|
||||
For those, fall back to `load_pdg_topn_map` to reproduce that exact
|
||||
behavior; a current checkpoint always has the key (possibly `null`, if
|
||||
`particle_type.target != "onehot"`), so this fallback never fires for one."""
|
||||
if "sec_type_topn_map" in ckpt:
|
||||
raw = ckpt["sec_type_topn_map"]
|
||||
return topnmap_from_json(raw, axis="pdg") if raw is not None else None
|
||||
return load_pdg_topn_map(ckpt)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class InferenceContext:
|
||||
"""Everything needed to run a trained checkpoint forward, resolved once."""
|
||||
|
||||
stage1: nn.Module | None
|
||||
stage2: nn.Module | None
|
||||
cond_norm: Normalizer
|
||||
tgt_norm: Normalizer
|
||||
sec_phys_norm: Normalizer
|
||||
pdg_map: dict[int, int]
|
||||
mat_map: dict[str, int]
|
||||
pdg_topn_map: TopNMap | None
|
||||
mat_topn_map: TopNMap | None
|
||||
sec_type_topn_map: TopNMap | None
|
||||
particle_conditioning: str
|
||||
material_conditioning: str
|
||||
k_max: int
|
||||
stage1_ddpm_steps: int
|
||||
stage2_ddpm_steps: int
|
||||
other_policy: str
|
||||
model_config: dict
|
||||
epoch: int | None
|
||||
best_val_loss: float | None
|
||||
|
||||
|
||||
def load_for_inference(
|
||||
checkpoint: Path,
|
||||
device: torch.device,
|
||||
command_name: str,
|
||||
weights: str = "raw",
|
||||
require_stage2: bool = True,
|
||||
) -> InferenceContext:
|
||||
"""Load *checkpoint* and reconstruct everything `predict`/`rollout` need
|
||||
to run it forward, on *device*, in `eval()` mode.
|
||||
|
||||
*command_name* (e.g. `"predict"`/`"rollout"`) only feeds the "needs both"
|
||||
error message below. *weights* is `"raw"` (the live training weights) or
|
||||
`"ema"` (the EMA shadow copy, see `--ema-decay`). *require_stage2*
|
||||
controls whether a checkpoint with an inactive stage 2
|
||||
(`stage2_model.active = false`) is an error (both current callers need
|
||||
both stages) or an acceptable `stage2 = None` result — kept as a real
|
||||
parameter since `stage{1,2}_model.active` is a real, if currently
|
||||
stage1+stage2-only-in-practice, config option.
|
||||
"""
|
||||
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
||||
for key in ("model_config", "sec_decoder"):
|
||||
if key not in ckpt:
|
||||
raise CheckpointCompatibilityError(f"checkpoint has no {key} — retrain with the current code")
|
||||
|
||||
if "sec_phys" not in ckpt.get("normalizer", {}):
|
||||
raise CheckpointCompatibilityError("checkpoint has no normalizer.sec_phys — retrain with the current code")
|
||||
|
||||
gconfig.warn_if_checkpoint_config_mismatch(checkpoint)
|
||||
|
||||
model_cfg = ckpt["model_config"]
|
||||
particle_conditioning, material_conditioning = conditioning_axes(model_cfg)
|
||||
pdg_topn_map = load_pdg_topn_map(ckpt)
|
||||
mat_topn_map = load_mat_topn_map(ckpt)
|
||||
if particle_conditioning == "onehot" and pdg_topn_map is None:
|
||||
raise CheckpointCompatibilityError(
|
||||
"checkpoint's conditioning.particle.type='onehot' but has no pdg_topn_map — retrain with the current code"
|
||||
)
|
||||
if material_conditioning == "onehot" and mat_topn_map is None:
|
||||
raise CheckpointCompatibilityError(
|
||||
"checkpoint's conditioning.material.type='onehot' but has no mat_topn_map — retrain with the current code"
|
||||
)
|
||||
sec_type_topn_map = load_sec_type_topn_map(ckpt)
|
||||
particle_type_target = stage_cfg(model_cfg, "stage2").get("particle_type", {}).get("target", "onehot")
|
||||
if particle_type_target == "onehot" and sec_type_topn_map is None:
|
||||
raise CheckpointCompatibilityError(
|
||||
"checkpoint's stage2_model.particle_type.target='onehot' but has no "
|
||||
"sec_type_topn_map — retrain with the current code"
|
||||
)
|
||||
other_policy = particle_type_other_policy(model_cfg)
|
||||
stage1_ddpm_steps = ddpm_steps(model_cfg, "stage1")
|
||||
stage2_ddpm_steps = ddpm_steps(model_cfg, "stage2")
|
||||
k_max = stage_cfg(model_cfg, "stage2").get("k_max", K_MAX)
|
||||
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_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
|
||||
tgt_norm = Normalizer.from_dict(ckpt["normalizer"]["target"])
|
||||
sec_phys_norm = Normalizer.from_dict(ckpt["normalizer"]["sec_phys"])
|
||||
|
||||
built = build_models(model_cfg)
|
||||
stage1, stage2 = built["stage1"], built["stage2"]
|
||||
if require_stage2 and (stage1 is None or stage2 is None):
|
||||
raise CheckpointCompatibilityError(
|
||||
f"checkpoint has an inactive stage1 or stage2 — {command_name} needs both (see stage{{1,2}}_model.active)"
|
||||
)
|
||||
|
||||
if weights == "raw":
|
||||
model_key, sec_key = "model", "sec_decoder"
|
||||
else:
|
||||
model_key, sec_key = "model_ema", "sec_decoder_ema"
|
||||
if model_key not in ckpt or sec_key not in ckpt:
|
||||
raise CheckpointCompatibilityError(
|
||||
f"{checkpoint} has no EMA weights (trained before --ema-decay, "
|
||||
"or with --ema-decay 0) — use --weights raw"
|
||||
)
|
||||
if stage1 is not None:
|
||||
stage1.load_state_dict(ckpt[model_key])
|
||||
stage1.to(device).eval()
|
||||
if stage2 is not None:
|
||||
stage2.load_state_dict(ckpt[sec_key])
|
||||
stage2.to(device).eval()
|
||||
|
||||
return InferenceContext(
|
||||
stage1=stage1,
|
||||
stage2=stage2,
|
||||
cond_norm=cond_norm,
|
||||
tgt_norm=tgt_norm,
|
||||
sec_phys_norm=sec_phys_norm,
|
||||
pdg_map=pdg_map,
|
||||
mat_map=mat_map,
|
||||
pdg_topn_map=pdg_topn_map,
|
||||
mat_topn_map=mat_topn_map,
|
||||
sec_type_topn_map=sec_type_topn_map,
|
||||
particle_conditioning=particle_conditioning,
|
||||
material_conditioning=material_conditioning,
|
||||
k_max=k_max,
|
||||
stage1_ddpm_steps=stage1_ddpm_steps,
|
||||
stage2_ddpm_steps=stage2_ddpm_steps,
|
||||
other_policy=other_policy,
|
||||
model_config=model_cfg,
|
||||
epoch=ckpt.get("epoch"),
|
||||
best_val_loss=ckpt.get("best_val_loss"),
|
||||
)
|
||||
+383
-453
File diff suppressed because it is too large
Load Diff
+254
-1293
File diff suppressed because it is too large
Load Diff
+42
-83
@@ -1,48 +1,15 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import IterableDataset
|
||||
|
||||
from giant.constants import K_MAX
|
||||
from giant.data.loader import event_id_offset, iter_file_chunks
|
||||
from giant.data.transforms import Normalizer, build_features, sorted_membership
|
||||
|
||||
|
||||
class StepBatch(NamedTuple):
|
||||
"""One training batch, as yielded by `StreamingStepsDataset`. Field order
|
||||
is load-bearing for existing positional unpacking elsewhere (`trainers.py`,
|
||||
`validate.py`, test fixtures) — append only, never insert or reorder.
|
||||
|
||||
cond_cont: (B, COND_DIM) float32
|
||||
cond_cat: (B, 2/3/4) int64 — width 2 unless conditioning="onehot"
|
||||
target_s1: (B, 9) float32 — normalised Stage-1 primary target
|
||||
n_sec: (B,) int64 — true secondary count per step
|
||||
sec_cont: (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit,
|
||||
local_dir, log_mass, charge] per slot (mass/charge
|
||||
normalised iff `sec_phys_normalizer` was given); always
|
||||
computed the same way regardless of
|
||||
stage2_model.particle_type.target, only actually used
|
||||
downstream under target="physical"
|
||||
proc_idx: (B,) int64 — process-class label (ProcessRouter supervision
|
||||
only; zeros when `proc_map` is None)
|
||||
sec_type_idx: (B, k_max) int64 — per-slot class index into
|
||||
`sec_type_class_map`, for particle_type.target in
|
||||
("onehot", "embedding"); zeros (unused) otherwise
|
||||
"""
|
||||
|
||||
cond_cont: torch.Tensor
|
||||
cond_cat: torch.Tensor
|
||||
target_s1: torch.Tensor
|
||||
n_sec: torch.Tensor
|
||||
sec_cont: torch.Tensor
|
||||
proc_idx: torch.Tensor
|
||||
sec_type_idx: torch.Tensor
|
||||
|
||||
|
||||
def make_event_split(
|
||||
all_event_ids: np.ndarray,
|
||||
val_fraction: float = 0.1,
|
||||
@@ -71,11 +38,18 @@ class StreamingStepsDataset(IterableDataset):
|
||||
rather than single rows, so the batch is assembled with vectorized
|
||||
numpy slicing instead of a per-row Python loop in the default collate.
|
||||
|
||||
Each batch is a `StepBatch` — see its docstring for field meanings.
|
||||
|
||||
`k_max` (constructor arg, default the module constant) should match
|
||||
`stage2_model.k_max` — it sets the padded
|
||||
width of `sec_cont`/`sec_type_idx` above.
|
||||
Each batch is a tuple:
|
||||
(cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx)
|
||||
where:
|
||||
cond_cont: (B, COND_DIM) float32
|
||||
cond_cat: (B, 2) int64
|
||||
target_s1: (B, 9) float32 — normalised Stage-1 primary target
|
||||
n_sec: (B,) int64 — true secondary count per step
|
||||
sec_cont: (B, K_MAX, SEC_SLOT_DIM) float32 — [stick_logit,
|
||||
local_dir, log_mass, charge] per slot (mass/charge
|
||||
normalised iff `sec_phys_normalizer` was given)
|
||||
proc_idx: (B,) int64 — process-class label (ProcessRouter supervision
|
||||
only; zeros when `proc_map` is None)
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -90,13 +64,8 @@ class StreamingStepsDataset(IterableDataset):
|
||||
shuffle_buffer: int = 65536,
|
||||
shuffle: bool = True,
|
||||
proc_map: dict[str, int] | None = None,
|
||||
particle_conditioning: str = "embedding",
|
||||
material_conditioning: str = "embedding",
|
||||
conditioning: str = "embedding",
|
||||
sec_phys_normalizer: Normalizer | None = None,
|
||||
pdg_topn_map: dict[int, int] | None = None,
|
||||
mat_topn_map: dict[str, int] | None = None,
|
||||
sec_type_class_map: dict | None = None,
|
||||
k_max: int = K_MAX,
|
||||
) -> None:
|
||||
self.files = list(files)
|
||||
self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)}
|
||||
@@ -110,13 +79,8 @@ class StreamingStepsDataset(IterableDataset):
|
||||
self.shuffle_buffer = max(shuffle_buffer, batch_size)
|
||||
self.shuffle = shuffle
|
||||
self.proc_map = proc_map
|
||||
self.particle_conditioning = particle_conditioning
|
||||
self.material_conditioning = material_conditioning
|
||||
self.conditioning = conditioning
|
||||
self.sec_phys_normalizer = sec_phys_normalizer
|
||||
self.pdg_topn_map = pdg_topn_map
|
||||
self.mat_topn_map = mat_topn_map
|
||||
self.sec_type_class_map = sec_type_class_map
|
||||
self.k_max = k_max
|
||||
|
||||
def __iter__(self):
|
||||
worker_info = torch.utils.data.get_worker_info()
|
||||
@@ -134,17 +98,25 @@ class StreamingStepsDataset(IterableDataset):
|
||||
buf_nsec: list[np.ndarray] = []
|
||||
buf_sec: list[np.ndarray] = []
|
||||
buf_proc: list[np.ndarray] = []
|
||||
buf_type: list[np.ndarray] = []
|
||||
buf_n = 0
|
||||
|
||||
for path in files:
|
||||
for chunk in iter_file_chunks(path, offset=self._offsets[path], k_max=self.k_max):
|
||||
for chunk in iter_file_chunks(path, offset=self._offsets[path]):
|
||||
mask = sorted_membership(chunk["event_id"], self._events_arr)
|
||||
if not mask.any():
|
||||
continue
|
||||
chunk = {k: v[mask] for k, v in chunk.items()}
|
||||
|
||||
feats = build_features(
|
||||
(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
target_s1,
|
||||
n_sec,
|
||||
sec_cont,
|
||||
proc_idx,
|
||||
_,
|
||||
_,
|
||||
) = build_features(
|
||||
chunk,
|
||||
self.pdg_map,
|
||||
self.mat_map,
|
||||
@@ -153,21 +125,15 @@ class StreamingStepsDataset(IterableDataset):
|
||||
sec_phys_normalizer=self.sec_phys_normalizer,
|
||||
proc_map=self.proc_map,
|
||||
require_secondaries=True,
|
||||
particle_conditioning=self.particle_conditioning,
|
||||
material_conditioning=self.material_conditioning,
|
||||
pdg_topn_map=self.pdg_topn_map,
|
||||
mat_topn_map=self.mat_topn_map,
|
||||
sec_type_class_map=self.sec_type_class_map,
|
||||
k_max=self.k_max,
|
||||
conditioning=self.conditioning,
|
||||
)
|
||||
buf_cont.append(feats.cond_cont)
|
||||
buf_cat.append(feats.cond_cat)
|
||||
buf_tgt.append(feats.target_s1)
|
||||
buf_nsec.append(feats.n_sec)
|
||||
buf_sec.append(feats.sec_cont)
|
||||
buf_proc.append(feats.proc_idx)
|
||||
buf_type.append(feats.sec_type_idx)
|
||||
buf_n += len(feats.cond_cont)
|
||||
buf_cont.append(cond_cont)
|
||||
buf_cat.append(cond_cat)
|
||||
buf_tgt.append(target_s1)
|
||||
buf_nsec.append(n_sec)
|
||||
buf_sec.append(sec_cont)
|
||||
buf_proc.append(proc_idx)
|
||||
buf_n += len(cond_cont)
|
||||
|
||||
if buf_n >= self.shuffle_buffer:
|
||||
(
|
||||
@@ -177,7 +143,6 @@ class StreamingStepsDataset(IterableDataset):
|
||||
buf_nsec,
|
||||
buf_sec,
|
||||
buf_proc,
|
||||
buf_type,
|
||||
buf_n,
|
||||
) = yield from self._flush(
|
||||
buf_cont,
|
||||
@@ -186,7 +151,6 @@ class StreamingStepsDataset(IterableDataset):
|
||||
buf_nsec,
|
||||
buf_sec,
|
||||
buf_proc,
|
||||
buf_type,
|
||||
final=False,
|
||||
)
|
||||
|
||||
@@ -198,7 +162,6 @@ class StreamingStepsDataset(IterableDataset):
|
||||
buf_nsec,
|
||||
buf_sec,
|
||||
buf_proc,
|
||||
buf_type,
|
||||
final=True,
|
||||
)
|
||||
|
||||
@@ -210,7 +173,6 @@ class StreamingStepsDataset(IterableDataset):
|
||||
buf_nsec: list[np.ndarray],
|
||||
buf_sec: list[np.ndarray],
|
||||
buf_proc: list[np.ndarray],
|
||||
buf_type: list[np.ndarray],
|
||||
final: bool,
|
||||
):
|
||||
cont = np.concatenate(buf_cont)
|
||||
@@ -219,30 +181,28 @@ class StreamingStepsDataset(IterableDataset):
|
||||
nsec = np.concatenate(buf_nsec)
|
||||
sec = np.concatenate(buf_sec)
|
||||
proc = np.concatenate(buf_proc)
|
||||
styp = np.concatenate(buf_type)
|
||||
|
||||
if self.shuffle:
|
||||
idx = np.random.permutation(len(cont))
|
||||
cont, cat, tgt = cont[idx], cat[idx], tgt[idx]
|
||||
nsec, sec, proc, styp = nsec[idx], sec[idx], proc[idx], styp[idx]
|
||||
nsec, sec, proc = nsec[idx], sec[idx], proc[idx]
|
||||
|
||||
bs = self.batch_size
|
||||
n = len(cont)
|
||||
n_full = n // bs if not final else (n + bs - 1) // bs
|
||||
for start in range(0, n_full * bs, bs):
|
||||
end = min(start + bs, n)
|
||||
yield StepBatch(
|
||||
cond_cont=torch.from_numpy(cont[start:end]).float(),
|
||||
cond_cat=torch.from_numpy(cat[start:end]).long(),
|
||||
target_s1=torch.from_numpy(tgt[start:end]).float(),
|
||||
n_sec=torch.from_numpy(nsec[start:end]).long(),
|
||||
sec_cont=torch.from_numpy(sec[start:end]).float(),
|
||||
proc_idx=torch.from_numpy(proc[start:end]).long(),
|
||||
sec_type_idx=torch.from_numpy(styp[start:end]).long(),
|
||||
yield (
|
||||
torch.from_numpy(cont[start:end]).float(),
|
||||
torch.from_numpy(cat[start:end]).long(),
|
||||
torch.from_numpy(tgt[start:end]).float(),
|
||||
torch.from_numpy(nsec[start:end]).long(),
|
||||
torch.from_numpy(sec[start:end]).float(),
|
||||
torch.from_numpy(proc[start:end]).long(),
|
||||
)
|
||||
|
||||
if final:
|
||||
return [], [], [], [], [], [], [], 0
|
||||
return [], [], [], [], [], [], 0
|
||||
rem = n_full * bs
|
||||
return (
|
||||
[cont[rem:]],
|
||||
@@ -251,6 +211,5 @@ class StreamingStepsDataset(IterableDataset):
|
||||
[nsec[rem:]],
|
||||
[sec[rem:]],
|
||||
[proc[rem:]],
|
||||
[styp[rem:]],
|
||||
n - rem,
|
||||
)
|
||||
|
||||
+37
-120
@@ -1,4 +1,3 @@
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Iterator
|
||||
|
||||
@@ -6,8 +5,6 @@ import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
from giant.constants import K_MAX
|
||||
|
||||
# A manifest is a plain text file listing one parquet path per line, used to
|
||||
# name a curated subset of files (e.g. a train/holdout pool) without copying
|
||||
# or symlinking the underlying parquet files. Lines are resolved relative to
|
||||
@@ -16,7 +13,7 @@ from giant.constants import K_MAX
|
||||
MANIFEST_SUFFIX = ".manifest"
|
||||
|
||||
# Each input parquet file is a separate Geant4 job converted 1:1 from its own
|
||||
# ROOT file (giant/tools/steps_to_parquet.py), and a job's event_id numbering
|
||||
# ROOT file (scripts/steps_to_parquet.py), and a job's event_id numbering
|
||||
# always restarts from 0 — so when multiple files are loaded together (a
|
||||
# directory or .manifest), raw event_id values collide across files even
|
||||
# though they refer to unrelated events. Every per-file event_id column gets
|
||||
@@ -115,7 +112,9 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar
|
||||
return out
|
||||
|
||||
|
||||
def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
|
||||
def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]:
|
||||
from giant.constants import K_MAX
|
||||
|
||||
has_sec_lists = "sec_E_list" in df.columns
|
||||
|
||||
d: dict[str, np.ndarray] = {
|
||||
@@ -134,7 +133,9 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[s
|
||||
# / ProcessRouter). Guarded like has_sec_lists: older parquet
|
||||
# conversions predating this column still load fine.
|
||||
"process": (
|
||||
df["process"].to_numpy(dtype=object) if "process" in df.columns else np.full(len(df), "", dtype=object)
|
||||
df["process"].to_numpy(dtype=object)
|
||||
if "process" in df.columns
|
||||
else np.full(len(df), "", dtype=object)
|
||||
),
|
||||
"step_length": df["step_length"].to_numpy(dtype=np.float32),
|
||||
"post_E": df["post_E"].to_numpy(dtype=np.float32),
|
||||
@@ -145,15 +146,17 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[s
|
||||
}
|
||||
|
||||
if has_sec_lists:
|
||||
d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max)
|
||||
d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max)
|
||||
d["sec_dir_list"] = _pad_dir_col(df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max)
|
||||
d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX)
|
||||
d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX)
|
||||
d["sec_dir_list"] = _pad_dir_col(
|
||||
df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX
|
||||
)
|
||||
|
||||
return d
|
||||
|
||||
|
||||
def load_steps(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]:
|
||||
return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max)
|
||||
def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]:
|
||||
return _df_to_dict(pd.read_parquet(path), offset=offset)
|
||||
|
||||
|
||||
def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
|
||||
@@ -162,15 +165,13 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray:
|
||||
return _offset_event_id(ids, offset)
|
||||
|
||||
|
||||
def iter_file_chunks(path: str | Path, offset: int = 0, k_max: int = K_MAX) -> Iterator[dict[str, np.ndarray]]:
|
||||
"""Yield one parquet row-group at a time so a large file never fully loads.
|
||||
|
||||
`k_max` sets the padded width of the sec_*_list columns (should match
|
||||
`stage2_model.k_max`); defaults to the
|
||||
module constant for callers that don't care (e.g. Stage-1-only reads)."""
|
||||
def iter_file_chunks(
|
||||
path: str | Path, offset: int = 0
|
||||
) -> Iterator[dict[str, np.ndarray]]:
|
||||
"""Yield one parquet row-group at a time so a large file never fully loads."""
|
||||
pf = pq.ParquetFile(path)
|
||||
for i in range(pf.num_row_groups):
|
||||
yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max)
|
||||
yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset)
|
||||
|
||||
|
||||
_COND_COLS = [
|
||||
@@ -204,11 +205,15 @@ def _cond_df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]
|
||||
}
|
||||
|
||||
|
||||
def iter_cond_chunks(path: str | Path, offset: int = 0) -> Iterator[dict[str, np.ndarray]]:
|
||||
def iter_cond_chunks(
|
||||
path: str | Path, offset: int = 0
|
||||
) -> Iterator[dict[str, np.ndarray]]:
|
||||
"""Yield conditioning-only row-groups (no post-step columns read from disk)."""
|
||||
pf = pq.ParquetFile(path)
|
||||
for i in range(pf.num_row_groups):
|
||||
yield _cond_df_to_dict(pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset)
|
||||
yield _cond_df_to_dict(
|
||||
pf.read_row_group(i, columns=_COND_COLS).to_pandas(), offset=offset
|
||||
)
|
||||
|
||||
|
||||
def build_index_maps(
|
||||
@@ -238,44 +243,6 @@ def build_index_maps_from_files(
|
||||
)
|
||||
|
||||
|
||||
def _accumulate_value_counts(counts: dict, series: pd.Series, cast) -> None:
|
||||
for name, count in series.value_counts().items():
|
||||
name = cast(name)
|
||||
counts[name] = counts.get(name, 0) + int(count)
|
||||
|
||||
|
||||
def _rank_by_frequency_from_files(files: list[Path], column: str, cast) -> dict:
|
||||
"""Scan `column` across `files` and return `{cast(value): total_count}`,
|
||||
accumulated in file order (see `fingerprint_files`'s docstring on why
|
||||
scan order — not a normalized/sorted order — is preserved: it drives
|
||||
tie-breaking in the frequency ranking below)."""
|
||||
counts: dict = {}
|
||||
for path in files:
|
||||
df = pd.read_parquet(path, columns=[column])
|
||||
_accumulate_value_counts(counts, df[column], cast)
|
||||
return counts
|
||||
|
||||
|
||||
def _topn_plus_other_map(counts: dict, n_classes: int) -> tuple[dict, dict]:
|
||||
"""Frequency-capped value->index map: the `n_classes - 1` most frequent
|
||||
keys get their own index; every rarer key is bucketed into a shared
|
||||
"other" index (`n_classes - 1`).
|
||||
|
||||
Returns `(class_map, other_members)` — `other_members` is `{key: count}`
|
||||
for every key bucketed into "other" (the empirical within-bucket
|
||||
distribution, for `other_policy = "sample"` at rollout).
|
||||
"""
|
||||
ranked = sorted(counts, key=lambda k: counts[k], reverse=True)
|
||||
keep = ranked[: max(n_classes - 1, 0)]
|
||||
class_map = {k: i for i, k in enumerate(keep)}
|
||||
other_idx = n_classes - 1
|
||||
other_members: dict = {}
|
||||
for k in ranked[len(keep) :]:
|
||||
class_map[k] = other_idx
|
||||
other_members[k] = counts[k]
|
||||
return class_map, other_members
|
||||
|
||||
|
||||
def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str, int]:
|
||||
"""Scan the `process` column and build a frequency-capped process->index map.
|
||||
|
||||
@@ -286,66 +253,16 @@ def build_process_map_from_files(files: list[Path], n_experts: int) -> dict[str,
|
||||
mirrors how `build_features` clamps the n_sec label to K_MAX for the
|
||||
fixed-width n_sec_head classifier.
|
||||
"""
|
||||
counts = _rank_by_frequency_from_files(files, "process", str)
|
||||
class_map, _ = _topn_plus_other_map(counts, n_experts)
|
||||
return class_map
|
||||
|
||||
|
||||
@dataclass
|
||||
class TopNMap:
|
||||
"""A frequency-capped value->index map for a conditioning/type axis (PDG
|
||||
or material), plus the empirical within-bucket distribution of whatever
|
||||
got folded into "other" — see `build_topn_map_from_files`."""
|
||||
|
||||
class_map: dict
|
||||
other_members: dict
|
||||
|
||||
|
||||
def build_topn_map_from_files(files: list[Path], column: str, n_classes: int, cast=str) -> TopNMap:
|
||||
"""Scan `column` and build a frequency-capped value->index map, structurally
|
||||
identical to `build_process_map_from_files` (shares its ranking core via
|
||||
`_topn_plus_other_map`), generalized over the source column and key type.
|
||||
|
||||
Used for the material axis (`column="material"`, `cast=str`, matching
|
||||
`mat_map`'s key type). The PDG axis uses
|
||||
`build_pdg_topn_map_from_files` instead (it needs to pool two columns,
|
||||
which this single-column form can't express). Also records
|
||||
`other_members` (the empirical within-"other" distribution), needed
|
||||
later for `other_policy = "sample"` at rollout — computed now since it's
|
||||
free during this same scan.
|
||||
"""
|
||||
counts = _rank_by_frequency_from_files(files, column, cast)
|
||||
class_map, other_members = _topn_plus_other_map(counts, n_classes)
|
||||
return TopNMap(class_map=class_map, other_members=other_members)
|
||||
|
||||
|
||||
def build_pdg_topn_map_from_files(files: list[Path], n_classes: int) -> TopNMap:
|
||||
"""PDG top-N-plus-other map, pooling counts from BOTH roles a PDG code
|
||||
plays in this dataset: a step's own primary particle (`pdg` column) and
|
||||
an emitted secondary's species (`sec_pdg_list`, exploded) — shared by
|
||||
`conditioning.particle.type = "onehot"` and
|
||||
`stage2_model.particle_type.target = "onehot"`. Pooling both is what
|
||||
keeps a species that's common as a secondary
|
||||
but rare as a primary (or vice versa) from being pushed into "other"
|
||||
just because one role's count alone looks small — the meeting's failure
|
||||
mode (zero photon secondaries, hallucinated antineutrinos) was
|
||||
specifically about secondary-species collapse, so the map this feeds
|
||||
needs to reflect secondary frequency, not just primary frequency.
|
||||
|
||||
`sec_pdg_list` is absent from parquet files predating the parent->child
|
||||
join (see `_df_to_dict`'s `has_sec_lists` guard) — silently skipped for
|
||||
those, same convention as elsewhere in this module.
|
||||
"""
|
||||
counts: dict = {}
|
||||
counts: dict[str, int] = {}
|
||||
for path in files:
|
||||
columns = ["pdg"]
|
||||
has_sec = "sec_pdg_list" in pq.ParquetFile(path).schema_arrow.names
|
||||
if has_sec:
|
||||
columns.append("sec_pdg_list")
|
||||
df = pd.read_parquet(path, columns=columns)
|
||||
_accumulate_value_counts(counts, df["pdg"], int)
|
||||
if has_sec:
|
||||
exploded = df["sec_pdg_list"].explode().dropna()
|
||||
_accumulate_value_counts(counts, exploded, int)
|
||||
class_map, other_members = _topn_plus_other_map(counts, n_classes)
|
||||
return TopNMap(class_map=class_map, other_members=other_members)
|
||||
df = pd.read_parquet(path, columns=["process"])
|
||||
for name, count in df["process"].value_counts().items():
|
||||
name = str(name)
|
||||
counts[name] = counts.get(name, 0) + int(count)
|
||||
ranked = sorted(counts, key=lambda name: counts[name], reverse=True)
|
||||
keep = ranked[: max(n_experts - 1, 0)]
|
||||
proc_map = {name: i for i, name in enumerate(keep)}
|
||||
other_idx = n_experts - 1
|
||||
for name in ranked[len(keep) :]:
|
||||
proc_map[name] = other_idx
|
||||
return proc_map
|
||||
|
||||
+31
-59
@@ -23,7 +23,7 @@ import numpy as np
|
||||
|
||||
from giant import config
|
||||
from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM, X_DIM
|
||||
from giant.data.loader import TopNMap, event_id_offset, load_event_ids
|
||||
from giant.data.loader import event_id_offset, load_event_ids
|
||||
from giant.data.transforms import Normalizer, sorted_membership
|
||||
|
||||
# Bump manually on a change to the data-encoding semantics (e.g. a future
|
||||
@@ -96,50 +96,10 @@ def fingerprint_files(files: list[Path]) -> list[list]:
|
||||
return out
|
||||
|
||||
|
||||
def normalizer_key(
|
||||
val_fraction: float,
|
||||
seed: int,
|
||||
particle_conditioning: str,
|
||||
material_conditioning: str,
|
||||
) -> str:
|
||||
def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str:
|
||||
# .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing
|
||||
# spurious cache misses between runs with the "same" val_fraction. The two
|
||||
# conditioning axes are independent and both
|
||||
# affect which cond_cont columns are computed for real vs. zero-filled
|
||||
# (giant.data.transforms._physical_cond_columns), so both must be part of
|
||||
# the key or two mixed-axis runs could collide on the same cache entry.
|
||||
return f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}_mcond={material_conditioning}"
|
||||
|
||||
|
||||
# Top-N-map axes: "pdg" keys match pdg_map's int
|
||||
# keys (shared by conditioning.particle.type="onehot" and
|
||||
# stage2_model.particle_type.target="onehot" — one map for both), "material"
|
||||
# keys match mat_map's str keys.
|
||||
_TOPN_AXIS_CASTS = {"pdg": int, "material": str}
|
||||
|
||||
|
||||
def topn_key(axis: str, n_classes: int) -> str:
|
||||
"""JSON-safe key for `SetupCache.topn_maps` — N is part of the key so the
|
||||
sidecar stays reusable across runs with different emb_dim (see the
|
||||
dict[int, dict] precedent `proc_maps` sets, keyed by n_experts)."""
|
||||
if axis not in _TOPN_AXIS_CASTS:
|
||||
raise ValueError(f"unknown top-N map axis {axis!r}, expected one of {sorted(_TOPN_AXIS_CASTS)}")
|
||||
return f"{axis}:{n_classes}"
|
||||
|
||||
|
||||
def topnmap_to_json(m: TopNMap) -> dict:
|
||||
return {
|
||||
"class_map": {str(k): v for k, v in m.class_map.items()},
|
||||
"other_members": {str(k): v for k, v in m.other_members.items()},
|
||||
}
|
||||
|
||||
|
||||
def topnmap_from_json(d: dict, axis: str) -> TopNMap:
|
||||
cast = _TOPN_AXIS_CASTS[axis]
|
||||
return TopNMap(
|
||||
class_map={cast(k): v for k, v in d["class_map"].items()},
|
||||
other_members={cast(k): v for k, v in d["other_members"].items()},
|
||||
)
|
||||
# spurious cache misses between runs with the "same" val_fraction.
|
||||
return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}"
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -159,7 +119,9 @@ class NormalizerEntry:
|
||||
"tgt_norm": self.tgt_norm.to_dict(),
|
||||
"sec_phys_norm": self.sec_phys_norm.to_dict(),
|
||||
"n_train_steps": self.n_train_steps,
|
||||
"energy_quantiles": np.asarray(self.energy_quantiles, dtype=np.float32).tolist(),
|
||||
"energy_quantiles": np.asarray(
|
||||
self.energy_quantiles, dtype=np.float32
|
||||
).tolist(),
|
||||
}
|
||||
|
||||
@classmethod
|
||||
@@ -181,8 +143,6 @@ class SetupCache:
|
||||
event_index: tuple[np.ndarray, np.ndarray] | None = None
|
||||
proc_maps: dict[int, dict[str, int]] = field(default_factory=dict)
|
||||
normalizers: dict[str, NormalizerEntry] = field(default_factory=dict)
|
||||
topn_maps: dict[str, TopNMap] = field(default_factory=dict)
|
||||
"""Keyed by `topn_key(axis, n_classes)`."""
|
||||
|
||||
@classmethod
|
||||
def empty(cls, files: list[Path]) -> "SetupCache":
|
||||
@@ -196,7 +156,6 @@ class SetupCache:
|
||||
"fingerprint": self.fingerprint,
|
||||
"proc_maps": {str(k): v for k, v in self.proc_maps.items()},
|
||||
"normalizers": {k: v.to_json() for k, v in self.normalizers.items()},
|
||||
"topn_maps": {k: topnmap_to_json(v) for k, v in self.topn_maps.items()},
|
||||
}
|
||||
if self.vocab is not None:
|
||||
pdg_map, mat_map = self.vocab
|
||||
@@ -226,8 +185,9 @@ class SetupCache:
|
||||
np.array(d["event_index"]["counts"], dtype=np.int64),
|
||||
)
|
||||
proc_maps = {int(k): v for k, v in d.get("proc_maps", {}).items()}
|
||||
normalizers = {k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()}
|
||||
topn_maps = {k: topnmap_from_json(v, axis=k.split(":", 1)[0]) for k, v in d.get("topn_maps", {}).items()}
|
||||
normalizers = {
|
||||
k: NormalizerEntry.from_json(v) for k, v in d.get("normalizers", {}).items()
|
||||
}
|
||||
return cls(
|
||||
fingerprint=d["fingerprint"],
|
||||
git_hash=d.get("git_hash", "unknown"),
|
||||
@@ -235,7 +195,6 @@ class SetupCache:
|
||||
event_index=event_index,
|
||||
proc_maps=proc_maps,
|
||||
normalizers=normalizers,
|
||||
topn_maps=topn_maps,
|
||||
)
|
||||
|
||||
def merge(self, other: "SetupCache") -> "SetupCache":
|
||||
@@ -250,14 +209,17 @@ class SetupCache:
|
||||
fingerprint=other.fingerprint,
|
||||
git_hash=other.git_hash,
|
||||
vocab=other.vocab if other.vocab is not None else self.vocab,
|
||||
event_index=(other.event_index if other.event_index is not None else self.event_index),
|
||||
event_index=(
|
||||
other.event_index if other.event_index is not None else self.event_index
|
||||
),
|
||||
proc_maps={**self.proc_maps, **other.proc_maps},
|
||||
normalizers={**self.normalizers, **other.normalizers},
|
||||
topn_maps={**self.topn_maps, **other.topn_maps},
|
||||
)
|
||||
|
||||
|
||||
def load(data: str | Path, files: list[Path], echo=lambda *a, **k: None) -> SetupCache | None:
|
||||
def load(
|
||||
data: str | Path, files: list[Path], echo=lambda *a, **k: None
|
||||
) -> SetupCache | None:
|
||||
"""Load and validate the sidecar for `data`; `None` on any miss (never raises).
|
||||
|
||||
A missing file, corrupt JSON, format-version mismatch, dimension-constant
|
||||
@@ -281,7 +243,9 @@ def load(data: str | Path, files: list[Path], echo=lambda *a, **k: None) -> Setu
|
||||
echo("setup cache: format version changed — ignoring stale cache")
|
||||
return None
|
||||
if raw.get("dims") != _DIMS:
|
||||
echo("setup cache: model dimension constants changed — ignoring stale cache")
|
||||
echo(
|
||||
"setup cache: model dimension constants changed — ignoring stale cache"
|
||||
)
|
||||
return None
|
||||
fp = fingerprint_files(files)
|
||||
if raw.get("fingerprint") != fp:
|
||||
@@ -324,7 +288,9 @@ def save(
|
||||
with open(lock_path, "a") as lock_file:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_EX)
|
||||
try:
|
||||
base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(files)
|
||||
base = load(data, files, echo=lambda *a, **k: None) or SetupCache.empty(
|
||||
files
|
||||
)
|
||||
merged = base.merge(sections)
|
||||
payload = json.dumps(merged.to_json(), separators=(",", ":"))
|
||||
tmp.write_text(payload)
|
||||
@@ -332,7 +298,9 @@ def save(
|
||||
finally:
|
||||
fcntl.flock(lock_file, fcntl.LOCK_UN)
|
||||
except OSError as exc:
|
||||
echo(f"setup cache: could not write {path} ({exc}) — continuing without caching")
|
||||
echo(
|
||||
f"setup cache: could not write {path} ({exc}) — continuing without caching"
|
||||
)
|
||||
try:
|
||||
tmp.unlink(missing_ok=True)
|
||||
except OSError:
|
||||
@@ -343,12 +311,16 @@ def compute_event_index_from_files(files: list[Path]) -> tuple[np.ndarray, np.nd
|
||||
"""Unique event ids + per-event row (step) counts, across all `files`."""
|
||||
if not files:
|
||||
return np.empty(0, dtype=np.int64), np.empty(0, dtype=np.int64)
|
||||
all_ids = np.concatenate([load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)])
|
||||
all_ids = np.concatenate(
|
||||
[load_event_ids(f, offset=event_id_offset(i)) for i, f in enumerate(files)]
|
||||
)
|
||||
unique_ids, counts = np.unique(all_ids, return_counts=True)
|
||||
return unique_ids, counts
|
||||
|
||||
|
||||
def n_train_steps_for_split(unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray) -> int:
|
||||
def n_train_steps_for_split(
|
||||
unique_ids: np.ndarray, counts: np.ndarray, train_events_arr: np.ndarray
|
||||
) -> int:
|
||||
"""Row (step) count summed over whichever `unique_ids` fall in `train_events_arr`.
|
||||
|
||||
`train_events_arr` must be ascending and duplicate-free (as produced by
|
||||
|
||||
+193
-333
@@ -1,10 +1,7 @@
|
||||
import warnings
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
|
||||
from giant.constants import K_MAX
|
||||
|
||||
_EPS = 1e-8
|
||||
|
||||
# Floor added to each energy fraction before taking log-ratios so the simplex
|
||||
@@ -13,6 +10,14 @@ _EPS = 1e-8
|
||||
# the conservation it slightly softens is physically negligible (~0.001%).
|
||||
_SIMPLEX_FLOOR = 1e-5
|
||||
|
||||
# Upper clip for a raw predicted log_mass before inv_log_transform: exp(y)
|
||||
# must stay well inside float32 range (~3.4e38, i.e. y < ~88.7) or it
|
||||
# overflows to inf, which — like the negative-mass case below — blows up the
|
||||
# next log_transform call once that mass is fed back in as conditioning.
|
||||
# 80.0 leaves comfortable headroom while still being far beyond any physical
|
||||
# particle mass a converged model would ever predict.
|
||||
_LOG_MASS_MAX = 80.0
|
||||
|
||||
|
||||
def log_transform(x: np.ndarray, eps: float = _EPS) -> np.ndarray:
|
||||
x = np.asarray(x, dtype=np.float32)
|
||||
@@ -85,7 +90,9 @@ def energy_simplex_encode(
|
||||
return z.astype(np.float32)
|
||||
|
||||
|
||||
def energy_simplex_decode(z: np.ndarray, pre_E: np.ndarray) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||||
def energy_simplex_decode(
|
||||
z: np.ndarray, pre_E: np.ndarray
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Inverse of `energy_simplex_encode`: ALR coords + pre_E → physical energies.
|
||||
|
||||
A softmax over `[z_edep, z_sec, 0]` recovers the three simplex fractions, so
|
||||
@@ -117,7 +124,9 @@ def _rodrigues_axis(pre_dir: np.ndarray) -> np.ndarray:
|
||||
arbitrary second operand) on every row; profiling on a 114M-row file
|
||||
showed `np.cross` as the single hottest call inside this rotation.
|
||||
"""
|
||||
axis = np.stack([pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1)
|
||||
axis = np.stack(
|
||||
[pre_dir[:, 1], -pre_dir[:, 0], np.zeros_like(pre_dir[:, 0])], axis=1
|
||||
)
|
||||
axis_norm = np.linalg.norm(axis, axis=1, keepdims=True) # (N,1)
|
||||
# axis_norm ~ 0 happens at BOTH poles: pre_dir ~ +ẑ (forward) and
|
||||
# pre_dir ~ -ẑ (near-exact backscatter) — ‖pre_dir × ẑ‖ = sin(angle to
|
||||
@@ -196,7 +205,9 @@ def local_frame_rotation(pre_dir: np.ndarray, post_dir: np.ndarray) -> np.ndarra
|
||||
kxv = _cross_with_z_axis(axis, post_dir) # (N,3)
|
||||
kdv = (axis * post_dir).sum(axis=1, keepdims=True) # (N,1)
|
||||
|
||||
return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
|
||||
return (post_dir * cos_t + kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
|
||||
np.float32
|
||||
)
|
||||
|
||||
|
||||
class Normalizer:
|
||||
@@ -336,21 +347,22 @@ def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray:
|
||||
return sorted_arr[idx] == values
|
||||
|
||||
|
||||
def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0) -> np.ndarray:
|
||||
def _vectorized_map_lookup(
|
||||
values: np.ndarray, mapping: dict, strict: bool = True
|
||||
) -> np.ndarray:
|
||||
"""Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`.
|
||||
|
||||
Replaces a per-element Python dict lookup with one `searchsorted` call.
|
||||
Raises `KeyError` if any value in `values` isn't a key of `mapping`,
|
||||
matching the dict-comprehension it replaces (never silently misassigns)
|
||||
— unless `strict=False`, in which case unmapped values get `default`
|
||||
instead. Only pass `strict=False` where the caller has independently
|
||||
verified the resulting index is either never actually read (e.g.
|
||||
— unless `strict=False`, in which case unmapped values get a dummy index
|
||||
of 0 instead. Only pass `strict=False` where the caller has independently
|
||||
verified the resulting index is never actually read (e.g.
|
||||
`build_cond_features` under `conditioning="physical"`, where
|
||||
`ConditionEncoder` ignores `cond_cat` entirely) or where `default` is a
|
||||
deliberate fallback class (e.g. a top-N map's "other" index for a raw
|
||||
value outside the training vocab). It exists so a rollout can be seeded
|
||||
with a species/material outside the training vocab without a spurious
|
||||
`KeyError`, which is the entire point of physical-property conditioning.
|
||||
`ConditionEncoder` ignores `cond_cat` entirely); it exists so a rollout
|
||||
can be seeded with a species/material outside the training vocab without
|
||||
a spurious `KeyError`, which is the entire point of physical-property
|
||||
conditioning.
|
||||
"""
|
||||
keys = np.asarray(list(mapping.keys()))
|
||||
vals = np.asarray(list(mapping.values()), dtype=np.int64)
|
||||
@@ -362,7 +374,7 @@ def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = Tru
|
||||
found = keys_sorted[pos] == values
|
||||
if not found.all():
|
||||
if not strict:
|
||||
out = np.full(values.shape, default, dtype=np.int64)
|
||||
out = np.zeros(values.shape, dtype=np.int64)
|
||||
out[found] = vals_sorted[pos[found]]
|
||||
return out
|
||||
missing = np.unique(values[~found])
|
||||
@@ -380,7 +392,9 @@ def travel_direction(pre_pos: np.ndarray, post_pos: np.ndarray) -> np.ndarray:
|
||||
disp = post_pos - pre_pos
|
||||
norm = np.linalg.norm(disp, axis=1, keepdims=True)
|
||||
safe_norm = np.where(norm < 1e-7, 1.0, norm)
|
||||
return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(np.float32)
|
||||
return np.where(norm < 1e-7, np.array([[0.0, 0.0, 1.0]]), disp / safe_norm).astype(
|
||||
np.float32
|
||||
)
|
||||
|
||||
|
||||
def reconstruct_post_pos(
|
||||
@@ -399,7 +413,9 @@ def reconstruct_post_pos(
|
||||
return (pre_pos + step_length.reshape(-1, 1) * travel_dir_world).astype(np.float32)
|
||||
|
||||
|
||||
def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) -> np.ndarray:
|
||||
def inv_local_frame_rotation(
|
||||
pre_dir: np.ndarray, post_dir_local: np.ndarray
|
||||
) -> np.ndarray:
|
||||
"""Inverse of local_frame_rotation: rotate from local frame back to world frame.
|
||||
|
||||
Applies R^T (same axis, negative angle) to post_dir_local.
|
||||
@@ -413,7 +429,9 @@ def inv_local_frame_rotation(pre_dir: np.ndarray, post_dir_local: np.ndarray) ->
|
||||
kdv = (axis * post_dir_local).sum(axis=1, keepdims=True)
|
||||
|
||||
# Negative angle: sin_t → -sin_t
|
||||
return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(np.float32)
|
||||
return (post_dir_local * cos_t - kxv * sin_t + axis * kdv * (1.0 - cos_t)).astype(
|
||||
np.float32
|
||||
)
|
||||
|
||||
|
||||
_STICK_LOGIT_CLIP = 10.0 # logit value used for the last valid secondary slot
|
||||
@@ -478,10 +496,14 @@ def encode_secondaries(
|
||||
remaining_raw = e_sec - cumsum[:, i - 1]
|
||||
shortfall_flagged |= sec_valid[:, i] & (remaining_raw < -_SHORTFALL_TOL)
|
||||
remaining = np.maximum(remaining_raw, _EPS)
|
||||
f = np.clip(sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS)
|
||||
f = np.clip(
|
||||
sec_E_list[:, i].astype(np.float64) / remaining, _EPS, 1.0 - _EPS
|
||||
)
|
||||
logit = np.log(f / (1.0 - f)).astype(np.float32)
|
||||
# Last valid slot: give it the full remaining budget
|
||||
is_last = sec_valid[:, i] & ~(sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool))
|
||||
is_last = sec_valid[:, i] & ~(
|
||||
sec_valid[:, i + 1] if i + 1 < K else np.zeros(N, dtype=bool)
|
||||
)
|
||||
logit = np.where(is_last, _STICK_LOGIT_CLIP, logit)
|
||||
logit = np.where(
|
||||
sec_valid[:, i],
|
||||
@@ -508,7 +530,9 @@ def encode_secondaries(
|
||||
# Only rotate valid slots; leave padded slots as (0,0,1) or whatever.
|
||||
valid_mask = sec_valid[:, i]
|
||||
if valid_mask.any():
|
||||
dir_local[valid_mask, i] = local_frame_rotation(pre_dir[valid_mask], sec_dir_list[valid_mask, i])
|
||||
dir_local[valid_mask, i] = local_frame_rotation(
|
||||
pre_dir[valid_mask], sec_dir_list[valid_mask, i]
|
||||
)
|
||||
|
||||
if sec_pdg_list is not None:
|
||||
from giant.particles import particle_phys_array
|
||||
@@ -534,69 +558,42 @@ def encode_secondaries(
|
||||
return sec_cont.astype(np.float32)
|
||||
|
||||
|
||||
def encode_secondary_type_idx(sec_pdg_list: np.ndarray, sec_valid: np.ndarray, class_map: dict) -> np.ndarray:
|
||||
"""Per-secondary-slot class index into `class_map` — (N, K_MAX) int64.
|
||||
|
||||
`class_map` is either a top-N-plus-other map's `class_map`
|
||||
(`stage2_model.particle_type.target = "onehot"`, see
|
||||
`giant.data.loader.build_pdg_topn_map_from_files`) or the dense `pdg_map`
|
||||
(`target = "embedding"`). Not used at all for `target = "physical"` —
|
||||
that target keeps using `encode_secondaries`'s (log_mass, charge)
|
||||
columns unchanged.
|
||||
|
||||
Padding slots get index 0 (their looked-up value is discarded downstream
|
||||
by the `sec_valid`/`n_sec` mask regardless, so any in-vocabulary dummy
|
||||
code works). A *real, valid* secondary whose code is missing from
|
||||
`class_map` raises `KeyError` (`strict=True`) rather than silently
|
||||
misassigning — for `target="onehot"` this should never actually
|
||||
trigger, since `build_pdg_topn_map_from_files` pools both primary and
|
||||
secondary occurrences precisely so every secondary species seen in
|
||||
these files has a key (in "other" at worst); for `target="embedding"`
|
||||
(which reuses the dense, primary-only `pdg_map`) it's a real signal
|
||||
that a secondary-only species exists with no primary-role counterpart.
|
||||
"""
|
||||
N, K = sec_pdg_list.shape
|
||||
# An arbitrary already-present key works as the padding-slot dummy code
|
||||
# (unlike encode_secondaries' physics-derived phys lookup, this is an
|
||||
# index into class_map's own vocabulary, so a fixed sentinel like 22
|
||||
# isn't guaranteed to be a key — an arbitrary present one always is).
|
||||
dummy = next(iter(class_map))
|
||||
safe_pdg = np.where(sec_valid, sec_pdg_list, dummy)
|
||||
idx = _vectorized_map_lookup(safe_pdg.reshape(-1), class_map, strict=True).reshape(N, K)
|
||||
return np.where(sec_valid, idx, 0).astype(np.int64)
|
||||
|
||||
|
||||
def decode_secondary_cont(
|
||||
def decode_secondaries(
|
||||
sec_cont: np.ndarray,
|
||||
n_sec: np.ndarray,
|
||||
e_sec: np.ndarray,
|
||||
pre_dir: np.ndarray,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Continuous-only half of `decode_secondaries`'s inverse: the
|
||||
stick-breaking energy split and local->world direction — generator/
|
||||
`particle_type.target`-independent, since every target (`"physical"`,
|
||||
`"onehot"`, `"embedding"`) shares the same `CONT_SLOT_DIM`-wide
|
||||
(stick_logit, dir) prefix and differs only
|
||||
in what follows it. `decode_secondaries` (target="physical") is the
|
||||
original all-in-one form built on top of this; `target` in `("onehot",
|
||||
"embedding")` decodes their type slice separately via
|
||||
`giant.particles.decode_topn_class`/`decode_embedding_nearest` and calls
|
||||
this directly instead — see `giant/rollout.py`.
|
||||
sec_phys_normalizer: "Normalizer | None" = None,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Inverse of encode_secondaries: continuous targets → physical secondary attrs.
|
||||
|
||||
sec_cont: (N, K, >=CONT_SLOT_DIM) — only columns `[:, :, :CONT_SLOT_DIM]`
|
||||
(stick_logit, local dir) are read; a caller may pass its full
|
||||
per-slot tensor (continuous + type) unsliced.
|
||||
sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y,
|
||||
local_dir_z, log_mass, charge] (log_mass/charge normalised iff
|
||||
`sec_phys_normalizer` was applied when this was produced — e.g. a
|
||||
raw model prediction; pass the same normalizer here to invert it)
|
||||
n_sec: (N,) integer secondary counts
|
||||
e_sec: (N,) total secondary energy budget [MeV]
|
||||
pre_dir: (N, 3) pre-step world-frame direction
|
||||
|
||||
Returns (sec_E, sec_dir_world, sec_valid), shapes (N, K), (N, K, 3),
|
||||
(N, K). The valid slots' energies (`sec_E[sec_valid]`, per row) always
|
||||
sum to exactly `e_sec` — see the rescaling below.
|
||||
Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each
|
||||
shape (N, K_MAX). The valid slots' energies (`sec_E[sec_valid]`, per row)
|
||||
always sum to exactly `e_sec` — see the rescaling below. mass/charge are
|
||||
the model's raw predicted physical identity for each secondary, used
|
||||
as-is (no snapping to a discrete PDG code) — see giant/particles.py for
|
||||
the separate, reporting-only nearest-PDG lookup callers may apply on top
|
||||
of this for display/bookkeeping purposes.
|
||||
"""
|
||||
N, K = sec_cont.shape[0], sec_cont.shape[1]
|
||||
if sec_phys_normalizer is not None:
|
||||
N_, K_, _ = sec_cont.shape
|
||||
phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2))
|
||||
sec_cont = sec_cont.copy()
|
||||
sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
|
||||
|
||||
N, K, _ = sec_cont.shape
|
||||
stick_logits = sec_cont[:, :, 0] # (N, K)
|
||||
dir_local = sec_cont[:, :, 1:4].copy() # (N, K, 3)
|
||||
log_mass = sec_cont[:, :, 4] # (N, K)
|
||||
charge = sec_cont[:, :, 5] # (N, K)
|
||||
|
||||
# Flow-matching output isn't guaranteed unit norm; normalise before the
|
||||
# rotation below, which preserves magnitude rather than fixing it up.
|
||||
@@ -636,57 +633,22 @@ def decode_secondary_cont(
|
||||
for i in range(K):
|
||||
valid = sec_valid[:, i]
|
||||
if valid.any():
|
||||
sec_dir_world[valid, i] = inv_local_frame_rotation(pre_dir[valid], dir_local[valid, i])
|
||||
sec_dir_world[valid, i] = inv_local_frame_rotation(
|
||||
pre_dir[valid], dir_local[valid, i]
|
||||
)
|
||||
|
||||
return sec_E, sec_dir_world, sec_valid
|
||||
|
||||
|
||||
def decode_secondaries(
|
||||
sec_cont: np.ndarray,
|
||||
n_sec: np.ndarray,
|
||||
e_sec: np.ndarray,
|
||||
pre_dir: np.ndarray,
|
||||
sec_phys_normalizer: "Normalizer | None" = None,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
|
||||
"""Inverse of encode_secondaries: continuous targets → physical secondary attrs.
|
||||
|
||||
`particle_type.target = "physical"` only (the type slice is a raw
|
||||
(log_mass, charge) regression target folded straight into `sec_cont`) —
|
||||
`"onehot"`/`"embedding"` decode through `decode_secondary_cont` +
|
||||
`giant.particles.decode_topn_class`/`decode_embedding_nearest` instead,
|
||||
since their type slice isn't (log_mass, charge) at all. See
|
||||
`decode_secondary_cont`'s docstring for why the two share the energy/
|
||||
direction logic below.
|
||||
|
||||
sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y,
|
||||
local_dir_z, log_mass, charge] (log_mass/charge normalised iff
|
||||
`sec_phys_normalizer` was applied when this was produced — e.g. a
|
||||
raw model prediction; pass the same normalizer here to invert it)
|
||||
n_sec: (N,) integer secondary counts
|
||||
e_sec: (N,) total secondary energy budget [MeV]
|
||||
pre_dir: (N, 3) pre-step world-frame direction
|
||||
|
||||
Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each
|
||||
shape (N, K_MAX). mass/charge are the model's raw predicted physical
|
||||
identity for each secondary, used as-is (no snapping to a discrete PDG
|
||||
code) — see giant/particles.py for the separate, reporting-only
|
||||
nearest-PDG lookup callers may apply on top of this for display/
|
||||
bookkeeping purposes.
|
||||
"""
|
||||
if sec_phys_normalizer is not None:
|
||||
N_, K_, _ = sec_cont.shape
|
||||
phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2))
|
||||
sec_cont = sec_cont.copy()
|
||||
sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2)
|
||||
|
||||
sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont, n_sec, e_sec, pre_dir)
|
||||
|
||||
log_mass = sec_cont[:, :, 4] # (N, K)
|
||||
charge = sec_cont[:, :, 5] # (N, K)
|
||||
|
||||
# mass is non-negative by construction (inv_log_transform of a real
|
||||
# number is always > 0); clip to 0 for padded/invalid slots rather than
|
||||
# leaving a spurious small positive floor from the log inverse.
|
||||
# log_mass is a raw model prediction, not itself the output of
|
||||
# log_transform, so it can land far outside the range that round-trips
|
||||
# cleanly through inv_log_transform: too negative and exp(log_mass)
|
||||
# undershoots _EPS, making inv_log_transform go slightly negative; too
|
||||
# positive and exp(log_mass) overflows float32 to inf. Either one then
|
||||
# blows up the next log_transform call on this track's mass once it's
|
||||
# fed back in as conditioning for a further rollout step
|
||||
# (giant/rollout.py -> build_cond_features -> _physical_cond_columns).
|
||||
# Clip log_mass to a range whose inverse is guaranteed finite and >= 0
|
||||
# before that can happen; clip to 0 separately for padded/invalid slots
|
||||
# rather than leaving a spurious small positive floor.
|
||||
log_mass = np.clip(log_mass, np.log(_EPS), _LOG_MASS_MAX)
|
||||
sec_mass = np.where(sec_valid, inv_log_transform(log_mass), 0.0).astype(np.float32)
|
||||
sec_charge = np.where(sec_valid, charge, 0.0).astype(np.float32)
|
||||
|
||||
@@ -694,66 +656,52 @@ def decode_secondaries(
|
||||
|
||||
|
||||
def _physical_cond_columns(
|
||||
data: dict[str, np.ndarray],
|
||||
particle_conditioning: str,
|
||||
material_conditioning: str,
|
||||
data: dict[str, np.ndarray], conditioning: str
|
||||
) -> np.ndarray:
|
||||
"""(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns.
|
||||
|
||||
The particle and material blocks are gated independently and may mix
|
||||
freely — e.g. material `physical` with particle `embedding` — so e.g.
|
||||
`particle_conditioning="embedding"` + `material_conditioning="physical"`
|
||||
zero-fills only the particle columns and computes the material ones for
|
||||
real.
|
||||
|
||||
"embedding"/"onehot" zero-fill their block (cheap, and ConditionEncoder
|
||||
never reads these columns in either mode — so an unfilled
|
||||
giant.materials table can never crash an "embedding"/"onehot"-mode run).
|
||||
"physical" computes it for real: particle columns come from
|
||||
`data["mass"]`/`data["charge"]` when the caller already knows them
|
||||
directly (rollout.py, for a track descended from a model-predicted
|
||||
secondary — see giant/rollout.py's "no snapping" design), else derived
|
||||
from `data["pdg"]` via giant.particles; material columns always come
|
||||
from `data["material"]` via giant.materials, since material is never
|
||||
itself a model prediction.
|
||||
"embedding" mode zero-fills (cheap, and ConditionEncoder never reads
|
||||
these columns in that mode — so an unfilled giant.materials table can
|
||||
never crash an "embedding"-mode run). "physical" mode computes them for
|
||||
real: particle columns come from `data["mass"]`/`data["charge"]` when the
|
||||
caller already knows them directly (rollout.py, for a track descended
|
||||
from a model-predicted secondary — see giant/rollout.py's "no snapping"
|
||||
design), else derived from `data["pdg"]` via giant.particles; material
|
||||
columns always come from `data["material"]` via giant.materials, since
|
||||
material is never itself a model prediction.
|
||||
"""
|
||||
from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
|
||||
|
||||
n = len(next(iter(data.values())))
|
||||
if conditioning == "embedding":
|
||||
n = len(next(iter(data.values())))
|
||||
return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32)
|
||||
if conditioning != "physical":
|
||||
raise ValueError(f"unknown conditioning mode {conditioning!r}")
|
||||
|
||||
if particle_conditioning == "physical":
|
||||
from giant.particles import particle_phys_array
|
||||
from giant.materials import material_properties_array
|
||||
from giant.particles import particle_phys_array
|
||||
|
||||
if "mass" in data and "charge" in data:
|
||||
mass = np.asarray(data["mass"], dtype=np.float32)
|
||||
charge = np.asarray(data["charge"], dtype=np.float32)
|
||||
else:
|
||||
mass, charge = particle_phys_array(data["pdg"]).T
|
||||
particle_cols = np.column_stack([log_transform(mass), charge])
|
||||
elif particle_conditioning in ("embedding", "onehot"):
|
||||
particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32)
|
||||
if "mass" in data and "charge" in data:
|
||||
mass = np.asarray(data["mass"], dtype=np.float32)
|
||||
charge = np.asarray(data["charge"], dtype=np.float32)
|
||||
else:
|
||||
raise ValueError(f"unknown conditioning.particle.type {particle_conditioning!r}")
|
||||
mass, charge = particle_phys_array(data["pdg"]).T
|
||||
|
||||
if material_conditioning == "physical":
|
||||
from giant.materials import material_properties_array
|
||||
z_eff, a_eff, density, x0, lambda_int = material_properties_array(
|
||||
data["material"]
|
||||
).T
|
||||
|
||||
z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T
|
||||
material_cols = np.column_stack(
|
||||
[
|
||||
z_eff,
|
||||
a_eff,
|
||||
log_transform(density),
|
||||
log_transform(x0),
|
||||
log_transform(lambda_int),
|
||||
]
|
||||
)
|
||||
elif material_conditioning in ("embedding", "onehot"):
|
||||
material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32)
|
||||
else:
|
||||
raise ValueError(f"unknown conditioning.material.type {material_conditioning!r}")
|
||||
|
||||
return np.column_stack([particle_cols, material_cols]).astype(np.float32)
|
||||
return np.column_stack(
|
||||
[
|
||||
log_transform(mass),
|
||||
charge,
|
||||
z_eff,
|
||||
a_eff,
|
||||
log_transform(density),
|
||||
log_transform(x0),
|
||||
log_transform(lambda_int),
|
||||
]
|
||||
).astype(np.float32)
|
||||
|
||||
|
||||
def build_cond_features(
|
||||
@@ -761,24 +709,9 @@ def build_cond_features(
|
||||
pdg_map: dict[int, int],
|
||||
mat_map: dict[str, int],
|
||||
cond_normalizer: "Normalizer | None" = None,
|
||||
particle_conditioning: str = "embedding",
|
||||
material_conditioning: str = "embedding",
|
||||
pdg_topn_map: dict[int, int] | None = None,
|
||||
mat_topn_map: dict[str, int] | None = None,
|
||||
conditioning: str = "embedding",
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Build conditioning arrays only — no target, no post-step variables.
|
||||
|
||||
`particle_conditioning`/`material_conditioning` are independent —
|
||||
e.g. `particle_conditioning="embedding"` +
|
||||
`material_conditioning="physical"` is a valid mix.
|
||||
|
||||
`pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see
|
||||
`giant.data.loader.build_topn_map_from_files`) append extra `cond_cat`
|
||||
columns read by `ConditionEncoder`'s `"onehot"` mode: pdg topN index at
|
||||
column 2 (iff `pdg_topn_map` given), material topN index at column 3
|
||||
(iff `mat_topn_map` given, after column 2 if both are). Only ever given when
|
||||
the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise.
|
||||
"""
|
||||
"""Build conditioning arrays only — no target, no post-step variables."""
|
||||
cond_cont = np.column_stack(
|
||||
[
|
||||
data["pre_pos"],
|
||||
@@ -788,69 +721,49 @@ def build_cond_features(
|
||||
]
|
||||
).astype(np.float32)
|
||||
cond_cont = np.column_stack(
|
||||
[
|
||||
cond_cont,
|
||||
_physical_cond_columns(data, particle_conditioning, material_conditioning),
|
||||
]
|
||||
[cond_cont, _physical_cond_columns(data, conditioning)]
|
||||
).astype(np.float32)
|
||||
|
||||
# In "physical" mode cond_cat's first two columns are only a
|
||||
# reporting/router convenience — ConditionEncoder never reads them
|
||||
# (giant/model/network.py) — so a species/material outside the training
|
||||
# vocab (the whole point of physical-property conditioning) gets a dummy
|
||||
# index instead of raising. In "embedding" mode those columns ARE the
|
||||
# conditioning signal, so an unmapped value must still raise loudly
|
||||
# rather than silently misassign. In "onehot" mode they again go unread
|
||||
# (the topN columns below are the real signal), so they're as permissive
|
||||
# as "physical". Each axis's strictness is independent.
|
||||
pdg_strict = particle_conditioning == "embedding"
|
||||
mat_strict = material_conditioning == "embedding"
|
||||
pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=pdg_strict)
|
||||
mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=mat_strict)
|
||||
cat_cols = [pdg_idx, mat_idx]
|
||||
if pdg_topn_map is not None:
|
||||
cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
|
||||
if mat_topn_map is not None:
|
||||
cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
|
||||
cond_cat = np.column_stack(cat_cols)
|
||||
# In "physical" mode cond_cat is only a reporting/router convenience —
|
||||
# ConditionEncoder never reads it (giant/model/network.py) — so a
|
||||
# species/material outside the training vocab (the whole point of
|
||||
# physical-property conditioning) gets a dummy index instead of raising.
|
||||
# In "embedding" mode cond_cat IS the conditioning signal, so an unmapped
|
||||
# value must still raise loudly rather than silently misassign.
|
||||
strict = conditioning == "embedding"
|
||||
pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict)
|
||||
mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict)
|
||||
cond_cat = np.column_stack([pdg_idx, mat_idx])
|
||||
|
||||
if cond_normalizer is not None:
|
||||
cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, particle_conditioning, material_conditioning)
|
||||
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",
|
||||
particle_conditioning: str,
|
||||
material_conditioning: str,
|
||||
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. When NEITHER axis is "physical" those columns are never read by
|
||||
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. If EITHER
|
||||
axis is "physical" its columns are load-bearing, so a mismatch there is a
|
||||
real incompatibility, not something to paper over.
|
||||
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:
|
||||
physical_load_bearing = "physical" in (
|
||||
particle_conditioning,
|
||||
material_conditioning,
|
||||
)
|
||||
if physical_load_bearing:
|
||||
if conditioning != "embedding":
|
||||
raise ValueError(
|
||||
f"cond normalizer has {mean.shape[-1]} columns, expected "
|
||||
f"{width}, and particle_conditioning={particle_conditioning!r}/"
|
||||
f"material_conditioning={material_conditioning!r} reads the "
|
||||
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."
|
||||
@@ -861,41 +774,6 @@ def _cond_normalizer_transform(
|
||||
return ((cond_cont - mean) / std).astype(np.float32)
|
||||
|
||||
|
||||
class StepFeatures(NamedTuple):
|
||||
"""Output of `build_features`. Field order is load-bearing for existing
|
||||
positional unpacking (tests, `StreamingStepsDataset`) — append only,
|
||||
never insert or reorder.
|
||||
|
||||
target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
|
||||
n_sec: (N,) integer secondary counts (target for n_sec head)
|
||||
sec_cont: (N, K_MAX, SEC_SLOT_DIM=6) continuous secondary targets
|
||||
[stick_logit, dir_local, log_mass, charge] — mass/charge are
|
||||
the secondary's real physical identity (from its ground-truth
|
||||
PDG code), a fixed regression target, not a learned/snapped one.
|
||||
Always computed the same way regardless of
|
||||
`stage2_model.particle_type.target` — only actually used
|
||||
downstream under `target = "physical"`.
|
||||
proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
|
||||
never conditioning). Zeros when `proc_map` is None or the loaded
|
||||
data has no "process" column (e.g. pre-conversion parquet files).
|
||||
sec_type_idx: (N, K_MAX) integer secondary class index into
|
||||
`sec_type_class_map`, for `stage2_model.particle_type.target`
|
||||
in `("onehot", "embedding")` — see `encode_secondary_type_idx`.
|
||||
Zero-filled (and unused) when `sec_type_class_map` is None
|
||||
(i.e. `target = "physical"`).
|
||||
"""
|
||||
|
||||
cond_cont: np.ndarray
|
||||
cond_cat: np.ndarray
|
||||
target_s1: np.ndarray
|
||||
n_sec: np.ndarray
|
||||
sec_cont: np.ndarray
|
||||
proc_idx: np.ndarray
|
||||
sec_type_idx: np.ndarray
|
||||
cond_normalizer: Normalizer | None
|
||||
target_normalizer: Normalizer | None
|
||||
|
||||
|
||||
def build_features(
|
||||
data: dict[str, np.ndarray],
|
||||
pdg_map: dict[int, int],
|
||||
@@ -906,17 +784,29 @@ def build_features(
|
||||
fit: bool = False,
|
||||
proc_map: dict[str, int] | None = None,
|
||||
require_secondaries: bool = False,
|
||||
particle_conditioning: str = "embedding",
|
||||
material_conditioning: str = "embedding",
|
||||
conditioning: str = "embedding",
|
||||
sec_phys_only: bool = False,
|
||||
pdg_topn_map: dict[int, int] | None = None,
|
||||
mat_topn_map: dict[str, int] | None = None,
|
||||
sec_type_class_map: dict | None = None,
|
||||
k_max: int = K_MAX,
|
||||
) -> StepFeatures:
|
||||
"""Assemble a `StepFeatures` of (cond_cont, cond_cat, target_s1, n_sec,
|
||||
sec_cont, proc_idx, sec_type_idx, cond_normalizer, target_normalizer) —
|
||||
see `StepFeatures` for field meanings.
|
||||
) -> tuple[
|
||||
np.ndarray,
|
||||
np.ndarray,
|
||||
np.ndarray,
|
||||
np.ndarray,
|
||||
np.ndarray,
|
||||
np.ndarray,
|
||||
Normalizer | None,
|
||||
Normalizer | None,
|
||||
]:
|
||||
"""Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) arrays.
|
||||
|
||||
target_s1: (N, 9) Stage-1 primary post-step target (unchanged from Phase 1)
|
||||
n_sec: (N,) integer secondary counts (target for n_sec head)
|
||||
sec_cont: (N, K_MAX, SEC_SLOT_DIM=6) continuous secondary targets
|
||||
[stick_logit, dir_local, log_mass, charge] — mass/charge are
|
||||
the secondary's real physical identity (from its ground-truth
|
||||
PDG code), a fixed regression target, not a learned/snapped one.
|
||||
proc_idx: (N,) integer process-class label (ProcessRouter supervision only —
|
||||
never conditioning). Zeros when `proc_map` is None or the loaded
|
||||
data has no "process" column (e.g. pre-conversion parquet files).
|
||||
|
||||
require_secondaries: when True, raise if any step has n_sec > 0 but the
|
||||
per-secondary list columns are absent (a mis-converted file that would
|
||||
@@ -927,27 +817,17 @@ def build_features(
|
||||
the stick-breaking/direction-rotation blocks of `sec_cont` (zero-filled
|
||||
instead) for callers (normalizer fitting) that only read
|
||||
`sec_cont[:, :, 4:6]` and would otherwise discard that work.
|
||||
|
||||
pdg_topn_map/mat_topn_map: appended `cond_cat` columns for
|
||||
`ConditionEncoder`'s `"onehot"` mode — see `build_cond_features`.
|
||||
|
||||
sec_type_class_map: the map `sec_type_idx` is looked up against — a
|
||||
top-N-plus-other map's `class_map` for `target = "onehot"`, or the
|
||||
dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself).
|
||||
`None` for `target = "physical"`.
|
||||
|
||||
k_max: should match `stage2_model.k_max` —
|
||||
overridden internally by `data["sec_E_list"]`'s own padded width when
|
||||
present (the loader already padded it to some k_max; that width is
|
||||
authoritative), so this only actually matters when secondary list
|
||||
columns are absent (Stage-1-only reads, or a pre-secondary-join
|
||||
file), where it sets `sec_cont`/`sec_type_idx`'s zero-filled width.
|
||||
"""
|
||||
from giant.constants import K_MAX
|
||||
|
||||
post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"])
|
||||
travel_dir_local = local_frame_rotation(data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"]))
|
||||
travel_dir_local = local_frame_rotation(
|
||||
data["pre_dir"], travel_direction(data["pre_pos"], data["post_pos"])
|
||||
)
|
||||
|
||||
energy_z = energy_simplex_encode(data["edep"], data["e_sec"], data["post_E"], data["pre_E"]) # (N, 2)
|
||||
energy_z = energy_simplex_encode(
|
||||
data["edep"], data["e_sec"], data["post_E"], data["pre_E"]
|
||||
) # (N, 2)
|
||||
|
||||
target_s1 = np.column_stack(
|
||||
[
|
||||
@@ -968,43 +848,30 @@ def build_features(
|
||||
]
|
||||
).astype(np.float32) # (N, COND_DIM_BASE=8)
|
||||
cond_cont = np.column_stack(
|
||||
[
|
||||
cond_cont,
|
||||
_physical_cond_columns(data, particle_conditioning, material_conditioning),
|
||||
]
|
||||
[cond_cont, _physical_cond_columns(data, conditioning)]
|
||||
).astype(np.float32) # (N, COND_DIM=15)
|
||||
|
||||
pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map)
|
||||
mat_idx = _vectorized_map_lookup(data["material"], mat_map)
|
||||
cat_cols = [pdg_idx, mat_idx]
|
||||
if pdg_topn_map is not None:
|
||||
cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map))
|
||||
if mat_topn_map is not None:
|
||||
cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map))
|
||||
cond_cat = np.column_stack(cat_cols) # (N, 2/3/4)
|
||||
cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2)
|
||||
|
||||
n_sec_raw = data["n_sec"].astype(np.int64) # (N,) unclamped, for the valid-slot mask
|
||||
n_sec_raw = data["n_sec"].astype(
|
||||
np.int64
|
||||
) # (N,) unclamped, for the valid-slot mask
|
||||
# Clamp the classification label to K_MAX: the head only has K_MAX+1 classes
|
||||
# (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already
|
||||
# applied to sec_cont by the loader's list padding. Without this, a rare
|
||||
# high-multiplicity step (real data goes up to ~37) hands cross_entropy
|
||||
# an out-of-range target and CUDA asserts.
|
||||
n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64) # (N,)
|
||||
|
||||
# Secondary continuous targets
|
||||
sec_E_list = data.get("sec_E_list")
|
||||
sec_dir_list = data.get("sec_dir_list")
|
||||
sec_pdg_list = data.get("sec_pdg_list")
|
||||
if sec_E_list is not None:
|
||||
# The loader already padded sec_*_list to some k_max (see
|
||||
# giant.data.loader.iter_file_chunks); that padded width is
|
||||
# authoritative over whatever this call happened to pass in, so the
|
||||
# two can never drift apart.
|
||||
k_max = sec_E_list.shape[1]
|
||||
|
||||
# Clamp the classification label to k_max: the head only has k_max+1
|
||||
# classes (0..k_max), and truncating here mirrors the k_max-slot
|
||||
# truncation already applied to sec_cont by the loader's list padding.
|
||||
# Without this, a rare high-multiplicity step (real data goes up to ~37)
|
||||
# hands cross_entropy an out-of-range target and CUDA asserts.
|
||||
n_sec = np.minimum(n_sec_raw, k_max).astype(np.int64) # (N,)
|
||||
|
||||
if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None:
|
||||
sec_valid = np.arange(k_max)[None, :] < n_sec_raw[:, None] # (N, k_max)
|
||||
sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None] # (N, K_MAX)
|
||||
sec_cont = encode_secondaries(
|
||||
sec_E_list,
|
||||
sec_dir_list,
|
||||
@@ -1013,12 +880,7 @@ def build_features(
|
||||
data["pre_dir"],
|
||||
sec_pdg_list=sec_pdg_list,
|
||||
phys_only=sec_phys_only,
|
||||
) # (N, k_max, 6)
|
||||
sec_type_idx = (
|
||||
encode_secondary_type_idx(sec_pdg_list, sec_valid, sec_type_class_map)
|
||||
if sec_type_class_map is not None
|
||||
else np.zeros((len(n_sec), k_max), dtype=np.int64)
|
||||
)
|
||||
) # (N, K_MAX, 6)
|
||||
else:
|
||||
# Guard against silently training Stage 2 on zeroed targets: if any step
|
||||
# actually spawned secondaries (n_sec > 0, from child_track_ids) but the
|
||||
@@ -1040,8 +902,7 @@ def build_features(
|
||||
"require_secondaries=False for Stage-1-only use."
|
||||
)
|
||||
N = len(n_sec)
|
||||
sec_cont = np.zeros((N, k_max, 6), dtype=np.float32)
|
||||
sec_type_idx = np.zeros((N, k_max), dtype=np.int64)
|
||||
sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
|
||||
|
||||
if fit:
|
||||
cond_normalizer = Normalizer().fit(cond_cont)
|
||||
@@ -1063,14 +924,13 @@ def build_features(
|
||||
else:
|
||||
proc_idx = np.zeros(len(cond_cat), dtype=np.int64)
|
||||
|
||||
return StepFeatures(
|
||||
cond_cont=cond_cont,
|
||||
cond_cat=cond_cat,
|
||||
target_s1=target_s1,
|
||||
n_sec=n_sec,
|
||||
sec_cont=sec_cont,
|
||||
proc_idx=proc_idx,
|
||||
sec_type_idx=sec_type_idx,
|
||||
cond_normalizer=cond_normalizer,
|
||||
target_normalizer=target_normalizer,
|
||||
return (
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
target_s1,
|
||||
n_sec,
|
||||
sec_cont,
|
||||
proc_idx,
|
||||
cond_normalizer,
|
||||
target_normalizer,
|
||||
)
|
||||
|
||||
+29
-7
@@ -31,7 +31,10 @@ import numpy as np
|
||||
import pandas as pd
|
||||
import pyarrow.parquet as pq
|
||||
|
||||
_INSTALL_HINT = "the geometry oracle needs scikit-learn — install it with `uv sync --extra cpu --extra geometry`"
|
||||
_INSTALL_HINT = (
|
||||
"the geometry oracle needs scikit-learn — install it with "
|
||||
"`uv sync --extra cpu --extra geometry`"
|
||||
)
|
||||
|
||||
|
||||
def _require_sklearn():
|
||||
@@ -60,7 +63,9 @@ class _SlabLookup:
|
||||
layer_ids: np.ndarray # (n_segments,) int64, layer_id of each segment
|
||||
radius_max: float # largest transverse radius seen in training data
|
||||
|
||||
def query(self, pos: np.ndarray, margin: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
def query(
|
||||
self, pos: np.ndarray, margin: float
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
|
||||
other = [i for i in range(3) if i != self.axis]
|
||||
z = pos[:, self.axis]
|
||||
radius = np.sqrt(pos[:, other[0]] ** 2 + pos[:, other[1]] ** 2)
|
||||
@@ -70,7 +75,11 @@ class _SlabLookup:
|
||||
material = self.materials[idx]
|
||||
layer_id = self.layer_ids[idx]
|
||||
|
||||
escaped = (z < self.z_edges[0] - margin) | (z > self.z_edges[-1] + margin) | (radius > self.radius_max + margin)
|
||||
escaped = (
|
||||
(z < self.z_edges[0] - margin)
|
||||
| (z > self.z_edges[-1] + margin)
|
||||
| (radius > self.radius_max + margin)
|
||||
)
|
||||
return material, layer_id, escaped
|
||||
|
||||
|
||||
@@ -291,12 +300,16 @@ def _fit_slab_lookup(
|
||||
"""
|
||||
other = [i for i in range(3) if i != axis]
|
||||
z = pos[:, axis].astype(np.float64)
|
||||
radius = np.sqrt(pos[:, other[0]].astype(np.float64) ** 2 + pos[:, other[1]].astype(np.float64) ** 2)
|
||||
radius = np.sqrt(
|
||||
pos[:, other[0]].astype(np.float64) ** 2
|
||||
+ pos[:, other[1]].astype(np.float64) ** 2
|
||||
)
|
||||
|
||||
z_min, z_max = float(z.min()), float(z.max())
|
||||
if z_min == z_max:
|
||||
raise ValueError(
|
||||
"all points share the same depth-axis coordinate — pick a different `depth_axis` or use method='knn'/'svm'"
|
||||
"all points share the same depth-axis coordinate — pick a "
|
||||
"different `depth_axis` or use method='knn'/'svm'"
|
||||
)
|
||||
edges = np.linspace(z_min, z_max, n_bins + 1)
|
||||
bin_idx = np.clip(np.searchsorted(edges, z, side="right") - 1, 0, n_bins - 1)
|
||||
@@ -331,7 +344,12 @@ def _fit_slab_lookup(
|
||||
bin_layer = bin_layer[fill_from]
|
||||
|
||||
# Run-length-encode consecutive bins sharing a label into segments.
|
||||
changed = np.flatnonzero((bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])) + 1
|
||||
changed = (
|
||||
np.flatnonzero(
|
||||
(bin_material[1:] != bin_material[:-1]) | (bin_layer[1:] != bin_layer[:-1])
|
||||
)
|
||||
+ 1
|
||||
)
|
||||
seg_starts = np.concatenate([[0], changed])
|
||||
z_edges = np.concatenate([edges[seg_starts], edges[-1:]])
|
||||
materials = bin_material[seg_starts]
|
||||
@@ -442,7 +460,11 @@ def build_geometry_oracle(
|
||||
# Escape threshold from the reference point spacing. Sample a subset for the
|
||||
# median 2-NN distance (the 1st neighbour of a training point is itself).
|
||||
nn = NearestNeighbors(n_neighbors=2).fit(X)
|
||||
probe = X if len(X) <= 20_000 else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
|
||||
probe = (
|
||||
X
|
||||
if len(X) <= 20_000
|
||||
else X[np.random.default_rng(seed).choice(len(X), 20_000, replace=False)]
|
||||
)
|
||||
d2, _ = nn.kneighbors(probe, n_neighbors=2)
|
||||
median_nn = float(np.median(d2[:, 1]))
|
||||
escape_threshold = escape_factor * median_nn
|
||||
|
||||
+26
-9
@@ -64,10 +64,18 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
|
||||
"G4_CESIUM_IODIDE": MaterialProperties(
|
||||
z_eff=54.0, a_eff=129.904539, density=4.51, x0=1.860288, lambda_int=39.305990
|
||||
),
|
||||
"G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950),
|
||||
"G4_W": MaterialProperties(z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580),
|
||||
"G4_Cu": MaterialProperties(z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940),
|
||||
"G4_Fe": MaterialProperties(z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300),
|
||||
"G4_Pb": MaterialProperties(
|
||||
z_eff=82.0, a_eff=207.216962, density=11.35, x0=0.561253, lambda_int=18.247950
|
||||
),
|
||||
"G4_W": MaterialProperties(
|
||||
z_eff=74.0, a_eff=183.841648, density=19.30, x0=0.350418, lambda_int=10.311580
|
||||
),
|
||||
"G4_Cu": MaterialProperties(
|
||||
z_eff=29.0, a_eff=63.545648, density=8.96, x0=1.435578, lambda_int=15.587940
|
||||
),
|
||||
"G4_Fe": MaterialProperties(
|
||||
z_eff=26.0, a_eff=55.845113, density=7.874, x0=1.757493, lambda_int=16.990300
|
||||
),
|
||||
"G4_BRASS": MaterialProperties(
|
||||
z_eff=30.939130,
|
||||
a_eff=68.500857,
|
||||
@@ -75,7 +83,9 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
|
||||
x0=1.367465,
|
||||
lambda_int=16.947420,
|
||||
),
|
||||
"G4_POLYSTYRENE": MaterialProperties(z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880),
|
||||
"G4_POLYSTYRENE": MaterialProperties(
|
||||
z_eff=3.5, a_eff=6.509339, density=1.06, x0=41.312510, lambda_int=68.749880
|
||||
),
|
||||
"G4_PLASTIC_SC_VINYLTOLUENE": MaterialProperties(
|
||||
z_eff=3.368421,
|
||||
a_eff=6.219791,
|
||||
@@ -98,15 +108,20 @@ MATERIAL_PROPERTIES: dict[str, MaterialProperties] = {
|
||||
x0=30392.070000,
|
||||
lambda_int=71009.500000,
|
||||
),
|
||||
"G4_lAr": MaterialProperties(z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400),
|
||||
"G4_lAr": MaterialProperties(
|
||||
z_eff=18.0, a_eff=39.947692, density=1.396, x0=14.003440, lambda_int=85.706400
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def get_material_properties(name: str, table: dict[str, MaterialProperties] | None = None) -> MaterialProperties:
|
||||
def get_material_properties(
|
||||
name: str, table: dict[str, MaterialProperties] | None = None
|
||||
) -> MaterialProperties:
|
||||
t = MATERIAL_PROPERTIES if table is None else table
|
||||
if name not in t:
|
||||
raise UnknownMaterialError(
|
||||
f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES -- add it (known: {sorted(t)})"
|
||||
f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES "
|
||||
f"-- add it (known: {sorted(t)})"
|
||||
)
|
||||
props = t[name]
|
||||
if any(v is None for v in props):
|
||||
@@ -119,7 +134,9 @@ def get_material_properties(name: str, table: dict[str, MaterialProperties] | No
|
||||
return props
|
||||
|
||||
|
||||
def material_properties_array(names: np.ndarray, table: dict[str, MaterialProperties] | None = None) -> np.ndarray:
|
||||
def material_properties_array(
|
||||
names: np.ndarray, table: dict[str, MaterialProperties] | None = None
|
||||
) -> np.ndarray:
|
||||
"""(N,) str material names -> (N, 5) float32 [z_eff, a_eff, density, x0, lambda_int]."""
|
||||
out = np.array(
|
||||
[get_material_properties(str(m), table) for m in np.asarray(names)],
|
||||
|
||||
@@ -1,114 +0,0 @@
|
||||
"""v0.2 -> v0.3 checkpoint migration: translates a v0.2 checkpoint's flat
|
||||
`model_config`/state dicts into the current nested shape (issues.md Issue 8;
|
||||
see also `giant._migration` and `giant.config.migrate_config`, the sibling
|
||||
config.toml migration surface — issues.md Issue 6)."""
|
||||
|
||||
from giant._migration import V02_FIXED_FACTS, reject_legacy_router_expert_sizing
|
||||
from giant.constants import EMB_DIM, K_MAX
|
||||
|
||||
|
||||
def _migrate_legacy_model_config(model_config: dict) -> dict:
|
||||
"""Translate a v0.2 checkpoint's flat `model_config` (giant/pipeline.py's
|
||||
old shape: `hidden_dim`/`n_blocks`/`emb_dim`/`dropout`/`conditioning`/
|
||||
`router`/`mode`/... all at one level) into the nested
|
||||
`{"pdg_vocab", "mat_vocab", "conditioning", "stage1_model",
|
||||
"stage2_model"}` shape `build_models` expects.
|
||||
|
||||
Sets `stage2_model.n_sec.owner = "stage1"` so the n_sec_head weights a v0.2
|
||||
checkpoint carries on its Stage-1 module keep loading there instead of the new
|
||||
default location (`Stage2OneShot`) — the n_sec head was trained against Stage 1's
|
||||
own `ConditionEncoder` output, so it has to stay attached to Stage 1's module, not
|
||||
just be labeled as such.
|
||||
|
||||
Only the monolithic (non-routed) trunk shape is exercised by the step-2
|
||||
migration test; a routed v0.2 checkpoint still builds correctly here
|
||||
(the router config passes through), but its
|
||||
state dict isn't covered by `migrate_legacy_state_dict` below.
|
||||
"""
|
||||
m = model_config
|
||||
conditioning_mode = m.get("conditioning", "embedding")
|
||||
generator = m.get("mode", "flow")
|
||||
hidden_dim = m.get("hidden_dim", 256)
|
||||
n_blocks = m.get("n_blocks", 6)
|
||||
emb_dim = m.get("emb_dim", EMB_DIM)
|
||||
dropout = m.get("dropout", 0.1)
|
||||
k_max = m.get("k_max", K_MAX)
|
||||
noise_dim = m.get("noise_dim", 64)
|
||||
router_cfg = dict(m.get("router") or {})
|
||||
reject_legacy_router_expert_sizing(router_cfg, source="this checkpoint's model_config.router")
|
||||
router_cfg.setdefault("enabled", False)
|
||||
|
||||
F = V02_FIXED_FACTS
|
||||
cond_n_layers = F["conditioning.particle.n_layers"] # same fact for both axes
|
||||
return {
|
||||
"pdg_vocab": m["pdg_vocab"],
|
||||
"mat_vocab": m["mat_vocab"],
|
||||
"conditioning": {
|
||||
"out_dim": F["conditioning.out_dim"],
|
||||
"share_stages": False,
|
||||
"particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
|
||||
"material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": cond_n_layers},
|
||||
},
|
||||
"stage1_model": {
|
||||
"active": F["stage1_model.active"],
|
||||
"generator": generator,
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_res_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
"flow": {"time_dim": F["stage1_model.flow.time_dim"]},
|
||||
"ddpm": {"time_dim": F["stage1_model.ddpm.time_dim"]},
|
||||
"wgan": {"noise_dim": noise_dim},
|
||||
"router": dict(router_cfg),
|
||||
},
|
||||
"stage2_model": {
|
||||
"active": F["stage2_model.active"],
|
||||
"decoder": F["stage2_model.decoder"],
|
||||
"generator": generator,
|
||||
"hidden_dim": hidden_dim,
|
||||
"n_res_blocks": n_blocks,
|
||||
"dropout": dropout,
|
||||
"k_max": k_max,
|
||||
"context_dim": F["stage2_model.context_dim"],
|
||||
"n_sec": {"mode": "head", "owner": "stage1"},
|
||||
"particle_type": {"target": F["stage2_model.particle_type.target"]},
|
||||
"flow": {"time_dim": F["stage2_model.flow.time_dim"]},
|
||||
"ddpm": {"time_dim": F["stage2_model.ddpm.time_dim"]},
|
||||
"wgan": {"noise_dim": noise_dim},
|
||||
"router": {**router_cfg, "tie_to_stage1": False},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def migrate_legacy_state_dict(old_stage1_sd: dict, old_stage2_sd: dict) -> tuple[dict, dict]:
|
||||
"""Remap a v0.2 checkpoint's (`DenoisingMLP`-or-`WGANGenerator`,
|
||||
`SecondaryDecoder`-or-`WGANSecondaryGenerator`) state dicts onto the new
|
||||
`(Stage1Model, Stage2OneShot)` module structure produced by
|
||||
`build_models(_migrate_legacy_model_config(model_config))`.
|
||||
|
||||
Only the monolithic (non-routed) trunk shape is handled.
|
||||
"""
|
||||
|
||||
def _trunk_prefix(k: str) -> str:
|
||||
if k.startswith(("input_proj.", "blocks.", "out_proj.")):
|
||||
return f"trunk.{k}"
|
||||
return k
|
||||
|
||||
new_stage1 = {}
|
||||
for k, v in old_stage1_sd.items():
|
||||
if k.startswith("n_sec_head."):
|
||||
new_stage1[k] = v # stays top-level (n_sec.owner="stage1")
|
||||
else:
|
||||
new_stage1[_trunk_prefix(k)] = v
|
||||
|
||||
new_stage2 = {}
|
||||
for k, v in old_stage2_sd.items():
|
||||
if k.startswith("cond_enc.base."):
|
||||
new_stage2["cond_enc." + k[len("cond_enc.base.") :]] = v
|
||||
elif k.startswith("cond_enc.stage1_proj."):
|
||||
new_stage2["context_adapter.proj." + k[len("cond_enc.stage1_proj.") :]] = v
|
||||
elif k.startswith("cond_enc.fuse."):
|
||||
new_stage2["fuse." + k[len("cond_enc.fuse.") :]] = v
|
||||
else:
|
||||
new_stage2[_trunk_prefix(k)] = v
|
||||
|
||||
return new_stage1, new_stage2
|
||||
@@ -1,211 +0,0 @@
|
||||
"""Factories: `build_models`/`build_critics` assemble the top-level stage
|
||||
models from a config dict (issues.md Issue 8)."""
|
||||
|
||||
import torch.nn as nn
|
||||
|
||||
from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig
|
||||
from giant.constants import X_DIM
|
||||
from giant.model._legacy import _migrate_legacy_model_config
|
||||
from giant.model.encoders import ConditionEncoder
|
||||
from giant.model.models import (
|
||||
CriticModel,
|
||||
Stage1Model,
|
||||
Stage2Autoregressive,
|
||||
Stage2OneShot,
|
||||
resolve_type_n_classes,
|
||||
stage2_trunk_sec_dim,
|
||||
)
|
||||
from giant.model.routers import Router, _build_router_from_cfg
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Factories
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def build_models(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
"""Construct `{"stage1": ..., "stage2": ...}` from a config dict — either
|
||||
the new nested shape (has a `"stage1_model"` key, plus `"pdg_vocab"`/
|
||||
`"mat_vocab"`/`"conditioning"` at the top level) or a v0.2 checkpoint's
|
||||
flat `model_config`, auto-migrated via `_migrate_legacy_model_config`.
|
||||
|
||||
A stage is `None` in the result when that stage's `active = False`.
|
||||
`stage2_model.router.tie_to_stage1` shares stage 1's literal `Router`
|
||||
instance rather than building a second, independently-parameterized one
|
||||
(v0.2's actual — probably accidental — behaviour: two routers built from
|
||||
one config with no semantic relationship between them).
|
||||
|
||||
`conditioning.share_stages = true` builds one `ConditionEncoder`
|
||||
instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/
|
||||
`Stage2Autoregressive`'s `cond_enc` param), instead of each stage
|
||||
building its own — halving the conditioning parameter count and forcing a
|
||||
common representation. `false` (default) keeps v0.2 behaviour:
|
||||
independent instances with identical config but independent weights.
|
||||
"""
|
||||
cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
|
||||
pdg_vocab = cfg["pdg_vocab"]
|
||||
mat_vocab = cfg["mat_vocab"]
|
||||
conditioning = cfg["conditioning"]
|
||||
particle_cfg = conditioning["particle"]
|
||||
material_cfg = conditioning["material"]
|
||||
particle_conditioning = particle_cfg["type"]
|
||||
conditioning_cfg = ConditioningConfig.from_dict(conditioning)
|
||||
s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
|
||||
s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
|
||||
cond_out_dim = conditioning_cfg.out_dim
|
||||
shared_cond_enc: ConditionEncoder | None = None
|
||||
if conditioning_cfg.share_stages:
|
||||
shared_cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
|
||||
|
||||
result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
|
||||
|
||||
stage1_router: Router | None = None
|
||||
if s1_spec.active:
|
||||
router_cfg = cfg["stage1_model"].get("router") or {}
|
||||
if s1_spec.router.enabled:
|
||||
stage1_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
|
||||
generator = s1_spec.generator
|
||||
# wgan has no time_dim concept (no diffusion/flow time variable) —
|
||||
# matches the pre-dataclass .get("time_dim", 64) fallback, which
|
||||
# always hit its default for a wgan sub-block too.
|
||||
time_dim = getattr(s1_spec, generator).time_dim if generator != "wgan" else 64
|
||||
n_sec_owner = s2_spec.n_sec.owner
|
||||
n_sec_head_k_max = s2_spec.k_max if n_sec_owner == "stage1" else None
|
||||
result["stage1"] = Stage1Model(
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=s1_spec.hidden_dim,
|
||||
n_res_blocks=s1_spec.n_res_blocks,
|
||||
cond_out_dim=cond_out_dim,
|
||||
dropout=s1_spec.dropout,
|
||||
generator=generator,
|
||||
time_dim=time_dim,
|
||||
noise_dim=s1_spec.wgan.noise_dim,
|
||||
router=stage1_router,
|
||||
n_sec_head_k_max=n_sec_head_k_max,
|
||||
cond_enc=shared_cond_enc,
|
||||
)
|
||||
|
||||
if s2_spec.active:
|
||||
decoder = s2_spec.decoder
|
||||
router_cfg = cfg["stage2_model"].get("router") or {}
|
||||
stage2_router: Router | None = None
|
||||
if s2_spec.router.enabled:
|
||||
if s2_spec.router.tie_to_stage1 and stage1_router is not None:
|
||||
stage2_router = stage1_router
|
||||
else:
|
||||
stage2_router = _build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, particle_conditioning)
|
||||
generator = s2_spec.generator
|
||||
# wgan has no time_dim concept — see the matching comment in stage 1
|
||||
# above.
|
||||
time_dim = getattr(s2_spec, generator).time_dim if generator != "wgan" else 64
|
||||
n_sec_owner = s2_spec.n_sec.owner
|
||||
k_max = s2_spec.k_max
|
||||
particle_type_cfg = s2_spec.particle_type.to_dict()
|
||||
|
||||
if decoder == "autoregressive":
|
||||
ar_cfg = s2_spec.autoregressive
|
||||
result["stage2"] = Stage2Autoregressive(
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=s2_spec.hidden_dim,
|
||||
n_res_blocks=s2_spec.n_res_blocks,
|
||||
cond_out_dim=cond_out_dim,
|
||||
context_dim=s2_spec.context_dim,
|
||||
dropout=s2_spec.dropout,
|
||||
generator=generator,
|
||||
time_dim=time_dim,
|
||||
noise_dim=s2_spec.wgan.noise_dim,
|
||||
k_max=k_max,
|
||||
router=stage2_router,
|
||||
build_n_sec_head=n_sec_owner != "stage1",
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
history=ar_cfg.history,
|
||||
attn_n_heads=ar_cfg.attn_n_heads,
|
||||
attn_n_layers=ar_cfg.attn_n_layers,
|
||||
cond_enc=shared_cond_enc,
|
||||
)
|
||||
else:
|
||||
sec_dim = stage2_trunk_sec_dim(
|
||||
particle_type_cfg, generator, k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
|
||||
)
|
||||
result["stage2"] = Stage2OneShot(
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=s2_spec.hidden_dim,
|
||||
n_res_blocks=s2_spec.n_res_blocks,
|
||||
cond_out_dim=cond_out_dim,
|
||||
context_dim=s2_spec.context_dim,
|
||||
sec_dim=sec_dim,
|
||||
dropout=s2_spec.dropout,
|
||||
generator=generator,
|
||||
time_dim=time_dim,
|
||||
noise_dim=s2_spec.wgan.noise_dim,
|
||||
k_max=k_max,
|
||||
router=stage2_router,
|
||||
build_n_sec_head=n_sec_owner != "stage1",
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
cond_enc=shared_cond_enc,
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def build_critics(model_config: dict) -> dict[str, nn.Module | None]:
|
||||
"""Construct `{"stage1": ..., "stage2": ...}` critics for `generator =
|
||||
"wgan"` training. Training-only — never persisted for inference the way
|
||||
`build_models`'s pair is. `None` for a stage that's inactive or not
|
||||
WGAN."""
|
||||
cfg = model_config if "stage1_model" in model_config else _migrate_legacy_model_config(model_config)
|
||||
pdg_vocab = cfg["pdg_vocab"]
|
||||
mat_vocab = cfg["mat_vocab"]
|
||||
conditioning = cfg["conditioning"]
|
||||
particle_cfg = conditioning["particle"]
|
||||
material_cfg = conditioning["material"]
|
||||
conditioning_cfg = ConditioningConfig.from_dict(conditioning)
|
||||
cond_out_dim = conditioning_cfg.out_dim
|
||||
s1_spec = Stage1ModelConfig.from_dict(cfg["stage1_model"])
|
||||
s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
|
||||
|
||||
result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None}
|
||||
|
||||
if s1_spec.active and s1_spec.generator == "wgan":
|
||||
result["stage1"] = CriticModel(
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
in_dim=X_DIM,
|
||||
hidden_dim=s1_spec.wgan.critic_hidden_dim or s1_spec.hidden_dim,
|
||||
n_res_blocks=s1_spec.wgan.critic_n_res_blocks or s1_spec.n_res_blocks,
|
||||
cond_out_dim=cond_out_dim,
|
||||
dropout=s1_spec.dropout,
|
||||
stage="stage1",
|
||||
)
|
||||
|
||||
if s2_spec.active and s2_spec.generator == "wgan":
|
||||
k_max = s2_spec.k_max
|
||||
particle_type_cfg = s2_spec.particle_type.to_dict()
|
||||
in_dim = stage2_trunk_sec_dim(
|
||||
particle_type_cfg, "wgan", k_max, resolve_type_n_classes(particle_type_cfg, particle_cfg["emb_dim"])
|
||||
)
|
||||
result["stage2"] = CriticModel(
|
||||
pdg_vocab=pdg_vocab,
|
||||
mat_vocab=mat_vocab,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
in_dim=in_dim,
|
||||
hidden_dim=s2_spec.wgan.critic_hidden_dim or s2_spec.hidden_dim,
|
||||
n_res_blocks=s2_spec.wgan.critic_n_res_blocks or s2_spec.n_res_blocks,
|
||||
cond_out_dim=cond_out_dim,
|
||||
dropout=s2_spec.dropout,
|
||||
stage="stage2",
|
||||
context_dim=s2_spec.context_dim,
|
||||
)
|
||||
|
||||
return result
|
||||
@@ -1,122 +0,0 @@
|
||||
"""Conditioning encoder — fuses continuous conditioning with particle/material
|
||||
identity (issues.md Issue 8)."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM
|
||||
from giant.model.layers import _make_axis_mlp
|
||||
|
||||
|
||||
def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]:
|
||||
"""`cond_cat` column indices for each axis's top-N-onehot index, or
|
||||
`None` if that axis isn't `"onehot"`.
|
||||
|
||||
Columns 0/1 are always the dense pdg/material vocab index. The particle
|
||||
top-N column (if any) comes next, then the material top-N column (if
|
||||
any) — `giant.data.transforms.build_cond_features`/`build_features`
|
||||
append columns in this same order, so the two sides must never drift
|
||||
apart.
|
||||
"""
|
||||
col = 2
|
||||
particle_col = None
|
||||
if particle_type == "onehot":
|
||||
particle_col = col
|
||||
col += 1
|
||||
material_col = None
|
||||
if material_type == "onehot":
|
||||
material_col = col
|
||||
col += 1
|
||||
return particle_col, material_col
|
||||
|
||||
|
||||
class ConditionEncoder(nn.Module):
|
||||
"""Fuses continuous conditioning with particle/material identity.
|
||||
|
||||
The particle and material axes are configured independently
|
||||
(`particle_cfg`/`material_cfg`, each `{"type", "emb_dim", "n_layers"}`)
|
||||
and may mix freely, e.g. material "physical" with particle "embedding".
|
||||
Three modes per axis:
|
||||
- "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s
|
||||
dense training-vocab index. Memorizes the training menu.
|
||||
- "physical": an `n_layers`-deep MLP over the axis's raw physical
|
||||
properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see
|
||||
giant.data.transforms.build_features), computable for any PDG code /
|
||||
material name rather than only ones seen in training.
|
||||
- "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other
|
||||
class map (`giant.data.loader.build_topn_map_from_files`/
|
||||
`build_pdg_topn_map_from_files`), read from `cond_cat`'s extra
|
||||
top-N-index column(s) — see `_cat_col_layout`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
particle_cfg: dict,
|
||||
material_cfg: dict,
|
||||
cont_dim: int = COND_DIM,
|
||||
out_dim: int = 128,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.particle_cfg = dict(particle_cfg)
|
||||
self.material_cfg = dict(material_cfg)
|
||||
self._particle_topn_col, self._material_topn_col = cat_col_layout(particle_cfg["type"], material_cfg["type"])
|
||||
|
||||
p_type = particle_cfg["type"]
|
||||
p_emb_dim = particle_cfg["emb_dim"]
|
||||
if p_type == "embedding":
|
||||
self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim)
|
||||
elif p_type == "physical":
|
||||
self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1))
|
||||
elif p_type != "onehot":
|
||||
raise ValueError(f"unknown conditioning.particle.type {p_type!r}")
|
||||
|
||||
m_type = material_cfg["type"]
|
||||
m_emb_dim = material_cfg["emb_dim"]
|
||||
if m_type == "embedding":
|
||||
self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim)
|
||||
elif m_type == "physical":
|
||||
self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1))
|
||||
elif m_type != "onehot":
|
||||
raise ValueError(f"unknown conditioning.material.type {m_type!r}")
|
||||
|
||||
in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(in_dim, out_dim),
|
||||
nn.SiLU(),
|
||||
nn.Linear(out_dim, out_dim),
|
||||
)
|
||||
|
||||
def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
|
||||
p_type = self.particle_cfg["type"]
|
||||
if p_type == "embedding":
|
||||
return self.pdg_emb(cond_cat[:, 0])
|
||||
if p_type == "physical":
|
||||
particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM]
|
||||
return self.particle_mlp(particle_phys)
|
||||
assert self._particle_topn_col is not None
|
||||
return F.one_hot(
|
||||
cond_cat[:, self._particle_topn_col],
|
||||
num_classes=self.particle_cfg["emb_dim"],
|
||||
).float()
|
||||
|
||||
def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor):
|
||||
m_type = self.material_cfg["type"]
|
||||
if m_type == "embedding":
|
||||
return self.mat_emb(cond_cat[:, 1])
|
||||
if m_type == "physical":
|
||||
material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :]
|
||||
return self.material_mlp(material_phys)
|
||||
assert self._material_topn_col is not None
|
||||
return F.one_hot(
|
||||
cond_cat[:, self._material_topn_col],
|
||||
num_classes=self.material_cfg["emb_dim"],
|
||||
).float()
|
||||
|
||||
def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
pdg_e = self._particle_embed(cond_cont, cond_cat)
|
||||
mat_e = self._material_embed(cond_cont, cond_cat)
|
||||
x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1)
|
||||
return self.mlp(x)
|
||||
@@ -1,154 +0,0 @@
|
||||
"""History encoders — stage-2 autoregressive only. Self-contained, no
|
||||
dependency on any other `giant.model` submodule (issues.md Issue 8)."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class HistoryEncoder(nn.Module):
|
||||
"""Interface for stage-2 autoregressive per-token history summaries:
|
||||
`forward(feat, has_prev) -> (B, K, out_dim)`, a single parallel pass over
|
||||
a full (teacher-forced) token sequence — used by training. `MarkovHistory`
|
||||
and `AttentionHistory` are the two implementations. Inference
|
||||
(`giant/sample.py`) generates one token at a
|
||||
time and cannot afford `forward`'s per-step cost to be O(K) (attention
|
||||
would then be O(K^2) over a rollout's k_max loop); encoders that need
|
||||
incremental state for that path additionally implement `init_cache`/
|
||||
`step` (see `AttentionHistory`) — `MarkovHistory` doesn't need to, since
|
||||
its per-step cost is already O(1) (it only ever looks at the previous
|
||||
token, not the full prefix)."""
|
||||
|
||||
def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class MarkovHistory(HistoryEncoder):
|
||||
"""Summarizes the previous secondary's own `(energy_fraction, direction,
|
||||
type_representation)` through one small MLP — the "markov" history:
|
||||
token i+1 only ever sees token i plus the running scalars
|
||||
(`remaining_frac`/`slot_idx`, fused in separately by
|
||||
`Stage2Autoregressive._token_cond`), not the full prefix.
|
||||
|
||||
At slot 0 (`has_prev` False) substitutes a learned start vector rather
|
||||
than zeros — a reasonable default.
|
||||
"""
|
||||
|
||||
def __init__(self, in_dim: int, out_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.start = nn.Parameter(torch.zeros(in_dim))
|
||||
self.mlp = nn.Sequential(nn.Linear(in_dim, out_dim), nn.SiLU())
|
||||
|
||||
def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
|
||||
start = self.start.view(1, 1, -1).expand_as(feat)
|
||||
x = torch.where(has_prev.unsqueeze(-1), feat, start)
|
||||
return self.mlp(x)
|
||||
|
||||
|
||||
class _CausalAttnBlock(nn.Module):
|
||||
"""One pre-norm causal self-attention block for `AttentionHistory`.
|
||||
|
||||
Exposes two forward paths that must agree (see
|
||||
`test_attention_history_step_matches_forward` in `tests/test_network.py`):
|
||||
`forward` — the full-sequence, causally-masked pass used for training;
|
||||
`step` — an incremental pass for inference, given the *pre-attention*
|
||||
normalized hidden states of every earlier position (`kv_cache`, i.e.
|
||||
`norm1(x)` for positions `< t`, not `x` itself). Caching `norm1(x)` rather
|
||||
than raw `x` is what makes `step` correct: this block's attention needs
|
||||
exactly that quantity as keys/values, and `LayerNorm` has no cross-position
|
||||
interaction, so recomputing it per position instead of caching it would
|
||||
still be correct but pointlessly repeat work. The *next* block's cache is
|
||||
built from a different sequence (this block's output), so each block owns
|
||||
an independent cache entry.
|
||||
"""
|
||||
|
||||
def __init__(self, dim: int, n_heads: int, dropout: float = 0.0) -> None:
|
||||
super().__init__()
|
||||
self.norm1 = nn.LayerNorm(dim)
|
||||
self.attn = nn.MultiheadAttention(dim, n_heads, dropout=dropout, batch_first=True)
|
||||
self.norm2 = nn.LayerNorm(dim)
|
||||
self.mlp = nn.Sequential(nn.Linear(dim, 4 * dim), nn.GELU(), nn.Linear(4 * dim, dim))
|
||||
|
||||
def forward(self, x: torch.Tensor, causal_mask: torch.Tensor) -> torch.Tensor:
|
||||
h = self.norm1(x)
|
||||
attn_out, _ = self.attn(h, h, h, attn_mask=causal_mask, need_weights=False)
|
||||
x = x + attn_out
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
return x
|
||||
|
||||
def step(self, x_new: torch.Tensor, kv_cache: torch.Tensor | None) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""`x_new`: `(B, 1, dim)`, this position's input. `kv_cache`: `None`
|
||||
(first position) or `(B, T, dim)` — `norm1(x)` of every earlier
|
||||
position at this same block. Returns `(out, new_kv_cache)`, `out`
|
||||
being this position's block output (`(B, 1, dim)`, to feed the next
|
||||
block's `step`), `new_kv_cache` the same cache extended by this
|
||||
position (to reuse at this block's *next* `step` call)."""
|
||||
h_new = self.norm1(x_new)
|
||||
kv = h_new if kv_cache is None else torch.cat([kv_cache, h_new], dim=1)
|
||||
attn_out, _ = self.attn(h_new, kv, kv, need_weights=False)
|
||||
x = x_new + attn_out
|
||||
x = x + self.mlp(self.norm2(x))
|
||||
return x, kv
|
||||
|
||||
|
||||
class AttentionHistory(HistoryEncoder):
|
||||
"""Causal self-attention over the emitted-token prefix — the more
|
||||
expressive alternative to `MarkovHistory`'s fixed previous-token-only
|
||||
summary. `feat`/`has_prev`
|
||||
follow the same shifted-by-one convention `MarkovHistory` and
|
||||
`Stage2Autoregressive._token_cond` use: `feat[:, i]` is token `i - 1`'s
|
||||
own `(energy_fraction, direction, type_representation)`, with a learned
|
||||
start vector substituted at `has_prev == False` positions (only slot 0 in
|
||||
practice — see `giant.training.stage2_inputs._ar_has_prev`). Causal masking then makes
|
||||
position `i`'s output a function of `feat[:, 1:i+1]` — i.e. tokens
|
||||
`0..i-1` — exactly the prefix available when predicting token `i`.
|
||||
|
||||
`forward` is the parallel training path (one pass over the whole
|
||||
teacher-forced sequence); `init_cache`/`step` are the incremental
|
||||
inference path `giant/sample.py` uses, one new token per call, to avoid
|
||||
re-encoding the whole prefix from scratch every slot — `step` must be
|
||||
called exactly once per slot (its cache-extension is not idempotent),
|
||||
so a slot's output must be reused for
|
||||
every model call within that slot (`forward`'s ODE substeps, or a separate
|
||||
`predict_type` call) rather than re-derived — see
|
||||
`Stage2Autoregressive.history_step`.
|
||||
"""
|
||||
|
||||
def __init__(self, in_dim: int, out_dim: int, n_heads: int = 4, n_layers: int = 2) -> None:
|
||||
super().__init__()
|
||||
self.start = nn.Parameter(torch.zeros(in_dim))
|
||||
self.in_proj = nn.Linear(in_dim, out_dim)
|
||||
self.blocks = nn.ModuleList([_CausalAttnBlock(out_dim, n_heads) for _ in range(n_layers)])
|
||||
|
||||
def _embed(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
|
||||
start = self.start.view(1, 1, -1).expand_as(feat)
|
||||
x = torch.where(has_prev.unsqueeze(-1), feat, start)
|
||||
return self.in_proj(x)
|
||||
|
||||
def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
|
||||
B, K, _ = feat.shape
|
||||
x = self._embed(feat, has_prev)
|
||||
mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
|
||||
for block in self.blocks:
|
||||
x = block(x, mask)
|
||||
return x
|
||||
|
||||
def init_cache(self) -> list[torch.Tensor | None]:
|
||||
return [None for _ in self.blocks]
|
||||
|
||||
def step(
|
||||
self,
|
||||
token_feat: torch.Tensor,
|
||||
has_prev: torch.Tensor,
|
||||
cache: list[torch.Tensor | None],
|
||||
) -> tuple[torch.Tensor, list[torch.Tensor | None]]:
|
||||
"""`token_feat`/`has_prev`: `(B, 1, in_dim)`/`(B, 1)` — the newest
|
||||
token's own features (what would be `feat[:, k]` in `forward`).
|
||||
Advances every block's cache by this position and returns this
|
||||
position's output (`(B, 1, out_dim)`, the correct history summary for
|
||||
the NEXT slot) plus the updated cache."""
|
||||
x = self._embed(token_feat, has_prev)
|
||||
new_cache: list[torch.Tensor | None] = []
|
||||
for block, kv in zip(self.blocks, cache):
|
||||
x, kv_new = block.step(x, kv)
|
||||
new_cache.append(kv_new)
|
||||
return x, new_cache
|
||||
@@ -1,76 +0,0 @@
|
||||
"""Small stateless-ish building blocks shared across encoders/trunks/models —
|
||||
no dependency on any other `giant.model` submodule (issues.md Issue 8)."""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class SinusoidalEmbedding(nn.Module):
|
||||
def __init__(self, dim: int) -> None:
|
||||
super().__init__()
|
||||
assert dim % 2 == 0, "dim must be even"
|
||||
half = dim // 2
|
||||
freqs = torch.exp(-math.log(10000) * torch.arange(half, dtype=torch.float32) / max(half - 1, 1))
|
||||
self.register_buffer("freqs", freqs)
|
||||
|
||||
def forward(self, t: torch.Tensor) -> torch.Tensor:
|
||||
t = t.reshape(-1, 1).float()
|
||||
args = t * self.freqs.unsqueeze(0) # (B, half)
|
||||
return torch.cat([args.sin(), args.cos()], dim=-1) # (B, dim)
|
||||
|
||||
|
||||
def _make_axis_mlp(in_dim: int, emb_dim: int, n_layers: int) -> nn.Sequential:
|
||||
"""`n_layers`-deep MLP producing an `emb_dim`-wide vector from `in_dim`
|
||||
physical properties (`conditioning.{particle,material}.n_layers`).
|
||||
|
||||
`n_layers=1` (the v0.3.0 default): a single `Linear`, no hidden
|
||||
activation. `n_layers=2` reproduces v0.2's hardcoded depth exactly —
|
||||
`Linear -> SiLU -> Linear` — which is why `migrate_config` back-fills
|
||||
`n_layers=2` for migrated configs rather than the v0.3 default of 1 (see
|
||||
its docstring).
|
||||
"""
|
||||
if n_layers < 1:
|
||||
raise ValueError(f"n_layers must be >= 1, got {n_layers}")
|
||||
if n_layers == 1:
|
||||
return nn.Sequential(nn.Linear(in_dim, emb_dim))
|
||||
layers: list[nn.Module] = [nn.Linear(in_dim, emb_dim), nn.SiLU()]
|
||||
for _ in range(n_layers - 2):
|
||||
layers += [nn.Linear(emb_dim, emb_dim), nn.SiLU()]
|
||||
layers.append(nn.Linear(emb_dim, emb_dim))
|
||||
return nn.Sequential(*layers)
|
||||
|
||||
|
||||
class ContextAdapter(nn.Module):
|
||||
"""Projects a stage's outcome (e.g. Stage 1's 9D target) down to a
|
||||
fixed-width context vector for a downstream stage's conditioning —
|
||||
`stage2_model.context_dim`. Was `SecondaryConditionEncoder.stage1_proj`
|
||||
(+ its `tanh`) in v0.2; pulled out as its own module in v0.3.0 since
|
||||
`SecondaryConditionEncoder` as a wrapper class disappears."""
|
||||
|
||||
def __init__(self, in_dim: int, context_dim: int) -> None:
|
||||
super().__init__()
|
||||
self.proj = nn.Linear(in_dim, context_dim)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return torch.tanh(self.proj(x))
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
def __init__(self, dim: int, cond_dim: int, dropout: float = 0.0) -> None:
|
||||
super().__init__()
|
||||
self.norm = nn.LayerNorm(dim)
|
||||
self.linear1 = nn.Linear(dim, dim)
|
||||
self.cond_proj = nn.Linear(cond_dim, dim, bias=False)
|
||||
self.act = nn.SiLU()
|
||||
self.dropout = nn.Dropout(dropout)
|
||||
self.linear2 = nn.Linear(dim, dim)
|
||||
|
||||
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
||||
h = self.norm(x)
|
||||
h = self.linear1(h) + self.cond_proj(cond)
|
||||
h = self.act(h)
|
||||
h = self.dropout(h)
|
||||
h = self.linear2(h)
|
||||
return x + h
|
||||
@@ -1,593 +0,0 @@
|
||||
"""Top-level stage models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`,
|
||||
`CriticModel` — composed from encoders/trunks/history (issues.md Issue 8)."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
|
||||
from giant.model.encoders import ConditionEncoder
|
||||
from giant.model.history import AttentionHistory, HistoryEncoder, MarkovHistory
|
||||
from giant.model.layers import ContextAdapter, ResBlock, SinusoidalEmbedding
|
||||
from giant.model.routers import Router
|
||||
from giant.model.trunks import build_trunk
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Stage models
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def resolve_type_n_classes(particle_type_cfg: dict, particle_emb_dim: int) -> int:
|
||||
"""Effective width fed to `stage2_type_dim`/`stage2_trunk_sec_dim` in
|
||||
place of a bare `conditioning.particle.emb_dim` read. Under
|
||||
`target = "onehot"` this is `stage2_model.particle_type.n_classes` (0 =
|
||||
inherit `conditioning.particle.emb_dim`) — see gitea #29, which decoupled
|
||||
the secondary-species vocabulary size from the unrelated
|
||||
physical-conditioning MLP's output width. Under `target = "embedding"`
|
||||
(or `"physical"`, which ignores this value entirely) `n_classes` doesn't
|
||||
apply — the width stays `conditioning.particle.emb_dim`, the embedding
|
||||
table's own dimensionality (`validate_config` requires
|
||||
`conditioning.particle.type = "embedding"` here)."""
|
||||
if particle_type_cfg.get("target", "physical") == "onehot":
|
||||
return particle_type_cfg.get("n_classes", 0) or particle_emb_dim
|
||||
return particle_emb_dim
|
||||
|
||||
|
||||
def stage2_type_dim(particle_type_cfg: dict, emb_dim: int) -> int:
|
||||
"""Width of a single secondary slot's type slice —
|
||||
`PARTICLE_PHYS_DIM` (log_mass, charge) for `target = "physical"`, else
|
||||
`emb_dim` (both `"onehot"` class logits and `"embedding"` vectors are
|
||||
this many classes/dims wide — callers resolve `emb_dim` via
|
||||
`resolve_type_n_classes` first)."""
|
||||
target = particle_type_cfg.get("target", "physical")
|
||||
return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
|
||||
|
||||
|
||||
def stage2_trunk_sec_dim(particle_type_cfg: dict, generator: str, k_max: int, emb_dim: int) -> int:
|
||||
"""`Stage2OneShot`'s trunk output width.
|
||||
|
||||
`target = "physical"` is untouched from v0.2/today:
|
||||
`k_max * SEC_SLOT_DIM`, the type slice folded into the same
|
||||
flow-matched/WGAN vector as the continuous stick/dir slots.
|
||||
|
||||
`target` in `("onehot", "embedding")`: under `generator == "wgan"` the
|
||||
type slice is still folded in (adversarial for onehot via ST-Gumbel,
|
||||
already-continuous for embedding), just `emb_dim` wide instead of
|
||||
`PARTICLE_PHYS_DIM` wide: `k_max * (CONT_SLOT_DIM + emb_dim)`. Under
|
||||
`generator in ("flow", "ddpm")` the type slice isn't part of this vector
|
||||
at all — it's `Stage2OneShot.type_head`'s job instead — so the trunk
|
||||
only covers `k_max * CONT_SLOT_DIM`.
|
||||
"""
|
||||
target = particle_type_cfg.get("target", "physical")
|
||||
if target == "physical":
|
||||
return k_max * SEC_SLOT_DIM
|
||||
if generator == "wgan":
|
||||
return k_max * (CONT_SLOT_DIM + emb_dim)
|
||||
return k_max * CONT_SLOT_DIM
|
||||
|
||||
|
||||
class Stage1Model(nn.Module):
|
||||
"""Predicts the 9D primary post-step vector. No `n_sec_head` — fresh runs
|
||||
move it to stage 2, except for a migrated v0.2 checkpoint
|
||||
(`n_sec_head_k_max` given), where it stays attached here
|
||||
since that's where its weights live and what conditioning it was trained
|
||||
against (see `_migrate_legacy_model_config`).
|
||||
|
||||
`cond_enc`, if given, is used in place of building a fresh
|
||||
`ConditionEncoder` — `conditioning.share_stages = true`: `build_models`
|
||||
constructs one shared instance and passes it to both stages, halving the
|
||||
conditioning parameter count and forcing a common representation."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
particle_cfg: dict,
|
||||
material_cfg: dict,
|
||||
hidden_dim: int = 256,
|
||||
n_res_blocks: int = 6,
|
||||
cond_out_dim: int = 128,
|
||||
x_dim: int = X_DIM,
|
||||
dropout: float = 0.0,
|
||||
generator: str = "flow",
|
||||
time_dim: int = 64,
|
||||
noise_dim: int = 64,
|
||||
router: Router | None = None,
|
||||
n_sec_head_k_max: int | None = None,
|
||||
cond_enc: ConditionEncoder | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.generator_kind = generator
|
||||
self.noise_dim = noise_dim
|
||||
self.cond_enc = (
|
||||
cond_enc
|
||||
if cond_enc is not None
|
||||
else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
|
||||
)
|
||||
has_time = generator in ("flow", "ddpm")
|
||||
self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
|
||||
merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
|
||||
in_dim = noise_dim if generator == "wgan" else x_dim
|
||||
self.trunk = build_trunk(router, in_dim, x_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
|
||||
self.n_sec_head = None
|
||||
if n_sec_head_k_max is not None:
|
||||
self.n_sec_head = nn.Sequential(
|
||||
nn.Linear(cond_out_dim, hidden_dim // 2),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_dim // 2, n_sec_head_k_max + 1),
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_t: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
t: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
c_emb = self.cond_enc(cond_cont, cond_cat)
|
||||
cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
|
||||
return self.trunk(x_t, cond, cond_cont, cond_cat)
|
||||
|
||||
def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
"""Return n_sec logits (B, K_MAX+1) from conditioning alone. Only
|
||||
valid on a migrated v0.2 checkpoint's Stage1Model — fresh v0.3.0
|
||||
configs predict n_sec from Stage2OneShot instead."""
|
||||
if self.n_sec_head is None:
|
||||
raise RuntimeError(
|
||||
"this Stage1Model has no n_sec_head — n_sec now lives on "
|
||||
"stage 2 by default; this method only exists "
|
||||
"for a migrated v0.2 checkpoint (n_sec.owner='stage1')"
|
||||
)
|
||||
c_emb = self.cond_enc(cond_cont, cond_cat)
|
||||
return self.n_sec_head(c_emb)
|
||||
|
||||
|
||||
class Stage2OneShot(nn.Module):
|
||||
"""Predicts all `k_max` secondary slots simultaneously — v0.2 behaviour,
|
||||
reproduced exactly (`decoder = "autoregressive"` is `Stage2Autoregressive`,
|
||||
step 4/5, not implemented yet).
|
||||
|
||||
Owns `n_sec_head` by default unless `build_n_sec_head=False`
|
||||
(a migrated v0.2 checkpoint, whose n_sec_head instead attaches to
|
||||
Stage1Model — see `_migrate_legacy_model_config`).
|
||||
|
||||
`particle_type_cfg["target"]` (default `"physical"`) selects the
|
||||
secondary-type mechanism: `"physical"` keeps the type slice folded into
|
||||
the trunk's own
|
||||
flow-matched/WGAN output, unchanged from v0.2 (`sec_dim` — computed by
|
||||
the caller via `stage2_trunk_sec_dim` — already reflects this). Under
|
||||
`"onehot"`/`"embedding"` with `generator in ("flow", "ddpm")`, the type
|
||||
slice is predicted by a separate `type_head` instead (same shape pattern
|
||||
as `n_sec_head`) — `sec_dim` then covers only the continuous
|
||||
stick/dir slots, `type_head` covers `k_max * emb_dim` type logits/vectors.
|
||||
Under `generator == "wgan"` the type slice stays folded into `sec_dim`
|
||||
(just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is
|
||||
unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation.
|
||||
|
||||
`cond_enc`, if given, is used in place of building a fresh
|
||||
`ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
particle_cfg: dict,
|
||||
material_cfg: dict,
|
||||
hidden_dim: int = 256,
|
||||
n_res_blocks: int = 6,
|
||||
cond_out_dim: int = 128,
|
||||
context_dim: int = 64,
|
||||
sec_dim: int = SEC_DIM,
|
||||
x_dim: int = X_DIM,
|
||||
dropout: float = 0.0,
|
||||
generator: str = "wgan",
|
||||
time_dim: int = 64,
|
||||
noise_dim: int = 64,
|
||||
k_max: int = K_MAX,
|
||||
router: Router | None = None,
|
||||
build_n_sec_head: bool = True,
|
||||
particle_type_cfg: dict | None = None,
|
||||
cond_enc: ConditionEncoder | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.generator_kind = generator
|
||||
self.noise_dim = noise_dim
|
||||
self.k_max = k_max
|
||||
self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
|
||||
self.type_dim = stage2_type_dim(
|
||||
self.particle_type_cfg, resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
|
||||
)
|
||||
self.cond_enc = (
|
||||
cond_enc
|
||||
if cond_enc is not None
|
||||
else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
|
||||
)
|
||||
self.context_adapter = ContextAdapter(x_dim, context_dim)
|
||||
self.fuse = nn.Sequential(
|
||||
nn.Linear(cond_out_dim + context_dim, cond_out_dim),
|
||||
nn.SiLU(),
|
||||
)
|
||||
has_time = generator in ("flow", "ddpm")
|
||||
self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
|
||||
merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
|
||||
in_dim = noise_dim if generator == "wgan" else sec_dim
|
||||
self.trunk = build_trunk(router, in_dim, sec_dim, hidden_dim, n_res_blocks, merged_cond_dim, dropout)
|
||||
self.n_sec_head = None
|
||||
if build_n_sec_head:
|
||||
self.n_sec_head = nn.Sequential(
|
||||
nn.Linear(cond_out_dim, hidden_dim // 2),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_dim // 2, k_max + 1),
|
||||
)
|
||||
self.type_head = None
|
||||
target = self.particle_type_cfg.get("target", "physical")
|
||||
if target != "physical" and generator in ("flow", "ddpm"):
|
||||
emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
|
||||
self.type_head = nn.Sequential(
|
||||
nn.Linear(cond_out_dim, hidden_dim // 2),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_dim // 2, k_max * emb_dim),
|
||||
)
|
||||
self._type_k_max = k_max
|
||||
self._type_emb_dim = emb_dim
|
||||
|
||||
def _cond_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
|
||||
base = self.cond_enc(cond_cont, cond_cat)
|
||||
ctx = self.context_adapter(stage1_out)
|
||||
return self.fuse(torch.cat([base, ctx], dim=-1))
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_t: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
t: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
|
||||
cond = torch.cat([self.time_emb(t), c_emb], dim=-1) if self.time_emb is not None else c_emb
|
||||
return self.trunk(x_t, cond, cond_cont, cond_cat)
|
||||
|
||||
def predict_n_sec(
|
||||
self,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
if self.n_sec_head is None:
|
||||
raise RuntimeError(
|
||||
"this Stage2OneShot has no n_sec_head — it belongs to a "
|
||||
"migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
|
||||
"stage1.predict_n_sec(cond_cont, cond_cat) instead"
|
||||
)
|
||||
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
|
||||
return self.n_sec_head(c_emb)
|
||||
|
||||
def predict_type(
|
||||
self,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""`(B, k_max, emb_dim)` per-slot type logits (`target="onehot"`) or
|
||||
vectors (`target="embedding"`) — only under `generator in ("flow",
|
||||
"ddpm")`; `generator == "wgan"` folds the type slice into `forward`'s
|
||||
own output instead (see class docstring)."""
|
||||
if self.type_head is None:
|
||||
raise RuntimeError(
|
||||
"this Stage2OneShot has no type_head — either "
|
||||
"particle_type.target='physical' (the type slice is part of "
|
||||
"forward()'s own output) or generator='wgan' (the WGAN "
|
||||
"trainer reads the type slice out of forward()'s output "
|
||||
"directly instead)"
|
||||
)
|
||||
c_emb = self._cond_embed(cond_cont, cond_cat, stage1_out)
|
||||
return self.type_head(c_emb).view(-1, self._type_k_max, self._type_emb_dim)
|
||||
|
||||
|
||||
class Stage2Autoregressive(nn.Module):
|
||||
"""Emits secondaries one at a time in descending-energy order, instead
|
||||
of `Stage2OneShot`'s simultaneous
|
||||
k_max-slot prediction. `history` selects `MarkovHistory` or
|
||||
`AttentionHistory` (`attn_n_heads`/`attn_n_layers`, attention only).
|
||||
`teacher_forcing` handling lives entirely in the trainer
|
||||
(`giant/train.py`), since it only affects how training inputs are
|
||||
assembled, not this module's architecture.
|
||||
|
||||
Under teacher forcing every token's conditioning is built from ground
|
||||
truth, so a whole K-token sequence trains in one parallel batched pass:
|
||||
`forward` accepts `(B, K, ...)` tensors for an arbitrary K (not hardcoded
|
||||
to `k_max`) — this also means a future one-token-at-a-time inference loop
|
||||
(`K=1` per call, step 6) needs no interface change here.
|
||||
|
||||
Two independent conditioning paths, mirroring `Stage2OneShot`'s
|
||||
`_cond_embed` but split in two: `_base_cond` (`cond_enc` +
|
||||
`context_adapter` only) feeds `predict_n_sec`, since n_sec doesn't depend
|
||||
on token position; `_token_cond` additionally fuses in the history
|
||||
encoding and two running scalars (remaining energy-budget fraction,
|
||||
normalized slot index), and feeds `forward`/`predict_type`/the trunk.
|
||||
|
||||
`cond_enc`, if given, is used in place of building a fresh
|
||||
`ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`).
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
particle_cfg: dict,
|
||||
material_cfg: dict,
|
||||
hidden_dim: int = 256,
|
||||
n_res_blocks: int = 6,
|
||||
cond_out_dim: int = 128,
|
||||
context_dim: int = 64,
|
||||
x_dim: int = X_DIM,
|
||||
dropout: float = 0.0,
|
||||
generator: str = "wgan",
|
||||
time_dim: int = 64,
|
||||
noise_dim: int = 64,
|
||||
k_max: int = K_MAX,
|
||||
router: Router | None = None,
|
||||
build_n_sec_head: bool = True,
|
||||
particle_type_cfg: dict | None = None,
|
||||
history: str = "markov",
|
||||
attn_n_heads: int = 4,
|
||||
attn_n_layers: int = 2,
|
||||
cond_enc: ConditionEncoder | None = None,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if history not in ("markov", "attention"):
|
||||
raise ValueError(f"stage2_model.autoregressive.history={history!r} — must be 'markov' or 'attention'")
|
||||
self.history_kind = history
|
||||
self.generator_kind = generator
|
||||
self.noise_dim = noise_dim
|
||||
self.k_max = k_max
|
||||
self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"})
|
||||
emb_dim = resolve_type_n_classes(self.particle_type_cfg, particle_cfg["emb_dim"])
|
||||
self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim)
|
||||
|
||||
self.cond_enc = (
|
||||
cond_enc
|
||||
if cond_enc is not None
|
||||
else ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
|
||||
)
|
||||
self.context_adapter = ContextAdapter(x_dim, context_dim)
|
||||
self.base_fuse = nn.Sequential(
|
||||
nn.Linear(cond_out_dim + context_dim, cond_out_dim),
|
||||
nn.SiLU(),
|
||||
)
|
||||
|
||||
# Reuses conditioning.out_dim for the history encoder's own output
|
||||
# width — there's no dedicated stage2_model.autoregressive key for
|
||||
# this, a reasonable default rather than a design-doc-specified value.
|
||||
history_dim = cond_out_dim
|
||||
hist_in_dim = CONT_SLOT_DIM + self.type_dim
|
||||
self.history_encoder: HistoryEncoder = (
|
||||
AttentionHistory(hist_in_dim, history_dim, n_heads=attn_n_heads, n_layers=attn_n_layers)
|
||||
if history == "attention"
|
||||
else MarkovHistory(hist_in_dim, history_dim)
|
||||
)
|
||||
token_fuse_in = cond_out_dim + context_dim + history_dim + 2 # +2: remaining_frac, slot_idx
|
||||
self.token_fuse = nn.Sequential(
|
||||
nn.Linear(token_fuse_in, cond_out_dim),
|
||||
nn.SiLU(),
|
||||
)
|
||||
|
||||
has_time = generator in ("flow", "ddpm")
|
||||
self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None
|
||||
merged_cond_dim = (time_dim if has_time else 0) + cond_out_dim
|
||||
token_dim = stage2_trunk_sec_dim(self.particle_type_cfg, generator, 1, emb_dim)
|
||||
in_dim = noise_dim if generator == "wgan" else token_dim
|
||||
self.trunk = build_trunk(
|
||||
router,
|
||||
in_dim,
|
||||
token_dim,
|
||||
hidden_dim,
|
||||
n_res_blocks,
|
||||
merged_cond_dim,
|
||||
dropout,
|
||||
)
|
||||
|
||||
self.n_sec_head = None
|
||||
if build_n_sec_head:
|
||||
self.n_sec_head = nn.Sequential(
|
||||
nn.Linear(cond_out_dim, hidden_dim // 2),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_dim // 2, k_max + 1),
|
||||
)
|
||||
self.type_head = None
|
||||
target = self.particle_type_cfg.get("target", "physical")
|
||||
if target != "physical" and generator in ("flow", "ddpm"):
|
||||
self.type_head = nn.Sequential(
|
||||
nn.Linear(cond_out_dim, hidden_dim // 2),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_dim // 2, self.type_dim),
|
||||
)
|
||||
|
||||
def _base_cond(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
|
||||
base = self.cond_enc(cond_cont, cond_cat)
|
||||
ctx = self.context_adapter(stage1_out)
|
||||
return self.base_fuse(torch.cat([base, ctx], dim=-1))
|
||||
|
||||
def _token_cond(
|
||||
self,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
history_feat: torch.Tensor,
|
||||
has_prev: torch.Tensor,
|
||||
remaining_frac: torch.Tensor,
|
||||
slot_idx: torch.Tensor,
|
||||
hist: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""`hist`, if given, overrides recomputing `self.history_encoder`
|
||||
from `history_feat`/`has_prev` — the inference-time KV-cache path
|
||||
(`Stage2Autoregressive.history_step`) precomputes it once per slot and
|
||||
passes it in here so a slot's (possibly several) model calls — an ODE
|
||||
loop's substeps, or a separate `predict_type` call — read the same
|
||||
cached history instead of each re-deriving (and, under attention,
|
||||
re-appending to the cache — see `AttentionHistory.step`'s docstring)."""
|
||||
K = history_feat.size(1)
|
||||
base = self.cond_enc(cond_cont, cond_cat).unsqueeze(1).expand(-1, K, -1)
|
||||
ctx = self.context_adapter(stage1_out).unsqueeze(1).expand(-1, K, -1)
|
||||
if hist is None:
|
||||
hist = self.history_encoder(history_feat, has_prev)
|
||||
scalars = torch.stack([remaining_frac, slot_idx], dim=-1)
|
||||
return self.token_fuse(torch.cat([base, ctx, hist, scalars], dim=-1))
|
||||
|
||||
def init_history_cache(self):
|
||||
"""Inference-only incremental-decoding state for `self.history_encoder`
|
||||
(`giant/sample.py`'s AR loop): `None` under `history="markov"` (its
|
||||
per-step cost is already O(1) — see `HistoryEncoder`'s docstring), or
|
||||
`AttentionHistory.init_cache()` under `history="attention"`."""
|
||||
if isinstance(self.history_encoder, AttentionHistory):
|
||||
return self.history_encoder.init_cache()
|
||||
return None
|
||||
|
||||
def history_step(self, token_feat: torch.Tensor, has_prev: torch.Tensor, cache) -> tuple[torch.Tensor, object]:
|
||||
"""One inference slot's worth of history encoding: advances `cache`
|
||||
(from `init_history_cache`, or a previous `history_step` call) by
|
||||
`token_feat`/`has_prev` (`(B, 1, ...)` — the just-emitted previous
|
||||
token, same convention `giant.sample.sample_secondaries_ar` already
|
||||
threads as `prev_repr`), and returns `(hist, new_cache)` — `hist` is
|
||||
this slot's history summary (pass it as `_token_cond`'s `hist=` to
|
||||
every model call made for this slot), `new_cache` is what to pass into
|
||||
the *next* slot's `history_step`. Must be called exactly once per
|
||||
slot — see `AttentionHistory.step`'s docstring."""
|
||||
if isinstance(self.history_encoder, AttentionHistory):
|
||||
return self.history_encoder.step(token_feat, has_prev, cache)
|
||||
return self.history_encoder(token_feat, has_prev), cache
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x_t: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
history_feat: torch.Tensor,
|
||||
has_prev: torch.Tensor,
|
||||
remaining_frac: torch.Tensor,
|
||||
slot_idx: torch.Tensor,
|
||||
t: torch.Tensor | None = None,
|
||||
hist: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
B, K = x_t.shape[0], x_t.shape[1]
|
||||
c_emb = self._token_cond(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
hist=hist,
|
||||
)
|
||||
if self.time_emb is not None:
|
||||
assert t is not None
|
||||
t_emb = self.time_emb(t.reshape(-1)).view(B, K, -1)
|
||||
cond = torch.cat([t_emb, c_emb], dim=-1)
|
||||
else:
|
||||
cond = c_emb
|
||||
x_flat = x_t.reshape(B * K, -1)
|
||||
cond_flat = cond.reshape(B * K, -1)
|
||||
cond_cont_flat = cond_cont.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
|
||||
cond_cat_flat = cond_cat.unsqueeze(1).expand(-1, K, -1).reshape(B * K, -1)
|
||||
out = self.trunk(x_flat, cond_flat, cond_cont_flat, cond_cat_flat)
|
||||
return out.view(B, K, -1)
|
||||
|
||||
def predict_n_sec(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, stage1_out: torch.Tensor) -> torch.Tensor:
|
||||
if self.n_sec_head is None:
|
||||
raise RuntimeError(
|
||||
"this Stage2Autoregressive has no n_sec_head — it belongs to "
|
||||
"a migrated v0.2 checkpoint (n_sec.owner='stage1'); call "
|
||||
"stage1.predict_n_sec(cond_cont, cond_cat) instead"
|
||||
)
|
||||
return self.n_sec_head(self._base_cond(cond_cont, cond_cat, stage1_out))
|
||||
|
||||
def predict_type(
|
||||
self,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
history_feat: torch.Tensor,
|
||||
has_prev: torch.Tensor,
|
||||
remaining_frac: torch.Tensor,
|
||||
slot_idx: torch.Tensor,
|
||||
hist: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
if self.type_head is None:
|
||||
raise RuntimeError(
|
||||
"this Stage2Autoregressive has no type_head — either "
|
||||
"particle_type.target='physical' (the type slice is part of "
|
||||
"forward()'s own output) or generator='wgan' (the WGAN "
|
||||
"trainer reads the type slice out of forward()'s output "
|
||||
"directly instead)"
|
||||
)
|
||||
c_emb = self._token_cond(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
hist=hist,
|
||||
)
|
||||
B, K, _ = c_emb.shape
|
||||
return self.type_head(c_emb.reshape(B * K, -1)).view(B, K, self.type_dim)
|
||||
|
||||
|
||||
class CriticModel(nn.Module):
|
||||
"""Generator-agnostic WGAN-GP critic body: a scalar realism score, for
|
||||
either stage (`stage="stage1"` mirrors v0.2 `Critic`; `stage="stage2"`
|
||||
mirrors v0.2 `SecondaryCritic`, adding the same context-fusion path as
|
||||
`Stage2OneShot`). Used only when that stage's `generator == "wgan"`."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
particle_cfg: dict,
|
||||
material_cfg: dict,
|
||||
in_dim: int,
|
||||
hidden_dim: int = 256,
|
||||
n_res_blocks: int = 6,
|
||||
cond_out_dim: int = 128,
|
||||
dropout: float = 0.0,
|
||||
stage: str = "stage1",
|
||||
context_dim: int = 64,
|
||||
context_in_dim: int = X_DIM,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
if stage not in ("stage1", "stage2"):
|
||||
raise ValueError(f"stage must be 'stage1' or 'stage2', got {stage!r}")
|
||||
self.stage = stage
|
||||
self.cond_enc = ConditionEncoder(pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim)
|
||||
if stage == "stage2":
|
||||
self.context_adapter = ContextAdapter(context_in_dim, context_dim)
|
||||
self.fuse = nn.Sequential(
|
||||
nn.Linear(cond_out_dim + context_dim, cond_out_dim),
|
||||
nn.SiLU(),
|
||||
)
|
||||
self.input_proj = nn.Linear(in_dim, hidden_dim)
|
||||
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_res_blocks)])
|
||||
self.out_norm = nn.LayerNorm(hidden_dim)
|
||||
self.out_proj = nn.Linear(hidden_dim, 1)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
base = self.cond_enc(cond_cont, cond_cat)
|
||||
if self.stage == "stage2":
|
||||
ctx = self.context_adapter(stage1_out)
|
||||
cond = self.fuse(torch.cat([base, ctx], dim=-1))
|
||||
else:
|
||||
cond = base
|
||||
h = self.input_proj(x)
|
||||
for block in self.blocks:
|
||||
h = block(h, cond)
|
||||
return self.out_proj(self.out_norm(h)).squeeze(-1)
|
||||
+1377
-83
File diff suppressed because it is too large
Load Diff
@@ -1,375 +0,0 @@
|
||||
"""Mixture-of-experts routing: `Router` base + registry, the four concrete
|
||||
router types, and composed/config-driven construction — self-contained, no
|
||||
dependency on any other `giant.model` submodule (issues.md Issue 8)."""
|
||||
|
||||
import inspect
|
||||
import math
|
||||
import re
|
||||
from collections.abc import Sequence
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from giant.constants import COND_DIM
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Routers — carried over unchanged from v0.2
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class Router(nn.Module):
|
||||
"""Contract for a pluggable mixture-of-experts routing axis.
|
||||
|
||||
Subclasses implement `gate` (soft partition-of-unity weights over
|
||||
experts, used in train mode for a fully differentiable mixture);
|
||||
`top1` and `balance_loss` have working defaults so a new routing axis
|
||||
is usually a one-method add. See `ROUTER_REGISTRY` / `build_router`.
|
||||
"""
|
||||
|
||||
def __init__(self, n_experts: int) -> None:
|
||||
super().__init__()
|
||||
self.n_experts = n_experts
|
||||
self.gumbel = False
|
||||
self.gumbel_tau = 1.0
|
||||
|
||||
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
"""(B, n_experts) soft weights, rows summing to 1."""
|
||||
raise NotImplementedError
|
||||
|
||||
def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
"""(B, n_experts) train-time expert-combination weights.
|
||||
|
||||
Default (`gumbel=False`): identical to `gate()`. Opt-in
|
||||
straight-through Gumbel-softmax (`gumbel=True`, train mode only):
|
||||
hardens the forward pass to a one-hot sample (matching eval-time
|
||||
top-1 dispatch) while keeping the soft sample's gradient on backward.
|
||||
"""
|
||||
probs = self.gate(cond_cont, cond_cat)
|
||||
if not (self.gumbel and self.training):
|
||||
return probs
|
||||
log_probs = torch.log(probs.clamp_min(1e-8))
|
||||
return F.gumbel_softmax(log_probs, tau=self.gumbel_tau, hard=True, dim=-1)
|
||||
|
||||
def top1(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
"""(B,) hard expert index, used for eval-time grouped dispatch."""
|
||||
return self.gate(cond_cont, cond_cat).argmax(dim=-1)
|
||||
|
||||
def balance_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
"""Importance CV^2 load-balancing auxiliary loss (Shazeer et al. 2017)."""
|
||||
importance = self.gate(cond_cont, cond_cat).sum(dim=0) # (n_experts,)
|
||||
return (importance.std() / (importance.mean() + 1e-8)) ** 2
|
||||
|
||||
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
||||
"""Optional supervised auxiliary loss shaping the router's own belief.
|
||||
|
||||
Default: none (a scalar 0). Routers gating on an unobservable
|
||||
pre-step quantity (e.g. ProcessRouter) override this.
|
||||
"""
|
||||
return torch.zeros((), device=cond_cont.device)
|
||||
|
||||
def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
"""Optional auxiliary loss rewarding sharper (lower-entropy) routing."""
|
||||
norm_entropy, _ = self.gate_stats(cond_cont, cond_cat)
|
||||
return norm_entropy
|
||||
|
||||
def gate_stats(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Diagnostics: `(norm_entropy, importance)` — see v0.2 docstring for
|
||||
the full explanation, unchanged in v0.3.0."""
|
||||
gate = self.gate(cond_cont, cond_cat) # (B, n_experts)
|
||||
row_entropy = -(gate * (gate + 1e-8).log()).sum(dim=-1) # (B,)
|
||||
norm_entropy = row_entropy.mean() / math.log(self.n_experts)
|
||||
importance = gate.sum(dim=0) # (n_experts,)
|
||||
return norm_entropy, importance
|
||||
|
||||
|
||||
ROUTER_REGISTRY: dict[str, type[Router]] = {}
|
||||
|
||||
|
||||
def register_router(name: str):
|
||||
def decorator(cls: type[Router]) -> type[Router]:
|
||||
ROUTER_REGISTRY[name] = cls
|
||||
return cls
|
||||
|
||||
return decorator
|
||||
|
||||
|
||||
def build_router(name: str, n_experts: int, **kwargs) -> Router:
|
||||
"""Factory: look up a `Router` subclass by name from the registry.
|
||||
|
||||
Every registered router type is fed the same `router` config dict;
|
||||
kwargs not declared by that type's constructor are silently dropped, so
|
||||
per-type hyperparameters (e.g. EnergyRouter's `temperature`) can coexist
|
||||
in one config without special-casing.
|
||||
"""
|
||||
if name not in ROUTER_REGISTRY:
|
||||
raise ValueError(f"unknown router type {name!r}; available: {sorted(ROUTER_REGISTRY)}")
|
||||
cls = ROUTER_REGISTRY[name]
|
||||
accepted = set(inspect.signature(cls.__init__).parameters) - {"self", "n_experts"}
|
||||
filtered = {k: v for k, v in kwargs.items() if k in accepted}
|
||||
return cls(n_experts=n_experts, **filtered)
|
||||
|
||||
|
||||
def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor:
|
||||
"""Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient
|
||||
bound used for EnergyRouter's `learn_width`/`learn_temperature` modes."""
|
||||
return lo + (hi - lo) * torch.sigmoid(raw)
|
||||
|
||||
|
||||
def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float:
|
||||
"""Inverse of `_bounded_interp`, used once at construction to warm-start
|
||||
`raw` so the initial effective width/temperature exactly equals `value`."""
|
||||
p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6)
|
||||
return math.log(p / (1 - p))
|
||||
|
||||
|
||||
@register_router("energy")
|
||||
class EnergyRouter(Router):
|
||||
"""Soft turn-on gate over normalized pre-step log-energy.
|
||||
|
||||
Reads `cond_cont[:, energy_idx]` (ignores cond_cat). `gate(e) =
|
||||
softmax_i(-(e - c_i)^2 / tau)`; as tau -> 0 this hardens to
|
||||
nearest-center (Voronoi) selection, exactly what `top1` uses at eval.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_experts: int = 4,
|
||||
temperature: float = 0.5,
|
||||
learn_centers: bool = True,
|
||||
energy_idx: int = 3,
|
||||
centers_init: Sequence[float] | None = None,
|
||||
learn_width: bool = False,
|
||||
learn_temperature: bool = False,
|
||||
width_min_ratio: float = 0.1,
|
||||
width_max_ratio: float = 10.0,
|
||||
) -> None:
|
||||
super().__init__(n_experts)
|
||||
if learn_width and learn_temperature:
|
||||
raise ValueError("learn_width and learn_temperature are mutually exclusive")
|
||||
self.temperature = temperature
|
||||
self.energy_idx = energy_idx
|
||||
self.learn_width = learn_width
|
||||
self.learn_temperature = learn_temperature
|
||||
if learn_width or learn_temperature:
|
||||
if not (width_min_ratio < 1.0 < width_max_ratio):
|
||||
raise ValueError(
|
||||
f"width_min_ratio ({width_min_ratio}) and width_max_ratio ({width_max_ratio}) must bracket 1.0"
|
||||
)
|
||||
self._width_lo = width_min_ratio * temperature
|
||||
self._width_hi = width_max_ratio * temperature
|
||||
raw0 = _inverse_bounded_interp(temperature, self._width_lo, self._width_hi)
|
||||
if learn_width:
|
||||
self.raw_width = nn.Parameter(torch.full((n_experts,), raw0))
|
||||
else:
|
||||
self.raw_temperature = nn.Parameter(torch.tensor(raw0))
|
||||
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, expected n_experts={n_experts}")
|
||||
centers = torch.tensor(list(centers_init), dtype=torch.float32)
|
||||
if learn_centers:
|
||||
self.centers = nn.Parameter(centers)
|
||||
else:
|
||||
self.register_buffer("centers", centers)
|
||||
|
||||
def effective_width(self) -> torch.Tensor | float:
|
||||
if self.learn_width:
|
||||
return _bounded_interp(self.raw_width, self._width_lo, self._width_hi)
|
||||
if self.learn_temperature:
|
||||
return _bounded_interp(self.raw_temperature, self._width_lo, self._width_hi)
|
||||
return self.temperature
|
||||
|
||||
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
e = cond_cont[:, self.energy_idx].unsqueeze(-1) # (B, 1)
|
||||
d2 = (e - self.centers.unsqueeze(0)) ** 2 # (B, n_experts)
|
||||
return torch.softmax(-d2 / self.effective_width(), dim=-1)
|
||||
|
||||
|
||||
@register_router("pdg")
|
||||
class PdgRouter(Router):
|
||||
"""Soft turn-on gate over a learned PDG embedding (own table, separate
|
||||
from the trunk's `ConditionEncoder`). No supervision needed — PDG code
|
||||
is already known at pre-step time."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_experts: int,
|
||||
pdg_vocab: int,
|
||||
emb_dim: int = 8,
|
||||
temperature: float = 0.5,
|
||||
learn_centers: bool = True,
|
||||
) -> None:
|
||||
super().__init__(n_experts)
|
||||
self.temperature = temperature
|
||||
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
|
||||
centers = torch.randn(n_experts, emb_dim) * 0.1
|
||||
if learn_centers:
|
||||
self.centers = nn.Parameter(centers)
|
||||
else:
|
||||
self.register_buffer("centers", centers)
|
||||
|
||||
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim)
|
||||
d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1) # (B, n_experts)
|
||||
return torch.softmax(-d2 / self.temperature, dim=-1)
|
||||
|
||||
|
||||
@register_router("process")
|
||||
class ProcessRouter(Router):
|
||||
"""Routes on the physics process expected to end the step — a post-step
|
||||
outcome, so a small classifier over pre-step conditioning predicts it
|
||||
(own pdg/material embeddings, separate from the trunk's ConditionEncoder).
|
||||
`n_experts` doubles as the number of process classes. Supervised via
|
||||
`classify_loss` against the true `process` label at train time only;
|
||||
`gate`/`top1` never see it."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
n_experts: int,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
emb_dim: int = 8,
|
||||
hidden_dim: int = 64,
|
||||
) -> None:
|
||||
super().__init__(n_experts)
|
||||
self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim)
|
||||
self.mat_emb = nn.Embedding(mat_vocab, emb_dim)
|
||||
self.classifier = nn.Sequential(
|
||||
nn.Linear(COND_DIM + 2 * emb_dim, hidden_dim),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_dim, n_experts),
|
||||
)
|
||||
|
||||
def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
pdg_e = self.pdg_emb(cond_cat[:, 0])
|
||||
mat_e = self.mat_emb(cond_cat[:, 1])
|
||||
h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1)
|
||||
return self.classifier(h)
|
||||
|
||||
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
return torch.softmax(self.logits(cond_cont, cond_cat), dim=-1)
|
||||
|
||||
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
||||
return F.cross_entropy(self.logits(cond_cont, cond_cat), labels)
|
||||
|
||||
|
||||
class ComposedRouter(Router):
|
||||
"""Joint router over independent axes (e.g. energy x pdg), outer-product
|
||||
gated. Not registered in `ROUTER_REGISTRY`; use `build_composed_router`."""
|
||||
|
||||
def __init__(self, routers: list[Router]) -> None:
|
||||
if not routers:
|
||||
raise ValueError("ComposedRouter needs at least one sub-router")
|
||||
n_experts = 1
|
||||
for r in routers:
|
||||
n_experts *= r.n_experts
|
||||
super().__init__(n_experts)
|
||||
self.routers = nn.ModuleList(routers)
|
||||
|
||||
def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor:
|
||||
joint = self.routers[0].gate(cond_cont, cond_cat) # (B, n_0)
|
||||
for router in self.routers[1:]:
|
||||
g = router.gate(cond_cont, cond_cat) # (B, n_i)
|
||||
joint = (joint.unsqueeze(-1) * g.unsqueeze(1)).flatten(1) # (B, prod so far)
|
||||
return joint
|
||||
|
||||
def classify_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
|
||||
total = torch.zeros((), device=cond_cont.device)
|
||||
for router in self.routers:
|
||||
total = total + router.classify_loss(cond_cont, cond_cat, labels)
|
||||
return total
|
||||
|
||||
|
||||
def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter:
|
||||
"""Build a `ComposedRouter` from a list of per-axis router specs — see
|
||||
`_parse_composed_axes`."""
|
||||
routers = [
|
||||
build_router(
|
||||
spec["type"],
|
||||
spec["n_experts"],
|
||||
**{
|
||||
**shared_kwargs,
|
||||
**{k: v for k, v in spec.items() if k not in ("type", "n_experts")},
|
||||
},
|
||||
)
|
||||
for spec in specs
|
||||
]
|
||||
return ComposedRouter(routers)
|
||||
|
||||
|
||||
_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$")
|
||||
|
||||
|
||||
def _parse_composed_axes(router_cfg: dict) -> list[dict]:
|
||||
"""Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts.
|
||||
|
||||
e.g. `axis0_type = "energy"`, `axis0_n_experts = 4`, `axis1_type = "pdg"`,
|
||||
`axis1_n_experts = 3`, `axis1_emb_dim = 8`. Axis indices must be
|
||||
contiguous from 0.
|
||||
"""
|
||||
axes: dict[int, dict] = {}
|
||||
for key, value in router_cfg.items():
|
||||
m = _AXIS_KEY_RE.match(key)
|
||||
if m is None:
|
||||
continue
|
||||
idx, field = int(m.group(1)), m.group(2)
|
||||
axes.setdefault(idx, {})[field] = value
|
||||
missing = set(range(len(axes))) - axes.keys()
|
||||
if missing:
|
||||
raise ValueError(f"composed router config has gaps at axis indices {missing}")
|
||||
return [axes[i] for i in range(len(axes))]
|
||||
|
||||
|
||||
# Router types that read cond_cat's pdg index through their own
|
||||
# nn.Embedding(pdg_vocab, ...), regardless of the trunk's particle
|
||||
# conditioning mode — see _check_router_conditioning_compat.
|
||||
_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process")
|
||||
|
||||
|
||||
def _check_router_conditioning_compat(router_types: list[str], particle_conditioning: str) -> None:
|
||||
"""Reject a router axis that reintroduces a training-vocab PDG lookup
|
||||
under `conditioning.particle.type = "physical"`.
|
||||
|
||||
`PdgRouter`/`ProcessRouter` always build their own dataset-scoped
|
||||
`nn.Embedding(pdg_vocab, ...)`, independent of `ConditionEncoder`'s
|
||||
particle mode. Pairing either with `"physical"` would silently
|
||||
reintroduce a training-menu-scoped lookup at the routing layer,
|
||||
defeating the point of physical-property conditioning. Raised loudly at
|
||||
model-build time.
|
||||
"""
|
||||
bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES))
|
||||
if bad and particle_conditioning == "physical":
|
||||
raise ValueError(
|
||||
f"router type(s) {bad} always use a training-vocab PDG embedding, "
|
||||
"which is incompatible with conditioning.particle.type='physical' "
|
||||
"(whose whole point is generalizing beyond that vocab) — pick a "
|
||||
"different router type (e.g. 'energy') or use "
|
||||
"conditioning.particle.type='embedding'."
|
||||
)
|
||||
|
||||
|
||||
def _build_router_from_cfg(
|
||||
router_cfg: dict,
|
||||
pdg_vocab: int,
|
||||
mat_vocab: int,
|
||||
particle_conditioning: str = "embedding",
|
||||
) -> Router:
|
||||
"""Resolve one stage's `router` config into a `Router`, single-axis or
|
||||
composed. `gumbel` is set as a post-construction attribute (shared by
|
||||
every router type, not a per-type constructor kwarg)."""
|
||||
shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
|
||||
if router_cfg["type"] == "composed":
|
||||
axes = _parse_composed_axes(router_cfg)
|
||||
_check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning)
|
||||
router = build_composed_router(axes, **shared_vocab)
|
||||
router.gumbel = bool(router_cfg.get("gumbel", False))
|
||||
return router
|
||||
_check_router_conditioning_compat([router_cfg["type"]], particle_conditioning)
|
||||
router_kwargs = {k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts")}
|
||||
router_kwargs.setdefault("pdg_vocab", pdg_vocab)
|
||||
router_kwargs.setdefault("mat_vocab", mat_vocab)
|
||||
router = build_router(router_cfg["type"], router_cfg["n_experts"], **router_kwargs)
|
||||
router.gumbel = bool(router_cfg.get("gumbel", False))
|
||||
return router
|
||||
+17
-101
@@ -11,7 +11,9 @@ class CosineSchedule:
|
||||
steps = np.arange(T + 1, dtype=np.float64)
|
||||
f = np.cos(((steps / T + s) / (1.0 + s)) * np.pi / 2.0) ** 2
|
||||
alpha_bars = (f / f[0]).astype(np.float32)
|
||||
betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)
|
||||
betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(
|
||||
np.float32
|
||||
)
|
||||
|
||||
self.betas = torch.from_numpy(betas)
|
||||
self.alphas = torch.from_numpy((1.0 - betas))
|
||||
@@ -46,7 +48,7 @@ class CosineSchedule:
|
||||
noise = torch.randn_like(x0)
|
||||
x_t = self.q_sample(x0, t, noise)
|
||||
t_norm = t.float() / self.T
|
||||
pred = model(x_t, cond_cont, cond_cat, t=t_norm)
|
||||
pred = model(x_t, t_norm, cond_cont, cond_cat)
|
||||
return F.mse_loss(pred, noise)
|
||||
|
||||
|
||||
@@ -65,7 +67,7 @@ def flow_matching_loss(
|
||||
x0 = torch.randn_like(x1)
|
||||
x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
|
||||
u_t = x1 - x0
|
||||
v_t = model(x_t, cond_cont, cond_cat, t=t)
|
||||
v_t = model(x_t, t, cond_cont, cond_cat)
|
||||
return F.mse_loss(v_t, u_t)
|
||||
|
||||
|
||||
@@ -76,125 +78,39 @@ def flow_matching_loss_secondary(
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
sec_mask: torch.Tensor,
|
||||
type_dim: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Flow matching loss for the secondary decoder with per-slot masking.
|
||||
|
||||
x1: (B, K_MAX * (CONT_SLOT_DIM + type_dim)) — flattened secondary target
|
||||
(stick_logit, dir, then a `type_dim`-wide type slice)
|
||||
x1: (B, SEC_DIM) — flattened secondary target (stick_logit, dir, log_mass, charge)
|
||||
sec_mask: (B, K_MAX) bool — True for valid secondary slots
|
||||
type_dim: width of the per-slot type slice folded into `x1` — defaults to
|
||||
`PARTICLE_PHYS_DIM` (log_mass, charge), `particle_type.target =
|
||||
"physical"`'s width and the only case this function handled before
|
||||
v0.3.0 step 4. `0` means no type slice is in `x1` at all (`target`
|
||||
in `("onehot", "embedding")` under `generator in ("flow", "ddpm")` —
|
||||
`Stage2OneShot.type_head` handles the type loss separately in that
|
||||
case).
|
||||
|
||||
Only valid-slot dimensions contribute to the loss; padded slots are zeroed
|
||||
before averaging, so the loss is not diluted by empty slots.
|
||||
|
||||
Each slot packs CONT_SLOT_DIM continuous dims (stick_logit, dir) followed
|
||||
by the `type_dim`-wide type slice — under `target = "physical"` (the
|
||||
default) that's the secondary's predicted physical identity, a fixed
|
||||
regression target (see giant.data.transforms.encode_secondaries); under
|
||||
`target = "embedding"` (folded in only for `generator = "wgan"`, so
|
||||
`type_dim > 0` here only ever means "physical") it would be the detached
|
||||
embedding-table row. Even though the two blocks are the same order of
|
||||
magnitude now (unlike the 16-wide learned embedding block "physical"
|
||||
replaced), they're still on different physical scales, so they're each
|
||||
averaged over their own width first and then combined with equal weight
|
||||
— this stays correct if type_dim/CONT_SLOT_DIM change.
|
||||
by PARTICLE_PHYS_DIM physical-identity dims (log_mass, charge) — the
|
||||
secondary's predicted physical identity, a fixed regression target (see
|
||||
giant.data.transforms.encode_secondaries). Even though the two blocks are
|
||||
the same order of magnitude now (unlike the 16-wide learned embedding
|
||||
block this replaced), they're still on different physical scales, so
|
||||
they're each averaged over their own width first and then combined with
|
||||
equal weight — this stays correct if PARTICLE_PHYS_DIM/CONT_SLOT_DIM change.
|
||||
"""
|
||||
from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
|
||||
|
||||
if type_dim is None:
|
||||
type_dim = PARTICLE_PHYS_DIM
|
||||
from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
|
||||
|
||||
B = x1.size(0)
|
||||
k_max = sec_mask.size(1)
|
||||
t = torch.rand(B, device=x1.device)
|
||||
x0 = torch.randn_like(x1)
|
||||
x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1
|
||||
u_t = x1 - x0
|
||||
v_t = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
||||
v_t = model(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
|
||||
slot_dim = CONT_SLOT_DIM + type_dim
|
||||
err = ((v_t - u_t) ** 2).view(B, k_max, slot_dim)
|
||||
cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1) # (B, k_max)
|
||||
|
||||
mask = sec_mask.float()
|
||||
denom = mask.sum().clamp(min=1)
|
||||
cont_loss = (cont_err * mask).sum() / denom
|
||||
if type_dim == 0:
|
||||
return cont_loss
|
||||
phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1)
|
||||
phys_loss = (phys_err * mask).sum() / denom
|
||||
return cont_loss + phys_loss
|
||||
|
||||
|
||||
def flow_matching_loss_secondary_ar(
|
||||
model: torch.nn.Module,
|
||||
x1: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
history_feat: torch.Tensor,
|
||||
has_prev: torch.Tensor,
|
||||
remaining_frac: torch.Tensor,
|
||||
slot_idx: torch.Tensor,
|
||||
sec_mask: torch.Tensor,
|
||||
type_dim: int | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""`Stage2Autoregressive` analogue of `flow_matching_loss_secondary`, same
|
||||
masked, per-block (continuous vs. type) loss recipe — but native to
|
||||
`Stage2Autoregressive`'s `(B, K_MAX, token_dim)` I/O and its extra
|
||||
per-token conditioning args, rather than a flattened `(B, K_MAX*token_dim)`
|
||||
vector. Kept as a sibling rather than unified with the flat version: the
|
||||
model call signature differs enough (four extra per-token conditioning
|
||||
tensors) that merging would need an awkward shape-flag + closure.
|
||||
|
||||
Under teacher forcing this is still a
|
||||
single parallel pass over all K_MAX tokens — `x1`/`history_feat`/etc. are
|
||||
already built from ground truth for every slot by the caller
|
||||
(`giant.training.stage2_inputs._assemble_stage2_ar_inputs`/`_assemble_stage2_ar_target`).
|
||||
|
||||
x1: (B, K_MAX, CONT_SLOT_DIM + type_dim) — per-token flattened target
|
||||
(stick_logit, dir, then a `type_dim`-wide type slice)
|
||||
sec_mask: (B, K_MAX) bool — True for valid secondary slots
|
||||
type_dim: as `flow_matching_loss_secondary` — defaults to
|
||||
`PARTICLE_PHYS_DIM`, `0` means no type slice is in `x1` at all.
|
||||
"""
|
||||
from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
|
||||
|
||||
if type_dim is None:
|
||||
type_dim = PARTICLE_PHYS_DIM
|
||||
|
||||
B, K, _ = x1.shape
|
||||
t = torch.rand(B, K, device=x1.device)
|
||||
x0 = torch.randn_like(x1)
|
||||
x_t = (1.0 - t.unsqueeze(-1)) * x0 + t.unsqueeze(-1) * x1
|
||||
u_t = x1 - x0
|
||||
v_t = model(
|
||||
x_t,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
t=t,
|
||||
)
|
||||
|
||||
err = (v_t - u_t) ** 2
|
||||
err = ((v_t - u_t) ** 2).view(B, K_MAX, SEC_SLOT_DIM)
|
||||
cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1) # (B, K_MAX)
|
||||
phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM].mean(dim=-1)
|
||||
|
||||
mask = sec_mask.float()
|
||||
denom = mask.sum().clamp(min=1)
|
||||
cont_loss = (cont_err * mask).sum() / denom
|
||||
if type_dim == 0:
|
||||
return cont_loss
|
||||
phys_err = err[..., CONT_SLOT_DIM : CONT_SLOT_DIM + type_dim].mean(dim=-1)
|
||||
phys_loss = (phys_err * mask).sum() / denom
|
||||
return cont_loss + phys_loss
|
||||
|
||||
@@ -1,155 +0,0 @@
|
||||
"""Trunks: everything downstream of the fused conditioning vector — monolithic
|
||||
or expert-routed (issues.md Issue 8)."""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from giant.model.layers import ResBlock
|
||||
from giant.model.routers import Router
|
||||
|
||||
|
||||
class ExpertTrunk(nn.Module):
|
||||
"""One small expert: `input_proj -> ResBlock stack -> out_proj`.
|
||||
|
||||
Unlike v0.2, `out_dim` is independent of `in_dim` — needed by stage-2 AR
|
||||
tokens later (`noise_dim` in, `4 + type_dim` out), even though every
|
||||
step-2/3 caller still has `in_dim == out_dim`.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
hidden_dim: int,
|
||||
n_blocks: int,
|
||||
cond_dim: int,
|
||||
dropout: float = 0.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.input_proj = nn.Linear(in_dim, hidden_dim)
|
||||
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_blocks)])
|
||||
self.out_proj = nn.Linear(hidden_dim, out_dim)
|
||||
|
||||
def forward(self, x: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
|
||||
x = self.input_proj(x)
|
||||
for block in self.blocks:
|
||||
x = block(x, cond)
|
||||
return self.out_proj(x)
|
||||
|
||||
|
||||
def _route_forward(
|
||||
experts: nn.ModuleList,
|
||||
router: Router,
|
||||
x: torch.Tensor,
|
||||
cond: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
training: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Shared dispatch for `RoutedTrunk`.
|
||||
|
||||
Train mode: full mixture `sum_i weight_i * expert_i(x)` — always
|
||||
N-expert dense compute, fully differentiable (`weight` is
|
||||
`router.combine_weights`). Eval mode: grouped top-1 dispatch — each row
|
||||
runs exactly one expert, the actual source of the per-call speedup.
|
||||
"""
|
||||
if training:
|
||||
weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts)
|
||||
out = torch.zeros(x.shape[0], experts[0].out_proj.out_features, device=x.device)
|
||||
for i, expert in enumerate(experts):
|
||||
out = out + weights[:, i : i + 1] * expert(x, cond)
|
||||
return out
|
||||
|
||||
idx = router.top1(cond_cont, cond_cat) # (B,)
|
||||
out_dim = experts[0].out_proj.out_features
|
||||
out = torch.zeros(x.shape[0], out_dim, device=x.device)
|
||||
for i, expert in enumerate(experts):
|
||||
mask = idx == i
|
||||
if mask.any():
|
||||
out[mask] = expert(x[mask], cond[mask])
|
||||
return out
|
||||
|
||||
|
||||
class Trunk(nn.Module):
|
||||
"""Interface implemented by `MonolithicTrunk`/`RoutedTrunk`: everything
|
||||
downstream of the fused conditioning vector, i.e. the actual generative
|
||||
trunk of a stage (`input_proj -> blocks -> out_proj`, monolithic or
|
||||
expert-routed)."""
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cond: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
raise NotImplementedError
|
||||
|
||||
|
||||
class MonolithicTrunk(Trunk):
|
||||
def __init__(
|
||||
self,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
hidden_dim: int,
|
||||
n_res_blocks: int,
|
||||
cond_dim: int,
|
||||
dropout: float = 0.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.input_proj = nn.Linear(in_dim, hidden_dim)
|
||||
self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_dim, dropout=dropout) for _ in range(n_res_blocks)])
|
||||
self.out_proj = nn.Linear(hidden_dim, out_dim)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cond: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
x = self.input_proj(x)
|
||||
for block in self.blocks:
|
||||
x = block(x, cond)
|
||||
return self.out_proj(x)
|
||||
|
||||
|
||||
class RoutedTrunk(Trunk):
|
||||
def __init__(
|
||||
self,
|
||||
router: Router,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
hidden_dim: int,
|
||||
n_res_blocks: int,
|
||||
cond_dim: int,
|
||||
dropout: float = 0.0,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.router = router
|
||||
self.experts = nn.ModuleList(
|
||||
[ExpertTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout) for _ in range(router.n_experts)]
|
||||
)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
x: torch.Tensor,
|
||||
cond: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
return _route_forward(self.experts, self.router, x, cond, cond_cont, cond_cat, self.training)
|
||||
|
||||
|
||||
def build_trunk(
|
||||
router: Router | None,
|
||||
in_dim: int,
|
||||
out_dim: int,
|
||||
hidden_dim: int,
|
||||
n_res_blocks: int,
|
||||
cond_dim: int,
|
||||
dropout: float = 0.0,
|
||||
) -> Trunk:
|
||||
if router is not None:
|
||||
return RoutedTrunk(router, in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
|
||||
return MonolithicTrunk(in_dim, out_dim, hidden_dim, n_res_blocks, cond_dim, dropout)
|
||||
+6
-123
@@ -7,28 +7,16 @@ table involved) for isomer/excited nuclear codes the package's ground-state-only
|
||||
nuclide table doesn't cover — confirmed necessary for ~32% of the nuclear codes
|
||||
actually present in the multi-material dataset
|
||||
(`0932fb02-f2ce-43ca-a4ef-60a2b1221bbc.parquet`).
|
||||
|
||||
Also holds the v0.3.0 stage-2 categorical-type rollout decode:
|
||||
`decode_topn_class`/`decode_embedding_nearest` turn
|
||||
`Stage2Autoregressive`/`Stage2OneShot`'s `"onehot"`/`"embedding"` type
|
||||
predictions back into concrete PDG codes, the one place a secondary's
|
||||
categorical/continuous type representation is ever discretized (its
|
||||
free-running history representation stays unsnapped — see
|
||||
`giant/sample.py`'s AR loop).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from functools import lru_cache
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import numpy as np
|
||||
from particle import InvalidParticle, Particle, ParticleNotFound
|
||||
from particle import pdgid as _pdgid
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from giant.data.loader import TopNMap
|
||||
|
||||
# First-pass nuclear mass approximation (A * atomic mass unit); no
|
||||
# binding-energy correction. Only used for codes missing from `particle`'s
|
||||
# ground-state nuclide table -- ground-state codes get the package's real
|
||||
@@ -58,7 +46,9 @@ def particle_mass_charge(pdg: int) -> tuple[float, float]:
|
||||
if _pdgid.is_nucleus(pdg):
|
||||
z, a = _pdgid.Z(pdg), _pdgid.A(pdg)
|
||||
if z is None or a is None:
|
||||
raise ValueError(f"PDG {pdg}: is_nucleus but Z/A decode failed") from None
|
||||
raise ValueError(
|
||||
f"PDG {pdg}: is_nucleus but Z/A decode failed"
|
||||
) from None
|
||||
return float(a) * _AMU_MEV, float(z)
|
||||
raise ValueError(
|
||||
f"PDG code {pdg} could not be resolved via the `particle` package "
|
||||
@@ -107,7 +97,9 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd
|
||||
if len(resolved) == 0:
|
||||
raise ValueError("nearest_known_pdg: no resolvable candidates")
|
||||
codes = np.array([r[0] for r in resolved], dtype=np.int64)
|
||||
table_log_mass = np.log(np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS)
|
||||
table_log_mass = np.log(
|
||||
np.array([r[1] for r in resolved], dtype=np.float64) + _LOG_EPS
|
||||
)
|
||||
table_charge = np.array([r[2] for r in resolved], dtype=np.float64)
|
||||
|
||||
mass = np.asarray(mass, dtype=np.float64)
|
||||
@@ -119,112 +111,3 @@ def nearest_known_pdg(mass: np.ndarray, charge: np.ndarray, candidates) -> np.nd
|
||||
) ** 2
|
||||
idx = d2.argmin(axis=1)
|
||||
return codes[idx]
|
||||
|
||||
|
||||
def invert_dense_map(m: dict[int, int]) -> dict[int, int]:
|
||||
"""index -> key, inverting a dense, bijective value->index map (`pdg_map`,
|
||||
or a `TopNMap.class_map`'s non-"other" entries — see `decode_topn_class`,
|
||||
which needs a *partial* inverse, not this general one, because its "other"
|
||||
index isn't unique-preimage). `pdg_map` itself is always a true bijection
|
||||
(`giant.data.loader.build_index_maps_from_files` enumerates the vocab), so
|
||||
a plain dict-comprehension inversion is exact here — used for
|
||||
`stage2_model.particle_type.target = "embedding"` decode, whose vocabulary
|
||||
is the full dense `pdg_map`, not a top-N-plus-other map."""
|
||||
return {v: k for k, v in m.items()}
|
||||
|
||||
|
||||
def decode_topn_class(
|
||||
class_idx: np.ndarray,
|
||||
topn_map: "TopNMap",
|
||||
n_classes: int,
|
||||
other_policy: str = "sample",
|
||||
rng: np.random.Generator | None = None,
|
||||
) -> np.ndarray:
|
||||
"""`stage2_model.particle_type.target = "onehot"` inference decode:
|
||||
per-row top-N class index -> concrete secondary-species PDG code.
|
||||
|
||||
class_idx: int array, any shape, values in `[0, n_classes)`.
|
||||
topn_map: the `TopNMap` (`giant.data.loader.build_pdg_topn_map_from_files`)
|
||||
this class index was built from — `class_map` (PDG -> class, injective
|
||||
except at the shared "other" index) plus `other_members` (the
|
||||
empirical within-"other" distribution, needed for `other_policy =
|
||||
"sample"`/`"modal"`).
|
||||
n_classes: the resolved secondary-species class count
|
||||
(`giant.model.models.resolve_type_n_classes` —
|
||||
`stage2_model.particle_type.n_classes`, 0 = inherit
|
||||
`conditioning.particle.emb_dim`; see gitea #29); the "other" bucket
|
||||
is index `n_classes - 1` by construction
|
||||
(`giant.data.loader._topn_plus_other_map`).
|
||||
other_policy: `"sample"` draws from `other_members`' empirical frequency;
|
||||
`"modal"` always the single most common "other" member; `"drop"`
|
||||
returns PDG `0` for those rows (not a valid PDG code — the caller
|
||||
must treat it as "no secondary", the same convention as
|
||||
`TERM_UNKNOWN_PDG` elsewhere in the rollout driver).
|
||||
|
||||
Every non-"other" class index has a unique inverse (the top `n_classes -
|
||||
1` keys each got their own index in `_topn_plus_other_map`), so those
|
||||
rows decode exactly; only "other" rows need `other_policy`.
|
||||
"""
|
||||
other_idx = n_classes - 1
|
||||
inv = np.zeros(n_classes, dtype=np.int64)
|
||||
for pdg, idx in topn_map.class_map.items():
|
||||
if idx != other_idx:
|
||||
inv[idx] = pdg
|
||||
|
||||
flat = np.asarray(class_idx, dtype=np.int64).reshape(-1)
|
||||
out = inv[np.clip(flat, 0, n_classes - 1)]
|
||||
|
||||
other_mask = flat == other_idx
|
||||
n_other = int(other_mask.sum())
|
||||
if n_other:
|
||||
if not topn_map.other_members:
|
||||
raise ValueError("decode_topn_class: 'other' class predicted but topn_map.other_members is empty")
|
||||
members = np.array(list(topn_map.other_members.keys()), dtype=np.int64)
|
||||
counts = np.array(list(topn_map.other_members.values()), dtype=np.float64)
|
||||
if other_policy == "drop":
|
||||
out[other_mask] = 0
|
||||
elif other_policy == "modal":
|
||||
out[other_mask] = members[counts.argmax()]
|
||||
elif other_policy == "sample":
|
||||
rng = rng if rng is not None else np.random.default_rng()
|
||||
probs = counts / counts.sum()
|
||||
out[other_mask] = rng.choice(members, size=n_other, p=probs)
|
||||
else:
|
||||
raise ValueError(f"unknown other_policy {other_policy!r}")
|
||||
|
||||
return out.reshape(np.asarray(class_idx).shape)
|
||||
|
||||
|
||||
def decode_embedding_nearest(
|
||||
vectors: np.ndarray,
|
||||
emb_weight: np.ndarray,
|
||||
idx_to_pdg: dict[int, int],
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""`stage2_model.particle_type.target = "embedding"` inference decode:
|
||||
L1-nearest row of the conditioning's own particle embedding table, since
|
||||
a generative model's continuous output
|
||||
essentially never lands within float tolerance of a table row (the exact-
|
||||
match form is only valid as a round-trip test assertion, never here).
|
||||
|
||||
vectors: `(..., emb_dim)` raw predicted vectors, any leading shape.
|
||||
emb_weight: `(vocab, emb_dim)` — `ConditionEncoder.pdg_emb.weight`,
|
||||
detached and moved to numpy by the caller. This is the SAME table
|
||||
`particle_type.target = "embedding"` was regressed against
|
||||
(`validate_config` requires `conditioning.particle.type =
|
||||
"embedding"` whenever this target is used — one table, not two).
|
||||
idx_to_pdg: `invert_dense_map(pdg_map)` — embedding row index -> PDG.
|
||||
|
||||
Returns `(pdg, l1_dist)`, both shaped like `vectors.shape[:-1]`. `l1_dist`
|
||||
is a diagnostic: a heavy tail means the decoder is emitting vectors off
|
||||
the embedding manifold, the direct analogue of the species-collapse
|
||||
symptom this redesign exists to fix.
|
||||
"""
|
||||
emb_dim = vectors.shape[-1]
|
||||
flat = np.asarray(vectors, dtype=np.float64).reshape(-1, emb_dim)
|
||||
table = np.asarray(emb_weight, dtype=np.float64)
|
||||
d = np.abs(flat[:, None, :] - table[None, :, :]).sum(axis=-1) # (N, vocab)
|
||||
nearest = d.argmin(axis=1)
|
||||
dist = d[np.arange(len(nearest)), nearest]
|
||||
pdg = np.array([idx_to_pdg[int(i)] for i in nearest], dtype=np.int64)
|
||||
lead_shape = vectors.shape[:-1]
|
||||
return pdg.reshape(lead_shape), dist.reshape(lead_shape).astype(np.float32)
|
||||
|
||||
+145
-202
@@ -9,19 +9,19 @@ from torch.utils.data import DataLoader
|
||||
from giant import config
|
||||
from giant.constants import (
|
||||
COND_DIM,
|
||||
EMB_DIM,
|
||||
K_MAX,
|
||||
PARTICLE_PHYS_DIM,
|
||||
SEC_SLOT_DIM,
|
||||
X_DIM,
|
||||
)
|
||||
from giant.data import setup_cache
|
||||
from giant.data.loader import (
|
||||
TopNMap,
|
||||
event_id_offset,
|
||||
find_parquet_files,
|
||||
iter_file_chunks,
|
||||
build_index_maps_from_files,
|
||||
build_pdg_topn_map_from_files,
|
||||
build_process_map_from_files,
|
||||
build_topn_map_from_files,
|
||||
)
|
||||
from giant.data.transforms import (
|
||||
Normalizer,
|
||||
@@ -31,8 +31,8 @@ from giant.data.transforms import (
|
||||
sorted_membership,
|
||||
)
|
||||
from giant.data.dataset import make_event_split, StreamingStepsDataset
|
||||
from giant.model.network import build_models, build_critics, resolve_type_n_classes
|
||||
from giant.training import train as run_training
|
||||
from giant.model.network import build_models, build_critics
|
||||
from giant.train import train as run_training
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -48,9 +48,6 @@ class SetupStageResult:
|
||||
pdg_map: dict[int, int]
|
||||
mat_map: dict[str, int]
|
||||
proc_map: dict[str, int] | None
|
||||
pdg_topn_map: TopNMap | None
|
||||
sec_type_topn_map: TopNMap | None
|
||||
mat_topn_map: TopNMap | None
|
||||
cond_norm: Normalizer
|
||||
tgt_norm: Normalizer
|
||||
sec_phys_norm: Normalizer
|
||||
@@ -59,62 +56,26 @@ class SetupStageResult:
|
||||
n_train_steps: int
|
||||
|
||||
|
||||
def _seed_energy_router(
|
||||
router_cfg: dict,
|
||||
cond_norm: Normalizer,
|
||||
energy_quantiles: np.ndarray,
|
||||
energy_idx: int,
|
||||
echo,
|
||||
) -> None:
|
||||
"""Mutate `router_cfg["centers_init"]` in place from real data quantiles,
|
||||
when this stage's router is an enabled EnergyRouter. Shared by both
|
||||
stages' router configs — each seeded independently, since v0.3.0 stages
|
||||
may have entirely different router configs."""
|
||||
active = router_cfg.get("enabled") and router_cfg.get("type") == "energy"
|
||||
if not active:
|
||||
return
|
||||
if energy_quantiles.size == 0:
|
||||
echo(" warning: no energy samples collected — EnergyRouter falls back to default centers")
|
||||
return
|
||||
assert cond_norm.mean is not None and cond_norm.std is not None
|
||||
levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
|
||||
raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
|
||||
centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
|
||||
router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
|
||||
echo(f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}")
|
||||
|
||||
|
||||
def run_setup_stage(
|
||||
data: str | Path,
|
||||
val_fraction: float,
|
||||
seed: int,
|
||||
cfg: dict,
|
||||
conditioning: str,
|
||||
router_cfg: dict,
|
||||
cache_setup: bool = True,
|
||||
rebuild_setup_cache: bool = False,
|
||||
echo=print,
|
||||
) -> SetupStageResult:
|
||||
"""Scan `data` for everything training needs before the epoch loop: the
|
||||
train/val event split, pdg/material vocab maps, an optional process map
|
||||
(needed if either stage's router is type="process"), and the Stage-1/
|
||||
Stage-2 normalizers.
|
||||
|
||||
`cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/
|
||||
`stage2_model`), already passed through `giant.config.validate_config`.
|
||||
`conditioning.particle.type` and `conditioning.material.type` are
|
||||
independent and may differ.
|
||||
(`router_cfg.type == "process"`), and the Stage-1/Stage-2 normalizers.
|
||||
|
||||
Reads from and writes to the `giant.data.setup_cache` sidecar when
|
||||
`cache_setup` is set (`rebuild_setup_cache` ignores — but still
|
||||
refreshes — any existing sidecar content). Each stage's `router` config
|
||||
is mutated in place: an active `EnergyRouter` (`router.type == "energy"`)
|
||||
refreshes — any existing sidecar content). `router_cfg` may be mutated
|
||||
in place: an active `EnergyRouter` (`router_cfg["type"] == "energy"`)
|
||||
gets its `centers_init` seeded from real data quantiles here.
|
||||
"""
|
||||
particle_conditioning = cfg["conditioning"]["particle"]["type"]
|
||||
material_conditioning = cfg["conditioning"]["material"]["type"]
|
||||
k_max = cfg["stage2_model"]["k_max"]
|
||||
stage1_router = cfg["stage1_model"].get("router") or {}
|
||||
stage2_router = cfg["stage2_model"].get("router") or {}
|
||||
|
||||
files = find_parquet_files(data)
|
||||
echo(f"found {len(files)} parquet file(s)")
|
||||
|
||||
@@ -136,14 +97,23 @@ def run_setup_stage(
|
||||
if cache is not None:
|
||||
cache.event_index = (unique_ids, counts)
|
||||
|
||||
train_events, val_events = make_event_split(unique_ids, val_fraction=val_fraction, seed=seed)
|
||||
train_events, val_events = make_event_split(
|
||||
unique_ids, val_fraction=val_fraction, seed=seed
|
||||
)
|
||||
events_arr = np.array(sorted(train_events))
|
||||
n_train_steps = setup_cache.n_train_steps_for_split(unique_ids, counts, events_arr)
|
||||
echo(f" {int(counts.sum()):,} steps | {len(train_events)} train events | {len(val_events)} val events")
|
||||
echo(
|
||||
f" {int(counts.sum()):,} steps | "
|
||||
f"{len(train_events)} train events | "
|
||||
f"{len(val_events)} val events"
|
||||
)
|
||||
|
||||
if cache is not None and cache.vocab is not None:
|
||||
pdg_map, mat_map = cache.vocab
|
||||
echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)")
|
||||
echo(
|
||||
f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, "
|
||||
f"{len(mat_map)} materials)"
|
||||
)
|
||||
else:
|
||||
echo("building vocabulary maps …")
|
||||
pdg_map, mat_map = build_index_maps_from_files(files)
|
||||
@@ -151,23 +121,16 @@ def run_setup_stage(
|
||||
if cache is not None:
|
||||
cache.vocab = (pdg_map, mat_map)
|
||||
|
||||
# A process map is needed if either stage's router reads the physics
|
||||
# process label (type="process"). Only one map is built even if both
|
||||
# stages want one — see the module-level note in giant/cli.py's
|
||||
# _router_total_experts for why composed-router n_experts isn't a plain
|
||||
# int; process routers are never composed in practice, so this doesn't
|
||||
# need that generality.
|
||||
proc_map: dict[str, int] | None = None
|
||||
process_router_cfg = next(
|
||||
(r for r in (stage1_router, stage2_router) if r.get("enabled") and r.get("type") == "process"),
|
||||
None,
|
||||
)
|
||||
if process_router_cfg is not None:
|
||||
n_experts = process_router_cfg["n_experts"]
|
||||
if router_cfg.get("enabled") and router_cfg.get("type") == "process":
|
||||
n_experts = router_cfg["n_experts"]
|
||||
cached_proc_map = cache.proc_maps.get(n_experts) if cache is not None else None
|
||||
if cached_proc_map is not None:
|
||||
proc_map = cached_proc_map
|
||||
echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)")
|
||||
echo(
|
||||
f"process vocabulary: cache hit ({len(proc_map)} labels, "
|
||||
f"{n_experts} experts)"
|
||||
)
|
||||
else:
|
||||
echo("building process vocabulary …")
|
||||
proc_map = build_process_map_from_files(files, n_experts=n_experts)
|
||||
@@ -175,62 +138,11 @@ def run_setup_stage(
|
||||
if cache is not None:
|
||||
cache.proc_maps[n_experts] = proc_map
|
||||
|
||||
# Top-N-plus-other maps for onehot conditioning/type axes.
|
||||
# The PDG axis is used independently by conditioning.particle.type="onehot"
|
||||
# (cond_cat's onehot feature) and stage2_model.particle_type.target="onehot"
|
||||
# (secondary-species decode) — their class counts can now differ (gitea
|
||||
# #29: stage2_model.particle_type.n_classes, 0 = inherit
|
||||
# conditioning.particle.emb_dim), so each is resolved and built
|
||||
# independently via _pdg_topn below. cache.topn_maps is keyed by
|
||||
# (axis, n_classes) (setup_cache.topn_key), so when the two resolve to
|
||||
# the same N the second call is a cache hit against the first — no extra
|
||||
# scan in the common case where they still match. The material axis is
|
||||
# independent of both.
|
||||
particle_cfg = cfg["conditioning"]["particle"]
|
||||
material_cfg = cfg["conditioning"]["material"]
|
||||
particle_type_cfg_dict = cfg["stage2_model"].get("particle_type") or {}
|
||||
particle_type_target = config.ParticleTypeConfig.from_dict(particle_type_cfg_dict).target
|
||||
|
||||
def _pdg_topn(n_classes: int) -> TopNMap:
|
||||
cache_key = setup_cache.topn_key("pdg", n_classes)
|
||||
cached = cache.topn_maps.get(cache_key) if cache is not None else None
|
||||
if cached is not None:
|
||||
echo(f"pdg top-N map: cache hit ({len(cached.class_map)} codes, {n_classes} classes)")
|
||||
return cached
|
||||
echo("building pdg top-N map …")
|
||||
topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes)
|
||||
echo(f" {len(topn_map.class_map)} pdg codes mapped to {n_classes} classes")
|
||||
if cache is not None:
|
||||
cache.topn_maps[cache_key] = topn_map
|
||||
return topn_map
|
||||
|
||||
pdg_topn_map: TopNMap | None = None
|
||||
if particle_cfg["type"] == "onehot":
|
||||
pdg_topn_map = _pdg_topn(particle_cfg["emb_dim"])
|
||||
|
||||
sec_type_topn_map: TopNMap | None = None
|
||||
if particle_type_target == "onehot":
|
||||
sec_type_n_classes = resolve_type_n_classes(particle_type_cfg_dict, particle_cfg["emb_dim"])
|
||||
sec_type_topn_map = _pdg_topn(sec_type_n_classes)
|
||||
|
||||
mat_topn_map: TopNMap | None = None
|
||||
if material_cfg["type"] == "onehot":
|
||||
n_classes = material_cfg["emb_dim"]
|
||||
cache_key = setup_cache.topn_key("material", n_classes)
|
||||
cached = cache.topn_maps.get(cache_key) if cache is not None else None
|
||||
if cached is not None:
|
||||
mat_topn_map = cached
|
||||
echo(f"material top-N map: cache hit ({len(mat_topn_map.class_map)} materials, {n_classes} classes)")
|
||||
else:
|
||||
echo("building material top-N map …")
|
||||
mat_topn_map = build_topn_map_from_files(files, "material", n_classes=n_classes, cast=str)
|
||||
echo(f" {len(mat_topn_map.class_map)} materials mapped to {n_classes} classes")
|
||||
if cache is not None:
|
||||
cache.topn_maps[cache_key] = mat_topn_map
|
||||
|
||||
energy_router_active = any(r.get("enabled") and r.get("type") == "energy" for r in (stage1_router, stage2_router))
|
||||
energy_idx = 3
|
||||
norm_key = setup_cache.normalizer_key(val_fraction, seed, particle_conditioning, material_conditioning)
|
||||
energy_router_active = (
|
||||
router_cfg.get("enabled") and router_cfg.get("type") == "energy"
|
||||
)
|
||||
energy_idx = router_cfg.get("energy_idx", 3)
|
||||
norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning)
|
||||
entry = cache.normalizers.get(norm_key) if cache is not None else None
|
||||
|
||||
if entry is not None:
|
||||
@@ -251,37 +163,33 @@ def run_setup_stage(
|
||||
# grid (setup_cache.energy_quantiles_from_sample) so centers can
|
||||
# instead be seeded from actual data quantiles below. Collected
|
||||
# whenever the setup cache is being populated, not only when *this*
|
||||
# run's router is energy-typed, so a later run enabling an energy
|
||||
# router against this same (val_fraction, seed, conditioning) key
|
||||
# never needs to rescan just to seed centers.
|
||||
# run's router is energy-typed, so a later run enabling
|
||||
# --router-type energy against this same (val_fraction, seed,
|
||||
# conditioning) key never needs to rescan just to seed centers.
|
||||
collect_energy_sample = energy_router_active or cache is not None
|
||||
energy_sampler = _ReservoirSampler(capacity=100_000) if collect_energy_sample else None
|
||||
energy_sampler = (
|
||||
_ReservoirSampler(capacity=100_000) if collect_energy_sample else None
|
||||
)
|
||||
for i, path in enumerate(files):
|
||||
for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max):
|
||||
for chunk in iter_file_chunks(path, offset=event_id_offset(i)):
|
||||
mask = sorted_membership(chunk["event_id"], events_arr)
|
||||
if not mask.any():
|
||||
continue
|
||||
chunk_tr = {k: v[mask] for k, v in chunk.items()}
|
||||
feats = build_features(
|
||||
cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = build_features(
|
||||
chunk_tr,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
proc_map=proc_map,
|
||||
require_secondaries=True,
|
||||
particle_conditioning=particle_conditioning,
|
||||
material_conditioning=material_conditioning,
|
||||
conditioning=conditioning,
|
||||
sec_phys_only=True,
|
||||
k_max=k_max,
|
||||
)
|
||||
cond_cont = feats.cond_cont
|
||||
target_s1 = feats.target_s1
|
||||
n_sec = feats.n_sec
|
||||
sec_cont = feats.sec_cont
|
||||
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(sec_cont.shape[1])[None, :] < n_sec[:, None]
|
||||
sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None]
|
||||
sec_phys = sec_cont[:, :, 4:6][sec_valid]
|
||||
if len(sec_phys) > 0:
|
||||
sec_phys_acc.update(sec_phys)
|
||||
@@ -298,8 +206,22 @@ def run_setup_stage(
|
||||
cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles
|
||||
)
|
||||
|
||||
_seed_energy_router(stage1_router, cond_norm, energy_quantiles, energy_idx, echo)
|
||||
_seed_energy_router(stage2_router, cond_norm, energy_quantiles, energy_idx, echo)
|
||||
if energy_router_active and energy_quantiles.size > 0:
|
||||
assert cond_norm.mean is not None and cond_norm.std is not None
|
||||
levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
|
||||
raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
|
||||
centers_init = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[
|
||||
energy_idx
|
||||
]
|
||||
router_cfg["centers_init"] = centers_init.astype(np.float32).tolist()
|
||||
echo(
|
||||
f" seeded EnergyRouter centers from data quantiles: {router_cfg['centers_init']}"
|
||||
)
|
||||
elif energy_router_active:
|
||||
echo(
|
||||
" warning: no energy samples collected — EnergyRouter falls back to "
|
||||
"default centers"
|
||||
)
|
||||
|
||||
if cache is not None:
|
||||
setup_cache.save(data, files, cache, echo=echo)
|
||||
@@ -309,9 +231,6 @@ def run_setup_stage(
|
||||
pdg_map=pdg_map,
|
||||
mat_map=mat_map,
|
||||
proc_map=proc_map,
|
||||
pdg_topn_map=pdg_topn_map,
|
||||
sec_type_topn_map=sec_type_topn_map,
|
||||
mat_topn_map=mat_topn_map,
|
||||
cond_norm=cond_norm,
|
||||
tgt_norm=tgt_norm,
|
||||
sec_phys_norm=sec_phys_norm,
|
||||
@@ -333,7 +252,7 @@ def run_train_job(
|
||||
rebuild_setup_cache: bool = False,
|
||||
echo=print,
|
||||
) -> None:
|
||||
t = cfg["train"]
|
||||
t, m = cfg["train"], cfg["model"]
|
||||
config.seed_everything(t["seed"])
|
||||
|
||||
out_dir = Path(out_dir)
|
||||
@@ -352,15 +271,20 @@ def run_train_job(
|
||||
"section)"
|
||||
)
|
||||
|
||||
config.validate_config(cfg)
|
||||
particle_conditioning = cfg["conditioning"]["particle"]["type"]
|
||||
material_conditioning = cfg["conditioning"]["material"]["type"]
|
||||
k_max = cfg["stage2_model"]["k_max"]
|
||||
router_cfg = m["router"]
|
||||
if t["mode"] == "wgan" and router_cfg.get("enabled"):
|
||||
raise ValueError(
|
||||
"--mode wgan does not support --router (no routed WGAN generator/"
|
||||
"critic exists) — disable one or the other"
|
||||
)
|
||||
|
||||
conditioning = m["conditioning"]
|
||||
setup = run_setup_stage(
|
||||
data,
|
||||
val_fraction=t["val_fraction"],
|
||||
seed=t["seed"],
|
||||
cfg=cfg,
|
||||
conditioning=conditioning,
|
||||
router_cfg=router_cfg,
|
||||
cache_setup=cache_setup,
|
||||
rebuild_setup_cache=rebuild_setup_cache,
|
||||
echo=echo,
|
||||
@@ -378,35 +302,6 @@ def run_train_job(
|
||||
setup.n_train_steps,
|
||||
)
|
||||
|
||||
# cond_cat's onehot columns are present per-axis, independently, under
|
||||
# that axis's own conditioning.{particle,material}.type == "onehot"
|
||||
# (the two axes may mix freely). run_setup_stage builds each map
|
||||
# whenever its own axis is "onehot" (see its own
|
||||
# particle_cfg["type"]/material_cfg["type"]
|
||||
# checks), so they're guaranteed non-None here — asserted, not just
|
||||
# assumed, so a future wiring bug fails loudly instead of silently
|
||||
# dropping the onehot columns.
|
||||
cond_pdg_topn = None
|
||||
cond_mat_topn = None
|
||||
if particle_conditioning == "onehot":
|
||||
assert setup.pdg_topn_map is not None
|
||||
cond_pdg_topn = setup.pdg_topn_map.class_map
|
||||
if material_conditioning == "onehot":
|
||||
assert setup.mat_topn_map is not None
|
||||
cond_mat_topn = setup.mat_topn_map.class_map
|
||||
|
||||
# The secondary type-index map depends on stage2_model.particle_type.target,
|
||||
# independently of conditioning's own onehot/embedding choice above
|
||||
# (physical stays untouched/None).
|
||||
particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target
|
||||
if particle_type_target == "onehot":
|
||||
assert setup.sec_type_topn_map is not None
|
||||
sec_type_class_map = setup.sec_type_topn_map.class_map
|
||||
elif particle_type_target == "embedding":
|
||||
sec_type_class_map = pdg_map
|
||||
else:
|
||||
sec_type_class_map = None
|
||||
|
||||
total_train_batches = n_train_steps // t["batch_size"]
|
||||
echo(f" ~{n_train_steps:,} train steps, ~{total_train_batches:,} batches")
|
||||
|
||||
@@ -421,13 +316,8 @@ def run_train_job(
|
||||
shuffle_buffer=shuffle_buffer,
|
||||
shuffle=True,
|
||||
proc_map=proc_map,
|
||||
particle_conditioning=particle_conditioning,
|
||||
material_conditioning=material_conditioning,
|
||||
conditioning=conditioning,
|
||||
sec_phys_normalizer=sec_phys_norm,
|
||||
pdg_topn_map=cond_pdg_topn,
|
||||
mat_topn_map=cond_mat_topn,
|
||||
sec_type_class_map=sec_type_class_map,
|
||||
k_max=k_max,
|
||||
)
|
||||
val_ds = StreamingStepsDataset(
|
||||
files=files,
|
||||
@@ -439,13 +329,8 @@ def run_train_job(
|
||||
batch_size=t["batch_size"],
|
||||
shuffle=False,
|
||||
proc_map=proc_map,
|
||||
particle_conditioning=particle_conditioning,
|
||||
material_conditioning=material_conditioning,
|
||||
conditioning=conditioning,
|
||||
sec_phys_normalizer=sec_phys_norm,
|
||||
pdg_topn_map=cond_pdg_topn,
|
||||
mat_topn_map=cond_mat_topn,
|
||||
sec_type_class_map=sec_type_class_map,
|
||||
k_max=k_max,
|
||||
)
|
||||
|
||||
pin = device.type == "cuda"
|
||||
@@ -462,19 +347,60 @@ def run_train_job(
|
||||
pin_memory=pin,
|
||||
)
|
||||
|
||||
emb_dim = m.get("emb_dim", EMB_DIM)
|
||||
expert_hidden_dim, expert_n_blocks = config.resolve_expert_dims(
|
||||
router_cfg, m["hidden_dim"], m["n_blocks"]
|
||||
)
|
||||
if router_cfg.get("enabled") and (expert_hidden_dim, expert_n_blocks) != (
|
||||
m["hidden_dim"],
|
||||
m["n_blocks"],
|
||||
):
|
||||
# Only reachable via an explicit router.expert_hidden_dim/n_blocks
|
||||
# override (the 0/"unset" sentinel always resolves to m["hidden_dim"]/
|
||||
# ["n_blocks"] — see resolve_expert_dims), so this is never a false
|
||||
# positive from inheritance, only a deliberate narrow/wide-experts
|
||||
# config the checkpoint dir name (_h{hidden_dim}_b{n_blocks}) won't
|
||||
# reflect.
|
||||
echo(
|
||||
f" warning: experts are {expert_hidden_dim}x{expert_n_blocks}, "
|
||||
f"different from model.hidden_dim/n_blocks ({m['hidden_dim']}x"
|
||||
f"{m['n_blocks']}) — the checkpoint dir name reflects the latter, "
|
||||
"not the experts actually being trained"
|
||||
)
|
||||
|
||||
model_config = {
|
||||
"pdg_vocab": len(pdg_map),
|
||||
"mat_vocab": len(mat_map),
|
||||
"conditioning": cfg["conditioning"],
|
||||
"stage1_model": cfg["stage1_model"],
|
||||
"stage2_model": cfg["stage2_model"],
|
||||
"hidden_dim": m["hidden_dim"],
|
||||
"n_blocks": m["n_blocks"],
|
||||
"emb_dim": emb_dim,
|
||||
"dropout": m["dropout"],
|
||||
"k_max": K_MAX,
|
||||
"sec_slot_dim": SEC_SLOT_DIM,
|
||||
"conditioning": conditioning,
|
||||
"router": dict(router_cfg),
|
||||
"expert_hidden_dim": expert_hidden_dim,
|
||||
"expert_n_blocks": expert_n_blocks,
|
||||
# Read by `predict`/`rollout` (which never receive their own --mode
|
||||
# flag) to auto-detect which sampler a checkpoint needs.
|
||||
"mode": t["mode"],
|
||||
"noise_dim": m.get("noise_dim", 64),
|
||||
}
|
||||
|
||||
models = build_models(model_config)
|
||||
critics = build_critics(model_config)
|
||||
for name, model in models.items():
|
||||
if model is not None:
|
||||
echo(f"{name}: {sum(p.numel() for p in model.parameters()):,} parameters")
|
||||
stage1_model, sec_decoder = build_models(model_config)
|
||||
echo(
|
||||
f"stage1: {sum(p.numel() for p in stage1_model.parameters()):,} parameters | "
|
||||
f"sec_decoder: {sum(p.numel() for p in sec_decoder.parameters()):,} parameters"
|
||||
)
|
||||
|
||||
critic = None
|
||||
sec_critic = None
|
||||
if t["mode"] == "wgan":
|
||||
critic, sec_critic = build_critics(model_config)
|
||||
echo(
|
||||
f"critic: {sum(p.numel() for p in critic.parameters()):,} parameters | "
|
||||
f"sec_critic: {sum(p.numel() for p in sec_critic.parameters()):,} parameters"
|
||||
)
|
||||
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
meta = config.build_run_meta(
|
||||
@@ -489,13 +415,25 @@ def run_train_job(
|
||||
config.save_config(cfg, out_dir, meta)
|
||||
|
||||
run_training(
|
||||
cfg=cfg,
|
||||
models=models,
|
||||
critics=critics,
|
||||
stage1_model=stage1_model,
|
||||
sec_decoder=sec_decoder,
|
||||
train_loader=train_loader,
|
||||
val_loader=val_loader,
|
||||
mode=t["mode"],
|
||||
epochs=t["epochs"],
|
||||
lr=t["lr"],
|
||||
weight_decay=t["weight_decay"],
|
||||
ema_decay=t["ema_decay"],
|
||||
warmup_epochs=t["warmup_epochs"],
|
||||
device=device,
|
||||
out_dir=out_dir,
|
||||
lambda_nsec=t.get("lambda_nsec", 0.1),
|
||||
lambda_s2=t.get("lambda_s2", 1.0),
|
||||
lambda_balance=router_cfg.get("lambda_balance", 0.0),
|
||||
lambda_proc=router_cfg.get("lambda_proc", 0.0),
|
||||
lambda_entropy=router_cfg.get("lambda_entropy", 0.0),
|
||||
gumbel_tau_start=router_cfg.get("gumbel_tau_start", 1.0),
|
||||
gumbel_tau_end=router_cfg.get("gumbel_tau_end", 0.1),
|
||||
normalizer_dict={
|
||||
"cond": cond_norm.to_dict(),
|
||||
"target": tgt_norm.to_dict(),
|
||||
@@ -504,13 +442,18 @@ def run_train_job(
|
||||
pdg_map={str(k): v for k, v in pdg_map.items()},
|
||||
mat_map={str(k): v for k, v in mat_map.items()},
|
||||
proc_map=proc_map,
|
||||
pdg_topn_map=setup.pdg_topn_map,
|
||||
sec_type_topn_map=setup.sec_type_topn_map,
|
||||
mat_topn_map=setup.mat_topn_map,
|
||||
model_config=model_config,
|
||||
resume_path=resume,
|
||||
validate_every=t["validate_every"],
|
||||
validate_steps=t["validate_steps"],
|
||||
max_val_batches=t["max_val_batches"],
|
||||
total_train_batches=total_train_batches,
|
||||
use_wandb=t.get("wandb", True),
|
||||
critic=critic,
|
||||
sec_critic=sec_critic,
|
||||
n_critic=t.get("n_critic", 5),
|
||||
gp_weight=t.get("gp_weight", 10.0),
|
||||
critic_lr=t.get("critic_lr") or None,
|
||||
use_wandb=t.get("wandb", False),
|
||||
wandb_project=t.get("wandb_project", "giant"),
|
||||
wandb_run_name=t.get("wandb_run_name", ""),
|
||||
wandb_log_every=t.get("wandb_log_every", 50),
|
||||
|
||||
+83
-262
@@ -17,7 +17,7 @@ treated as detector leakage and not deposited.
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import Counter
|
||||
from typing import TYPE_CHECKING, Callable, TypedDict
|
||||
from typing import Callable, TypedDict
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
@@ -33,156 +33,18 @@ from giant.data.transforms import (
|
||||
Normalizer,
|
||||
build_cond_features,
|
||||
decode_secondaries,
|
||||
decode_secondary_cont,
|
||||
energy_simplex_decode,
|
||||
inv_local_frame_rotation,
|
||||
inv_log_transform,
|
||||
reconstruct_post_pos,
|
||||
)
|
||||
from giant.particles import (
|
||||
decode_embedding_nearest,
|
||||
decode_topn_class,
|
||||
invert_dense_map,
|
||||
nearest_known_pdg,
|
||||
particle_phys_array,
|
||||
from giant.particles import nearest_known_pdg, particle_phys_array
|
||||
from giant.sample import (
|
||||
sample_flow,
|
||||
sample_secondaries,
|
||||
sample_wgan,
|
||||
sample_secondaries_wgan,
|
||||
)
|
||||
from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from giant.data.loader import TopNMap
|
||||
|
||||
|
||||
class L1DistCollector:
|
||||
"""Accumulates the L1-distance diagnostic across a whole rollout
|
||||
run: the L1 distance between each emitted secondary's raw predicted
|
||||
embedding vector and the nearest table row it snapped to (only
|
||||
meaningful under `particle_type.target = "embedding"` —
|
||||
`giant.particles.decode_embedding_nearest`). A heavy tail means the
|
||||
decoder is emitting vectors off the embedding manifold — the direct
|
||||
analogue of the species-collapse symptom the v0.3.0 redesign exists to
|
||||
fix.
|
||||
|
||||
Not folded into `rollout()`'s own return value (which is shape-typed as
|
||||
step records, see `_RECORD_KEYS`/`RolloutSummary`) — passed in and read
|
||||
back by the caller instead, mirroring the existing `on_chunk` pattern.
|
||||
O(1) memory via a fixed log-spaced histogram rather than raw samples,
|
||||
since a heavy right tail is exactly what this diagnostic watches for.
|
||||
"""
|
||||
|
||||
def __init__(self, n_bins: int = 50, lo: float = 1e-3, hi: float = 1e3) -> None:
|
||||
self.n = 0
|
||||
self.total = 0.0
|
||||
self.total_sq = 0.0
|
||||
self.minimum = float("inf")
|
||||
self.maximum = 0.0
|
||||
self.hist_edges = np.geomspace(lo, hi, n_bins + 1)
|
||||
self.hist_counts = np.zeros(n_bins, dtype=np.int64)
|
||||
|
||||
def add(self, dist: np.ndarray, valid: np.ndarray) -> None:
|
||||
vals = np.asarray(dist)[np.asarray(valid)]
|
||||
if vals.size == 0:
|
||||
return
|
||||
self.n += int(vals.size)
|
||||
self.total += float(vals.sum())
|
||||
self.total_sq += float(np.square(vals).sum())
|
||||
self.minimum = min(self.minimum, float(vals.min()))
|
||||
self.maximum = max(self.maximum, float(vals.max()))
|
||||
self.hist_counts += np.histogram(vals, bins=self.hist_edges)[0]
|
||||
|
||||
def summary(self) -> dict | None:
|
||||
"""`None` if nothing was ever added (target != "embedding", or a
|
||||
run with zero secondaries) — the caller should omit the diagnostic
|
||||
entirely rather than write a degenerate summary."""
|
||||
if self.n == 0:
|
||||
return None
|
||||
mean = self.total / self.n
|
||||
variance = max(self.total_sq / self.n - mean**2, 0.0)
|
||||
return {
|
||||
"n": self.n,
|
||||
"mean": mean,
|
||||
"std": variance**0.5,
|
||||
"min": self.minimum,
|
||||
"max": self.maximum,
|
||||
"hist_edges": self.hist_edges.tolist(),
|
||||
"hist_counts": self.hist_counts.tolist(),
|
||||
}
|
||||
|
||||
|
||||
def decode_secondary_identity(
|
||||
sec_decoder: torch.nn.Module,
|
||||
sec_cont: torch.Tensor,
|
||||
sec_type: torch.Tensor,
|
||||
n_sec_np: np.ndarray,
|
||||
e_sec: np.ndarray,
|
||||
pre_dir: np.ndarray,
|
||||
sec_phys_norm: Normalizer,
|
||||
pdg_map: dict[int, int],
|
||||
sec_type_topn_map: "TopNMap | None",
|
||||
other_policy: str,
|
||||
rng: np.random.Generator | None,
|
||||
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
|
||||
"""Decode Stage 2's raw (sec_cont, sec_type) output into physical
|
||||
secondary attributes, branching on `sec_decoder.particle_type_cfg`:
|
||||
|
||||
- `"physical"`: unchanged v0.2 path — `sec_type` already *is* (log_mass,
|
||||
charge), used as the secondary's identity as-is (no snapping).
|
||||
- `"onehot"`: `sec_type` is per-slot class logits — argmax, then
|
||||
`giant.particles.decode_topn_class` (+ `other_policy`) resolves a
|
||||
concrete PDG, whose real physics (log_mass, charge) then come from
|
||||
`giant.particles.particle_phys_array` — unlike "physical", the PDG
|
||||
resolution IS the secondary's identity here, not just a reporting
|
||||
label.
|
||||
- `"embedding"`: `sec_type` is a raw vector in the conditioning's own
|
||||
embedding space — `giant.particles.decode_embedding_nearest` L1-snaps
|
||||
it to the nearest table row for the PDG (+ physics via
|
||||
`particle_phys_array`), and also returns the L1 distance (see this
|
||||
module's `L1DistCollector`).
|
||||
|
||||
Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg,
|
||||
sec_type_l1_dist) — the last is `None` except under `"embedding"`.
|
||||
"""
|
||||
target = sec_decoder.particle_type_cfg.get("target", "physical")
|
||||
|
||||
if target == "physical":
|
||||
sec_full = torch.cat([sec_cont, sec_type], dim=-1).cpu().numpy()
|
||||
sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
|
||||
sec_full, n_sec_np, e_sec, pre_dir, sec_phys_normalizer=sec_phys_norm
|
||||
)
|
||||
sec_pdg = nearest_known_pdg(sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()).reshape(
|
||||
sec_mass.shape
|
||||
)
|
||||
return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, None
|
||||
|
||||
sec_E, sec_dir_world, sec_valid = decode_secondary_cont(sec_cont.cpu().numpy(), n_sec_np, e_sec, pre_dir)
|
||||
sec_type_np = sec_type.cpu().numpy()
|
||||
l1_dist = None
|
||||
|
||||
if target == "onehot":
|
||||
if sec_type_topn_map is None:
|
||||
raise RuntimeError(
|
||||
"particle_type.target='onehot' rollout needs sec_type_topn_map "
|
||||
"(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
|
||||
)
|
||||
class_idx = sec_type_np.argmax(axis=-1)
|
||||
sec_pdg = decode_topn_class(
|
||||
class_idx,
|
||||
sec_type_topn_map,
|
||||
n_classes=sec_decoder.type_dim,
|
||||
other_policy=other_policy,
|
||||
rng=rng,
|
||||
)
|
||||
else: # "embedding"
|
||||
idx_to_pdg = invert_dense_map(pdg_map)
|
||||
emb_weight = sec_decoder.cond_enc.pdg_emb.weight.detach().cpu().numpy()
|
||||
sec_pdg, l1_dist = decode_embedding_nearest(sec_type_np, emb_weight, idx_to_pdg)
|
||||
l1_dist = np.where(sec_valid, l1_dist, 0.0).astype(np.float32)
|
||||
|
||||
sec_mass, sec_charge = particle_phys_array(sec_pdg.reshape(-1)).T
|
||||
sec_mass = np.where(sec_valid, sec_mass.reshape(sec_pdg.shape), 0.0).astype(np.float32)
|
||||
sec_charge = np.where(sec_valid, sec_charge.reshape(sec_pdg.shape), 0.0).astype(np.float32)
|
||||
sec_pdg = np.where(sec_valid, sec_pdg, 0).astype(np.int64)
|
||||
return sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg, l1_dist
|
||||
|
||||
|
||||
# Record columns produced per step / per terminal marker.
|
||||
_RECORD_KEYS = [
|
||||
@@ -291,9 +153,13 @@ class _Recorder:
|
||||
RAM until the very end.
|
||||
"""
|
||||
|
||||
def __init__(self, sink: Callable[[dict[str, np.ndarray]], None] | None = None) -> None:
|
||||
def __init__(
|
||||
self, sink: Callable[[dict[str, np.ndarray]], None] | None = None
|
||||
) -> None:
|
||||
self._sink = sink
|
||||
self._cols: dict[str, list] | None = None if sink is not None else {k: [] for k in _RECORD_KEYS}
|
||||
self._cols: dict[str, list] | None = (
|
||||
None if sink is not None else {k: [] for k in _RECORD_KEYS}
|
||||
)
|
||||
self.n_rows = 0
|
||||
self.termination_reason_counts: Counter[str] = Counter()
|
||||
|
||||
@@ -301,7 +167,10 @@ class _Recorder:
|
||||
n = len(cols["event_id"])
|
||||
if n == 0:
|
||||
return
|
||||
row = {k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n) for k in _RECORD_KEYS}
|
||||
row = {
|
||||
k: np.asarray(cols[k], dtype=_RECORD_DTYPES[k]).reshape(n)
|
||||
for k in _RECORD_KEYS
|
||||
}
|
||||
self.n_rows += n
|
||||
reasons = row["termination_reason"]
|
||||
nonempty = reasons[reasons != ""]
|
||||
@@ -318,7 +187,8 @@ class _Recorder:
|
||||
|
||||
def to_dict(self) -> dict[str, np.ndarray]:
|
||||
assert self._cols is not None, (
|
||||
"to_dict() is unavailable when streaming to a sink — use n_rows/termination_reason_counts instead"
|
||||
"to_dict() is unavailable when streaming to a sink — use "
|
||||
"n_rows/termination_reason_counts instead"
|
||||
)
|
||||
out = {}
|
||||
for k, chunks in self._cols.items():
|
||||
@@ -335,7 +205,7 @@ def make_seed_frontier(
|
||||
pre_pos: np.ndarray,
|
||||
pre_E: np.ndarray,
|
||||
pre_dir: np.ndarray,
|
||||
particle_conditioning: str = "embedding",
|
||||
conditioning: str = "embedding",
|
||||
) -> tuple[dict[str, np.ndarray], dict[int, int]]:
|
||||
"""Build the initial frontier from primary entry states.
|
||||
|
||||
@@ -355,17 +225,17 @@ def make_seed_frontier(
|
||||
dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None)
|
||||
|
||||
pdg_arr = np.asarray(pdg, dtype=np.int64)
|
||||
if particle_conditioning == "physical":
|
||||
if conditioning == "physical":
|
||||
# Real primaries always have a genuine ground-truth PDG code, looked
|
||||
# up once here and carried forward unchanged for the track's lifetime
|
||||
# (its species never changes mid-track) — same lifecycle as "pdg"
|
||||
# itself.
|
||||
mass, charge = particle_phys_array(pdg_arr).T
|
||||
else:
|
||||
# "embedding"/"onehot" never read mass/charge (see
|
||||
# "embedding" mode never reads mass/charge (see
|
||||
# _physical_cond_columns), so resolving them here would only risk
|
||||
# crashing a rollout on a PDG code giant.particles can't resolve, for
|
||||
# a value that's never used.
|
||||
# crashing an embedding-mode rollout on a PDG code giant.particles
|
||||
# can't resolve, for a value that's never used.
|
||||
mass = np.zeros(n, dtype=np.float64)
|
||||
charge = np.zeros(n, dtype=np.float64)
|
||||
|
||||
@@ -440,16 +310,8 @@ def rollout(
|
||||
max_tracks_per_event: int | None = None,
|
||||
escape_threshold: float | None = None,
|
||||
on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
|
||||
particle_conditioning: str = "embedding",
|
||||
material_conditioning: str = "embedding",
|
||||
pdg_topn_map: "TopNMap | None" = None,
|
||||
mat_topn_map: "TopNMap | None" = None,
|
||||
sec_type_topn_map: "TopNMap | None" = None,
|
||||
other_policy: str = "sample",
|
||||
seed: int | None = None,
|
||||
stage1_ddpm_steps: int = 1000,
|
||||
stage2_ddpm_steps: int = 1000,
|
||||
l1_dist_collector: "L1DistCollector | None" = None,
|
||||
conditioning: str = "embedding",
|
||||
mode: str = "flow",
|
||||
) -> dict[str, np.ndarray] | RolloutSummary:
|
||||
"""Run showers to completion.
|
||||
|
||||
@@ -462,51 +324,12 @@ def rollout(
|
||||
dict[str, int]}`. Use this for large `--n-events`/`--max-steps` runs,
|
||||
where the full record set would otherwise scale with
|
||||
`n_events * max_steps * avg_tracks_per_event`.
|
||||
|
||||
There is no `mode` parameter — each stage's generative objective is read
|
||||
directly off the model instance's own `generator_kind` (stage 1 and
|
||||
stage 2 objectives are independent, e.g. `stage1_model.generator="flow"`
|
||||
+ `stage2_model.generator="wgan"`), and the decoder (one-shot vs
|
||||
autoregressive) is inferred from `sec_decoder`'s own class — see
|
||||
`sample_stage1`/`sample_stage2` (giant.sample).
|
||||
|
||||
`pdg_topn_map`/`mat_topn_map`/`sec_type_topn_map` serve three independent
|
||||
purposes, no longer required to share one map (see gitea #29):
|
||||
`pdg_topn_map`/`mat_topn_map` are required whenever
|
||||
`particle_conditioning`/`material_conditioning` is `"onehot"` (feeds
|
||||
`build_cond_features`'s extra `cond_cat` top-N columns); `sec_type_topn_map`/
|
||||
`other_policy` are required instead under
|
||||
`stage2_model.particle_type.target = "onehot"` (secondary-species
|
||||
decode) — its class count (`stage2_model.particle_type.n_classes`) may
|
||||
differ from `pdg_topn_map`'s. `seed` seeds the `other_policy = "sample"`
|
||||
draw only (torch/numpy sampling itself is seeded by the caller, same as
|
||||
today).
|
||||
|
||||
`l1_dist_collector`, if given, accumulates the embedding-distance
|
||||
diagnostic across the whole run — see `L1DistCollector`. Only populated
|
||||
under `particle_type.target = "embedding"`; a no-op otherwise.
|
||||
"""
|
||||
if particle_conditioning == "onehot" and pdg_topn_map is None:
|
||||
raise RuntimeError(
|
||||
"conditioning.particle.type='onehot' rollout needs pdg_topn_map "
|
||||
"(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']"
|
||||
)
|
||||
if sec_decoder.particle_type_cfg.get("target") == "onehot" and sec_type_topn_map is None:
|
||||
raise RuntimeError(
|
||||
"stage2_model.particle_type.target='onehot' rollout needs sec_type_topn_map "
|
||||
"(the checkpoint's saved top-N map) — see ckpt['sec_type_topn_map']"
|
||||
)
|
||||
if material_conditioning == "onehot" and mat_topn_map is None:
|
||||
raise RuntimeError(
|
||||
"conditioning.material.type='onehot' rollout needs mat_topn_map "
|
||||
"(the checkpoint's saved top-N map) — see ckpt['mat_topn_map']"
|
||||
)
|
||||
device = device or torch.device("cpu")
|
||||
stage1_model.eval()
|
||||
sec_decoder.eval()
|
||||
if escape_threshold is not None:
|
||||
oracle.escape_threshold = float(escape_threshold)
|
||||
rng = np.random.default_rng(seed)
|
||||
|
||||
frontier, counts = make_seed_frontier(
|
||||
seeds["event_id"],
|
||||
@@ -514,7 +337,7 @@ def rollout(
|
||||
seeds["pre_pos"],
|
||||
seeds["pre_E"],
|
||||
seeds["pre_dir"],
|
||||
particle_conditioning=particle_conditioning,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
rec = _Recorder(sink=on_chunk)
|
||||
|
||||
@@ -541,16 +364,8 @@ def rollout(
|
||||
steps,
|
||||
device,
|
||||
max_tracks_per_event,
|
||||
particle_conditioning,
|
||||
material_conditioning,
|
||||
pdg_topn_map,
|
||||
mat_topn_map,
|
||||
sec_type_topn_map,
|
||||
other_policy,
|
||||
rng,
|
||||
stage1_ddpm_steps,
|
||||
stage2_ddpm_steps,
|
||||
l1_dist_collector,
|
||||
conditioning,
|
||||
mode,
|
||||
)
|
||||
)
|
||||
frontier = _concat_frontiers(next_parts)
|
||||
@@ -580,16 +395,8 @@ def _step_chunk(
|
||||
steps,
|
||||
device,
|
||||
max_tracks_per_event,
|
||||
particle_conditioning,
|
||||
material_conditioning,
|
||||
pdg_topn_map,
|
||||
mat_topn_map,
|
||||
sec_type_topn_map,
|
||||
other_policy,
|
||||
rng,
|
||||
stage1_ddpm_steps,
|
||||
stage2_ddpm_steps,
|
||||
l1_dist_collector,
|
||||
conditioning,
|
||||
mode="flow",
|
||||
) -> dict[str, np.ndarray]:
|
||||
"""Advance one chunk of tracks by a single step; return the next frontier."""
|
||||
n = len(tr["event_id"])
|
||||
@@ -600,7 +407,7 @@ def _step_chunk(
|
||||
tr["_material"] = material
|
||||
tr["_layer_id"] = layer_id
|
||||
|
||||
if particle_conditioning == "physical":
|
||||
if conditioning == "physical":
|
||||
# Under physical-property conditioning, mass/charge (already resolved
|
||||
# on every track — see the cond_dict comment below) drive the model,
|
||||
# not a training-vocab PDG embedding — build_cond_features passes
|
||||
@@ -614,19 +421,33 @@ def _step_chunk(
|
||||
# --- Pre-step termination gates (in priority order; each track picks one) ---
|
||||
stop = np.zeros(n, dtype=bool)
|
||||
escaped_sel = escaped & ~stop
|
||||
rec.add(**_terminal_rows(tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))))
|
||||
rec.add(
|
||||
**_terminal_rows(
|
||||
tr, escaped_sel, TERM_ESCAPED, edep=np.zeros(int(escaped_sel.sum()))
|
||||
)
|
||||
)
|
||||
stop |= escaped_sel
|
||||
|
||||
unknown_sel = ~known_pdg & ~stop
|
||||
rec.add(**_terminal_rows(tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]))
|
||||
rec.add(
|
||||
**_terminal_rows(
|
||||
tr, unknown_sel, TERM_UNKNOWN_PDG, edep=tr["pre_E"][unknown_sel]
|
||||
)
|
||||
)
|
||||
stop |= unknown_sel
|
||||
|
||||
cutoff_sel = (tr["pre_E"] < energy_cutoff) & ~stop
|
||||
rec.add(**_terminal_rows(tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]))
|
||||
rec.add(
|
||||
**_terminal_rows(
|
||||
tr, cutoff_sel, TERM_ENERGY_CUTOFF, edep=tr["pre_E"][cutoff_sel]
|
||||
)
|
||||
)
|
||||
stop |= cutoff_sel
|
||||
|
||||
maxstep_sel = (tr["step_in_track"] >= max_steps) & ~stop
|
||||
rec.add(**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel]))
|
||||
rec.add(
|
||||
**_terminal_rows(tr, maxstep_sel, TERM_MAX_STEPS, edep=tr["pre_E"][maxstep_sel])
|
||||
)
|
||||
stop |= maxstep_sel
|
||||
|
||||
active = ~stop
|
||||
@@ -654,60 +475,60 @@ def _step_chunk(
|
||||
"charge": tr["charge"],
|
||||
}
|
||||
cond_cont, cond_cat = build_cond_features(
|
||||
cond_dict,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
cond_norm,
|
||||
particle_conditioning=particle_conditioning,
|
||||
material_conditioning=material_conditioning,
|
||||
pdg_topn_map=pdg_topn_map.class_map if particle_conditioning == "onehot" else None,
|
||||
mat_topn_map=mat_topn_map.class_map if material_conditioning == "onehot" else None,
|
||||
cond_dict, pdg_map, mat_map, cond_norm, conditioning=conditioning
|
||||
)
|
||||
cc = torch.from_numpy(cond_cont).float().to(device)
|
||||
ck = torch.from_numpy(cond_cat).long().to(device)
|
||||
|
||||
stage1_norm, n_sec_pred_stage1 = sample_stage1(stage1_model, cc, ck, steps, stage1_ddpm_steps)
|
||||
if mode == "wgan":
|
||||
stage1_norm, n_sec_pred = sample_wgan(stage1_model, cc, ck)
|
||||
else:
|
||||
stage1_norm, n_sec_pred = sample_flow(stage1_model, cc, ck, steps=steps)
|
||||
raw = tgt_norm.inverse_transform(stage1_norm.cpu().numpy())
|
||||
|
||||
step_length = inv_log_transform(raw[:, 0])
|
||||
edep, e_sec, post_E, _delta = energy_simplex_decode(raw[:, 1:3], tr["pre_E"])
|
||||
|
||||
post_dir_local = raw[:, 3:6].copy()
|
||||
post_dir_local /= np.clip(np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None)
|
||||
post_dir_local /= np.clip(
|
||||
np.linalg.norm(post_dir_local, axis=1, keepdims=True), 1e-8, None
|
||||
)
|
||||
post_dir_world = inv_local_frame_rotation(tr["pre_dir"], post_dir_local)
|
||||
|
||||
travel_dir_local = raw[:, 6:9].copy()
|
||||
travel_dir_local /= np.clip(np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None)
|
||||
post_pos = reconstruct_post_pos(tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local)
|
||||
travel_dir_local /= np.clip(
|
||||
np.linalg.norm(travel_dir_local, axis=1, keepdims=True), 1e-8, None
|
||||
)
|
||||
post_pos = reconstruct_post_pos(
|
||||
tr["pre_pos"], tr["pre_dir"], step_length, travel_dir_local
|
||||
)
|
||||
|
||||
n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cc, ck, stage1_norm, n_sec_pred_stage1)
|
||||
n_sec_np = n_sec_pred.cpu().numpy().astype(np.int64)
|
||||
|
||||
# --- Secondaries ---
|
||||
# No snapping for "physical"/history-facing state elsewhere in the
|
||||
# pipeline: sec_mass/sec_charge (or, for "onehot"/"embedding", the
|
||||
# resolved sec_pdg -> real physics) are the secondary's identity, used
|
||||
# as-is for the spawned track's own future conditioning — see
|
||||
# decode_secondary_identity's docstring for how each
|
||||
# particle_type.target differs on whether PDG resolution is a real
|
||||
# identity decision or just a reporting label.
|
||||
sec_cont, sec_type, _valid = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps)
|
||||
sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, sec_type_l1_dist = decode_secondary_identity(
|
||||
sec_decoder,
|
||||
sec_cont,
|
||||
sec_type,
|
||||
# No snapping: sec_mass/sec_charge are the model's raw predicted physical
|
||||
# identity, used as-is for the spawned track's own future conditioning.
|
||||
# sec_pdg_code below is a *separate*, reporting-only nearest-known-PDG
|
||||
# label (never fed back into the model) — see giant/particles.py.
|
||||
if mode == "wgan":
|
||||
sec_cont, sec_phys, _valid = sample_secondaries_wgan(
|
||||
sec_decoder, cc, ck, stage1_norm, n_sec_pred
|
||||
)
|
||||
else:
|
||||
sec_cont, sec_phys, _valid = sample_secondaries(
|
||||
sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps
|
||||
)
|
||||
sec_full = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy()
|
||||
sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid = decode_secondaries(
|
||||
sec_full,
|
||||
n_sec_np,
|
||||
e_sec,
|
||||
tr["pre_dir"],
|
||||
sec_phys_norm,
|
||||
pdg_map,
|
||||
sec_type_topn_map,
|
||||
other_policy,
|
||||
rng,
|
||||
sec_phys_normalizer=sec_phys_norm,
|
||||
)
|
||||
sec_valid = np.arange(sec_E.shape[1])[None, :] < n_sec_np[:, None]
|
||||
if l1_dist_collector is not None and sec_type_l1_dist is not None:
|
||||
l1_dist_collector.add(sec_type_l1_dist, sec_valid)
|
||||
sec_pdg_code = nearest_known_pdg(
|
||||
sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys()
|
||||
).reshape(sec_mass.shape)
|
||||
|
||||
edep = edep.astype(np.float64)
|
||||
post_E = post_E.astype(np.float64)
|
||||
|
||||
+109
-370
@@ -1,22 +1,26 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from giant.constants import CONT_SLOT_DIM, X_DIM
|
||||
from giant.model.network import Stage2Autoregressive, stage2_trunk_sec_dim
|
||||
from giant.model.schedule import CosineSchedule
|
||||
from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
|
||||
|
||||
|
||||
def _predict_n_sec_if_owned(
|
||||
model: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor
|
||||
) -> torch.Tensor | None:
|
||||
"""Stage-1 `n_sec_head` is only present on a migrated v0.2 checkpoint
|
||||
(fresh runs move it to stage 2 — see `Stage1Model`'s docstring). `None`
|
||||
here means "ask stage 2 instead", which every caller (`giant/rollout.py`,
|
||||
`giant/cli.py`) must do for a fresh checkpoint."""
|
||||
if getattr(model, "n_sec_head", None) is None:
|
||||
return None
|
||||
logits = model.predict_n_sec(cond_cont, cond_cat)
|
||||
return logits.argmax(dim=-1)
|
||||
def _slots_from_flat(
|
||||
x: torch.Tensor, n_sec_pred: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Reshape a flat (B, SEC_DIM) decoder output into per-slot tensors.
|
||||
|
||||
Returns (sec_cont, sec_phys, sec_valid) — see `sample_secondaries`'s
|
||||
docstring for their shapes/meaning. Shared by both the flow-matching and
|
||||
WGAN Stage-2 samplers, which differ only in how `x` was produced.
|
||||
"""
|
||||
B = x.size(0)
|
||||
device = x.device
|
||||
x_slots = x.view(B, K_MAX, SEC_SLOT_DIM)
|
||||
sec_cont = x_slots[:, :, :4]
|
||||
sec_phys = x_slots[:, :, 4:]
|
||||
sec_valid = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(
|
||||
1
|
||||
)
|
||||
return sec_cont, sec_phys, sec_valid
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -25,14 +29,12 @@ def sample_flow(
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
steps: int = 10,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Euler integration of the Stage-1 vector field from t=0 to t=1.
|
||||
|
||||
Returns (primary_sample, n_sec_pred):
|
||||
primary_sample: (B, X_DIM) — normalised 9D primary post-step output
|
||||
n_sec_pred: (B,) int64 — predicted secondary count, or `None` if
|
||||
`model` has no `n_sec_head` (a fresh v0.3.0 Stage1Model — see
|
||||
`_predict_n_sec_if_owned`).
|
||||
n_sec_pred: (B,) int64 — predicted secondary count
|
||||
"""
|
||||
model.eval()
|
||||
B = cond_cont.size(0)
|
||||
@@ -41,138 +43,11 @@ def sample_flow(
|
||||
dt = 1.0 / steps
|
||||
for i in range(steps):
|
||||
t = torch.full((B,), i * dt, device=device)
|
||||
v = model(x, cond_cont, cond_cat, t=t)
|
||||
v = model(x, t, cond_cont, cond_cat)
|
||||
x = x + v * dt
|
||||
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_ddpm(
|
||||
model: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
schedule,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)
|
||||
— see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
|
||||
model.eval()
|
||||
B = cond_cont.size(0)
|
||||
device = cond_cont.device
|
||||
x = torch.randn(B, X_DIM, device=device)
|
||||
T = schedule.T
|
||||
for i in reversed(range(T)):
|
||||
t_norm = torch.full((B,), i / T, device=device)
|
||||
eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
|
||||
beta = schedule.betas[i]
|
||||
alpha = schedule.alphas[i]
|
||||
alpha_bar = schedule.alpha_bars[i]
|
||||
z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
|
||||
x = (1.0 / alpha.sqrt()) * (x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred) + beta.sqrt() * z
|
||||
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_ddim(
|
||||
model: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
schedule,
|
||||
steps: int = 50,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)
|
||||
— see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
|
||||
model.eval()
|
||||
B = cond_cont.size(0)
|
||||
device = cond_cont.device
|
||||
T = schedule.T
|
||||
timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
|
||||
x = torch.randn(B, X_DIM, device=device)
|
||||
for step_idx, ts in enumerate(timesteps):
|
||||
t_idx = int(ts.item())
|
||||
t_norm = torch.full((B,), t_idx / T, device=device)
|
||||
eps_pred = model(x, cond_cont, cond_cat, t=t_norm)
|
||||
ab_t = schedule.alpha_bars[t_idx]
|
||||
if step_idx + 1 < len(timesteps):
|
||||
ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
|
||||
else:
|
||||
ab_prev = torch.ones(1, device=device)
|
||||
x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
|
||||
x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
|
||||
return x, _predict_n_sec_if_owned(model, cond_cont, cond_cat)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_wgan(
|
||||
generator: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)
|
||||
— see `sample_flow`'s docstring for the `n_sec_pred` `None` case."""
|
||||
generator.eval()
|
||||
B = cond_cont.size(0)
|
||||
z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
|
||||
x = generator(z, cond_cont, cond_cat)
|
||||
return x, _predict_n_sec_if_owned(generator, cond_cont, cond_cat)
|
||||
|
||||
|
||||
def _stage2_flat_width(sec_decoder: torch.nn.Module) -> int:
|
||||
"""The width of `sec_decoder`'s own trunk in/out vector — folded
|
||||
(continuous + type) under `particle_type.target = "physical"` or
|
||||
`generator = "wgan"`, continuous-only otherwise (the type slice then
|
||||
comes from `predict_type` instead — see `stage2_trunk_sec_dim`'s
|
||||
docstring)."""
|
||||
return stage2_trunk_sec_dim(
|
||||
sec_decoder.particle_type_cfg,
|
||||
sec_decoder.generator_kind,
|
||||
sec_decoder.k_max,
|
||||
sec_decoder.type_dim,
|
||||
)
|
||||
|
||||
|
||||
def _type_folded(sec_decoder: torch.nn.Module) -> bool:
|
||||
target = sec_decoder.particle_type_cfg.get("target", "physical")
|
||||
return target == "physical" or sec_decoder.generator_kind == "wgan"
|
||||
|
||||
|
||||
def _decode_stage2_flat(
|
||||
sec_decoder: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
n_sec_pred: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Reshape a flat `(B, flat_width)` `Stage2OneShot` output into per-slot
|
||||
tensors, generator/`particle_type.target`-agnostic: shared by
|
||||
`sample_secondaries`/`sample_secondaries_wgan`, which differ only in how
|
||||
`x` was produced.
|
||||
|
||||
Returns (sec_cont, sec_type, sec_valid):
|
||||
sec_cont: (B, k_max, CONT_SLOT_DIM) — [stick_logit, local_dir]
|
||||
sec_type: (B, k_max, type_dim) — under `target="physical"` this is
|
||||
[log_mass, charge] (normalised iff the checkpoint's sec_phys
|
||||
normalizer was applied at training time — denormalize before
|
||||
treating as physical units; see
|
||||
giant.data.transforms.decode_secondaries); under `"onehot"` /
|
||||
`"embedding"` it is raw class logits / an embedding-space vector
|
||||
— decode via giant.particles.decode_topn_class /
|
||||
decode_embedding_nearest (see giant/rollout.py).
|
||||
sec_valid: (B, k_max) bool — True for slots i < n_sec_pred
|
||||
"""
|
||||
B = x.size(0)
|
||||
device = x.device
|
||||
k_max = sec_decoder.k_max
|
||||
type_dim = sec_decoder.type_dim
|
||||
if _type_folded(sec_decoder):
|
||||
x_slots = x.view(B, k_max, CONT_SLOT_DIM + type_dim)
|
||||
sec_cont = x_slots[:, :, :CONT_SLOT_DIM]
|
||||
sec_type = x_slots[:, :, CONT_SLOT_DIM:]
|
||||
else:
|
||||
sec_cont = x.view(B, k_max, CONT_SLOT_DIM)
|
||||
sec_type = sec_decoder.predict_type(cond_cont, cond_cat, stage1_out)
|
||||
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
|
||||
return sec_cont, sec_type, sec_valid
|
||||
n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
|
||||
n_sec_pred = n_sec_logits.argmax(dim=-1)
|
||||
return x, n_sec_pred
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -184,25 +59,76 @@ def sample_secondaries(
|
||||
n_sec_pred: torch.Tensor,
|
||||
steps: int = 10,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Euler integration of `Stage2OneShot`'s (flow/ddpm) vector field; return
|
||||
raw slot outputs — see `_decode_stage2_flat`'s docstring for the returned
|
||||
(sec_cont, sec_type, sec_valid) shapes/meaning.
|
||||
"""Euler integration of the Stage-2 vector field; return raw slot outputs.
|
||||
|
||||
n_sec_pred: (B,) int64 — number of valid secondaries per step
|
||||
|
||||
Returns (sec_cont, sec_phys, sec_valid):
|
||||
sec_cont: (B, K_MAX, 4) — [stick_logit, local_dir_x, local_dir_y, local_dir_z]
|
||||
sec_phys: (B, K_MAX, PARTICLE_PHYS_DIM) — predicted [log_mass, charge]
|
||||
per slot (normalised iff the checkpoint's sec_phys
|
||||
normalizer was applied at training time — denormalize
|
||||
before treating as physical units; see
|
||||
giant.data.transforms.decode_secondaries). Used as-is —
|
||||
no snapping to a discrete PDG code.
|
||||
sec_valid: (B, K_MAX) bool — True for slots i < n_sec_pred
|
||||
"""
|
||||
sec_decoder.eval()
|
||||
B = cond_cont.size(0)
|
||||
device = cond_cont.device
|
||||
flat_width = _stage2_flat_width(sec_decoder)
|
||||
|
||||
x = torch.randn(B, flat_width, device=device)
|
||||
x = torch.randn(B, SEC_DIM, device=device)
|
||||
dt = 1.0 / steps
|
||||
for i in range(steps):
|
||||
t = torch.full((B,), i * dt, device=device)
|
||||
v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t)
|
||||
v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out)
|
||||
x = x + v * dt
|
||||
|
||||
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
||||
return _slots_from_flat(x, n_sec_pred)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_ddpm(
|
||||
model: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
schedule,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Full DDPM ancestral sampling (T reverse steps). Returns (sample, n_sec_pred)."""
|
||||
model.eval()
|
||||
B = cond_cont.size(0)
|
||||
device = cond_cont.device
|
||||
x = torch.randn(B, X_DIM, device=device)
|
||||
T = schedule.T
|
||||
for i in reversed(range(T)):
|
||||
t_norm = torch.full((B,), i / T, device=device)
|
||||
eps_pred = model(x, t_norm, cond_cont, cond_cat)
|
||||
beta = schedule.betas[i]
|
||||
alpha = schedule.alphas[i]
|
||||
alpha_bar = schedule.alpha_bars[i]
|
||||
z = torch.randn_like(x) if i > 0 else torch.zeros_like(x)
|
||||
x = (1.0 / alpha.sqrt()) * (
|
||||
x - (1.0 - alpha) / (1.0 - alpha_bar).sqrt() * eps_pred
|
||||
) + beta.sqrt() * z
|
||||
n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
|
||||
n_sec_pred = n_sec_logits.argmax(dim=-1)
|
||||
return x, n_sec_pred
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_wgan(
|
||||
generator: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Single-pass Stage-1 WGAN generator sample. Returns (sample, n_sec_pred)."""
|
||||
generator.eval()
|
||||
B = cond_cont.size(0)
|
||||
z = torch.randn(B, generator.noise_dim, device=cond_cont.device)
|
||||
x = generator(z, cond_cont, cond_cat)
|
||||
n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat)
|
||||
n_sec_pred = n_sec_logits.argmax(dim=-1)
|
||||
return x, n_sec_pred
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
@@ -213,228 +139,41 @@ def sample_secondaries_wgan(
|
||||
stage1_out: torch.Tensor,
|
||||
n_sec_pred: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Single-pass `Stage2OneShot` WGAN generator sample; see
|
||||
`_decode_stage2_flat`'s docstring for the returned (sec_cont, sec_type,
|
||||
sec_valid) shapes/meaning."""
|
||||
"""Single-pass Stage-2 WGAN generator sample; see `sample_secondaries`'s
|
||||
docstring for the returned (sec_cont, sec_phys, sec_valid) shapes."""
|
||||
sec_decoder.eval()
|
||||
B = cond_cont.size(0)
|
||||
z = torch.randn(B, sec_decoder.noise_dim, device=cond_cont.device)
|
||||
x = sec_decoder(z, cond_cont, cond_cat, stage1_out)
|
||||
return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
||||
return _slots_from_flat(x, n_sec_pred)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def sample_secondaries_ar(
|
||||
sec_decoder: torch.nn.Module,
|
||||
def sample_ddim(
|
||||
model: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
n_sec_pred: torch.Tensor,
|
||||
steps: int = 10,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""`Stage2Autoregressive` inference loop: one token at a time, in
|
||||
descending-energy slot order, `k_max` sequential calls. Unlike training
|
||||
(teacher forcing — a single parallel pass over ground-truth tokens, see
|
||||
`giant.training.stage2_inputs._assemble_stage2_ar_inputs`), there is no
|
||||
ground truth at inference: each token's conditioning is built
|
||||
free-running, from the PREVIOUS TOKEN'S OWN just-generated output — the
|
||||
train/inference gap that is the cost of markov history's
|
||||
expressiveness.
|
||||
|
||||
A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
|
||||
the "K sequential forwards" cost applies per-token here, not
|
||||
once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
|
||||
physics step.
|
||||
|
||||
Under `history="attention"` the history encoding is computed once per
|
||||
slot via `Stage2Autoregressive.history_step` (a KV-cache append)
|
||||
rather than re-derived by every model call inside that slot — so an ODE
|
||||
loop's `steps` substeps, and the separate `predict_type` call when the
|
||||
type slice isn't folded into the trunk output, all reuse the SAME `hist`
|
||||
tensor for a given `k`. Recomputing per call instead would be merely
|
||||
wasteful under markov (its per-call cost is already O(1)) but wrong under
|
||||
attention: `AttentionHistory.step` mutates the cache by appending, so
|
||||
calling it more than once per slot would double-count that slot's own
|
||||
(not-yet-existing) predecessor.
|
||||
|
||||
The free-running history feature stays UNSNAPPED (mirrors the
|
||||
established "no snapping" precedent for `particle_type.target =
|
||||
"physical"` secondaries feeding their own future conditioning):
|
||||
`"physical"` carries the raw (log_mass, charge) forward as-is;
|
||||
`"embedding"` carries the raw predicted vector as-is; `"onehot"` is the
|
||||
one exception — its history slot must be a probability-simplex-shaped
|
||||
vector (that's what `MarkovHistory`/`AttentionHistory` were trained on,
|
||||
`_type_repr`'s `F.one_hot` ground truth), so it's the hard one-hot of
|
||||
`argmax(logits)`, not the raw logits themselves. Discretizing further,
|
||||
into a concrete PDG code, only ever happens once — at secondary-spawn
|
||||
time in `giant/rollout.py` — never inside this loop.
|
||||
|
||||
Returns (sec_cont, sec_type, sec_valid) — same shapes/meaning as
|
||||
`sample_secondaries`/`sample_secondaries_wgan`'s (see
|
||||
`_decode_stage2_flat`'s docstring); `sec_type` is raw per-slot output in
|
||||
all three `particle_type.target` cases (never one-hot-collapsed), so the
|
||||
caller decodes it exactly the same way regardless of which decoder
|
||||
produced it.
|
||||
"""
|
||||
sec_decoder.eval()
|
||||
schedule,
|
||||
steps: int = 50,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred)."""
|
||||
model.eval()
|
||||
B = cond_cont.size(0)
|
||||
device = cond_cont.device
|
||||
k_max = sec_decoder.k_max
|
||||
type_dim = sec_decoder.type_dim
|
||||
generator = sec_decoder.generator_kind
|
||||
target = sec_decoder.particle_type_cfg.get("target", "physical")
|
||||
type_folded = _type_folded(sec_decoder)
|
||||
token_dim = CONT_SLOT_DIM + type_dim if type_folded else CONT_SLOT_DIM
|
||||
|
||||
sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
|
||||
sec_type = torch.zeros(B, k_max, type_dim, device=device)
|
||||
|
||||
# Running per-token state, threaded from one slot to the next.
|
||||
prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
|
||||
remaining = torch.ones(B, device=device)
|
||||
history_cache = sec_decoder.init_history_cache()
|
||||
|
||||
for k in range(k_max):
|
||||
has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
|
||||
history_feat = prev_repr.unsqueeze(1) # (B, 1, CONT_SLOT_DIM + type_dim)
|
||||
remaining_frac = remaining.unsqueeze(1) # (B, 1)
|
||||
slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
|
||||
hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
|
||||
|
||||
if generator == "wgan":
|
||||
z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
|
||||
token = sec_decoder(
|
||||
z,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
hist=hist,
|
||||
)
|
||||
T = schedule.T
|
||||
timesteps = torch.linspace(T - 1, 0, steps, dtype=torch.long, device=device)
|
||||
x = torch.randn(B, X_DIM, device=device)
|
||||
for step_idx, ts in enumerate(timesteps):
|
||||
t_idx = int(ts.item())
|
||||
t_norm = torch.full((B,), t_idx / T, device=device)
|
||||
eps_pred = model(x, t_norm, cond_cont, cond_cat)
|
||||
ab_t = schedule.alpha_bars[t_idx]
|
||||
if step_idx + 1 < len(timesteps):
|
||||
ab_prev = schedule.alpha_bars[int(timesteps[step_idx + 1].item())]
|
||||
else:
|
||||
x = torch.randn(B, 1, token_dim, device=device)
|
||||
dt = 1.0 / steps
|
||||
for i in range(steps):
|
||||
t = torch.full((B, 1), i * dt, device=device)
|
||||
v = sec_decoder(
|
||||
x,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
t=t,
|
||||
hist=hist,
|
||||
)
|
||||
x = x + v * dt
|
||||
token = x
|
||||
|
||||
token = token.squeeze(1) # (B, token_dim)
|
||||
cont_k = token[:, :CONT_SLOT_DIM]
|
||||
if type_folded:
|
||||
type_k = token[:, CONT_SLOT_DIM:]
|
||||
else:
|
||||
type_k = sec_decoder.predict_type(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
hist=hist,
|
||||
).squeeze(1)
|
||||
|
||||
sec_cont[:, k] = cont_k
|
||||
sec_type[:, k] = type_k
|
||||
|
||||
if target == "onehot":
|
||||
type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
|
||||
else:
|
||||
type_for_history = type_k
|
||||
|
||||
stick_fraction = torch.sigmoid(cont_k[:, 0])
|
||||
prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
|
||||
remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
|
||||
|
||||
sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < n_sec_pred.unsqueeze(1)
|
||||
return sec_cont, sec_type, sec_valid
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Per-stage dispatch — shared by giant/rollout.py and giant/cli.py's
|
||||
# `predict` command, since both need "given a stage model, produce a
|
||||
# sample" without hand-picking the sampler themselves (each stage's
|
||||
# generative objective is independent, read off the model's own
|
||||
# `generator_kind`, not a caller-supplied `mode` string).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def sample_stage1(
|
||||
stage1_model: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
steps: int,
|
||||
ddpm_steps: int = 1000,
|
||||
) -> tuple[torch.Tensor, torch.Tensor | None]:
|
||||
"""Dispatches on `stage1_model.generator_kind`."""
|
||||
kind = stage1_model.generator_kind
|
||||
if kind == "wgan":
|
||||
return sample_wgan(stage1_model, cond_cont, cond_cat)
|
||||
if kind == "ddpm":
|
||||
schedule = CosineSchedule(T=ddpm_steps).to(cond_cont.device)
|
||||
return sample_ddpm(stage1_model, cond_cont, cond_cat, schedule)
|
||||
return sample_flow(stage1_model, cond_cont, cond_cat, steps=steps)
|
||||
|
||||
|
||||
def sample_stage2(
|
||||
sec_decoder: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
n_sec_pred: torch.Tensor,
|
||||
steps: int,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""Dispatches on `decoder` (one-shot vs autoregressive — the class
|
||||
itself, via `isinstance`) and `sec_decoder.generator_kind` (flow/ddpm/
|
||||
wgan). DDPM secondaries aren't supported — no `Stage2*` class was ever
|
||||
built with `generator="ddpm"` in practice and `flow_matching_loss_secondary*`
|
||||
is the only stage-2 training path that exists for the non-adversarial
|
||||
case, so there's nothing to dispatch to here.
|
||||
"""
|
||||
if isinstance(sec_decoder, Stage2Autoregressive):
|
||||
return sample_secondaries_ar(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
|
||||
if sec_decoder.generator_kind == "wgan":
|
||||
return sample_secondaries_wgan(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
||||
return sample_secondaries(sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=steps)
|
||||
|
||||
|
||||
def resolve_n_sec(
|
||||
stage1_model: torch.nn.Module,
|
||||
sec_decoder: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_out: torch.Tensor,
|
||||
n_sec_pred: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
"""`n_sec_pred` is already populated when `stage1_model` owns a legacy
|
||||
`n_sec_head` (a migrated v0.2 checkpoint — see `Stage1Model`'s
|
||||
docstring); otherwise ask stage 2, which owns it by default. Raises if
|
||||
neither stage owns a head at all — the only way that happens is
|
||||
`stage2_model.n_sec.mode` other than `"head"` (`"truth"`/`"stop_token"`),
|
||||
neither of which is a valid rollout-/predict-capable checkpoint."""
|
||||
if n_sec_pred is not None:
|
||||
return n_sec_pred
|
||||
if getattr(sec_decoder, "n_sec_head", None) is None:
|
||||
raise RuntimeError(
|
||||
"checkpoint has no n_sec_head on either stage — needs "
|
||||
"stage2_model.n_sec.mode = 'head' (the default); 'truth' is "
|
||||
"standalone-evaluation-only and 'stop_token' isn't implemented"
|
||||
)
|
||||
logits = sec_decoder.predict_n_sec(cond_cont, cond_cat, stage1_out)
|
||||
return logits.argmax(dim=-1)
|
||||
ab_prev = torch.ones(1, device=device)
|
||||
x0_pred = (x - (1.0 - ab_t).sqrt() * eps_pred) / ab_t.sqrt()
|
||||
x = ab_prev.sqrt() * x0_pred + (1.0 - ab_prev).sqrt() * eps_pred
|
||||
n_sec_logits = model.predict_n_sec(cond_cont, cond_cat)
|
||||
n_sec_pred = n_sec_logits.argmax(dim=-1)
|
||||
return x, n_sec_pred
|
||||
|
||||
@@ -1,161 +0,0 @@
|
||||
"""Portal-machine follow-up for v0.3.0 step 2: diff a real v0.2 checkpoint's
|
||||
outputs against the new `build_models` on the same input batch.
|
||||
|
||||
`tests/test_migration_v02_v03.py` already proves this bit-identical with
|
||||
synthetic random weights, but that test can't run where it matters (no
|
||||
`/ceph` on local dev machines — see CLAUDE.md's Compute environment
|
||||
section). This script is the real-checkpoint counterpart: run it on a portal
|
||||
machine against an actual trained checkpoint before merging
|
||||
`v0.3.0-stage2-autoregressive` to `master`.
|
||||
|
||||
Usage (from the repo root, on a portal machine):
|
||||
|
||||
uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt
|
||||
uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --ema
|
||||
uv run python giant/tools/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --batch 32 --seed 1
|
||||
|
||||
Run it once against a flow (or ddpm) checkpoint and once against a wgan
|
||||
checkpoint ("one flow checkpoint and one WGAN checkpoint").
|
||||
A routed checkpoint (`model_config["router"]["enabled"]`) is only checked for
|
||||
successful construction — `giant.model.network.migrate_legacy_state_dict`
|
||||
doesn't yet remap routed (Expert-per-router) state dicts, so the
|
||||
bit-identical assertion is skipped with a clear warning in that case (see the
|
||||
function's own docstring for why).
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
|
||||
|
||||
from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM # noqa: E402
|
||||
from giant.model import network as net # noqa: E402
|
||||
from tests.legacy import network_v02_snapshot as legacy # noqa: E402
|
||||
|
||||
|
||||
def _random_batch(model_config: dict, batch: int, seed: int):
|
||||
g = torch.Generator().manual_seed(seed)
|
||||
pdg_vocab = model_config["pdg_vocab"]
|
||||
mat_vocab = model_config["mat_vocab"]
|
||||
k_max = model_config.get("k_max", 15)
|
||||
noise_dim = model_config.get("noise_dim", 64)
|
||||
|
||||
cond_cont = torch.randn(batch, COND_DIM, generator=g)
|
||||
cond_cat = torch.stack(
|
||||
[
|
||||
torch.randint(0, pdg_vocab, (batch,), generator=g),
|
||||
torch.randint(0, mat_vocab, (batch,), generator=g),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
x1 = torch.randn(batch, X_DIM, generator=g)
|
||||
x2 = torch.randn(batch, k_max * SEC_SLOT_DIM, generator=g)
|
||||
t = torch.rand(batch, generator=g)
|
||||
z1 = torch.randn(batch, noise_dim, generator=g)
|
||||
z2 = torch.randn(batch, noise_dim, generator=g)
|
||||
return cond_cont, cond_cat, x1, x2, t, z1, z2
|
||||
|
||||
|
||||
def _max_abs_diff(a: torch.Tensor, b: torch.Tensor) -> float:
|
||||
return (a - b).abs().max().item()
|
||||
|
||||
|
||||
def main() -> int:
|
||||
p = argparse.ArgumentParser(description=__doc__)
|
||||
p.add_argument("checkpoint", type=Path, help="Path to a v0.2 best.pt/last.pt")
|
||||
p.add_argument(
|
||||
"--ema",
|
||||
action="store_true",
|
||||
help="Use the checkpoint's EMA weights (model_ema/sec_decoder_ema) — "
|
||||
"what predict/rollout actually sample from — instead of raw weights.",
|
||||
)
|
||||
p.add_argument("--batch", type=int, default=16)
|
||||
p.add_argument("--seed", type=int, default=0)
|
||||
args = p.parse_args()
|
||||
|
||||
ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
|
||||
if "model_config" not in ckpt:
|
||||
print(f"FAIL: {args.checkpoint} has no 'model_config' key — can't migrate it")
|
||||
return 1
|
||||
model_config = ckpt["model_config"]
|
||||
mode = model_config.get("mode", "flow")
|
||||
routed = bool((model_config.get("router") or {}).get("enabled"))
|
||||
print(f"checkpoint: {args.checkpoint}")
|
||||
print(f" mode={mode!r} conditioning={model_config.get('conditioning')!r} routed={routed} ema={args.ema}")
|
||||
|
||||
stage1_key = "model_ema" if args.ema and "model_ema" in ckpt else "model"
|
||||
stage2_key = "sec_decoder_ema" if args.ema and "sec_decoder_ema" in ckpt else "sec_decoder"
|
||||
if args.ema and stage1_key == "model":
|
||||
print(" warning: --ema requested but no model_ema in checkpoint, using raw weights")
|
||||
|
||||
# --- old side: the frozen v0.2 snapshot, loaded with the checkpoint's own weights ---
|
||||
old_stage1, old_stage2 = legacy.build_models(model_config)
|
||||
old_stage1.load_state_dict(ckpt[stage1_key])
|
||||
old_stage2.load_state_dict(ckpt[stage2_key])
|
||||
old_stage1.eval()
|
||||
old_stage2.eval()
|
||||
|
||||
# --- new side: migrated config + remapped state dict, through the new build_models ---
|
||||
new_models = net.build_models(model_config)
|
||||
new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"]
|
||||
assert new_stage1 is not None and new_stage2 is not None
|
||||
|
||||
if routed:
|
||||
print(
|
||||
" routed checkpoint: migrate_legacy_state_dict only handles the "
|
||||
"monolithic trunk shape — verifying construction only, skipping "
|
||||
"the bit-identical weight/output comparison."
|
||||
)
|
||||
print("PASS (construction only, routed checkpoint)")
|
||||
return 0
|
||||
|
||||
remapped1, remapped2 = net.migrate_legacy_state_dict(ckpt[stage1_key], ckpt[stage2_key])
|
||||
missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
|
||||
missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
|
||||
if missing1 or unexpected1 or missing2 or unexpected2:
|
||||
print("FAIL: state dict mismatch after remap")
|
||||
print(f" stage1 missing={missing1} unexpected={unexpected1}")
|
||||
print(f" stage2 missing={missing2} unexpected={unexpected2}")
|
||||
return 1
|
||||
new_stage1.eval()
|
||||
new_stage2.eval()
|
||||
|
||||
cond_cont, cond_cat, x1, x2, t, z1, z2 = _random_batch(model_config, args.batch, args.seed)
|
||||
|
||||
ok = True
|
||||
with torch.no_grad():
|
||||
if mode == "wgan":
|
||||
old_out1 = old_stage1(z1, cond_cont, cond_cat)
|
||||
new_out1 = new_stage1(z1, cond_cont, cond_cat)
|
||||
else:
|
||||
old_out1 = old_stage1(x1, t, cond_cont, cond_cat)
|
||||
new_out1 = new_stage1(x1, cond_cont, cond_cat, t=t)
|
||||
old_n_sec = old_stage1.predict_n_sec(cond_cont, cond_cat)
|
||||
new_n_sec = new_stage1.predict_n_sec(cond_cont, cond_cat)
|
||||
if mode == "wgan":
|
||||
old_out2 = old_stage2(z2, cond_cont, cond_cat, old_out1)
|
||||
new_out2 = new_stage2(z2, cond_cont, cond_cat, new_out1)
|
||||
else:
|
||||
old_out2 = old_stage2(x2, t, cond_cont, cond_cat, old_out1)
|
||||
new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t)
|
||||
|
||||
for label, old_out, new_out in [
|
||||
("stage1 output", old_out1, new_out1),
|
||||
("n_sec logits", old_n_sec, new_n_sec),
|
||||
("stage2 output", old_out2, new_out2),
|
||||
]:
|
||||
identical = torch.equal(old_out, new_out)
|
||||
diff = _max_abs_diff(old_out, new_out)
|
||||
status = "OK" if identical else "MISMATCH"
|
||||
print(f" {label}: {status} (max abs diff = {diff:.3e})")
|
||||
ok = ok and identical
|
||||
|
||||
print("PASS" if ok else "FAIL")
|
||||
return 0 if ok else 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
+1155
File diff suppressed because it is too large
Load Diff
@@ -1,30 +0,0 @@
|
||||
"""Training: per-stage trainers, metric collection, checkpointing, the loop.
|
||||
|
||||
Split out of the former single-module `giant/train.py`. The public surface is
|
||||
`train` (the entry point `giant.pipeline` calls) plus the trainer/spec types
|
||||
that tests and tooling construct directly.
|
||||
"""
|
||||
|
||||
from giant.training.checkpoint import build_checkpoint, load_checkpoint
|
||||
from giant.training.metrics import MetricsCollector, MetricSpec
|
||||
from giant.training.loop import train
|
||||
from giant.training.trainers import (
|
||||
FlowDDPMStageTrainer,
|
||||
StageSpec,
|
||||
StageTrainer,
|
||||
WGANStageTrainer,
|
||||
build_stage_trainers,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"FlowDDPMStageTrainer",
|
||||
"MetricSpec",
|
||||
"MetricsCollector",
|
||||
"StageSpec",
|
||||
"StageTrainer",
|
||||
"WGANStageTrainer",
|
||||
"build_checkpoint",
|
||||
"build_stage_trainers",
|
||||
"load_checkpoint",
|
||||
"train",
|
||||
]
|
||||
@@ -1,68 +0,0 @@
|
||||
"""Checkpoint assembly and restore.
|
||||
|
||||
The on-disk layout is unchanged from v0.2/v0.3.0 and is read by
|
||||
`giant/cli.py`, `giant/rollout.py`, `giant/sample.py` and
|
||||
`giant/analysis/router_gating.py` — stage 1's weights live under `model`,
|
||||
stage 2's under `sec_decoder`, with `_ema`/`critic`/`sec_critic` companions
|
||||
and per-stage `optimizer_<stage>` / `optimizer_d_<stage>` / `lr_sched_<stage>`
|
||||
entries.
|
||||
"""
|
||||
|
||||
from giant.training.trainers import StageTrainer
|
||||
|
||||
#: Stage name -> the checkpoint key its weights live under. Historical: stage
|
||||
#: 1 predates the two-stage split, so it kept the bare "model" key.
|
||||
_STAGE_KEY = {"stage1": "model", "stage2": "sec_decoder"}
|
||||
_CRITIC_KEY = {"stage1": "critic", "stage2": "sec_critic"}
|
||||
|
||||
|
||||
def build_checkpoint(
|
||||
trainers: dict[str, StageTrainer],
|
||||
epoch: int,
|
||||
global_step: int,
|
||||
best_val_loss: float,
|
||||
extras: dict,
|
||||
) -> dict:
|
||||
"""`extras` carries the dataset-level sidecars (normalizer, vocab maps,
|
||||
model_config) that `train()` receives as arguments; `None` values are
|
||||
omitted so an absent sidecar leaves no key behind."""
|
||||
ckpt: dict = {
|
||||
"epoch": epoch,
|
||||
"best_val_loss": best_val_loss,
|
||||
"global_step": global_step,
|
||||
}
|
||||
for name, trainer in trainers.items():
|
||||
sd = trainer.state_dict()
|
||||
key = _STAGE_KEY[name]
|
||||
ckpt[key] = sd["model"]
|
||||
if "model_ema" in sd:
|
||||
ckpt[f"{key}_ema"] = sd["model_ema"]
|
||||
if "critic" in sd:
|
||||
ckpt[_CRITIC_KEY[name]] = sd["critic"]
|
||||
ckpt[f"optimizer_d_{name}"] = sd["optimizer_d"]
|
||||
ckpt[f"optimizer_{name}"] = sd["optimizer"]
|
||||
ckpt[f"lr_sched_{name}"] = sd["lr_sched"]
|
||||
ckpt.update({k: v for k, v in extras.items() if v is not None})
|
||||
return ckpt
|
||||
|
||||
|
||||
def load_checkpoint(trainers: dict[str, StageTrainer], ckpt: dict, lr: float) -> None:
|
||||
"""Restore every active stage, then hand `lr`'s authority back to the
|
||||
config — `load_state_dict` would otherwise leave the checkpoint's own
|
||||
base LR in place, silently ignoring `--lr` on resume."""
|
||||
for name, trainer in trainers.items():
|
||||
key = _STAGE_KEY[name]
|
||||
sd = {
|
||||
"model": ckpt[key],
|
||||
"optimizer": ckpt[f"optimizer_{name}"],
|
||||
"lr_sched": ckpt[f"lr_sched_{name}"],
|
||||
}
|
||||
ema_key = f"{key}_ema"
|
||||
if ema_key in ckpt:
|
||||
sd["model_ema"] = ckpt[ema_key]
|
||||
crit_key = _CRITIC_KEY[name]
|
||||
if crit_key in ckpt:
|
||||
sd["critic"] = ckpt[crit_key]
|
||||
sd["optimizer_d"] = ckpt[f"optimizer_d_{name}"]
|
||||
trainer.load_state_dict(sd)
|
||||
trainer.resume_lr(lr)
|
||||
@@ -1,297 +0,0 @@
|
||||
"""The training loop.
|
||||
|
||||
`train()` owns the epoch structure and nothing else: the per-stage step is
|
||||
`giant.training.trainers`' job, every number reported is
|
||||
`giant.training.metrics`' job, and the on-disk checkpoint is
|
||||
`giant.training.checkpoint`'s.
|
||||
"""
|
||||
|
||||
import os
|
||||
import signal
|
||||
import time
|
||||
from pathlib import Path
|
||||
from types import FrameType
|
||||
from typing import Callable
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from tqdm import tqdm
|
||||
|
||||
from giant.data.loader import TopNMap
|
||||
from giant.data.setup_cache import topnmap_to_json
|
||||
from giant.training.checkpoint import build_checkpoint, load_checkpoint
|
||||
from giant.training.metrics import MetricsCollector
|
||||
from giant.training.trainers import (
|
||||
FlowDDPMStageTrainer,
|
||||
StageTrainer,
|
||||
build_stage_trainers,
|
||||
)
|
||||
from giant.validate import validate_marginals
|
||||
|
||||
_CATCHABLE_SIGNALS = (signal.SIGINT, signal.SIGTERM)
|
||||
|
||||
|
||||
class _GracefulShutdown:
|
||||
"""Turns SIGINT/SIGTERM into a flag check instead of an immediate crash.
|
||||
|
||||
A second signal while already shutting down restores the default
|
||||
handler and re-sends the signal, so an unresponsive run can still be
|
||||
force-killed.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.requested = False
|
||||
self._previous: dict[
|
||||
int,
|
||||
Callable[[int, FrameType | None], object] | signal.Handlers | int | None,
|
||||
] = {}
|
||||
|
||||
def __enter__(self) -> "_GracefulShutdown":
|
||||
for sig in _CATCHABLE_SIGNALS:
|
||||
self._previous[sig] = signal.getsignal(sig)
|
||||
signal.signal(sig, self._handle)
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc_info) -> None:
|
||||
for sig, handler in self._previous.items():
|
||||
signal.signal(sig, handler)
|
||||
|
||||
def _handle(self, signum: int, frame) -> None:
|
||||
if self.requested:
|
||||
signal.signal(signum, self._previous[signum])
|
||||
os.kill(os.getpid(), signum)
|
||||
return
|
||||
self.requested = True
|
||||
print(
|
||||
f"\nreceived {signal.Signals(signum).name} — finishing the current "
|
||||
"batch, then saving a checkpoint and exiting (send again to force-quit)"
|
||||
)
|
||||
|
||||
|
||||
def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs):
|
||||
"""Runs `validate_marginals` on `trainer`'s sampling model (EMA model if
|
||||
present, else the raw model). `validate_marginals` itself dispatches
|
||||
through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec`, so
|
||||
this is generator- and one-shot-vs-autoregressive-agnostic."""
|
||||
model = trainer.sampling_model()
|
||||
return validate_marginals(model, val_loader, device=device, **kwargs)
|
||||
|
||||
|
||||
def _marginal_kl(trainers: dict[str, StageTrainer], val_loader, device, **kwargs) -> float:
|
||||
"""Mean marginal KL over the stage-1 sampling chain, or NaN when stage 1
|
||||
is inactive or `validate_marginals` declined to produce a result."""
|
||||
stage1 = trainers.get("stage1")
|
||||
if stage1 is None:
|
||||
return float("nan")
|
||||
result = _try_validate_marginals(
|
||||
stage1,
|
||||
val_loader,
|
||||
device,
|
||||
sec_decoder=trainers["stage2"].sampling_model() if "stage2" in trainers else None,
|
||||
**kwargs,
|
||||
)
|
||||
if result is None:
|
||||
return float("nan")
|
||||
return float(np.mean(result["kl_divergence"]))
|
||||
|
||||
|
||||
def train(
|
||||
cfg: dict,
|
||||
models: dict[str, torch.nn.Module | None],
|
||||
critics: dict[str, torch.nn.Module | None],
|
||||
train_loader: DataLoader,
|
||||
val_loader: DataLoader,
|
||||
device: torch.device,
|
||||
out_dir: str | Path,
|
||||
normalizer_dict: dict | None = None,
|
||||
pdg_map: dict | None = None,
|
||||
mat_map: dict | None = None,
|
||||
proc_map: dict | None = None,
|
||||
pdg_topn_map: TopNMap | None = None,
|
||||
sec_type_topn_map: TopNMap | None = None,
|
||||
mat_topn_map: TopNMap | None = None,
|
||||
model_config: dict | None = None,
|
||||
resume_path: str | Path | None = None,
|
||||
total_train_batches: int = 0,
|
||||
use_wandb: bool = False,
|
||||
wandb_project: str = "giant",
|
||||
wandb_run_name: str = "",
|
||||
wandb_log_every: int = 50,
|
||||
) -> None:
|
||||
"""Train whichever of stage1/stage2 are active, each through its own
|
||||
`StageTrainer`. `models`/`critics` are the dicts
|
||||
`giant.model.network.build_models`/`build_critics` return — a `None`
|
||||
entry means that stage is `active = false`.
|
||||
"""
|
||||
out_dir = Path(out_dir)
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
t = cfg["train"]
|
||||
epochs = t["epochs"]
|
||||
validate_every = t.get("validate_every", 0)
|
||||
validate_steps = t.get("validate_steps", 10)
|
||||
max_val_batches = t.get("max_val_batches", 0)
|
||||
|
||||
trainers = build_stage_trainers(cfg, models, critics, device, total_train_batches)
|
||||
if not trainers:
|
||||
raise ValueError("no active stage — stage1_model.active and stage2_model.active are both false")
|
||||
has_adversarial = any(not tr.supports_val_loss for tr in trainers.values())
|
||||
|
||||
checkpoint_extras = {
|
||||
"normalizer": normalizer_dict,
|
||||
"pdg_map": pdg_map,
|
||||
"mat_map": mat_map,
|
||||
"proc_map": proc_map,
|
||||
"pdg_topn_map": topnmap_to_json(pdg_topn_map) if pdg_topn_map is not None else None,
|
||||
"sec_type_topn_map": topnmap_to_json(sec_type_topn_map) if sec_type_topn_map is not None else None,
|
||||
"mat_topn_map": topnmap_to_json(mat_topn_map) if mat_topn_map is not None else None,
|
||||
"model_config": model_config,
|
||||
}
|
||||
|
||||
start_epoch = 1
|
||||
best_val_loss = float("inf")
|
||||
global_step = 0
|
||||
if resume_path is not None:
|
||||
ckpt = torch.load(resume_path, map_location=device, weights_only=False)
|
||||
load_checkpoint(trainers, ckpt, t["lr"])
|
||||
start_epoch = ckpt.get("epoch", 0) + 1
|
||||
best_val_loss = ckpt.get("best_val_loss", float("inf"))
|
||||
global_step = ckpt.get("global_step", 0)
|
||||
if start_epoch > epochs:
|
||||
print(f"checkpoint already completed epoch {start_epoch - 1} (>= --epochs {epochs}) — nothing to train")
|
||||
return
|
||||
|
||||
collector = MetricsCollector.create(
|
||||
trainers,
|
||||
out_dir,
|
||||
cfg,
|
||||
model_config,
|
||||
resume=resume_path is not None,
|
||||
use_wandb=use_wandb,
|
||||
wandb_project=wandb_project,
|
||||
wandb_run_name=wandb_run_name,
|
||||
wandb_log_every=wandb_log_every,
|
||||
)
|
||||
|
||||
epoch_w = len(str(epochs))
|
||||
last_completed_epoch = start_epoch - 1
|
||||
with _GracefulShutdown() as shutdown:
|
||||
for epoch in range(start_epoch, epochs + 1):
|
||||
epoch_start = time.monotonic()
|
||||
if device.type == "cuda":
|
||||
torch.cuda.reset_peak_memory_stats(device)
|
||||
collector.start_epoch(epoch)
|
||||
for trainer in trainers.values():
|
||||
trainer.train_mode()
|
||||
|
||||
bar = tqdm(
|
||||
train_loader,
|
||||
desc=f" epoch {epoch:{epoch_w}d}/{epochs}",
|
||||
total=total_train_batches or None,
|
||||
leave=False,
|
||||
unit="batch",
|
||||
dynamic_ncols=True,
|
||||
)
|
||||
for batch in bar:
|
||||
B = batch[0].size(0)
|
||||
collector.add_train_batch(
|
||||
{name: trainer.step(batch, device, global_step) for name, trainer in trainers.items()},
|
||||
B,
|
||||
)
|
||||
bar.set_postfix_str(collector.postfix(), refresh=False)
|
||||
global_step += 1
|
||||
collector.log_batch(global_step, batch, device)
|
||||
if shutdown.requested:
|
||||
break
|
||||
bar.close()
|
||||
|
||||
if shutdown.requested:
|
||||
ckpt = build_checkpoint(trainers, epoch - 1, global_step, best_val_loss, checkpoint_extras)
|
||||
torch.save(ckpt, out_dir / "last.pt")
|
||||
last_completed_epoch = epoch - 1
|
||||
print(
|
||||
f"saved in-progress weights from partway through epoch "
|
||||
f"{epoch} to {out_dir / 'last.pt'} "
|
||||
f"(resume will restart epoch {epoch})"
|
||||
)
|
||||
break
|
||||
|
||||
for trainer in trainers.values():
|
||||
trainer.eval_mode()
|
||||
|
||||
# --- per-stage validation ---
|
||||
scored = {name: tr for name, tr in trainers.items() if tr.supports_val_loss}
|
||||
if scored:
|
||||
with torch.no_grad():
|
||||
for val_batch_idx, batch in enumerate(val_loader):
|
||||
if max_val_batches > 0 and val_batch_idx >= max_val_batches:
|
||||
break
|
||||
B = batch[0].size(0)
|
||||
collector.add_val_batch(
|
||||
{name: tr.val_loss(batch, device) for name, tr in scored.items()},
|
||||
B,
|
||||
)
|
||||
collector.observe_routers(batch[0].to(device), batch[1].to(device), B)
|
||||
|
||||
# An adversarial stage has no averageable validation loss, so it
|
||||
# needs the marginal-KL signal every epoch to pick a best
|
||||
# checkpoint at all; a purely non-adversarial run only pays for
|
||||
# it every `validate_every` epochs.
|
||||
marginal_kl = float("nan")
|
||||
if has_adversarial:
|
||||
marginal_kl = _marginal_kl(trainers, val_loader, device)
|
||||
elif validate_every > 0 and epoch % validate_every == 0:
|
||||
stage1 = trainers.get("stage1")
|
||||
ddpm_steps = 1000
|
||||
if isinstance(stage1, FlowDDPMStageTrainer) and stage1.ddpm_schedule is not None:
|
||||
ddpm_steps = stage1.ddpm_schedule.T
|
||||
marginal_kl = _marginal_kl(
|
||||
trainers,
|
||||
val_loader,
|
||||
device,
|
||||
steps=validate_steps,
|
||||
ddpm_steps=ddpm_steps,
|
||||
)
|
||||
|
||||
val_loss = sum(
|
||||
trainer.val_objective(
|
||||
collector.train_means(name),
|
||||
collector.val_means(name),
|
||||
marginal_kl,
|
||||
)
|
||||
for name, trainer in trainers.items()
|
||||
)
|
||||
|
||||
epoch_time = time.monotonic() - epoch_start
|
||||
is_best = val_loss < best_val_loss
|
||||
collector.set("val/loss", val_loss)
|
||||
collector.set("val/marginal_kl", marginal_kl)
|
||||
collector.set(
|
||||
"gpu_mem_mb",
|
||||
torch.cuda.max_memory_allocated(device) / (1024 * 1024) if device.type == "cuda" else 0.0,
|
||||
)
|
||||
collector.set("samples_per_sec", collector.train_samples / max(epoch_time, 1e-8))
|
||||
collector.set("is_best", int(is_best))
|
||||
collector.set("epoch_time_s", epoch_time)
|
||||
|
||||
print(collector.summary_line(val_loss, epoch_time, is_best))
|
||||
collector.write_epoch(global_step)
|
||||
|
||||
ckpt = build_checkpoint(trainers, epoch, global_step, best_val_loss, checkpoint_extras)
|
||||
if is_best:
|
||||
best_val_loss = val_loss
|
||||
ckpt["best_val_loss"] = best_val_loss
|
||||
torch.save(ckpt, out_dir / "best.pt")
|
||||
torch.save(ckpt, out_dir / "last.pt")
|
||||
last_completed_epoch = epoch
|
||||
|
||||
if shutdown.requested:
|
||||
break
|
||||
|
||||
collector.close()
|
||||
|
||||
if shutdown.requested:
|
||||
print(
|
||||
f"stopped after epoch {last_completed_epoch} due to shutdown signal — "
|
||||
f"resume with --resume {out_dir / 'last.pt'}"
|
||||
)
|
||||
@@ -1,384 +0,0 @@
|
||||
"""Per-epoch metric accumulation, `metrics.csv`, and W&B logging.
|
||||
|
||||
Every scalar a training run reports is declared exactly once, as a
|
||||
`MetricSpec` on the `StageTrainer` that computes it (see
|
||||
`giant.training.trainers`). `MetricsCollector` derives the CSV/W&B column set
|
||||
from those declarations, so adding a metric means adding one line next to the
|
||||
code that produces it — there is no second list to keep in sync.
|
||||
|
||||
Column naming is uniform: `<stage>/train/<key>`, `<stage>/val/<key>`,
|
||||
`<stage>/<key>` for point-in-time values (`lr`, `critic_lr`),
|
||||
`<stage>/router/<key>` for routing diagnostics, and an unprefixed run-level
|
||||
tail (`val/loss`, `grad_norm`, `epoch_time_s`, ...). W&B groups panels on
|
||||
`/`, so the same names read well there.
|
||||
"""
|
||||
|
||||
import csv
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
|
||||
_ROUTER_KEYS = ("entropy", "util_min", "util_max", "util_std")
|
||||
|
||||
# Written after every stage's columns, by `MetricsCollector` itself rather
|
||||
# than by any one trainer — these describe the run, not a stage.
|
||||
_RUN_COLUMNS = (
|
||||
"val/loss",
|
||||
"val/marginal_kl",
|
||||
"grad_norm",
|
||||
"gpu_mem_mb",
|
||||
"samples_per_sec",
|
||||
"is_best",
|
||||
"epoch_time_s",
|
||||
)
|
||||
|
||||
# tqdm/W&B batch-granularity smoothing, matching v0.2/v0.3.0's inline EMA.
|
||||
_EMA_ALPHA = 0.05
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MetricSpec:
|
||||
"""One scalar a trainer emits per batch, and how it is reported.
|
||||
|
||||
`key` indexes the dict `StageTrainer.step()` / `.val_loss()` returns;
|
||||
`column` is the CSV/W&B column suffix, joined to the stage name with
|
||||
"/". `reduce` is either "mean" (batch-size-weighted average over the
|
||||
epoch) or "last" (the most recent value — for point-in-time quantities
|
||||
like the learning rate, which is a schedule readout, not a statistic).
|
||||
"""
|
||||
|
||||
key: str
|
||||
column: str
|
||||
reduce: str = "mean"
|
||||
|
||||
|
||||
def train_metric(key: str, column: str | None = None) -> MetricSpec:
|
||||
return MetricSpec(key, column or f"train/{key}")
|
||||
|
||||
|
||||
def val_metric(key: str, column: str | None = None) -> MetricSpec:
|
||||
return MetricSpec(key, column or f"val/{key}")
|
||||
|
||||
|
||||
def stage_metric(key: str, column: str | None = None) -> MetricSpec:
|
||||
"""A point-in-time stage-level readout (`lr`, `critic_lr`) — reported
|
||||
unprefixed by split, as `<stage>/<key>`."""
|
||||
return MetricSpec(key, column or key, reduce="last")
|
||||
|
||||
|
||||
def _wandb_run_config(cfg: dict, model_config: dict | None, param_counts: dict) -> dict:
|
||||
return {
|
||||
"train": cfg["train"],
|
||||
"conditioning": cfg["conditioning"],
|
||||
"stage1_model": cfg["stage1_model"],
|
||||
"stage2_model": cfg["stage2_model"],
|
||||
"model_config": model_config or {},
|
||||
"param_counts": param_counts,
|
||||
}
|
||||
|
||||
|
||||
class _Accumulator:
|
||||
"""Batch-size-weighted sums for one stage and one split."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.sums: dict[str, float] = {}
|
||||
self.n = 0
|
||||
self.last: dict[str, float] = {}
|
||||
|
||||
def add(self, stats: dict, keys: set[str], batch_size: int) -> None:
|
||||
for key in keys:
|
||||
if key in stats:
|
||||
self.sums[key] = self.sums.get(key, 0.0) + stats[key] * batch_size
|
||||
self.last.update(stats)
|
||||
self.n += batch_size
|
||||
|
||||
def mean(self, key: str) -> float:
|
||||
return self.sums.get(key, 0.0) / max(self.n, 1)
|
||||
|
||||
def means(self) -> dict[str, float]:
|
||||
return {key: self.mean(key) for key in self.sums}
|
||||
|
||||
def reset(self) -> None:
|
||||
self.sums.clear()
|
||||
self.last.clear()
|
||||
self.n = 0
|
||||
|
||||
|
||||
class _RouterAccumulator:
|
||||
"""Gate-diagnostic sums for one routed stage."""
|
||||
|
||||
def __init__(self, n_experts: int) -> None:
|
||||
self.n_experts = n_experts
|
||||
self.entropy = 0.0
|
||||
self.importance: torch.Tensor | None = None
|
||||
self.n = 0
|
||||
|
||||
def add(self, entropy: torch.Tensor, importance: torch.Tensor, n: int) -> None:
|
||||
self.entropy += entropy.item() * n
|
||||
self.importance = importance.clone() if self.importance is None else self.importance + importance
|
||||
self.n += n
|
||||
|
||||
def stats(self) -> dict[str, float]:
|
||||
if self.importance is None or self.n == 0:
|
||||
return dict.fromkeys(_ROUTER_KEYS, 0.0)
|
||||
util = self.importance / self.importance.sum().clamp_min(1e-8)
|
||||
return {
|
||||
"entropy": self.entropy / self.n,
|
||||
"util_min": util.min().item(),
|
||||
"util_max": util.max().item(),
|
||||
"util_std": util.std().item() if self.n_experts > 1 else 0.0,
|
||||
}
|
||||
|
||||
def reset(self) -> None:
|
||||
self.entropy = 0.0
|
||||
self.importance = None
|
||||
self.n = 0
|
||||
|
||||
|
||||
class MetricsCollector:
|
||||
"""Owns every number a training run reports.
|
||||
|
||||
Accumulates per-batch stats from each stage, writes one `metrics.csv` row
|
||||
per epoch, mirrors it to W&B, and formats the tqdm postfix and the epoch
|
||||
summary line — so `giant.training.loop.train` never carries a running
|
||||
sum, a column name, or a W&B call of its own.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
trainers: dict,
|
||||
out_dir: Path,
|
||||
*,
|
||||
epochs: int,
|
||||
resume: bool = False,
|
||||
wandb_run=None,
|
||||
wandb_log_every: int = 50,
|
||||
) -> None:
|
||||
self.trainers = trainers
|
||||
self.epochs = epochs
|
||||
self.wandb_run = wandb_run
|
||||
self.wandb_log_every = wandb_log_every
|
||||
self.epoch_width = len(str(epochs))
|
||||
|
||||
self._train = {name: _Accumulator() for name in trainers}
|
||||
self._val = {name: _Accumulator() for name in trainers}
|
||||
self._routers = {
|
||||
name: _RouterAccumulator(tr.router.n_experts) for name, tr in trainers.items() if tr.router is not None
|
||||
}
|
||||
# Only "mean" specs need summing; "last" specs are read straight off
|
||||
# the accumulator's most recent stats dict. "grad_norm" is always
|
||||
# summed — it feeds the run-level `grad_norm` column whether or not
|
||||
# a trainer reports it per stage.
|
||||
self._train_keys = {
|
||||
name: {spec.key for spec in tr.train_metrics if spec.reduce == "mean"} | {"grad_norm"}
|
||||
for name, tr in trainers.items()
|
||||
}
|
||||
self._val_keys = {
|
||||
name: {spec.key for spec in tr.val_metrics if spec.reduce == "mean"} for name, tr in trainers.items()
|
||||
}
|
||||
self._run_values: dict[str, float] = {}
|
||||
self._epoch = 0
|
||||
self._ema_loss = 0.0
|
||||
self._ema_grad_norm = 0.0
|
||||
self._ema_seeded = False
|
||||
self._batch_loss = 0.0
|
||||
self._batch_grad_norm = 0.0
|
||||
|
||||
self.fieldnames = self._build_fieldnames()
|
||||
metrics_path = out_dir / "metrics.csv"
|
||||
append = resume and metrics_path.exists()
|
||||
self._file = open(metrics_path, "a" if append else "w", newline="")
|
||||
self._writer = csv.DictWriter(self._file, fieldnames=self.fieldnames)
|
||||
if not append:
|
||||
self._writer.writeheader()
|
||||
|
||||
# --- construction ---------------------------------------------------
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls,
|
||||
trainers: dict,
|
||||
out_dir: Path,
|
||||
cfg: dict,
|
||||
model_config: dict | None,
|
||||
*,
|
||||
resume: bool = False,
|
||||
use_wandb: bool = False,
|
||||
wandb_project: str = "giant",
|
||||
wandb_run_name: str = "",
|
||||
wandb_log_every: int = 50,
|
||||
) -> "MetricsCollector":
|
||||
"""Build the collector, starting a W&B run first when enabled."""
|
||||
wandb_run = None
|
||||
if use_wandb:
|
||||
try:
|
||||
import wandb
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"train.wandb = true (--wandb) requires the 'wandb' package — install it via `uv sync --extra wandb`"
|
||||
) from exc
|
||||
param_counts = {name: sum(p.numel() for p in tr.model.parameters()) for name, tr in trainers.items()}
|
||||
param_counts["total"] = sum(param_counts.values())
|
||||
wandb_run = wandb.init(
|
||||
project=wandb_project,
|
||||
name=wandb_run_name or out_dir.name,
|
||||
id=out_dir.name,
|
||||
resume="allow",
|
||||
config=_wandb_run_config(cfg, model_config, param_counts),
|
||||
)
|
||||
return cls(
|
||||
trainers,
|
||||
out_dir,
|
||||
epochs=cfg["train"]["epochs"],
|
||||
resume=resume,
|
||||
wandb_run=wandb_run,
|
||||
wandb_log_every=wandb_log_every,
|
||||
)
|
||||
|
||||
def _build_fieldnames(self) -> list[str]:
|
||||
fields = ["epoch"]
|
||||
for name, trainer in self.trainers.items():
|
||||
for spec in trainer.train_metrics:
|
||||
fields.append(f"{name}/{spec.column}")
|
||||
for spec in trainer.val_metrics:
|
||||
fields.append(f"{name}/{spec.column}")
|
||||
if trainer.router is not None:
|
||||
fields += [f"{name}/router/{key}" for key in _ROUTER_KEYS]
|
||||
for spec in trainer.stage_metrics:
|
||||
fields.append(f"{name}/{spec.column}")
|
||||
fields += list(_RUN_COLUMNS)
|
||||
return fields
|
||||
|
||||
def close(self) -> None:
|
||||
self._file.close()
|
||||
if self.wandb_run is not None:
|
||||
self.wandb_run.finish()
|
||||
|
||||
# --- per-batch ------------------------------------------------------
|
||||
|
||||
def start_epoch(self, epoch: int) -> None:
|
||||
self._epoch = epoch
|
||||
for acc in self._train.values():
|
||||
acc.reset()
|
||||
for acc in self._val.values():
|
||||
acc.reset()
|
||||
for acc in self._routers.values():
|
||||
acc.reset()
|
||||
self._run_values.clear()
|
||||
self._ema_seeded = False
|
||||
|
||||
def add_train_batch(self, stats: dict[str, dict], batch_size: int) -> None:
|
||||
"""`stats` maps stage name -> the dict that stage's `step()` returned."""
|
||||
self._batch_loss = 0.0
|
||||
self._batch_grad_norm = 0.0
|
||||
for name, stage_stats in stats.items():
|
||||
self._train[name].add(stage_stats, self._train_keys[name], batch_size)
|
||||
self._batch_loss += self.trainers[name].batch_loss(stage_stats)
|
||||
self._batch_grad_norm += stage_stats.get("grad_norm", 0.0)
|
||||
if self._ema_seeded:
|
||||
self._ema_loss += _EMA_ALPHA * (self._batch_loss - self._ema_loss)
|
||||
self._ema_grad_norm += _EMA_ALPHA * (self._batch_grad_norm - self._ema_grad_norm)
|
||||
else:
|
||||
self._ema_loss = self._batch_loss
|
||||
self._ema_grad_norm = self._batch_grad_norm
|
||||
self._ema_seeded = True
|
||||
|
||||
def add_val_batch(self, stats: dict[str, dict], batch_size: int) -> None:
|
||||
for name, stage_stats in stats.items():
|
||||
self._val[name].add(stage_stats, self._val_keys[name], batch_size)
|
||||
|
||||
@torch.no_grad()
|
||||
def observe_routers(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor, batch_size: int) -> None:
|
||||
"""Record gate diagnostics for every routed stage on this batch.
|
||||
|
||||
Called from the validation pass only (as in v0.2/v0.3.0), so a stage
|
||||
whose trainer has no validation pass — i.e. WGAN — reports zeros.
|
||||
"""
|
||||
for name, acc in self._routers.items():
|
||||
router = self.trainers[name].router
|
||||
entropy, importance = router.gate_stats(cond_cont, cond_cat)
|
||||
acc.add(entropy, importance, batch_size)
|
||||
|
||||
def postfix(self) -> str:
|
||||
"""tqdm postfix for the training bar."""
|
||||
return f"loss={self._ema_loss:.4f} gnorm={self._ema_grad_norm:.3f}"
|
||||
|
||||
def log_batch(self, global_step: int, batch: tuple, device: torch.device) -> None:
|
||||
"""Batch-granularity W&B log, throttled to every `wandb_log_every`
|
||||
optimizer steps (a single epoch can be tens of thousands). `batch` is
|
||||
the raw training batch, needed only to re-derive routing entropy for
|
||||
routed stages — it is never moved to `device` otherwise."""
|
||||
if self.wandb_run is None or self.wandb_log_every <= 0:
|
||||
return
|
||||
if global_step % self.wandb_log_every != 0:
|
||||
return
|
||||
payload = {
|
||||
"batch/epoch": self._epoch,
|
||||
"batch/loss": self._batch_loss,
|
||||
"batch/loss_ema": self._ema_loss,
|
||||
"batch/grad_norm": self._batch_grad_norm,
|
||||
}
|
||||
for name, trainer in self.trainers.items():
|
||||
payload[f"batch/{name}/lr"] = trainer.optimizer.param_groups[0]["lr"]
|
||||
if trainer.router is not None:
|
||||
with torch.no_grad():
|
||||
entropy, _ = trainer.router.gate_stats(batch[0].to(device), batch[1].to(device))
|
||||
payload[f"batch/{name}/router/entropy"] = entropy.item()
|
||||
self.wandb_run.log(payload, step=global_step)
|
||||
|
||||
# --- per-epoch ------------------------------------------------------
|
||||
|
||||
def train_means(self, stage: str) -> dict[str, float]:
|
||||
return self._train[stage].means()
|
||||
|
||||
def val_means(self, stage: str) -> dict[str, float]:
|
||||
return self._val[stage].means()
|
||||
|
||||
@property
|
||||
def train_samples(self) -> int:
|
||||
"""Samples seen this epoch — identical across stages (every stage
|
||||
steps on every batch), so any one accumulator's count will do."""
|
||||
return max((acc.n for acc in self._train.values()), default=0)
|
||||
|
||||
def set(self, column: str, value: float) -> None:
|
||||
"""Record a run-level value for this epoch's row (`val/loss`,
|
||||
`gpu_mem_mb`, ...). Must name a column in `_RUN_COLUMNS`."""
|
||||
if column not in _RUN_COLUMNS:
|
||||
raise KeyError(f"{column!r} is not a run-level metrics column")
|
||||
self._run_values[column] = value
|
||||
|
||||
def summary_line(self, val_loss: float, epoch_time: float, is_best: bool) -> str:
|
||||
bits = [trainer.summary(self.train_means(name)) for name, trainer in self.trainers.items()]
|
||||
marker = " [best]" if is_best else ""
|
||||
return (
|
||||
f"epoch {self._epoch:{self.epoch_width}d}/{self.epochs} "
|
||||
+ " ".join(bits)
|
||||
+ f" val {val_loss:.4f} {epoch_time:.1f}s{marker}"
|
||||
)
|
||||
|
||||
def write_epoch(self, global_step: int) -> None:
|
||||
"""Assemble, write, and flush this epoch's row; mirror it to W&B."""
|
||||
row: dict = {"epoch": self._epoch}
|
||||
grad_norm_total = 0.0
|
||||
for name, trainer in self.trainers.items():
|
||||
train_acc, val_acc = self._train[name], self._val[name]
|
||||
for spec in trainer.train_metrics:
|
||||
row[f"{name}/{spec.column}"] = train_acc.mean(spec.key)
|
||||
for spec in trainer.val_metrics:
|
||||
row[f"{name}/{spec.column}"] = val_acc.mean(spec.key)
|
||||
if trainer.router is not None:
|
||||
for key, value in self._routers[name].stats().items():
|
||||
row[f"{name}/router/{key}"] = value
|
||||
for spec in trainer.stage_metrics:
|
||||
row[f"{name}/{spec.column}"] = train_acc.last.get(spec.key, 0.0)
|
||||
grad_norm_total += train_acc.mean("grad_norm")
|
||||
|
||||
for column in _RUN_COLUMNS:
|
||||
row[column] = self._run_values.get(column, float("nan"))
|
||||
row["grad_norm"] = grad_norm_total
|
||||
|
||||
self._writer.writerow(row)
|
||||
self._file.flush()
|
||||
if self.wandb_run is not None:
|
||||
self.wandb_run.log(row, step=global_step)
|
||||
@@ -1,318 +0,0 @@
|
||||
"""Ground-truth tensor assembly for stage-2 training.
|
||||
|
||||
Pure functions, no optimizer/model state: they turn a batch's ground-truth
|
||||
secondary tensors into the per-token targets and autoregressive conditioning
|
||||
inputs `giant.training.trainers` feeds to `Stage2OneShot` /
|
||||
`Stage2Autoregressive`. Split out of the trainers so the (target, generator,
|
||||
decoder) width rules — the fiddliest part of this codebase —
|
||||
live in one place and stay unit-testable on their own.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
|
||||
from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM
|
||||
from giant.sample import sample_secondaries_ar
|
||||
|
||||
|
||||
def _gumbel_tau(step: int, total_steps: int, tau_start: float, tau_end: float) -> float:
|
||||
"""Linear anneal of the straight-through Gumbel-softmax temperature.
|
||||
|
||||
Deterministic in `step`/`total_steps` alone (no extra state), so it
|
||||
recomputes correctly on `--resume` from a checkpoint's saved `global_step`
|
||||
without needing to persist anything new (see
|
||||
giant.model.network.Router.combine_weights).
|
||||
"""
|
||||
progress = min(step / max(total_steps, 1), 1.0)
|
||||
return tau_start + (tau_end - tau_start) * progress
|
||||
|
||||
|
||||
def _type_repr(
|
||||
sec_type_idx: torch.Tensor,
|
||||
sec_cont: torch.Tensor,
|
||||
particle_type_cfg: dict,
|
||||
cond_enc: torch.nn.Module,
|
||||
emb_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""(B, K_MAX, type_dim) ground-truth type representation, generator-
|
||||
independent (unlike `_assemble_stage2_ar_target`'s training *target*,
|
||||
which varies by generator/objective — see its docstring): `"physical"` ->
|
||||
`(log_mass, charge)`; `"onehot"` -> one-hot of the true class;
|
||||
`"embedding"` -> the conditioning's own detached embedding-table row.
|
||||
|
||||
Used both to build `_assemble_stage2_ar_target`'s wgan+onehot/embedding
|
||||
branch and as the AR history features' previous-secondary identity — the
|
||||
latter must always reflect the true physical secondary that came before,
|
||||
regardless of what the *current* token's own training objective is.
|
||||
"""
|
||||
target = particle_type_cfg.get("target", "physical")
|
||||
if target == "physical":
|
||||
return sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM]
|
||||
if target == "onehot":
|
||||
return F.one_hot(sec_type_idx, num_classes=emb_dim).float()
|
||||
return cond_enc.pdg_emb(sec_type_idx).detach()
|
||||
|
||||
|
||||
def _assemble_stage2_ar_target(
|
||||
sec_cont: torch.Tensor,
|
||||
sec_type_idx: torch.Tensor,
|
||||
particle_type_cfg: dict,
|
||||
generator: str,
|
||||
cond_enc: torch.nn.Module,
|
||||
emb_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""(B, K_MAX, token_dim) ground-truth per-token target — the unflattened
|
||||
analogue of `_assemble_stage2_real` (defined below in terms of this),
|
||||
matching whatever width `Stage2Autoregressive`'s (or `Stage2OneShot`'s)
|
||||
own trunk produces for this (target, generator) combination
|
||||
(`giant.model.network.stage2_trunk_sec_dim`):
|
||||
|
||||
- `target = "physical"`: unchanged from v0.2 — `sec_cont` (stick_logit,
|
||||
dir, log_mass, charge) as-is.
|
||||
- `target` in `("onehot", "embedding")` + `generator in ("flow", "ddpm")`:
|
||||
just the continuous stick/dir slots — the type slice isn't part of
|
||||
this tensor at all (`type_head` handles it separately).
|
||||
- `target` in `("onehot", "embedding")` + `generator == "wgan"`: stick/dir
|
||||
slots concatenated with the per-slot type representation (a one-hot of
|
||||
the true class, relaxed on the *generated* side only, by the caller;
|
||||
or the conditioning's own detached embedding-table row).
|
||||
"""
|
||||
target = particle_type_cfg.get("target", "physical")
|
||||
if target == "physical":
|
||||
return sec_cont
|
||||
cont = sec_cont[..., :CONT_SLOT_DIM]
|
||||
if generator != "wgan":
|
||||
return cont
|
||||
type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
|
||||
return torch.cat([cont, type_repr], dim=-1)
|
||||
|
||||
|
||||
def _assemble_stage2_real(
|
||||
sec_cont: torch.Tensor,
|
||||
sec_type_idx: torch.Tensor,
|
||||
particle_type_cfg: dict,
|
||||
generator: str,
|
||||
cond_enc: torch.nn.Module,
|
||||
emb_dim: int,
|
||||
) -> torch.Tensor:
|
||||
"""Ground-truth flattened stage-2 vector for `Stage2OneShot` — the
|
||||
flattened form of `_assemble_stage2_ar_target`, which
|
||||
`Stage2Autoregressive`'s per-token target also uses; the two must stay in
|
||||
lockstep. See `_assemble_stage2_ar_target`'s docstring for the
|
||||
(target, generator) width rules."""
|
||||
return _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim).flatten(
|
||||
1
|
||||
)
|
||||
|
||||
|
||||
def _stick_fraction(sec_cont: torch.Tensor) -> torch.Tensor:
|
||||
"""(B, K_MAX) — sigmoid of each slot's own stick-breaking logit
|
||||
(`sec_cont[...,0]`); scale-free (see `giant.data.transforms.
|
||||
encode_secondaries`), so this needs no absolute `e_sec`."""
|
||||
return torch.sigmoid(sec_cont[..., 0])
|
||||
|
||||
|
||||
def _remaining_energy_fraction(fraction: torch.Tensor) -> torch.Tensor:
|
||||
"""(B, K_MAX) — fraction of the original e_sec budget unclaimed entering
|
||||
slot i: `1.0` at `i=0`, `prod_{j<i}(1-fraction_j)` for `i>=1`
|
||||
("no re-derivation needed": the existing
|
||||
stick-breaking encoding is already scale-free, so this is derivable from
|
||||
the batch's ground-truth stick logits alone, no `e_sec` required)."""
|
||||
cumprod = torch.cumprod(1.0 - fraction, dim=1)
|
||||
return torch.cat([torch.ones_like(cumprod[:, :1]), cumprod[:, :-1]], dim=1)
|
||||
|
||||
|
||||
def _shift_prev(x: torch.Tensor) -> torch.Tensor:
|
||||
"""`(B, K, ...)` -> same shape, slot i holds slot i-1's value; slot 0 gets
|
||||
an arbitrary zero placeholder (never read as-is — see `_ar_has_prev`;
|
||||
`MarkovHistory` substitutes its own learned start vector there instead)."""
|
||||
return torch.cat([torch.zeros_like(x[:, :1]), x[:, :-1]], dim=1)
|
||||
|
||||
|
||||
def _ar_has_prev(k_max: int, device: torch.device) -> torch.Tensor:
|
||||
"""`(1, K_MAX)` bool: True for slot index `>= 1`. Correct without
|
||||
`n_sec`: `sec_mask` is a prefix mask, so any *valid* token at `k>=1`
|
||||
always has a valid predecessor at `k-1`; the only wrong cases are tokens
|
||||
that are themselves padding, already masked out of every loss."""
|
||||
return (torch.arange(k_max, device=device) >= 1).unsqueeze(0)
|
||||
|
||||
|
||||
def _ar_meta(k_max: int, batch: int, device: torch.device, fraction: torch.Tensor) -> dict[str, torch.Tensor]:
|
||||
"""`has_prev`/`remaining_frac`/`slot_idx` — the three per-token AR
|
||||
conditioning tensors that don't depend on *which* history representation
|
||||
(ground truth vs. the scheduled-sampling mix) produced `fraction`.
|
||||
Shared by `_assemble_stage2_ar_inputs` and
|
||||
`_assemble_stage2_ar_inputs_scheduled`, which differ only in
|
||||
`history_feat`."""
|
||||
slot_idx = (torch.arange(k_max, device=device).float() / max(k_max - 1, 1)).unsqueeze(0)
|
||||
return {
|
||||
"has_prev": _ar_has_prev(k_max, device).expand(batch, -1),
|
||||
"remaining_frac": _remaining_energy_fraction(fraction),
|
||||
"slot_idx": slot_idx.expand(batch, -1),
|
||||
}
|
||||
|
||||
|
||||
def _assemble_stage2_ar_inputs(
|
||||
sec_cont: torch.Tensor,
|
||||
sec_type_idx: torch.Tensor,
|
||||
particle_type_cfg: dict,
|
||||
cond_enc: torch.nn.Module,
|
||||
emb_dim: int,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Ground-truth per-token AR conditioning tensors — all `(B, K_MAX, ...)`
|
||||
or `(B, K_MAX)`, built in one vectorized pass (teacher forcing means
|
||||
every token's input is ground truth).
|
||||
Keys match `Stage2Autoregressive.forward`'s trailing kwargs."""
|
||||
device = sec_cont.device
|
||||
B, K = sec_cont.shape[0], sec_cont.shape[1]
|
||||
fraction = _stick_fraction(sec_cont)
|
||||
type_repr = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
|
||||
history_feat = torch.cat(
|
||||
[
|
||||
_shift_prev(fraction).unsqueeze(-1),
|
||||
_shift_prev(sec_cont[..., 1:CONT_SLOT_DIM]),
|
||||
_shift_prev(type_repr),
|
||||
],
|
||||
dim=-1,
|
||||
)
|
||||
return {"history_feat": history_feat, **_ar_meta(K, B, device, fraction)}
|
||||
|
||||
|
||||
def _stage2_tf_prob(mode: str, p_start: float, p_end: float, epoch: int, total_epochs: int) -> float:
|
||||
"""P(condition slot k+1 on the TRUE token k rather than the model's own
|
||||
prediction), for the current epoch
|
||||
(`stage2_model.autoregressive.teacher_forcing`).
|
||||
`"always"`/`"never"` are the two degenerate constants; `"scheduled"`
|
||||
linearly interpolates
|
||||
`p_start` (epoch 0) to `p_end` (the final epoch) — standard scheduled
|
||||
sampling (Bengio et al. 2015)."""
|
||||
if mode == "always":
|
||||
return 1.0
|
||||
if mode == "never":
|
||||
return 0.0
|
||||
frac = epoch / max(total_epochs - 1, 1)
|
||||
frac = min(max(frac, 0.0), 1.0)
|
||||
return p_start + (p_end - p_start) * frac
|
||||
|
||||
|
||||
def _history_repr_from_ar_sample(
|
||||
sec_cont_pred: torch.Tensor,
|
||||
sec_type_pred: torch.Tensor,
|
||||
particle_type_cfg: dict,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
||||
"""`(fraction, direction, type_repr)` — the same triple `_type_repr` /
|
||||
`_stick_fraction` derive from ground truth, but from a free-running
|
||||
`sample_secondaries_ar` self-sample instead, so the two can be mixed
|
||||
slot-by-slot under scheduled sampling (`_assemble_stage2_ar_inputs_scheduled`).
|
||||
`target="onehot"` collapses the raw per-slot type logits to a hard
|
||||
one-hot of `argmax` — `sample_secondaries_ar`'s own history convention
|
||||
(see its docstring), matching what `MarkovHistory`/`AttentionHistory`
|
||||
were trained on; the other two targets are already the right
|
||||
representation."""
|
||||
fraction = torch.sigmoid(sec_cont_pred[..., 0])
|
||||
direction = sec_cont_pred[..., 1:4]
|
||||
if particle_type_cfg.get("target", "physical") == "onehot":
|
||||
type_dim = sec_type_pred.size(-1)
|
||||
type_repr = F.one_hot(sec_type_pred.argmax(-1), num_classes=type_dim).float()
|
||||
else:
|
||||
type_repr = sec_type_pred
|
||||
return fraction, direction, type_repr
|
||||
|
||||
|
||||
def _assemble_stage2_ar_inputs_scheduled(
|
||||
model: torch.nn.Module,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_ctx: torch.Tensor,
|
||||
sec_cont: torch.Tensor,
|
||||
sec_type_idx: torch.Tensor,
|
||||
n_sec: torch.Tensor,
|
||||
particle_type_cfg: dict,
|
||||
cond_enc: torch.nn.Module,
|
||||
emb_dim: int,
|
||||
p_tf: float,
|
||||
sample_steps: int,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Scheduled-sampling counterpart of `_assemble_stage2_ar_inputs`
|
||||
(`teacher_forcing` = "scheduled"/"never"):
|
||||
each slot's history is the TRUE previous token with probability `p_tf`
|
||||
(an independent per-example, per-slot Bernoulli draw) and the model's own
|
||||
free-running prediction otherwise — closing the train/inference gap that
|
||||
`teacher_forcing="always"` (ground truth throughout training) never sees.
|
||||
`p_tf >= 1.0` degenerates exactly to `_assemble_stage2_ar_inputs` (and
|
||||
skips self-sampling entirely), so callers can call this unconditionally.
|
||||
|
||||
The free-running estimate is a REAL autoregressive self-sample —
|
||||
`giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
|
||||
cheap one-step proxy, so building it costs the same `k_max` (`* steps`
|
||||
for flow) sequential forwards `sample.py` pays at inference, EVERY batch
|
||||
this is called on (paid at train time too whenever teacher_forcing !=
|
||||
"always"). Fully detached: gradient only ever flows
|
||||
through the "real" target path each stage trainer already uses
|
||||
(`_assemble_stage2_ar_target`), never through this self-sample.
|
||||
"""
|
||||
device = sec_cont.device
|
||||
B, K = sec_cont.shape[0], sec_cont.shape[1]
|
||||
if p_tf >= 1.0:
|
||||
return _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, particle_type_cfg, cond_enc, emb_dim)
|
||||
|
||||
was_training = model.training
|
||||
sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
|
||||
model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
|
||||
)
|
||||
if was_training:
|
||||
model.train()
|
||||
|
||||
fraction_gt = _stick_fraction(sec_cont)
|
||||
dir_gt = sec_cont[..., 1:CONT_SLOT_DIM]
|
||||
type_repr_gt = _type_repr(sec_type_idx, sec_cont, particle_type_cfg, cond_enc, emb_dim)
|
||||
fraction_pred, dir_pred, type_repr_pred = _history_repr_from_ar_sample(
|
||||
sec_cont_pred, sec_type_pred, particle_type_cfg
|
||||
)
|
||||
|
||||
use_gt = torch.rand(B, K, device=device) < p_tf
|
||||
fraction = torch.where(use_gt, fraction_gt, fraction_pred)
|
||||
direction = torch.where(use_gt.unsqueeze(-1), dir_gt, dir_pred)
|
||||
type_repr = torch.where(use_gt.unsqueeze(-1), type_repr_gt, type_repr_pred)
|
||||
|
||||
own_feat = torch.cat([fraction.unsqueeze(-1), direction, type_repr], dim=-1)
|
||||
return {
|
||||
"history_feat": _shift_prev(own_feat),
|
||||
**_ar_meta(K, B, device, fraction),
|
||||
}
|
||||
|
||||
|
||||
def _relax_onehot_type_slice(
|
||||
x_flat: torch.Tensor,
|
||||
k_max: int,
|
||||
cont_dim: int,
|
||||
type_dim: int,
|
||||
tau: float,
|
||||
grad_probe: dict[str, float] | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Straight-through Gumbel-softmax relaxation of the per-slot type slice
|
||||
inside a flattened `(B, k_max * (cont_dim + type_dim))` WGAN generator
|
||||
output: the forward pass is a
|
||||
hard one-hot (matching what the critic sees from real data), the
|
||||
backward pass flows smooth gradient. Continuous slots (stick/dir, and
|
||||
the type slice itself under `target = "embedding"`, which never calls
|
||||
this) pass through unchanged.
|
||||
|
||||
`grad_probe`, if given, gets `["cont"]`/`["type"]` populated with the L2
|
||||
norm of the gradient reaching this split point during the next
|
||||
`.backward()` call that touches it — a backward hook, not a second
|
||||
backward pass. This is the differentiability validation-obligation
|
||||
instrumentation: the trunk-gradient contribution
|
||||
from the type slice vs. the continuous slices, for
|
||||
`particle_type.target="onehot"` + `generator="wgan"`. Only ever populated
|
||||
on a `did_g_step` batch — the critic step backprops through
|
||||
`fake.detach()`, which never reaches these hooks — so it stays empty
|
||||
(callers default to `0.0`) otherwise."""
|
||||
B = x_flat.size(0)
|
||||
x = x_flat.view(B, k_max, cont_dim + type_dim)
|
||||
cont, type_logits = x[..., :cont_dim], x[..., cont_dim:]
|
||||
if grad_probe is not None:
|
||||
cont.register_hook(lambda g: grad_probe.__setitem__("cont", g.norm().item()))
|
||||
type_logits.register_hook(lambda g: grad_probe.__setitem__("type", g.norm().item()))
|
||||
type_soft = F.gumbel_softmax(type_logits, tau=tau, hard=True, dim=-1)
|
||||
return torch.cat([cont, type_soft], dim=-1).reshape(B, -1)
|
||||
@@ -1,968 +0,0 @@
|
||||
"""Per-stage trainers: optimizer(s), EMA, LR schedule, and the per-batch step.
|
||||
|
||||
`StageSpec` resolves one stage's slice of the config once, so the two
|
||||
concrete trainers share a single constructor shape instead of ~24 keyword
|
||||
arguments each, and `StageTrainer` carries every piece that used to be
|
||||
copy-pasted between them (cosine warmup, EMA, checkpoint state, LR resume,
|
||||
train/eval toggling).
|
||||
|
||||
Each trainer also *declares* the metrics it emits, as `MetricSpec` lists —
|
||||
that declaration is the single source of truth for `metrics.csv` and W&B
|
||||
columns (see `giant.training.metrics`) — and exposes the three small hooks
|
||||
(`batch_loss`, `summary`, `val_objective`) that let the epoch loop treat
|
||||
adversarial and non-adversarial stages identically.
|
||||
"""
|
||||
|
||||
import copy
|
||||
import math
|
||||
from dataclasses import dataclass, field
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import torch.optim as optim
|
||||
|
||||
from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
|
||||
from giant.constants import CONT_SLOT_DIM
|
||||
from giant.data.dataset import StepBatch
|
||||
from giant.model.network import Router, resolve_type_n_classes, stage2_type_dim
|
||||
from giant.model.schedule import (
|
||||
CosineSchedule,
|
||||
flow_matching_loss,
|
||||
flow_matching_loss_secondary,
|
||||
flow_matching_loss_secondary_ar,
|
||||
)
|
||||
from giant.model.wgan import generator_loss, gradient_penalty
|
||||
from giant.training.metrics import MetricSpec, stage_metric, train_metric, val_metric
|
||||
from giant.training.stage2_inputs import (
|
||||
_assemble_stage2_ar_inputs_scheduled,
|
||||
_assemble_stage2_ar_target,
|
||||
_gumbel_tau,
|
||||
_relax_onehot_type_slice,
|
||||
_stage2_tf_prob,
|
||||
)
|
||||
|
||||
|
||||
@torch.no_grad()
|
||||
def _update_ema(ema_model: torch.nn.Module, model: torch.nn.Module, decay: float) -> None:
|
||||
for ema_p, p in zip(ema_model.parameters(), model.parameters()):
|
||||
ema_p.mul_(decay).add_(p, alpha=1 - decay)
|
||||
|
||||
|
||||
def _stage_router(model: torch.nn.Module) -> Router | None:
|
||||
"""A stage model's Router, if its trunk is routed — else None.
|
||||
|
||||
Post-step-2 refactor the router lives at `model.trunk.router`
|
||||
(`giant.model.network.RoutedTrunk`), not `model.router` directly.
|
||||
"""
|
||||
trunk = getattr(model, "trunk", None)
|
||||
return getattr(trunk, "router", None)
|
||||
|
||||
|
||||
def _cosine_warmup_lambda(warmup_steps: int, total_steps: int):
|
||||
"""Linear warmup for `warmup_steps`, then cosine decay to zero over the
|
||||
remainder — the LR schedule both trainers use, in their own step units
|
||||
(optimizer steps for flow/ddpm, generator steps for WGAN)."""
|
||||
|
||||
def _lr_lambda(step: int) -> float:
|
||||
if warmup_steps > 0 and step < warmup_steps:
|
||||
return (step + 1) / warmup_steps
|
||||
t = step - warmup_steps
|
||||
T = max(total_steps - warmup_steps, 1)
|
||||
return 0.5 * (1.0 + math.cos(math.pi * min(t, T) / T))
|
||||
|
||||
return _lr_lambda
|
||||
|
||||
|
||||
def _batch_to_device(batch: StepBatch, device: torch.device) -> StepBatch:
|
||||
return type(batch)(*(t.to(device) for t in batch))
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StageSpec:
|
||||
"""One stage's resolved training configuration.
|
||||
|
||||
Built once by `StageSpec.from_config`, which is the only place that reads
|
||||
the `cfg` dict — so a new config key means one new field and one new read,
|
||||
not another argument threaded through two constructors.
|
||||
"""
|
||||
|
||||
name: str
|
||||
is_stage2: bool
|
||||
generator: str
|
||||
decoder: str = "one_shot"
|
||||
|
||||
# loss weights
|
||||
lambda_weight: float = 1.0
|
||||
n_sec_lambda: float = 0.1
|
||||
|
||||
# particle-type target (stage 2 only)
|
||||
particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig)
|
||||
particle_type_n_classes: int = 16
|
||||
|
||||
# optimization
|
||||
lr: float = 3e-4
|
||||
weight_decay: float = 0.01
|
||||
ema_decay: float = 0.9999
|
||||
warmup_epochs: int = 0
|
||||
epochs: int = 1
|
||||
steps_per_epoch: int = 1
|
||||
|
||||
# routing auxiliaries
|
||||
lambda_balance: float = 0.0
|
||||
lambda_proc: float = 0.0
|
||||
lambda_entropy: float = 0.0
|
||||
gumbel_tau_start: float = 1.0
|
||||
gumbel_tau_end: float = 0.1
|
||||
|
||||
# autoregressive stage 2
|
||||
teacher_forcing: str = "always"
|
||||
tf_p_start: float = 1.0
|
||||
tf_p_end: float = 1.0
|
||||
ar_sample_steps: int = 10
|
||||
|
||||
# generator-specific
|
||||
ddpm_n_steps: int = 1000
|
||||
n_critic: int = 5
|
||||
gp_weight: float = 10.0
|
||||
critic_lr: float = 0.0
|
||||
type_gumbel_tau_start: float = 1.0
|
||||
type_gumbel_tau_end: float = 0.1
|
||||
|
||||
@classmethod
|
||||
def from_config(cls, cfg: dict, name: str, is_stage2: bool, steps_per_epoch: int) -> "StageSpec":
|
||||
t = TrainConfig.from_dict(cfg["train"])
|
||||
# n_sec/particle_type/decoder/autoregressive/wgan's gumbel_tau_* are
|
||||
# stage-2-only concepts, always read off s2_spec (guarded by
|
||||
# is_stage2 where the stage-1 StageSpec needs a different value) —
|
||||
# historically n_sec/particle_type were read from stage2_model
|
||||
# unconditionally even for the stage-1 StageSpec, preserved here for
|
||||
# behavioral parity. stage_spec covers the fields both stage configs
|
||||
# share structurally (generator, lambda, router, ddpm, and wgan's
|
||||
# base fields — Stage2ModelConfig's sub-configs all subclass
|
||||
# stage 1's).
|
||||
s2_spec = Stage2ModelConfig.from_dict(cfg["stage2_model"])
|
||||
stage_spec = s2_spec if is_stage2 else Stage1ModelConfig.from_dict(cfg["stage1_model"])
|
||||
return cls(
|
||||
name=name,
|
||||
is_stage2=is_stage2,
|
||||
generator=stage_spec.generator,
|
||||
decoder=s2_spec.decoder if is_stage2 else "one_shot",
|
||||
lambda_weight=stage_spec.lambda_weight,
|
||||
n_sec_lambda=s2_spec.n_sec.lambda_weight,
|
||||
particle_type=s2_spec.particle_type,
|
||||
particle_type_n_classes=resolve_type_n_classes(
|
||||
s2_spec.particle_type.to_dict(), cfg["conditioning"]["particle"]["emb_dim"]
|
||||
),
|
||||
# train.* keys are all guaranteed by DEFAULT_CONFIG's deep-merge
|
||||
# (giant/config.py), so TrainConfig.from_dict never has to fall
|
||||
# back to a literal here; the field defaults below exist only
|
||||
# for tests that construct StageSpec by hand.
|
||||
lr=t.lr,
|
||||
weight_decay=t.weight_decay,
|
||||
ema_decay=t.ema_decay,
|
||||
warmup_epochs=t.warmup_epochs,
|
||||
epochs=t.epochs,
|
||||
steps_per_epoch=max(steps_per_epoch, 1),
|
||||
lambda_balance=stage_spec.router.lambda_balance,
|
||||
lambda_proc=stage_spec.router.lambda_proc,
|
||||
lambda_entropy=stage_spec.router.lambda_entropy,
|
||||
gumbel_tau_start=stage_spec.router.gumbel_tau_start,
|
||||
gumbel_tau_end=stage_spec.router.gumbel_tau_end,
|
||||
teacher_forcing=s2_spec.autoregressive.teacher_forcing if is_stage2 else cls.teacher_forcing,
|
||||
tf_p_start=s2_spec.autoregressive.tf_p_start if is_stage2 else cls.tf_p_start,
|
||||
tf_p_end=s2_spec.autoregressive.tf_p_end if is_stage2 else cls.tf_p_end,
|
||||
# AR self-sampling under scheduled/never teacher forcing reuses
|
||||
# train.validate_steps as its flow-matching ODE step count — no
|
||||
# dedicated config key for this (the autoregressive config lists
|
||||
# tf_p_start/tf_p_end/attn_n_heads/attn_n_layers only).
|
||||
ar_sample_steps=t.validate_steps,
|
||||
ddpm_n_steps=stage_spec.ddpm.n_steps,
|
||||
n_critic=stage_spec.wgan.n_critic,
|
||||
gp_weight=stage_spec.wgan.gp_weight,
|
||||
critic_lr=stage_spec.wgan.critic_lr,
|
||||
type_gumbel_tau_start=s2_spec.wgan.gumbel_tau_start if is_stage2 else cls.type_gumbel_tau_start,
|
||||
type_gumbel_tau_end=s2_spec.wgan.gumbel_tau_end if is_stage2 else cls.type_gumbel_tau_end,
|
||||
)
|
||||
|
||||
|
||||
class StageTrainer:
|
||||
"""One active stage's optimizer(s), EMA, and per-batch step.
|
||||
|
||||
Reads only the shared `StepBatch` (`giant.data.dataset`) — stage 2 always
|
||||
conditions on the ground-truth `x1_s1` (`stage2_model.stage1_context =
|
||||
"truth"`, stage-level teacher forcing; `"sampled"` is not implemented),
|
||||
so stage trainers never need each other's output at train time. This means
|
||||
"stage-2-only training is a cheap ablation, not new plumbing" falls out
|
||||
for free: a trainer only exists for active stages, and inactive stages
|
||||
are simply never constructed.
|
||||
|
||||
Grad-norm clipping is per-stage here — v0.2's single shared optimizer
|
||||
clipped both stages' gradients jointly; splitting per stage is a small,
|
||||
disclosed behavior change. It doesn't affect Adam's per-parameter update
|
||||
math itself (no cross-parameter coupling), only the clip threshold's
|
||||
scope.
|
||||
"""
|
||||
|
||||
#: Metrics this trainer emits, declared once — `giant.training.metrics`
|
||||
#: derives every CSV/W&B column from these. Instance attributes rather
|
||||
#: than class constants because some are conditional on the stage's own
|
||||
#: configuration (see `WGANStageTrainer.__init__`).
|
||||
train_metrics: list[MetricSpec]
|
||||
val_metrics: list[MetricSpec]
|
||||
stage_metrics: list[MetricSpec]
|
||||
|
||||
#: False for adversarial stages, which have no monotone per-batch
|
||||
#: validation loss worth averaging (see `val_objective`).
|
||||
supports_val_loss: bool = True
|
||||
|
||||
#: Built by the subclass (the optimizer flavour differs) and wired to the
|
||||
#: schedule via `_init_lr_schedule`.
|
||||
optimizer: optim.Optimizer
|
||||
lr_sched: optim.lr_scheduler.LambdaLR
|
||||
total_steps: int
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
spec: StageSpec,
|
||||
model: torch.nn.Module,
|
||||
device: torch.device,
|
||||
extra_modules: tuple[torch.nn.Module, ...] = (),
|
||||
) -> None:
|
||||
self.spec = spec
|
||||
self.name = spec.name
|
||||
self.is_stage2 = spec.is_stage2
|
||||
self.generator = spec.generator
|
||||
self.decoder = spec.decoder
|
||||
self.device = device
|
||||
self.model = model.to(device)
|
||||
self.router = _stage_router(self.model)
|
||||
self._modules = (self.model, *extra_modules)
|
||||
|
||||
self.particle_type_cfg = spec.particle_type.to_dict()
|
||||
self.particle_type_n_classes = spec.particle_type_n_classes
|
||||
self.ema_decay = spec.ema_decay
|
||||
|
||||
self.ema_model: torch.nn.Module | None = None
|
||||
if spec.ema_decay > 0:
|
||||
self.ema_model = copy.deepcopy(self.model).eval()
|
||||
for p in self.ema_model.parameters():
|
||||
p.requires_grad_(False)
|
||||
|
||||
# --- schedule -------------------------------------------------------
|
||||
|
||||
def _init_lr_schedule(self, optimizer: optim.Optimizer, warmup_steps: int, total_steps: int) -> None:
|
||||
self._lr_lambda = _cosine_warmup_lambda(warmup_steps, total_steps)
|
||||
self.total_steps = total_steps
|
||||
self.lr_sched = optim.lr_scheduler.LambdaLR(optimizer, self._lr_lambda)
|
||||
|
||||
# --- per-batch (subclass responsibility) ----------------------------
|
||||
|
||||
def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
|
||||
raise NotImplementedError
|
||||
|
||||
def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
|
||||
raise NotImplementedError
|
||||
|
||||
# --- reporting hooks ------------------------------------------------
|
||||
|
||||
def batch_loss(self, stats: dict) -> float:
|
||||
"""The single number this stage contributes to the progress bar's
|
||||
smoothed loss."""
|
||||
raise NotImplementedError
|
||||
|
||||
def summary(self, means: dict) -> str:
|
||||
"""This stage's fragment of the end-of-epoch console line."""
|
||||
raise NotImplementedError
|
||||
|
||||
def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
|
||||
"""This stage's contribution to the best-checkpoint selection score."""
|
||||
raise NotImplementedError
|
||||
|
||||
# --- mode / state ---------------------------------------------------
|
||||
|
||||
def sampling_model(self) -> torch.nn.Module:
|
||||
return self.ema_model if self.ema_model is not None else self.model
|
||||
|
||||
def train_mode(self) -> None:
|
||||
for module in self._modules:
|
||||
module.train()
|
||||
|
||||
def eval_mode(self) -> None:
|
||||
for module in self._modules:
|
||||
module.eval()
|
||||
|
||||
# --- stage-2 secondary assembly (shared by both trainer subclasses) ---
|
||||
|
||||
def _ar_inputs(
|
||||
self,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_ctx: torch.Tensor,
|
||||
sec_cont: torch.Tensor,
|
||||
sec_type_idx: torch.Tensor,
|
||||
n_sec: torch.Tensor,
|
||||
epoch: int | None,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
"""Per-token AR conditioning for this stage's secondary decoder.
|
||||
|
||||
`epoch=None` means full teacher forcing (`p_tf=1.0`) regardless of
|
||||
`spec.teacher_forcing` — the val-loss convention, kept in this one
|
||||
place so both trainer subclasses honor it identically.
|
||||
"""
|
||||
p_tf = (
|
||||
1.0
|
||||
if epoch is None
|
||||
else _stage2_tf_prob(
|
||||
self.spec.teacher_forcing,
|
||||
self.spec.tf_p_start,
|
||||
self.spec.tf_p_end,
|
||||
epoch,
|
||||
self.spec.epochs,
|
||||
)
|
||||
)
|
||||
return _assemble_stage2_ar_inputs_scheduled(
|
||||
self.model,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_ctx,
|
||||
sec_cont,
|
||||
sec_type_idx,
|
||||
n_sec,
|
||||
self.particle_type_cfg,
|
||||
self.model.cond_enc,
|
||||
self.particle_type_n_classes,
|
||||
p_tf,
|
||||
self.spec.ar_sample_steps,
|
||||
)
|
||||
|
||||
def _sec_target(
|
||||
self,
|
||||
sec_cont: torch.Tensor,
|
||||
sec_type_idx: torch.Tensor,
|
||||
generator: str,
|
||||
*,
|
||||
flatten: bool,
|
||||
) -> torch.Tensor:
|
||||
"""Ground-truth stage-2 target for this stage's secondary decoder,
|
||||
per the (particle-type target, generator) width rules in
|
||||
`_assemble_stage2_ar_target`. `flatten=True` gives `Stage2OneShot`'s
|
||||
flattened `(B, K*token_dim)` form (the old `_real`); `flatten=False`
|
||||
gives `Stage2Autoregressive`'s per-token `(B, K, token_dim)` form (the
|
||||
old `_ar_target`) — the two are the same tensor modulo `.flatten(1)`,
|
||||
so the width rules live in one place (`stage2_inputs.py`)."""
|
||||
target = _assemble_stage2_ar_target(
|
||||
sec_cont,
|
||||
sec_type_idx,
|
||||
self.particle_type_cfg,
|
||||
generator,
|
||||
self.model.cond_enc,
|
||||
self.particle_type_n_classes,
|
||||
)
|
||||
return target.flatten(1) if flatten else target
|
||||
|
||||
@staticmethod
|
||||
def _sec_mask(n_sec: torch.Tensor, k_max: int, device: torch.device) -> torch.Tensor:
|
||||
"""`(B, K_MAX)` bool prefix mask: slot k is valid iff `k < n_sec`."""
|
||||
return torch.arange(k_max, device=device).unsqueeze(0) < n_sec.unsqueeze(1)
|
||||
|
||||
def _n_sec_loss(
|
||||
self,
|
||||
cond_cont: torch.Tensor,
|
||||
cond_cat: torch.Tensor,
|
||||
stage1_ctx: torch.Tensor,
|
||||
n_sec: torch.Tensor,
|
||||
device: torch.device,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""`(l_nsec, nsec_acc)` for this stage's multiplicity classifier —
|
||||
zeros when the stage owns no `n_sec_head` (stage 1 now that n_sec
|
||||
defaults to stage 2, or any stage without the head). Owns the only
|
||||
stage1-vs-stage2 `predict_n_sec` signature split, shared by the
|
||||
non-adversarial and WGAN trainers.
|
||||
|
||||
Gated on `n_sec_head is None`, not on `n_sec.mode`: a future
|
||||
`mode="stop_token"` model (currently rejected in
|
||||
`validate_config`) carries no head and would train its EOS signal in
|
||||
the generator/AR loss path instead, so this correctly stays zero.
|
||||
"""
|
||||
if self.model.n_sec_head is None:
|
||||
zero = torch.zeros((), device=device)
|
||||
return zero, zero
|
||||
logits = (
|
||||
self.model.predict_n_sec(cond_cont, cond_cat, stage1_ctx)
|
||||
if self.is_stage2
|
||||
else self.model.predict_n_sec(cond_cont, cond_cat)
|
||||
)
|
||||
l_nsec = F.cross_entropy(logits, n_sec)
|
||||
nsec_acc = (logits.argmax(dim=-1) == n_sec).float().mean()
|
||||
return l_nsec, nsec_acc
|
||||
|
||||
@staticmethod
|
||||
def _step_optimizer(optimizer: optim.Optimizer, loss: torch.Tensor, params: list) -> float:
|
||||
"""`zero_grad -> backward -> clip_grad_norm_(1.0) -> step`, returning
|
||||
the pre-clip grad norm. The one place the grad-clip constant lives."""
|
||||
optimizer.zero_grad()
|
||||
loss.backward()
|
||||
grad_norm = torch.nn.utils.clip_grad_norm_(params, 1.0)
|
||||
optimizer.step()
|
||||
return grad_norm.item()
|
||||
|
||||
def _extra_state(self) -> dict:
|
||||
"""Subclass state beyond model/optimizer/lr_sched/EMA."""
|
||||
return {}
|
||||
|
||||
def _load_extra_state(self, sd: dict) -> None:
|
||||
return None
|
||||
|
||||
def state_dict(self) -> dict:
|
||||
sd = {
|
||||
"model": self.model.state_dict(),
|
||||
"optimizer": self.optimizer.state_dict(),
|
||||
"lr_sched": self.lr_sched.state_dict(),
|
||||
}
|
||||
if self.ema_model is not None:
|
||||
sd["model_ema"] = self.ema_model.state_dict()
|
||||
sd.update(self._extra_state())
|
||||
return sd
|
||||
|
||||
def load_state_dict(self, sd: dict) -> None:
|
||||
self.model.load_state_dict(sd["model"])
|
||||
self.optimizer.load_state_dict(sd["optimizer"])
|
||||
self.lr_sched.load_state_dict(sd["lr_sched"])
|
||||
if self.ema_model is not None:
|
||||
self.ema_model.load_state_dict(sd.get("model_ema", sd["model"]))
|
||||
self._load_extra_state(sd)
|
||||
|
||||
def _resume_extra_lr(self, lr: float) -> None:
|
||||
return None
|
||||
|
||||
def resume_lr(self, lr: float) -> None:
|
||||
"""Restore the configured `lr`'s authority after `load_state_dict`
|
||||
restored the checkpoint's own base LR."""
|
||||
self.lr_sched.base_lrs = [lr for _ in self.lr_sched.base_lrs]
|
||||
resumed_lr = lr * self._lr_lambda(self.lr_sched.last_epoch)
|
||||
for group in self.optimizer.param_groups:
|
||||
group["lr"] = resumed_lr
|
||||
self._resume_extra_lr(lr)
|
||||
|
||||
|
||||
class FlowDDPMStageTrainer(StageTrainer):
|
||||
"""flow or ddpm generator for a single stage."""
|
||||
|
||||
def __init__(self, spec: StageSpec, model: torch.nn.Module, device: torch.device) -> None:
|
||||
if spec.is_stage2 and spec.generator not in ("flow",):
|
||||
raise NotImplementedError(
|
||||
f"stage2_model.generator={spec.generator!r} is accepted by the "
|
||||
"schema but not implemented in v0.3.0 for stage 2 (only "
|
||||
"'flow' and 'wgan' have a stage-2 secondary-decoder loss)"
|
||||
)
|
||||
super().__init__(spec, model, device)
|
||||
self.particle_type_lambda = self.particle_type_cfg.get("lambda", 1.0)
|
||||
# Width of the type slice actually folded into x1_s2 by _sec_target,
|
||||
# under this trainer's generator (flow/ddpm only — see the
|
||||
# NotImplementedError above): "physical" keeps it folded in
|
||||
# (PARTICLE_PHYS_DIM wide, unchanged from v0.2); "onehot"/"embedding"
|
||||
# pull it out into model.type_head instead (0 here).
|
||||
self._flow_type_dim = None if self.particle_type_cfg.get("target", "physical") == "physical" else 0
|
||||
|
||||
self.params = list(self.model.parameters())
|
||||
self.optimizer = optim.AdamW(self.params, lr=spec.lr, weight_decay=spec.weight_decay)
|
||||
self._init_lr_schedule(
|
||||
self.optimizer,
|
||||
warmup_steps=spec.warmup_epochs * spec.steps_per_epoch,
|
||||
total_steps=max(spec.epochs * spec.steps_per_epoch, 1),
|
||||
)
|
||||
self.ddpm_schedule = CosineSchedule(T=spec.ddpm_n_steps).to(device) if spec.generator == "ddpm" else None
|
||||
|
||||
self.train_metrics = [
|
||||
train_metric(key)
|
||||
for key in (
|
||||
"loss",
|
||||
"loss_gen",
|
||||
"loss_nsec",
|
||||
"loss_balance",
|
||||
"loss_proc",
|
||||
"loss_entropy",
|
||||
"nsec_acc",
|
||||
"loss_type",
|
||||
"type_acc",
|
||||
"grad_norm",
|
||||
)
|
||||
]
|
||||
self.val_metrics = [
|
||||
val_metric(key)
|
||||
for key in (
|
||||
"loss",
|
||||
"loss_gen",
|
||||
"loss_nsec",
|
||||
"nsec_acc",
|
||||
"loss_type",
|
||||
"type_acc",
|
||||
)
|
||||
]
|
||||
self.stage_metrics = [stage_metric("lr")]
|
||||
|
||||
def _generator_loss(self, cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=None):
|
||||
if not self.is_stage2:
|
||||
if self.generator == "flow":
|
||||
return flow_matching_loss(self.model, x1_s1, cond_cont, cond_cat)
|
||||
assert self.ddpm_schedule is not None
|
||||
return self.ddpm_schedule.loss(self.model, x1_s1, cond_cont, cond_cat)
|
||||
if self.decoder == "autoregressive":
|
||||
assert ar_inputs is not None
|
||||
return flow_matching_loss_secondary_ar(
|
||||
self.model,
|
||||
x1_s2,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_ctx,
|
||||
ar_inputs["history_feat"],
|
||||
ar_inputs["has_prev"],
|
||||
ar_inputs["remaining_frac"],
|
||||
ar_inputs["slot_idx"],
|
||||
sec_mask,
|
||||
type_dim=self._flow_type_dim,
|
||||
)
|
||||
return flow_matching_loss_secondary(
|
||||
self.model,
|
||||
x1_s2,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_ctx,
|
||||
sec_mask,
|
||||
type_dim=self._flow_type_dim,
|
||||
)
|
||||
|
||||
def _type_loss(
|
||||
self,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_ctx,
|
||||
sec_type_idx,
|
||||
sec_mask,
|
||||
device,
|
||||
ar_inputs=None,
|
||||
):
|
||||
"""CE (`target="onehot"`) or MSE (`target="embedding"`) loss for the
|
||||
stage-2 model's `type_head` — the non-adversarial counterpart to
|
||||
WGANStageTrainer's ST-Gumbel-into-the-critic path.
|
||||
Zero when this stage has no `type_head` (stage 1, or
|
||||
`particle_type.target = "physical"`)."""
|
||||
l_type = torch.zeros((), device=device)
|
||||
type_acc = torch.zeros((), device=device)
|
||||
type_head = getattr(self.model, "type_head", None)
|
||||
if not self.is_stage2 or type_head is None:
|
||||
return l_type, type_acc
|
||||
if self.decoder == "autoregressive":
|
||||
assert ar_inputs is not None
|
||||
type_out = self.model.predict_type(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_ctx,
|
||||
ar_inputs["history_feat"],
|
||||
ar_inputs["has_prev"],
|
||||
ar_inputs["remaining_frac"],
|
||||
ar_inputs["slot_idx"],
|
||||
)
|
||||
else:
|
||||
type_out = self.model.predict_type(cond_cont, cond_cat, stage1_ctx)
|
||||
mask = sec_mask.float()
|
||||
denom = mask.sum().clamp(min=1)
|
||||
if self.particle_type_cfg.get("target") == "onehot":
|
||||
ce = F.cross_entropy(type_out.transpose(1, 2), sec_type_idx, reduction="none")
|
||||
l_type = (ce * mask).sum() / denom
|
||||
type_acc = ((type_out.argmax(-1) == sec_type_idx).float() * mask).sum() / denom
|
||||
else: # "embedding"
|
||||
target_vec = self.model.cond_enc.pdg_emb(sec_type_idx).detach()
|
||||
se = ((type_out - target_vec) ** 2).mean(-1)
|
||||
l_type = (se * mask).sum() / denom
|
||||
return l_type, type_acc
|
||||
|
||||
def _compute(self, batch: StepBatch, device: torch.device, epoch: int | None = None) -> dict:
|
||||
"""`epoch=None` (the `val_loss` path) always uses full teacher
|
||||
forcing (`p_tf=1.0`) regardless of `spec.teacher_forcing` — validation
|
||||
should stay a stable, non-stochastic ground-truth comparison; only
|
||||
the training `step` path schedules `p_tf` by epoch."""
|
||||
(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
x1_s1,
|
||||
n_sec,
|
||||
sec_cont,
|
||||
proc_idx,
|
||||
sec_type_idx,
|
||||
) = _batch_to_device(batch, device)
|
||||
sec_mask = self._sec_mask(n_sec, sec_cont.size(1), device)
|
||||
stage1_ctx = x1_s1.detach()
|
||||
|
||||
x1_s2 = None
|
||||
ar_inputs = None
|
||||
if self.is_stage2 and self.decoder == "autoregressive":
|
||||
ar_inputs = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
|
||||
x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=False)
|
||||
elif self.is_stage2:
|
||||
x1_s2 = self._sec_target(sec_cont, sec_type_idx, self.generator, flatten=True)
|
||||
|
||||
l_gen = self._generator_loss(cond_cont, cond_cat, x1_s1, x1_s2, sec_mask, stage1_ctx, ar_inputs=ar_inputs)
|
||||
l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
|
||||
|
||||
l_type, type_acc = self._type_loss(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_ctx,
|
||||
sec_type_idx,
|
||||
sec_mask,
|
||||
device,
|
||||
ar_inputs=ar_inputs,
|
||||
)
|
||||
|
||||
l_balance = l_proc = l_entropy = torch.zeros((), device=device)
|
||||
if self.router is not None:
|
||||
if self.spec.lambda_balance > 0:
|
||||
l_balance = self.router.balance_loss(cond_cont, cond_cat)
|
||||
if self.spec.lambda_proc > 0:
|
||||
l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx)
|
||||
if self.spec.lambda_entropy > 0:
|
||||
l_entropy = self.router.entropy_loss(cond_cont, cond_cat)
|
||||
|
||||
total = self.spec.lambda_weight * l_gen + self.spec.n_sec_lambda * l_nsec + self.particle_type_lambda * l_type
|
||||
if self.spec.lambda_balance > 0:
|
||||
total = total + self.spec.lambda_balance * l_balance
|
||||
if self.spec.lambda_proc > 0:
|
||||
total = total + self.spec.lambda_proc * l_proc
|
||||
if self.spec.lambda_entropy > 0:
|
||||
total = total + self.spec.lambda_entropy * l_entropy
|
||||
|
||||
return {
|
||||
"loss": total,
|
||||
"loss_gen": l_gen,
|
||||
"loss_nsec": l_nsec,
|
||||
"loss_type": l_type,
|
||||
"type_acc": type_acc,
|
||||
"loss_balance": l_balance,
|
||||
"loss_proc": l_proc,
|
||||
"loss_entropy": l_entropy,
|
||||
"nsec_acc": nsec_acc,
|
||||
}
|
||||
|
||||
def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
|
||||
if self.router is not None:
|
||||
self.router.gumbel_tau = _gumbel_tau(
|
||||
global_step,
|
||||
self.total_steps,
|
||||
self.spec.gumbel_tau_start,
|
||||
self.spec.gumbel_tau_end,
|
||||
)
|
||||
epoch = global_step // self.spec.steps_per_epoch
|
||||
out = self._compute(batch, device, epoch=epoch)
|
||||
grad_norm = self._step_optimizer(self.optimizer, out["loss"], self.params)
|
||||
self.lr_sched.step()
|
||||
if self.ema_model is not None:
|
||||
_update_ema(self.ema_model, self.model, self.ema_decay)
|
||||
stats = {key: value.item() for key, value in out.items()}
|
||||
stats["grad_norm"] = grad_norm
|
||||
stats["lr"] = self.optimizer.param_groups[0]["lr"]
|
||||
return stats
|
||||
|
||||
@torch.no_grad()
|
||||
def val_loss(self, batch: StepBatch, device: torch.device) -> dict:
|
||||
return {key: value.item() for key, value in self._compute(batch, device).items()}
|
||||
|
||||
# --- reporting ------------------------------------------------------
|
||||
|
||||
def batch_loss(self, stats: dict) -> float:
|
||||
return stats["loss"]
|
||||
|
||||
def summary(self, means: dict) -> str:
|
||||
return f"{self.name}[loss={means.get('loss', 0.0):.3f}]"
|
||||
|
||||
def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
|
||||
return val_means.get("loss", 0.0)
|
||||
|
||||
|
||||
class _Stage2RealFakeBatch(NamedTuple):
|
||||
"""Subset of `StepBatch` that `_stage2_real_and_fake` needs."""
|
||||
|
||||
cond_cont: torch.Tensor
|
||||
cond_cat: torch.Tensor
|
||||
n_sec: torch.Tensor
|
||||
sec_cont: torch.Tensor
|
||||
sec_type_idx: torch.Tensor
|
||||
|
||||
|
||||
class WGANStageTrainer(StageTrainer):
|
||||
"""WGAN-GP generator+critic for a single stage (see giant/model/wgan.py).
|
||||
|
||||
Ports `_wgan_train_step` to operate on one stage instead of two fused
|
||||
together — the critic updates every batch; every `n_critic`-th batch
|
||||
additionally updates the generator (`did_g_step`). The (non-adversarial)
|
||||
n_sec classifier, when this stage's model owns it, updates every batch
|
||||
regardless — folded into whichever generator optimizer step happens this
|
||||
batch, same precedent as v0.2.
|
||||
"""
|
||||
|
||||
supports_val_loss = False
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
spec: StageSpec,
|
||||
model: torch.nn.Module,
|
||||
critic: torch.nn.Module,
|
||||
device: torch.device,
|
||||
) -> None:
|
||||
self.critic = critic.to(device)
|
||||
super().__init__(spec, model, device, extra_modules=(self.critic,))
|
||||
self.n_critic = max(spec.n_critic, 1)
|
||||
self.gp_weight = spec.gp_weight
|
||||
self.critic_lr = spec.critic_lr
|
||||
|
||||
self.g_params = list(self.model.parameters())
|
||||
self.d_params = list(self.critic.parameters())
|
||||
# WGAN-GP recipe (Gulrajani et al. 2017): Adam, beta1=0, no weight decay.
|
||||
self.optimizer = optim.Adam(self.g_params, lr=spec.lr, betas=(0.0, 0.9))
|
||||
self.optimizer_d = optim.Adam(
|
||||
self.d_params,
|
||||
lr=spec.critic_lr if spec.critic_lr > 0 else spec.lr,
|
||||
betas=(0.0, 0.9),
|
||||
)
|
||||
|
||||
# Generator steps fire every n_critic-th batch, so warmup/decay must
|
||||
# be counted in those units, matching v0.2.
|
||||
gen_steps_per_epoch = max(spec.steps_per_epoch // self.n_critic, 1)
|
||||
self._init_lr_schedule(
|
||||
self.optimizer,
|
||||
warmup_steps=spec.warmup_epochs * gen_steps_per_epoch,
|
||||
total_steps=max(spec.epochs * gen_steps_per_epoch, 1),
|
||||
)
|
||||
|
||||
train_keys = [
|
||||
"d_loss",
|
||||
"g_loss",
|
||||
"wasserstein",
|
||||
"gp_loss",
|
||||
"loss_nsec",
|
||||
"nsec_acc",
|
||||
"grad_norm_d",
|
||||
"grad_norm_g",
|
||||
]
|
||||
if self.is_stage2 and self.particle_type_cfg.get("target") == "onehot":
|
||||
# Differentiability instrumentation — only meaningful when the
|
||||
# type slice is a straight-through Gumbel relaxation.
|
||||
train_keys += ["grad_norm_type_slice", "grad_norm_cont_slice"]
|
||||
self.train_metrics = [train_metric(key) for key in train_keys]
|
||||
self.val_metrics = []
|
||||
self.stage_metrics = [stage_metric("lr"), stage_metric("critic_lr")]
|
||||
|
||||
def _stage2_real_and_fake(self, batch_tensors: _Stage2RealFakeBatch, stage1_ctx, global_step, device):
|
||||
"""Build `(real, fake_raw, mask, critic_fn)` for stage 2, covering
|
||||
both decoders and all three particle-type targets. `fake_raw` still
|
||||
needs the caller's straight-through relaxation under
|
||||
`particle_type.target = "onehot"`, and neither tensor is masked-and-
|
||||
multiplied on the fake side yet."""
|
||||
cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
|
||||
B = cond_cont.size(0)
|
||||
type_dim = stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes)
|
||||
slot_width = CONT_SLOT_DIM + type_dim
|
||||
k_max = sec_cont.size(1)
|
||||
|
||||
sec_mask = self._sec_mask(n_sec, k_max, device)
|
||||
mask = sec_mask.unsqueeze(-1).expand(-1, -1, slot_width).reshape(B, -1).float()
|
||||
|
||||
def critic_fn(x):
|
||||
return self.critic(x, cond_cont, cond_cat, stage1_ctx)
|
||||
|
||||
if self.decoder == "autoregressive":
|
||||
epoch = global_step // self.spec.steps_per_epoch
|
||||
ar = self._ar_inputs(cond_cont, cond_cat, stage1_ctx, sec_cont, sec_type_idx, n_sec, epoch)
|
||||
real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=False).reshape(B, -1) * mask
|
||||
z = torch.randn(B, k_max, self.model.noise_dim, device=device)
|
||||
fake_raw = self.model(
|
||||
z,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_ctx,
|
||||
ar["history_feat"],
|
||||
ar["has_prev"],
|
||||
ar["remaining_frac"],
|
||||
ar["slot_idx"],
|
||||
).reshape(B, -1)
|
||||
else:
|
||||
real = self._sec_target(sec_cont, sec_type_idx, "wgan", flatten=True) * mask
|
||||
z = torch.randn(B, self.model.noise_dim, device=device)
|
||||
fake_raw = self.model(z, cond_cont, cond_cat, stage1_ctx)
|
||||
|
||||
return real, fake_raw, mask, critic_fn
|
||||
|
||||
def step(self, batch: StepBatch, device: torch.device, global_step: int) -> dict:
|
||||
(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
x1_s1,
|
||||
n_sec,
|
||||
sec_cont,
|
||||
_proc_idx,
|
||||
sec_type_idx,
|
||||
) = _batch_to_device(batch, device)
|
||||
B = cond_cont.size(0)
|
||||
stage1_ctx = x1_s1.detach()
|
||||
grad_probe: dict[str, float] = {}
|
||||
|
||||
if not self.is_stage2:
|
||||
real = x1_s1
|
||||
|
||||
def critic_fn(x):
|
||||
return self.critic(x, cond_cont, cond_cat)
|
||||
|
||||
z = torch.randn(B, self.model.noise_dim, device=device)
|
||||
fake = self.model(z, cond_cont, cond_cat)
|
||||
mask = None
|
||||
else:
|
||||
real, fake_raw, mask, critic_fn = self._stage2_real_and_fake(
|
||||
_Stage2RealFakeBatch(cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx),
|
||||
stage1_ctx,
|
||||
global_step,
|
||||
device,
|
||||
)
|
||||
if self.particle_type_cfg.get("target", "physical") == "onehot":
|
||||
# Straight-through Gumbel-softmax relaxation of the type
|
||||
# slice only — the critic must see a hard one-hot forward
|
||||
# (matching what "real" data looks like) while gradient
|
||||
# still flows smoothly to the generator. grad_probe captures
|
||||
# the gradient-magnitude instrumentation — see
|
||||
# _relax_onehot_type_slice's docstring.
|
||||
tau = _gumbel_tau(
|
||||
global_step,
|
||||
self.total_steps,
|
||||
self.spec.type_gumbel_tau_start,
|
||||
self.spec.type_gumbel_tau_end,
|
||||
)
|
||||
fake_raw = _relax_onehot_type_slice(
|
||||
fake_raw,
|
||||
sec_cont.size(1),
|
||||
CONT_SLOT_DIM,
|
||||
stage2_type_dim(self.particle_type_cfg, self.particle_type_n_classes),
|
||||
tau,
|
||||
grad_probe=grad_probe,
|
||||
)
|
||||
fake = fake_raw * mask
|
||||
|
||||
# --- critic step (every batch) ---
|
||||
fake_detached = fake.detach()
|
||||
real_score = critic_fn(real)
|
||||
fake_score = critic_fn(fake_detached)
|
||||
gp = gradient_penalty(critic_fn, real, fake_detached, mask=mask)
|
||||
d_loss = fake_score.mean() - real_score.mean() + self.gp_weight * gp
|
||||
wasserstein = (real_score.mean() - fake_score.mean()).detach()
|
||||
|
||||
grad_norm_d = self._step_optimizer(self.optimizer_d, d_loss, self.d_params)
|
||||
|
||||
# --- generator (+ n_sec) step ---
|
||||
did_g_step = global_step % self.n_critic == 0
|
||||
l_nsec, nsec_acc = self._n_sec_loss(cond_cont, cond_cat, stage1_ctx, n_sec, device)
|
||||
|
||||
# On a non-generator-step batch with no n_sec_head on this stage
|
||||
# (n_sec now defaults to stage 2), there's nothing for
|
||||
# the generator optimizer to do this batch — g_loss would otherwise
|
||||
# be a graph-less zero tensor, which .backward() rejects outright.
|
||||
skip_g_step = not did_g_step and self.model.n_sec_head is None
|
||||
if did_g_step:
|
||||
g_loss_adv = generator_loss(critic_fn, fake)
|
||||
g_loss = self.spec.lambda_weight * g_loss_adv + self.spec.n_sec_lambda * l_nsec
|
||||
else:
|
||||
g_loss_adv = torch.zeros((), device=device)
|
||||
g_loss = self.spec.n_sec_lambda * l_nsec
|
||||
if skip_g_step:
|
||||
grad_norm_g = 0.0
|
||||
else:
|
||||
grad_norm_g = self._step_optimizer(self.optimizer, g_loss, self.g_params)
|
||||
|
||||
if did_g_step:
|
||||
self.lr_sched.step()
|
||||
if self.ema_model is not None:
|
||||
_update_ema(self.ema_model, self.model, self.ema_decay)
|
||||
|
||||
return {
|
||||
"d_loss": d_loss.item(),
|
||||
"g_loss": g_loss_adv.item(),
|
||||
"wasserstein": wasserstein.item(),
|
||||
"gp_loss": gp.item(),
|
||||
"loss_nsec": l_nsec.item(),
|
||||
"nsec_acc": nsec_acc.item(),
|
||||
"did_g_step": did_g_step,
|
||||
"grad_norm": grad_norm_d + grad_norm_g,
|
||||
"grad_norm_d": grad_norm_d,
|
||||
"grad_norm_g": grad_norm_g,
|
||||
"grad_norm_type_slice": grad_probe.get("type", 0.0),
|
||||
"grad_norm_cont_slice": grad_probe.get("cont", 0.0),
|
||||
"lr": self.optimizer.param_groups[0]["lr"],
|
||||
"critic_lr": self.optimizer_d.param_groups[0]["lr"],
|
||||
}
|
||||
|
||||
# --- reporting ------------------------------------------------------
|
||||
|
||||
def batch_loss(self, stats: dict) -> float:
|
||||
return stats["d_loss"] + stats["g_loss"]
|
||||
|
||||
def summary(self, means: dict) -> str:
|
||||
return f"{self.name}[d={means.get('d_loss', 0.0):.3f} g={means.get('g_loss', 0.0):.3f}]"
|
||||
|
||||
def val_objective(self, train_means: dict, val_means: dict, marginal_kl: float) -> float:
|
||||
"""No monotone per-batch WGAN loss fit for averaging, so
|
||||
best-checkpoint selection uses the real marginal-KL signal when
|
||||
`validate_marginals` produced one, and falls back to this epoch's own
|
||||
Wasserstein-distance magnitude otherwise.
|
||||
|
||||
Behavior change vs. every run up to v0.3.0: the pre-refactor code
|
||||
meant to do exactly this, but its guard
|
||||
(`{n: kl for n in wgan_names if n not in val_loss_per_stage}`) could
|
||||
never fire — `val_loss_per_stage` was pre-seeded with `0.0` for every
|
||||
stage, so a WGAN stage contributed a flat `0.0` and the marginal KL
|
||||
was recorded in the metrics row without ever influencing `best.pt`.
|
||||
Runs from before this commit therefore selected their best checkpoint
|
||||
on the non-adversarial stages alone."""
|
||||
if math.isfinite(marginal_kl):
|
||||
return marginal_kl
|
||||
return abs(train_means.get("wasserstein", 0.0))
|
||||
|
||||
# --- state ----------------------------------------------------------
|
||||
|
||||
def _extra_state(self) -> dict:
|
||||
return {
|
||||
"critic": self.critic.state_dict(),
|
||||
"optimizer_d": self.optimizer_d.state_dict(),
|
||||
}
|
||||
|
||||
def _load_extra_state(self, sd: dict) -> None:
|
||||
self.critic.load_state_dict(sd["critic"])
|
||||
self.optimizer_d.load_state_dict(sd["optimizer_d"])
|
||||
|
||||
def _resume_extra_lr(self, lr: float) -> None:
|
||||
resumed_critic_lr = self.critic_lr if self.critic_lr > 0 else lr
|
||||
for group in self.optimizer_d.param_groups:
|
||||
group["lr"] = resumed_critic_lr
|
||||
|
||||
|
||||
def build_stage_trainers(
|
||||
cfg: dict,
|
||||
models: dict[str, torch.nn.Module | None],
|
||||
critics: dict[str, torch.nn.Module | None],
|
||||
device: torch.device,
|
||||
total_train_batches: int,
|
||||
) -> dict[str, StageTrainer]:
|
||||
"""One trainer per active stage — `models[name] is None` means that stage
|
||||
is `active = false` and is simply never constructed."""
|
||||
trainers: dict[str, StageTrainer] = {}
|
||||
for name, is_stage2 in (("stage1", False), ("stage2", True)):
|
||||
model = models.get(name)
|
||||
if model is None:
|
||||
continue
|
||||
spec = StageSpec.from_config(cfg, name, is_stage2, max(total_train_batches, 1))
|
||||
if spec.generator == "wgan":
|
||||
critic = critics.get(name)
|
||||
assert critic is not None, (
|
||||
f"{name}_model.generator='wgan' requires a critic (see giant.model.network.build_critics)"
|
||||
)
|
||||
trainers[name] = WGANStageTrainer(spec, model, critic, device)
|
||||
else:
|
||||
trainers[name] = FlowDDPMStageTrainer(spec, model, device)
|
||||
return trainers
|
||||
+121
-137
@@ -2,13 +2,26 @@ import numpy as np
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
|
||||
from giant.constants import LOCAL_TARGET_NAMES
|
||||
from giant.sample import resolve_n_sec, sample_stage1, sample_stage2
|
||||
from giant.constants import K_MAX, LOCAL_TARGET_NAMES
|
||||
from giant.sample import (
|
||||
sample_flow,
|
||||
sample_ddpm,
|
||||
sample_ddim,
|
||||
sample_secondaries,
|
||||
sample_wgan,
|
||||
sample_secondaries_wgan,
|
||||
)
|
||||
|
||||
_SEC_PHYS_NAMES = ["log_mass", "charge"]
|
||||
|
||||
|
||||
def _histogram_kl(p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8) -> float:
|
||||
def _kw(steps: int | None) -> dict[str, int]:
|
||||
return {} if steps is None else {"steps": steps}
|
||||
|
||||
|
||||
def _histogram_kl(
|
||||
p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
|
||||
) -> float:
|
||||
"""KL(P || Q) between two 1D samples, estimated via a shared histogram."""
|
||||
lo = min(p_samples.min(), q_samples.min())
|
||||
hi = max(p_samples.max(), q_samples.max())
|
||||
@@ -31,40 +44,15 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray:
|
||||
return counts / total if total > 0 else counts
|
||||
|
||||
|
||||
def _categorical_kl(real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8) -> float:
|
||||
"""KL(P_real || Q_gen) between two class-index samples over `n_classes`
|
||||
categories, estimated from bincount fractions. NaN if either side has no
|
||||
valid samples (mirrors `_histogram_kl`'s empty-input handling)."""
|
||||
if len(real_idx) == 0 or len(gen_idx) == 0:
|
||||
return float("nan")
|
||||
p = _bincount_frac(real_idx, n_classes) + eps
|
||||
q = _bincount_frac(gen_idx, n_classes) + eps
|
||||
p /= p.sum()
|
||||
q /= q.sum()
|
||||
return float(np.sum(p * np.log(p / q)))
|
||||
|
||||
|
||||
def _embedding_nearest_class(vectors: torch.Tensor, emb_weight: torch.Tensor) -> np.ndarray:
|
||||
"""Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight`
|
||||
(vocab, emb_dim) — same computation as
|
||||
`giant.particles.decode_embedding_nearest`, but returning the raw class
|
||||
index instead of a decoded PDG code: validate.py only needs a
|
||||
real-vs-generated class-distribution comparison, not a rollout-usable
|
||||
identity, so there's no need for the pdg_map inversion here."""
|
||||
flat = vectors.reshape(-1, vectors.size(-1))
|
||||
dist = (flat.unsqueeze(1) - emb_weight.detach().unsqueeze(0)).abs().sum(-1)
|
||||
nearest = dist.argmin(dim=1)
|
||||
return nearest.reshape(vectors.shape[:-1]).cpu().numpy()
|
||||
|
||||
|
||||
def validate_marginals(
|
||||
stage1_model: torch.nn.Module,
|
||||
model: torch.nn.Module,
|
||||
val_loader: DataLoader,
|
||||
mode: str = "flow",
|
||||
schedule=None,
|
||||
device: torch.device | None = None,
|
||||
n_batches: int | None = None,
|
||||
kl_bins: int = 50,
|
||||
steps: int = 10,
|
||||
ddpm_steps: int = 1000,
|
||||
steps: int | None = None,
|
||||
sec_decoder: torch.nn.Module | None = None,
|
||||
) -> dict[str, np.ndarray | float]:
|
||||
"""Compare per-dimension marginals of generated vs. real steps.
|
||||
@@ -73,54 +61,51 @@ def validate_marginals(
|
||||
normalised space. `kl_divergence[j]` is KL(real || generated) for
|
||||
dimension j, estimated from a shared histogram over both samples.
|
||||
|
||||
Stage 1 is sampled via `giant.sample.sample_stage1`, which dispatches on
|
||||
`stage1_model.generator_kind` — `steps`/`ddpm_steps` are forwarded but
|
||||
only one of them is actually read, depending on that dispatch.
|
||||
`steps` overrides the number of sampler steps (flow ODE steps or DDIM
|
||||
substeps); `None` keeps each sampler's own default. Unused in "ddpm"
|
||||
mode, which always runs the full schedule.
|
||||
|
||||
When `sec_decoder` is given, also validates Stage 2 via
|
||||
`giant.sample.sample_stage2`/`resolve_n_sec` (generator- and
|
||||
one-shot-vs-autoregressive-agnostic): n_sec
|
||||
distribution (+ classification accuracy), per-slot energy-fraction
|
||||
marginals, and a particle-type marginal whose shape depends on
|
||||
`sec_decoder.particle_type_cfg["target"]` — restricted to each side's own
|
||||
valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since
|
||||
the two need not agree on how many slots are valid. Adds {"n_sec_real",
|
||||
"n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under
|
||||
`target = "physical"`, {"phys_real", "phys_generated", "phys_kl"}
|
||||
(continuous log_mass/charge marginals — v0.2 behaviour), or under
|
||||
`target` in `("onehot", "embedding")`, {"type_class_real",
|
||||
"type_class_gen", "type_class_kl"} (categorical class-index marginal:
|
||||
argmax for "onehot", L1-nearest conditioning-embedding row for
|
||||
"embedding" — see `_embedding_nearest_class`). Compared directly in
|
||||
normalised space (no denormalising — KL estimated from a shared
|
||||
per-sample histogram/bincount is invariant to a shared affine rescaling
|
||||
of both sides).
|
||||
When `sec_decoder` is given, also validates Stage 2: n_sec distribution
|
||||
(+ classification accuracy), predicted secondary physical-identity
|
||||
(log_mass, charge) marginals, and per-slot energy-fraction marginals —
|
||||
restricted to each side's own valid slots (real: `n_sec`; generated: the
|
||||
Stage-1 head's argmax), since the two need not agree on how many slots
|
||||
are valid. Compared directly in normalised space (no denormalising —
|
||||
KL estimated from a shared per-sample histogram is invariant to a shared
|
||||
affine rescaling of both sides). Adds {"n_sec_real", "n_sec_pred",
|
||||
"n_sec_accuracy", "phys_real", "phys_generated", "phys_kl",
|
||||
"energy_fraction_kl"} to the returned dict.
|
||||
"""
|
||||
if device is None:
|
||||
device = next(stage1_model.parameters()).device
|
||||
stage1_model.eval()
|
||||
device = next(model.parameters()).device
|
||||
model.eval()
|
||||
if sec_decoder is not None:
|
||||
sec_decoder.eval()
|
||||
|
||||
k_max = sec_decoder.k_max if sec_decoder is not None else 0
|
||||
target = sec_decoder.particle_type_cfg.get("target", "physical") if sec_decoder is not None else "physical"
|
||||
|
||||
all_real, all_gen = [], []
|
||||
all_n_sec_real, all_n_sec_pred = [], []
|
||||
all_phys_real, all_phys_gen = [], []
|
||||
all_type_class_real, all_type_class_gen = [], []
|
||||
all_frac_real: list[list[np.ndarray]] = [[] for _ in range(k_max)]
|
||||
all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(k_max)]
|
||||
all_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
|
||||
all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)]
|
||||
|
||||
for i, batch in enumerate(val_loader):
|
||||
if n_batches is not None and i >= n_batches:
|
||||
break
|
||||
# batch is a StepBatch (giant.data.dataset).
|
||||
x1, n_sec, sec_cont, sec_type_idx = batch.target_s1, batch.n_sec, batch.sec_cont, batch.sec_type_idx
|
||||
cond_cont = batch.cond_cont.to(device)
|
||||
cond_cat = batch.cond_cat.to(device)
|
||||
# Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx).
|
||||
cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx = batch
|
||||
cond_cont = cond_cont.to(device)
|
||||
cond_cat = cond_cat.to(device)
|
||||
|
||||
gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps)
|
||||
if mode == "flow":
|
||||
gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
|
||||
elif mode == "ddpm":
|
||||
gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule)
|
||||
elif mode == "wgan":
|
||||
gen, n_sec_pred = sample_wgan(model, cond_cont, cond_cat)
|
||||
else:
|
||||
gen, n_sec_pred = sample_ddim(
|
||||
model, cond_cont, cond_cat, schedule, **_kw(steps)
|
||||
)
|
||||
|
||||
all_real.append(x1.numpy())
|
||||
all_gen.append(gen.cpu().numpy())
|
||||
@@ -128,46 +113,54 @@ def validate_marginals(
|
||||
if sec_decoder is None:
|
||||
continue
|
||||
|
||||
n_sec_pred = resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred)
|
||||
n_sec_pred_np = n_sec_pred.cpu().numpy()
|
||||
n_sec_np = n_sec.numpy()
|
||||
all_n_sec_real.append(n_sec_np)
|
||||
all_n_sec_pred.append(n_sec_pred_np)
|
||||
|
||||
real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None] # (B, k_max)
|
||||
real_valid = np.arange(K_MAX)[None, :] < n_sec_np[:, None] # (B, K_MAX)
|
||||
real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64)))
|
||||
real_phys = sec_cont[:, :, 4:6].numpy() # (B, K_MAX, 2) [log_mass, charge]
|
||||
|
||||
sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2(
|
||||
sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps
|
||||
if mode == "wgan":
|
||||
sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries_wgan(
|
||||
sec_decoder, cond_cont, cond_cat, gen, n_sec_pred
|
||||
)
|
||||
else:
|
||||
sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries(
|
||||
sec_decoder,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
gen,
|
||||
n_sec_pred,
|
||||
steps=steps if steps is not None else 10,
|
||||
)
|
||||
gen_frac = 1.0 / (
|
||||
1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))
|
||||
)
|
||||
gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)))
|
||||
gen_phys = sec_phys_pred.cpu().numpy()
|
||||
gen_valid = sec_valid_pred.cpu().numpy()
|
||||
|
||||
if target == "physical":
|
||||
real_phys = sec_cont[:, :, 4:6].numpy() # (B, k_max, 2) [log_mass, charge]
|
||||
gen_phys = sec_type_pred.cpu().numpy()
|
||||
all_phys_real.append(real_phys[real_valid])
|
||||
all_phys_gen.append(gen_phys[gen_valid])
|
||||
else:
|
||||
sec_type_idx_np = sec_type_idx.numpy()
|
||||
all_type_class_real.append(sec_type_idx_np[real_valid])
|
||||
if target == "onehot":
|
||||
gen_class = sec_type_pred.argmax(dim=-1).cpu().numpy()
|
||||
else: # "embedding"
|
||||
emb_weight = sec_decoder.cond_enc.pdg_emb.weight
|
||||
gen_class = _embedding_nearest_class(sec_type_pred, emb_weight)
|
||||
all_type_class_gen.append(gen_class[gen_valid])
|
||||
|
||||
for j in range(k_max):
|
||||
all_phys_real.append(real_phys[real_valid])
|
||||
all_phys_gen.append(gen_phys[gen_valid])
|
||||
for j in range(K_MAX):
|
||||
all_frac_real[j].append(real_frac[real_valid[:, j], j])
|
||||
all_frac_gen[j].append(gen_frac[gen_valid[:, j], j])
|
||||
|
||||
real = np.concatenate(all_real, axis=0)
|
||||
generated = np.concatenate(all_gen, axis=0)
|
||||
|
||||
kl_divergence = np.array([_histogram_kl(real[:, j], generated[:, j], bins=kl_bins) for j in range(real.shape[1])])
|
||||
kl_divergence = np.array(
|
||||
[
|
||||
_histogram_kl(real[:, j], generated[:, j], bins=kl_bins)
|
||||
for j in range(real.shape[1])
|
||||
]
|
||||
)
|
||||
|
||||
header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
||||
header = (
|
||||
f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} "
|
||||
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
||||
)
|
||||
print(f"\n{header}")
|
||||
print("-" * len(header))
|
||||
for j, name in enumerate(LOCAL_TARGET_NAMES):
|
||||
@@ -188,9 +181,24 @@ def validate_marginals(
|
||||
n_sec_real = np.concatenate(all_n_sec_real, axis=0)
|
||||
n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0)
|
||||
n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean())
|
||||
phys_real = np.concatenate(all_phys_real, axis=0) # (M, 2)
|
||||
phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2)
|
||||
|
||||
energy_fraction_kl = np.full(k_max, np.nan)
|
||||
print(f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}")
|
||||
if len(phys_real) > 0 and len(phys_gen) > 0:
|
||||
phys_kl = np.array(
|
||||
[
|
||||
_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins)
|
||||
for j in range(2)
|
||||
]
|
||||
)
|
||||
else:
|
||||
phys_kl = np.full(2, np.nan)
|
||||
|
||||
energy_fraction_kl = np.full(K_MAX, np.nan)
|
||||
print(
|
||||
f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} "
|
||||
f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}"
|
||||
)
|
||||
n_sec_dist_header = f"{'n_sec value':<20} {'real_frac':>10} {'gen_frac':>10}"
|
||||
print(n_sec_dist_header)
|
||||
print("-" * len(n_sec_dist_header))
|
||||
@@ -200,70 +208,46 @@ def validate_marginals(
|
||||
for v in range(max_n_sec):
|
||||
print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}")
|
||||
|
||||
print(
|
||||
f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
|
||||
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
||||
)
|
||||
print("-" * 68)
|
||||
for j, name in enumerate(_SEC_PHYS_NAMES):
|
||||
r, g = phys_real[:, j], phys_gen[:, j]
|
||||
if len(r) == 0 or len(g) == 0:
|
||||
continue
|
||||
print(
|
||||
f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
|
||||
f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
|
||||
)
|
||||
|
||||
print(
|
||||
f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} "
|
||||
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
||||
)
|
||||
print("-" * 90)
|
||||
for j in range(k_max):
|
||||
for j in range(K_MAX):
|
||||
r = np.concatenate(all_frac_real[j]) if all_frac_real[j] else np.array([])
|
||||
g = np.concatenate(all_frac_gen[j]) if all_frac_gen[j] else np.array([])
|
||||
if len(r) == 0 or len(g) == 0:
|
||||
continue
|
||||
kl = _histogram_kl(r, g, bins=kl_bins)
|
||||
energy_fraction_kl[j] = kl
|
||||
print(f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}")
|
||||
print(
|
||||
f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} "
|
||||
f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}"
|
||||
)
|
||||
|
||||
result.update(
|
||||
{
|
||||
"n_sec_real": n_sec_real,
|
||||
"n_sec_pred": n_sec_pred_all,
|
||||
"n_sec_accuracy": n_sec_accuracy,
|
||||
"phys_real": phys_real,
|
||||
"phys_generated": phys_gen,
|
||||
"phys_kl": phys_kl,
|
||||
"energy_fraction_kl": energy_fraction_kl,
|
||||
}
|
||||
)
|
||||
|
||||
if target == "physical":
|
||||
phys_real = np.concatenate(all_phys_real, axis=0) # (M, 2)
|
||||
phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2)
|
||||
|
||||
if len(phys_real) > 0 and len(phys_gen) > 0:
|
||||
phys_kl = np.array([_histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) for j in range(2)])
|
||||
else:
|
||||
phys_kl = np.full(2, np.nan)
|
||||
|
||||
print(
|
||||
f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} "
|
||||
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
||||
)
|
||||
print("-" * 68)
|
||||
for j, name in enumerate(_SEC_PHYS_NAMES):
|
||||
r, g = phys_real[:, j], phys_gen[:, j]
|
||||
if len(r) == 0 or len(g) == 0:
|
||||
continue
|
||||
print(
|
||||
f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} {r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}"
|
||||
)
|
||||
|
||||
result.update({"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl})
|
||||
else:
|
||||
type_class_real = np.concatenate(all_type_class_real, axis=0)
|
||||
type_class_gen = np.concatenate(all_type_class_gen, axis=0)
|
||||
n_classes = sec_decoder.type_dim if target == "onehot" else sec_decoder.cond_enc.pdg_emb.weight.size(0)
|
||||
type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes)
|
||||
|
||||
print(
|
||||
f"\n{'sec type class (' + target + ')':<24} "
|
||||
f"n={len(type_class_real)}/{len(type_class_gen)} "
|
||||
f"KL(real||gen)={type_class_kl:.4f}"
|
||||
)
|
||||
|
||||
result.update(
|
||||
{
|
||||
"type_class_real": type_class_real,
|
||||
"type_class_gen": type_class_gen,
|
||||
"type_class_kl": type_class_kl,
|
||||
}
|
||||
)
|
||||
|
||||
return result
|
||||
|
||||
+3
-17
@@ -1,6 +1,6 @@
|
||||
[project]
|
||||
name = "giant"
|
||||
version = "0.3.0"
|
||||
version = "0.1.0"
|
||||
description = "Geant4 step-function surrogate via conditional flow matching"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.12"
|
||||
@@ -23,7 +23,6 @@ cuda = [
|
||||
]
|
||||
dev = [
|
||||
"pytest>=8,<10",
|
||||
"pytest-cov>=5,<8",
|
||||
"ruff>=0.15,<1",
|
||||
"ty>=0.0.50,<0.1",
|
||||
"giant[convert,analysis,geometry,wandb]",
|
||||
@@ -50,27 +49,14 @@ analysis = [
|
||||
|
||||
[project.scripts]
|
||||
giant = "giant.cli:app"
|
||||
dwarf = "giant.tools.dwarf:app"
|
||||
|
||||
[tool.ruff]
|
||||
line-length = 120
|
||||
|
||||
[tool.coverage.run]
|
||||
source = ["giant"]
|
||||
omit = ["*/legacy/*"]
|
||||
|
||||
[tool.coverage.report]
|
||||
exclude_also = [
|
||||
"if TYPE_CHECKING:",
|
||||
"raise NotImplementedError",
|
||||
]
|
||||
dwarf = "scripts.dwarf:app"
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[tool.hatch.build.targets.wheel]
|
||||
packages = ["giant"]
|
||||
packages = ["giant", "scripts"]
|
||||
|
||||
[tool.uv]
|
||||
conflicts = [
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
"""Cut a new raw generation or processed schema version for the geant_steps
|
||||
dataset tree (see giant/tools/migrate_geant_steps.py for the layout):
|
||||
dataset tree (see scripts/migrate_geant_steps.py for the layout):
|
||||
|
||||
raw/<kind>/<gen>/<detector>/shard-NNN.root
|
||||
processed/<kind>/<gen>/<schema>/<detector>/shard-NNN.parquet
|
||||
@@ -82,7 +82,9 @@ def _max_index(parent: Path, pattern: re.Pattern) -> int:
|
||||
|
||||
def _git_user_name() -> str | None:
|
||||
try:
|
||||
out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2)
|
||||
out = subprocess.run(
|
||||
["git", "config", "user.name"], capture_output=True, text=True, timeout=2
|
||||
)
|
||||
except (OSError, subprocess.SubprocessError):
|
||||
# OSError (e.g. git not on PATH) and subprocess.SubprocessError
|
||||
# (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
|
||||
@@ -140,7 +142,9 @@ def plan_bump_schema(
|
||||
raw_gen_dir = root / "raw" / kind / gen_tag
|
||||
processed_gen_dir = root / "processed" / kind / gen_tag
|
||||
if not raw_gen_dir.is_dir() and not processed_gen_dir.is_dir():
|
||||
raise SystemExit(f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first")
|
||||
raise SystemExit(
|
||||
f"error: {gen_tag} doesn't exist yet for kind={kind} — run bump-gen first"
|
||||
)
|
||||
if target is not None:
|
||||
if not SCHEMA_RE.match(target):
|
||||
raise SystemExit(f"error: --to must look like 'schemaN', got {target!r}")
|
||||
@@ -150,7 +154,9 @@ def plan_bump_schema(
|
||||
schema_tag = f"schema{next_schema}"
|
||||
new_dirs = [processed_gen_dir / schema_tag]
|
||||
by_suffix = f" ({by})" if by else ""
|
||||
log_line = f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
|
||||
log_line = (
|
||||
f"- `{gen_tag}`/`{schema_tag}` (kind={kind}) — {date} — {reason}{by_suffix}"
|
||||
)
|
||||
return new_dirs, log_line
|
||||
|
||||
|
||||
@@ -206,7 +212,9 @@ def _manifest_referenced_files(pools_root: Path) -> set[Path]:
|
||||
return referenced
|
||||
|
||||
|
||||
def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[int, int]:
|
||||
def _referenced_root_count(
|
||||
raw_gen_dir: Path, processed_gen_dir: Path
|
||||
) -> tuple[int, int]:
|
||||
"""(total .root files, count with a same-named .parquet under any schema) for one gen."""
|
||||
if not raw_gen_dir.is_dir():
|
||||
return 0, 0
|
||||
@@ -235,7 +243,9 @@ def _referenced_root_count(raw_gen_dir: Path, processed_gen_dir: Path) -> tuple[
|
||||
return total, referenced
|
||||
|
||||
|
||||
def _referenced_parquet_count(schema_dir: Path, manifest_referenced: set[Path]) -> tuple[int, int]:
|
||||
def _referenced_parquet_count(
|
||||
schema_dir: Path, manifest_referenced: set[Path]
|
||||
) -> tuple[int, int]:
|
||||
"""(total .parquet files, count listed in at least one manifest) for one schema dir."""
|
||||
if not schema_dir.is_dir():
|
||||
return 0, 0
|
||||
@@ -332,7 +342,11 @@ def print_status(root: Path) -> None:
|
||||
grand_files = 0
|
||||
for kind_dir in sorted(p for p in raw_root.iterdir() if p.is_dir()):
|
||||
kind = kind_dir.name
|
||||
gens = sorted(int(m.group(1)) for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir()) if m)
|
||||
gens = sorted(
|
||||
int(m.group(1))
|
||||
for m in (GEN_RE.match(p.name) for p in kind_dir.iterdir() if p.is_dir())
|
||||
if m
|
||||
)
|
||||
print(_colorize(f"{kind}/", "kind"))
|
||||
kind_total = 0
|
||||
kind_files = 0
|
||||
@@ -345,18 +359,25 @@ def print_status(root: Path) -> None:
|
||||
schemas = sorted(
|
||||
int(m.group(1))
|
||||
for m in (
|
||||
SCHEMA_RE.match(p.name) for p in (schema_dir.iterdir() if schema_dir.is_dir() else []) if p.is_dir()
|
||||
SCHEMA_RE.match(p.name)
|
||||
for p in (schema_dir.iterdir() if schema_dir.is_dir() else [])
|
||||
if p.is_dir()
|
||||
)
|
||||
if m
|
||||
)
|
||||
schema_sizes = {s: _du(schema_dir / f"schema{s}") for s in schemas}
|
||||
schema_counts = {
|
||||
s: _referenced_parquet_count(schema_dir / f"schema{s}", manifest_referenced) for s in schemas
|
||||
s: _referenced_parquet_count(
|
||||
schema_dir / f"schema{s}", manifest_referenced
|
||||
)
|
||||
for s in schemas
|
||||
}
|
||||
processed_size = sum(schema_sizes.values())
|
||||
processed_files = sum(c[0] for c in schema_counts.values())
|
||||
processed_referenced = sum(c[1] for c in schema_counts.values())
|
||||
raw_files, raw_referenced = _referenced_root_count(raw_gen_dir, processed_gen_dir)
|
||||
raw_files, raw_referenced = _referenced_root_count(
|
||||
raw_gen_dir, processed_gen_dir
|
||||
)
|
||||
gen_total = raw_size + processed_size
|
||||
gen_files = raw_files + processed_files
|
||||
kind_total += gen_total
|
||||
@@ -404,7 +425,9 @@ def print_status(root: Path) -> None:
|
||||
print(_reason_line(schema_reason, indent=4))
|
||||
else:
|
||||
print(_colorize(" (none)", "schema"))
|
||||
print(_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files))
|
||||
print(
|
||||
_row(f"{kind} total", kind_total, indent=1, level="gen", count=kind_files)
|
||||
)
|
||||
print()
|
||||
grand_total += kind_total
|
||||
grand_files += kind_files
|
||||
@@ -503,8 +526,13 @@ def plan_update_manifest(
|
||||
return result, missing
|
||||
|
||||
|
||||
def apply_update_manifest(manifest_path: Path, lines: list[tuple[str, str | None]]) -> None:
|
||||
out = [replacement if replacement is not None else original for original, replacement in lines]
|
||||
def apply_update_manifest(
|
||||
manifest_path: Path, lines: list[tuple[str, str | None]]
|
||||
) -> None:
|
||||
out = [
|
||||
replacement if replacement is not None else original
|
||||
for original, replacement in lines
|
||||
]
|
||||
manifest_path.write_text("\n".join(out) + "\n")
|
||||
|
||||
|
||||
@@ -524,7 +552,9 @@ def _resolve_manifest_files(manifest_path: Path) -> list[Path]:
|
||||
return files
|
||||
|
||||
|
||||
def plan_create_manifest(output_path: Path, parquet_files: list[Path]) -> tuple[list[str], list[Path], list[Path]]:
|
||||
def plan_create_manifest(
|
||||
output_path: Path, parquet_files: list[Path]
|
||||
) -> tuple[list[str], list[Path], list[Path]]:
|
||||
"""Return (relative_lines, missing_files, resolved_abs_paths)."""
|
||||
manifest_dir = output_path.resolve().parent
|
||||
lines: list[str] = []
|
||||
@@ -539,7 +569,9 @@ def plan_create_manifest(output_path: Path, parquet_files: list[Path]) -> tuple[
|
||||
return lines, missing, resolved
|
||||
|
||||
|
||||
def check_holdout_overlap(output_path: Path, resolved_new_files: list[Path]) -> list[tuple[str, Path]]:
|
||||
def check_holdout_overlap(
|
||||
output_path: Path, resolved_new_files: list[Path]
|
||||
) -> list[tuple[str, Path]]:
|
||||
"""Return (other_manifest_name, file) pairs where new files clash with existing manifests.
|
||||
|
||||
The check is triggered when output_path is (or will be) holdout.manifest, or when a
|
||||
@@ -579,7 +611,7 @@ def apply_create_manifest(output_path: Path, lines: list[str]) -> None:
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI entry points (called from giant/tools/dwarf.py)
|
||||
# CLI entry points (called from scripts/dwarf.py)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@@ -609,7 +641,9 @@ def _run_bump(
|
||||
if gen is None:
|
||||
new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
|
||||
else:
|
||||
new_dirs, log_line = plan_bump_schema(root_path, kind, gen, reason, by, date, to)
|
||||
new_dirs, log_line = plan_bump_schema(
|
||||
root_path, kind, gen, reason, by, date, to
|
||||
)
|
||||
|
||||
print(f"=== {'EXECUTING' if execute else 'DRY RUN'} ===")
|
||||
print("new directories:")
|
||||
@@ -32,7 +32,7 @@ from dataclasses import dataclass
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
|
||||
# Must match giant/tools/bump_dataset_version.py's GEN_RE.
|
||||
# Must match scripts/bump_dataset_version.py's GEN_RE.
|
||||
GEN_RE = re.compile(r"^gen\d+$")
|
||||
SHARD_RE = re.compile(r"^shard-(\d+)\.root$")
|
||||
|
||||
@@ -64,7 +64,9 @@ def parse_detector_spec(spec: str) -> tuple[str, str | None]:
|
||||
if ":" in spec:
|
||||
label, config = spec.split(":", 1)
|
||||
if not label or not config:
|
||||
raise PlanError(f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG")
|
||||
raise PlanError(
|
||||
f"invalid --detector spec {spec!r}: expected NAME or NAME:CONFIG"
|
||||
)
|
||||
return label, config
|
||||
return spec, None
|
||||
|
||||
@@ -92,7 +94,9 @@ def plan_jobs(
|
||||
raise PlanError(f"--gen must look like 'genN', got {gen!r}")
|
||||
gen_dir = dataset_root / "raw" / kind / gen
|
||||
if not gen_dir.is_dir():
|
||||
raise PlanError(f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first")
|
||||
raise PlanError(
|
||||
f"{gen_dir} doesn't exist — run bump_dataset_version.py bump-gen first"
|
||||
)
|
||||
|
||||
jobs = []
|
||||
for spec in detector_specs:
|
||||
@@ -120,7 +124,9 @@ def job_seed(kind: str, gen: str, job: SimJob, energy_gev: float | None) -> int:
|
||||
return zlib.crc32(key.encode()) & 0x7FFFFFFF
|
||||
|
||||
|
||||
def build_cmd(executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None) -> list[str]:
|
||||
def build_cmd(
|
||||
executable: Path, job: SimJob, events_per_file: int, energy_gev: float | None
|
||||
) -> list[str]:
|
||||
"""minicalosim executables take positional `[configName] nEvents [energy_GeV]`."""
|
||||
cmd = [str(executable)]
|
||||
if job.config:
|
||||
@@ -141,7 +147,10 @@ def run_job(
|
||||
gen: str,
|
||||
tmp_root: Path,
|
||||
) -> JobResult:
|
||||
workdir = tmp_root / f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
|
||||
workdir = (
|
||||
tmp_root
|
||||
/ f"{kind}-{gen}-{job.detector}-{job.shard_index:03d}-{uuid.uuid4().hex[:8]}"
|
||||
)
|
||||
workdir.mkdir(parents=True)
|
||||
|
||||
cmd = build_cmd(executable, job, events_per_file, energy_gev)
|
||||
@@ -165,12 +174,20 @@ def run_job(
|
||||
job,
|
||||
False,
|
||||
None,
|
||||
f"expected exactly one .root output in {workdir}, found {len(produced)}: {[p.name for p in produced]}",
|
||||
f"expected exactly one .root output in {workdir}, found {len(produced)}: "
|
||||
f"{[p.name for p in produced]}",
|
||||
result.stdout,
|
||||
result.stderr,
|
||||
)
|
||||
|
||||
dest = dataset_root / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
|
||||
dest = (
|
||||
dataset_root
|
||||
/ "raw"
|
||||
/ kind
|
||||
/ gen
|
||||
/ job.detector
|
||||
/ f"shard-{job.shard_index:03d}.root"
|
||||
)
|
||||
if dest.exists():
|
||||
return JobResult(
|
||||
job,
|
||||
@@ -262,7 +279,14 @@ def run_make_root(
|
||||
print(f"executable: {executable}")
|
||||
for job in planned_jobs:
|
||||
cmd = build_cmd(executable, job, events_per_file, energy_gev)
|
||||
dest = dataset_root_path / "raw" / kind / gen / job.detector / f"shard-{job.shard_index:03d}.root"
|
||||
dest = (
|
||||
dataset_root_path
|
||||
/ "raw"
|
||||
/ kind
|
||||
/ gen
|
||||
/ job.detector
|
||||
/ f"shard-{job.shard_index:03d}.root"
|
||||
)
|
||||
seed = job_seed(kind, gen, job, energy_gev)
|
||||
print(f" MINICALOSIM_SEED={seed} {' '.join(cmd)} -> {dest}")
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
"""dwarf — little helper to `giant`: dataset/tooling CLI for the geant_steps pipeline.
|
||||
|
||||
Unifies the standalone giant/tools/*.py conversion, migration, versioning, and
|
||||
Unifies the standalone scripts/*.py conversion, migration, versioning, and
|
||||
simulation-fanout tools into one Typer app so there's a single command name
|
||||
(and `--help`) to remember instead of five differently-hyphenated ones.
|
||||
"""
|
||||
@@ -14,20 +14,20 @@ import typer
|
||||
from typing_extensions import Annotated
|
||||
|
||||
from giant.config import Conditioning
|
||||
from giant.tools.bump_dataset_version import (
|
||||
from scripts.bump_dataset_version import (
|
||||
run_bump_gen,
|
||||
run_bump_schema,
|
||||
run_create_manifest,
|
||||
run_status,
|
||||
run_update_manifest,
|
||||
)
|
||||
from giant.tools.create_root_files import run_make_root
|
||||
from giant.tools.geometry_oracle import run_build_geometry_oracle
|
||||
from giant.tools.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
|
||||
from giant.tools.migrate_geant_steps import run_migration
|
||||
from giant.tools.steps_to_parquet import convert_steps_to_parquet
|
||||
from giant.tools.steps_to_parquet_parallel import run_parallel_job
|
||||
from giant.tools.warm_setup_cache import run_warm_setup_cache
|
||||
from scripts.create_root_files import run_make_root
|
||||
from scripts.geometry_oracle import run_build_geometry_oracle
|
||||
from scripts.hparam_scan import DATA_DEFAULT, SCAN_DIR_DEFAULT, run_hparam_scan
|
||||
from scripts.migrate_geant_steps import run_migration
|
||||
from scripts.steps_to_parquet import convert_steps_to_parquet
|
||||
from scripts.steps_to_parquet_parallel import run_parallel_job
|
||||
from scripts.warm_setup_cache import run_warm_setup_cache
|
||||
|
||||
app = typer.Typer(no_args_is_help=True)
|
||||
|
||||
@@ -80,7 +80,8 @@ def convert(
|
||||
typer.Option(
|
||||
"--output",
|
||||
"-o",
|
||||
help="Output Parquet file (default: <input>.parquet). Only valid with a single input file and --jobs 1.",
|
||||
help="Output Parquet file (default: <input>.parquet). Only valid "
|
||||
"with a single input file and --jobs 1.",
|
||||
),
|
||||
] = None,
|
||||
batch_size: Annotated[
|
||||
@@ -90,7 +91,9 @@ def convert(
|
||||
help="Uproot read batch size, e.g. '100 MB' or '500000' (rows)",
|
||||
),
|
||||
] = "100 MB",
|
||||
tree: Annotated[str, typer.Option("--tree", help="Tree name inside the ROOT file")] = "Steps",
|
||||
tree: Annotated[
|
||||
str, typer.Option("--tree", help="Tree name inside the ROOT file")
|
||||
] = "Steps",
|
||||
compression: Annotated[
|
||||
Compression, typer.Option("--compression", help="Parquet compression codec")
|
||||
] = Compression.snappy,
|
||||
@@ -126,11 +129,15 @@ def convert(
|
||||
raise typer.Exit(1)
|
||||
_warn_if_exceeds_shared_quota(jobs, "--jobs")
|
||||
|
||||
compression_value = "uncompressed" if compression is Compression.none else compression.value
|
||||
compression_value = (
|
||||
"uncompressed" if compression is Compression.none else compression.value
|
||||
)
|
||||
|
||||
if jobs == 1:
|
||||
if output is not None and len(root_files) > 1:
|
||||
typer.echo("error: --output can only be used with a single input file", err=True)
|
||||
typer.echo(
|
||||
"error: --output can only be used with a single input file", err=True
|
||||
)
|
||||
raise typer.Exit(1)
|
||||
total_orphaned = 0
|
||||
for root_file in root_files:
|
||||
@@ -143,7 +150,10 @@ def convert(
|
||||
)
|
||||
total_orphaned += n_orphaned
|
||||
if total_orphaned:
|
||||
typer.echo(f"\n{total_orphaned} orphaned child track(s) dropped across {len(root_files)} file(s).")
|
||||
typer.echo(
|
||||
f"\n{total_orphaned} orphaned child track(s) dropped across "
|
||||
f"{len(root_files)} file(s)."
|
||||
)
|
||||
return
|
||||
|
||||
if output is not None:
|
||||
@@ -166,7 +176,9 @@ def convert(
|
||||
|
||||
@app.command()
|
||||
def migrate(
|
||||
root: Annotated[Path, typer.Argument(help="Dataset root to migrate in place")] = _DATASET_ROOT_DEFAULT,
|
||||
root: Annotated[
|
||||
Path, typer.Argument(help="Dataset root to migrate in place")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
@@ -178,7 +190,8 @@ def migrate(
|
||||
bool,
|
||||
typer.Option(
|
||||
"--copy",
|
||||
help="Copy instead of move, leaving the originals in place (e.g. if another process is still reading them)",
|
||||
help="Copy instead of move, leaving the originals in place "
|
||||
"(e.g. if another process is still reading them)",
|
||||
),
|
||||
] = False,
|
||||
) -> None:
|
||||
@@ -189,8 +202,12 @@ def migrate(
|
||||
@app.command("bump-gen")
|
||||
def bump_gen(
|
||||
reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
|
||||
kind: Annotated[
|
||||
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
|
||||
] = "steps",
|
||||
by: Annotated[
|
||||
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
|
||||
] = None,
|
||||
date: Annotated[
|
||||
Optional[str],
|
||||
typer.Option("--date", help="Override date (default: today, ISO)"),
|
||||
@@ -203,8 +220,12 @@ def bump_gen(
|
||||
help="Target gen tag (default: one past the current highest)",
|
||||
),
|
||||
] = None,
|
||||
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Apply (default: dry run)")
|
||||
] = False,
|
||||
root: Annotated[
|
||||
Path, typer.Option("--root", help="Dataset root")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
) -> None:
|
||||
"""Cut a new raw generation."""
|
||||
run_bump_gen(
|
||||
@@ -222,8 +243,12 @@ def bump_gen(
|
||||
def bump_schema(
|
||||
gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
|
||||
reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
|
||||
kind: Annotated[
|
||||
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
|
||||
] = "steps",
|
||||
by: Annotated[
|
||||
Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")
|
||||
] = None,
|
||||
date: Annotated[
|
||||
Optional[str],
|
||||
typer.Option("--date", help="Override date (default: today, ISO)"),
|
||||
@@ -236,8 +261,12 @@ def bump_schema(
|
||||
help="Target schema tag (default: one past the current highest)",
|
||||
),
|
||||
] = None,
|
||||
execute: Annotated[bool, typer.Option("--execute", help="Apply (default: dry run)")] = False,
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Apply (default: dry run)")
|
||||
] = False,
|
||||
root: Annotated[
|
||||
Path, typer.Option("--root", help="Dataset root")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
) -> None:
|
||||
"""Cut a new schema within a gen."""
|
||||
run_bump_schema(
|
||||
@@ -254,7 +283,9 @@ def bump_schema(
|
||||
|
||||
@app.command()
|
||||
def status(
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
root: Annotated[
|
||||
Path, typer.Option("--root", help="Dataset root")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
) -> None:
|
||||
"""List existing gens/schemas per kind."""
|
||||
run_status(str(root))
|
||||
@@ -262,7 +293,9 @@ def status(
|
||||
|
||||
@app.command("update-manifest")
|
||||
def update_manifest(
|
||||
manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
|
||||
manifests: Annotated[
|
||||
list[Path], typer.Argument(help="One or more .manifest files to update")
|
||||
],
|
||||
schema: Annotated[
|
||||
Optional[str],
|
||||
typer.Option(
|
||||
@@ -273,7 +306,9 @@ def update_manifest(
|
||||
] = None,
|
||||
gen: Annotated[
|
||||
Optional[str],
|
||||
typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
|
||||
typer.Option(
|
||||
"--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"
|
||||
),
|
||||
] = None,
|
||||
execute: Annotated[
|
||||
bool,
|
||||
@@ -281,7 +316,9 @@ def update_manifest(
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Repoint manifest(s) to a new gen and/or schema, verifying all target files exist."""
|
||||
run_update_manifest([str(m) for m in manifests], schema=schema, execute=execute, gen=gen)
|
||||
run_update_manifest(
|
||||
[str(m) for m in manifests], schema=schema, execute=execute, gen=gen
|
||||
)
|
||||
|
||||
|
||||
@app.command("create-manifest")
|
||||
@@ -296,15 +333,22 @@ def create_manifest(
|
||||
typer.Option(
|
||||
"--pool",
|
||||
metavar="DETECTOR",
|
||||
help="Detector name; combined with --type and --root to form <root>/pools/<detector>/<type>.manifest",
|
||||
help="Detector name; combined with --type and --root to form "
|
||||
"<root>/pools/<detector>/<type>.manifest",
|
||||
),
|
||||
] = None,
|
||||
type_: Annotated[
|
||||
Optional[PoolType],
|
||||
typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
|
||||
typer.Option(
|
||||
"--type", help="Pool type — full, holdout, or dev (required with --pool)"
|
||||
),
|
||||
] = None,
|
||||
root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[bool, typer.Option("--execute", help="Write the manifest (default: dry run)")] = False,
|
||||
root: Annotated[
|
||||
Path, typer.Option("--root", help="Dataset root (used with --pool)")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
execute: Annotated[
|
||||
bool, typer.Option("--execute", help="Write the manifest (default: dry run)")
|
||||
] = False,
|
||||
force: Annotated[
|
||||
bool,
|
||||
typer.Option("--force", help="Overwrite the manifest if it already exists"),
|
||||
@@ -324,7 +368,9 @@ def create_manifest(
|
||||
|
||||
@app.command("make-root")
|
||||
def make_root(
|
||||
executable: Annotated[Path, typer.Option("--executable", help="Built minicalosim run_* executable")],
|
||||
executable: Annotated[
|
||||
Path, typer.Option("--executable", help="Built minicalosim run_* executable")
|
||||
],
|
||||
detector: Annotated[
|
||||
list[str],
|
||||
typer.Option(
|
||||
@@ -336,9 +382,15 @@ def make_root(
|
||||
"Repeatable.",
|
||||
),
|
||||
],
|
||||
num_files: Annotated[int, typer.Option("--num-files", help="New shards to create per detector")],
|
||||
events_per_file: Annotated[int, typer.Option("--events-per-file", help="nEvents passed to the executable")],
|
||||
gen: Annotated[str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")],
|
||||
num_files: Annotated[
|
||||
int, typer.Option("--num-files", help="New shards to create per detector")
|
||||
],
|
||||
events_per_file: Annotated[
|
||||
int, typer.Option("--events-per-file", help="nEvents passed to the executable")
|
||||
],
|
||||
gen: Annotated[
|
||||
str, typer.Option("--gen", help="Existing gen tag under raw/<kind>/, e.g. gen1")
|
||||
],
|
||||
energy_gev: Annotated[
|
||||
float | None,
|
||||
typer.Option(
|
||||
@@ -349,12 +401,20 @@ def make_root(
|
||||
"to name the dataset accordingly.",
|
||||
),
|
||||
] = None,
|
||||
kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
|
||||
dataset_root: Annotated[Path, typer.Option("--dataset-root", help="Dataset root")] = _DATASET_ROOT_DEFAULT,
|
||||
jobs: Annotated[int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")] = 4,
|
||||
kind: Annotated[
|
||||
str, typer.Option("--kind", help="steps | hits | ... (default: steps)")
|
||||
] = "steps",
|
||||
dataset_root: Annotated[
|
||||
Path, typer.Option("--dataset-root", help="Dataset root")
|
||||
] = _DATASET_ROOT_DEFAULT,
|
||||
jobs: Annotated[
|
||||
int, typer.Option("--jobs", "-j", help="Parallel simulation runs (default: 4)")
|
||||
] = 4,
|
||||
execute: Annotated[
|
||||
bool,
|
||||
typer.Option("--execute", help="Actually run jobs (default: dry run / print plan)"),
|
||||
typer.Option(
|
||||
"--execute", help="Actually run jobs (default: dry run / print plan)"
|
||||
),
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Generate new ROOT shards via a minicalosim executable."""
|
||||
@@ -381,7 +441,9 @@ class OracleMethod(str, Enum):
|
||||
|
||||
@app.command("build-geometry-oracle")
|
||||
def build_geometry_oracle(
|
||||
data: Annotated[Path, typer.Argument(help="Steps parquet file or directory of steps files")],
|
||||
data: Annotated[
|
||||
Path, typer.Argument(help="Steps parquet file or directory of steps files")
|
||||
],
|
||||
out: Annotated[Path, typer.Option("--out", "-o", help="Output oracle .pkl path")],
|
||||
method: Annotated[
|
||||
OracleMethod,
|
||||
@@ -394,7 +456,9 @@ def build_geometry_oracle(
|
||||
),
|
||||
),
|
||||
] = OracleMethod.slab,
|
||||
k: Annotated[int, typer.Option("--k", help="Neighbours for the knn classifier")] = 1,
|
||||
k: Annotated[
|
||||
int, typer.Option("--k", help="Neighbours for the knn classifier")
|
||||
] = 1,
|
||||
subsample: Annotated[
|
||||
int,
|
||||
typer.Option("--subsample", help="Max reference points sampled from the data"),
|
||||
@@ -440,7 +504,9 @@ def build_geometry_oracle(
|
||||
def warm_cache(
|
||||
data: Annotated[
|
||||
Path,
|
||||
typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
|
||||
typer.Argument(
|
||||
help="Parquet file, directory, or .manifest — same as `giant train`'s"
|
||||
),
|
||||
],
|
||||
val_fraction: Annotated[
|
||||
float,
|
||||
@@ -452,52 +518,49 @@ def warm_cache(
|
||||
] = 0.1,
|
||||
seed: Annotated[
|
||||
int,
|
||||
typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"),
|
||||
] = 0,
|
||||
particle_conditioning: Annotated[
|
||||
Conditioning,
|
||||
typer.Option(
|
||||
"--particle-conditioning",
|
||||
help="Must match the `giant train` run(s)' conditioning.particle.type to warm for",
|
||||
"--seed", "-s", help="Must match the `giant train` run(s) to warm for"
|
||||
),
|
||||
] = Conditioning.physical,
|
||||
material_conditioning: Annotated[
|
||||
] = 0,
|
||||
conditioning: Annotated[
|
||||
Conditioning,
|
||||
typer.Option(
|
||||
"--material-conditioning",
|
||||
help="Must match the `giant train` run(s)' conditioning.material.type "
|
||||
"to warm for — independent of --particle-conditioning "
|
||||
"(the two axes may differ)",
|
||||
"--conditioning", help="Must match the `giant train` run(s) to warm for"
|
||||
),
|
||||
] = Conditioning.physical,
|
||||
router: Annotated[
|
||||
bool,
|
||||
typer.Option(
|
||||
"--router/--no-router",
|
||||
help="Warm the process vocabulary too (only takes effect with --router-type process)",
|
||||
help="Warm the process vocabulary too (only takes effect with "
|
||||
"--router-type process)",
|
||||
),
|
||||
] = False,
|
||||
router_type: Annotated[str, typer.Option("--router-type", help="Router implementation name")] = "energy",
|
||||
n_experts: Annotated[int, typer.Option("--n-experts", help="Number of routed experts")] = 4,
|
||||
router_type: Annotated[
|
||||
str, typer.Option("--router-type", help="Router implementation name")
|
||||
] = "energy",
|
||||
n_experts: Annotated[
|
||||
int, typer.Option("--n-experts", help="Number of routed experts")
|
||||
] = 4,
|
||||
rebuild: Annotated[
|
||||
bool,
|
||||
typer.Option("--rebuild", help="Ignore any existing sidecar and recompute every section"),
|
||||
typer.Option(
|
||||
"--rebuild", help="Ignore any existing sidecar and recompute every section"
|
||||
),
|
||||
] = False,
|
||||
) -> None:
|
||||
"""Precompute `giant train`'s setup-stage sidecar for `data` ahead of time.
|
||||
|
||||
Warms the vocab maps, event-id split index, and the normalizer entry for
|
||||
the given --val-fraction/--seed/--particle-conditioning/
|
||||
--material-conditioning, so a later `giant train` run (or a `dwarf
|
||||
hparam-scan` sweep, which shares one such entry across every run) skips
|
||||
straight to training. See giant/data/setup_cache.py.
|
||||
the given --val-fraction/--seed/--conditioning, so a later `giant train`
|
||||
run (or a `dwarf hparam-scan` sweep, which shares one such entry across
|
||||
every run) skips straight to training. See giant/data/setup_cache.py.
|
||||
"""
|
||||
run_warm_setup_cache(
|
||||
data=str(data),
|
||||
val_fraction=val_fraction,
|
||||
seed=seed,
|
||||
particle_conditioning=particle_conditioning.value,
|
||||
material_conditioning=material_conditioning.value,
|
||||
conditioning=conditioning.value,
|
||||
router_enabled=router,
|
||||
router_type=router_type,
|
||||
n_experts=n_experts,
|
||||
@@ -38,7 +38,9 @@ def run_build_geometry_oracle(
|
||||
n_bins=n_bins,
|
||||
)
|
||||
|
||||
print(f"method: {method} reference points: {oracle.metadata['n_reference_points']:,}")
|
||||
print(
|
||||
f"method: {method} reference points: {oracle.metadata['n_reference_points']:,}"
|
||||
)
|
||||
print("classes (material, layer_id):")
|
||||
for material, layer_id in oracle.classes:
|
||||
print(f" {material:<12} layer_id={layer_id}")
|
||||
@@ -68,8 +68,8 @@ def final_metrics(metrics_path: Path) -> tuple[int, float, float]:
|
||||
with open(metrics_path, newline="") as f:
|
||||
rows = list(csv.DictReader(f))
|
||||
epochs_completed = int(rows[-1]["epoch"])
|
||||
final_val_loss = float(rows[-1]["val/loss"])
|
||||
best_val_loss = min(float(r["val/loss"]) for r in rows)
|
||||
final_val_loss = float(rows[-1]["val_loss"])
|
||||
best_val_loss = min(float(r["val_loss"]) for r in rows)
|
||||
return epochs_completed, final_val_loss, best_val_loss
|
||||
|
||||
|
||||
@@ -154,7 +154,9 @@ def run_hparam_scan(
|
||||
wall_time_s = time.monotonic() - start
|
||||
|
||||
if metrics_path.exists():
|
||||
epochs_completed, final_val_loss, best_val_loss = final_metrics(metrics_path)
|
||||
epochs_completed, final_val_loss, best_val_loss = final_metrics(
|
||||
metrics_path
|
||||
)
|
||||
append_summary(
|
||||
summary_path,
|
||||
{
|
||||
@@ -169,6 +171,11 @@ def run_hparam_scan(
|
||||
"wall_time_s": round(wall_time_s, 1),
|
||||
},
|
||||
)
|
||||
print(f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} ({wall_time_s:.1f}s)")
|
||||
print(
|
||||
f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} "
|
||||
f"({wall_time_s:.1f}s)"
|
||||
)
|
||||
else:
|
||||
print(f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log")
|
||||
print(
|
||||
f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
|
||||
)
|
||||
@@ -50,8 +50,12 @@ PREDICTED_RE = re.compile(
|
||||
r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)"
|
||||
r"_predicted(?P<local>_local)?\.parquet$"
|
||||
)
|
||||
SHARD_RE = re.compile(r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$")
|
||||
LEGACY_PREDICTED_RE = re.compile(r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$")
|
||||
SHARD_RE = re.compile(
|
||||
r"^(?P<detector>[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P<shard>\d+)\.(?P<ext>root|parquet)$"
|
||||
)
|
||||
LEGACY_PREDICTED_RE = re.compile(
|
||||
r"^pbwo4_10000events_hits_predicted(?P<local>_local)?\.parquet$"
|
||||
)
|
||||
LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?P<ext>root|parquet)$")
|
||||
|
||||
|
||||
@@ -106,9 +110,24 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
|
||||
if m:
|
||||
detector, shard, ext = m["detector"], int(m["shard"]), m["ext"]
|
||||
if ext == "root":
|
||||
dst = src_root / "raw" / "steps" / GEN / detector / f"shard-{shard:03d}.root"
|
||||
dst = (
|
||||
src_root
|
||||
/ "raw"
|
||||
/ "steps"
|
||||
/ GEN
|
||||
/ detector
|
||||
/ f"shard-{shard:03d}.root"
|
||||
)
|
||||
else:
|
||||
dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
|
||||
dst = (
|
||||
src_root
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ GEN
|
||||
/ SCHEMA
|
||||
/ detector
|
||||
/ f"shard-{shard:03d}.parquet"
|
||||
)
|
||||
moves.append((path, dst))
|
||||
continue
|
||||
|
||||
@@ -116,9 +135,19 @@ def plan_moves(src_root: Path) -> tuple[list[tuple[Path, Path]], list[Path]]:
|
||||
if m:
|
||||
ext = m["ext"]
|
||||
if ext == "root":
|
||||
dst = src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
|
||||
dst = (
|
||||
src_root / "raw" / "hits" / LEGACY_GEN / "pbwo4" / "shard-000.root"
|
||||
)
|
||||
else:
|
||||
dst = src_root / "processed" / "hits" / LEGACY_GEN / LEGACY_SCHEMA / "pbwo4" / "shard-000.parquet"
|
||||
dst = (
|
||||
src_root
|
||||
/ "processed"
|
||||
/ "hits"
|
||||
/ LEGACY_GEN
|
||||
/ LEGACY_SCHEMA
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
moves.append((path, dst))
|
||||
continue
|
||||
|
||||
@@ -136,9 +165,20 @@ def plan_manifests(src_root: Path) -> dict[Path, list[str]]:
|
||||
for pool, shards in rules.items():
|
||||
manifest_path = manifest_dir / f"{pool}{MANIFEST_SUFFIX}"
|
||||
for shard in shards:
|
||||
dst = src_root / "processed" / "steps" / GEN / SCHEMA / detector / f"shard-{shard:03d}.parquet"
|
||||
dst = (
|
||||
src_root
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ GEN
|
||||
/ SCHEMA
|
||||
/ detector
|
||||
/ f"shard-{shard:03d}.parquet"
|
||||
)
|
||||
manifests[manifest_path].append((shard, dst))
|
||||
return {k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)] for k, v in manifests.items()}
|
||||
return {
|
||||
k: [os.path.relpath(dst, start=k.parent) for _, dst in sorted(v)]
|
||||
for k, v in manifests.items()
|
||||
}
|
||||
|
||||
|
||||
def run_migration(root: str, execute: bool, copy: bool) -> None:
|
||||
@@ -13,7 +13,7 @@ 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 giant/tools/profile_analysis_costs.py``
|
||||
Usage: ``uv run python scripts/profile_analysis_costs.py``
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -94,7 +94,9 @@ def _make_rollout(n: int, n_events: int, seed: int) -> pl.DataFrame:
|
||||
"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),
|
||||
"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),
|
||||
@@ -161,7 +163,9 @@ def _make_reference(n: int, n_events: int, seed: int) -> pl.DataFrame:
|
||||
)
|
||||
|
||||
|
||||
def _time(spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path) -> float:
|
||||
def _time(
|
||||
spec_id: str, rollout: Path, reference: Path, shared: Path, out: Path
|
||||
) -> float:
|
||||
t0 = time.perf_counter()
|
||||
compute_reduced(
|
||||
spec_id,
|
||||
@@ -2,11 +2,11 @@
|
||||
|
||||
A single `dwarf convert` call converts a list of files one at a time; this
|
||||
module runs up to --jobs conversions concurrently, each as its own `dwarf
|
||||
convert` subprocess (invoked via `python -m giant.tools.dwarf`, so it picks up
|
||||
convert` subprocess (invoked via `python -m scripts.dwarf`, so it picks up
|
||||
the active venv/uv environment automatically).
|
||||
|
||||
Inputs must live under <dataset-root>/raw/<kind>/<gen>/<detector>/<file>.root
|
||||
(see giant/tools/migrate_geant_steps.py) — each is written to the matching
|
||||
(see scripts/migrate_geant_steps.py) — each is written to the matching
|
||||
processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet, where <schema>
|
||||
defaults to the highest schemaN already under processed/<kind>/<gen>/ (pass
|
||||
--schema to pick a specific one, e.g. one just created by `dwarf bump-schema`).
|
||||
@@ -22,7 +22,7 @@ import sys
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from pathlib import Path
|
||||
|
||||
# Must match giant/tools/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
|
||||
# Must match scripts/bump_dataset_version.py's GEN_RE / SCHEMA_RE.
|
||||
GEN_RE = re.compile(r"^gen\d+$")
|
||||
SCHEMA_RE = re.compile(r"^schema(\d+)$")
|
||||
|
||||
@@ -49,7 +49,9 @@ def latest_schema_tag(processed_gen_dir: Path) -> str | None:
|
||||
return best_tag
|
||||
|
||||
|
||||
def resolve_destination(root_file: Path, dataset_root: Path, schema_override: str | None) -> Path:
|
||||
def resolve_destination(
|
||||
root_file: Path, dataset_root: Path, schema_override: str | None
|
||||
) -> Path:
|
||||
"""Map raw/<kind>/<gen>/<detector>/<file>.root (relative to *dataset_root*)
|
||||
to processed/<kind>/<gen>/<schema>/<detector>/<file>.parquet.
|
||||
|
||||
@@ -64,7 +66,12 @@ def resolve_destination(root_file: Path, dataset_root: Path, schema_override: st
|
||||
raise DestinationError(f"{root_file} is not under dataset root {dataset_root}")
|
||||
|
||||
parts = rel.parts
|
||||
if len(parts) != 5 or parts[0] != "raw" or not GEN_RE.match(parts[2]) or not parts[4].endswith(".root"):
|
||||
if (
|
||||
len(parts) != 5
|
||||
or parts[0] != "raw"
|
||||
or not GEN_RE.match(parts[2])
|
||||
or not parts[4].endswith(".root")
|
||||
):
|
||||
raise DestinationError(
|
||||
f"{root_file} does not match raw/<kind>/<gen>/<detector>/<file>.root "
|
||||
f"under {dataset_root} (got relative path: {rel})"
|
||||
@@ -82,7 +89,7 @@ def resolve_destination(root_file: Path, dataset_root: Path, schema_override: st
|
||||
return processed_gen_dir / schema_tag / detector / f"{shard_stem}.parquet"
|
||||
|
||||
|
||||
_DWARF_CONVERT_CMD = [sys.executable, "-m", "giant.tools.dwarf", "convert"]
|
||||
_DWARF_CONVERT_CMD = [sys.executable, "-m", "scripts.dwarf", "convert"]
|
||||
|
||||
|
||||
def _convert_one(
|
||||
@@ -127,7 +134,7 @@ def run_parallel(
|
||||
written next to the input .root).
|
||||
|
||||
*cmd_prefix* overrides the subprocess command run per file (defaults to
|
||||
`python -m giant.tools.dwarf convert`) — used by tests to substitute a fake
|
||||
`python -m scripts.dwarf convert`) — used by tests to substitute a fake
|
||||
conversion script.
|
||||
|
||||
Returns one (root_file, returncode, stdout, stderr) tuple per file, in
|
||||
@@ -209,7 +216,13 @@ def run_parallel_job(
|
||||
print(f" {root_file}", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
|
||||
total_orphaned = sum(int(m.group(1)) for _, _, stdout, _ in results for m in _ORPHAN_RE.finditer(stdout))
|
||||
total_orphaned = sum(
|
||||
int(m.group(1))
|
||||
for _, _, stdout, _ in results
|
||||
for m in _ORPHAN_RE.finditer(stdout)
|
||||
)
|
||||
if total_orphaned:
|
||||
print(f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s).")
|
||||
print(
|
||||
f"\n{total_orphaned} orphaned child track(s) dropped across {len(results)} file(s)."
|
||||
)
|
||||
print(f"\nAll {len(results)} conversion(s) completed.")
|
||||
@@ -9,8 +9,6 @@ for the sidecar itself.
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from giant import config as gconfig
|
||||
from giant.constants import K_MAX
|
||||
from giant.pipeline import run_setup_stage
|
||||
|
||||
|
||||
@@ -18,8 +16,7 @@ def run_warm_setup_cache(
|
||||
data: str,
|
||||
val_fraction: float = 0.1,
|
||||
seed: int = 0,
|
||||
particle_conditioning: str = "physical",
|
||||
material_conditioning: str = "physical",
|
||||
conditioning: str = "physical",
|
||||
router_enabled: bool = False,
|
||||
router_type: str = "energy",
|
||||
n_experts: int = 4,
|
||||
@@ -28,11 +25,9 @@ def run_warm_setup_cache(
|
||||
) -> None:
|
||||
"""Populate (or refresh) the setup cache sidecar for `data`.
|
||||
|
||||
`val_fraction`/`seed`/`particle_conditioning`/`material_conditioning`
|
||||
select the normalizer cache entry
|
||||
`val_fraction`/`seed`/`conditioning` select the normalizer cache entry
|
||||
(`giant.data.setup_cache.normalizer_key`) — pass the same values a later
|
||||
`giant train` invocation will use so it hits this warmed entry. The two
|
||||
conditioning axes are independent and may differ.
|
||||
`giant train` invocation will use so it hits this warmed entry.
|
||||
`router_enabled`/`router_type`/`n_experts` only matter for
|
||||
`router_type == "process"` (warms that `n_experts`'s process map); the
|
||||
energy-router quantile summary is always collected regardless, so a
|
||||
@@ -44,30 +39,12 @@ def run_warm_setup_cache(
|
||||
"type": router_type,
|
||||
"n_experts": n_experts,
|
||||
}
|
||||
# Merged against DEFAULT_CONFIG (not a hand-rolled partial dict) so
|
||||
# run_setup_stage always sees every key it might read (e.g.
|
||||
# conditioning.particle.emb_dim, stage2_model.particle_type.target) at
|
||||
# its real default, not silently missing/None — see issues.md Issue 1.
|
||||
# This CLI only ever configures one router (matching today's single
|
||||
# --router-type flag), so it's placed on stage1_model; stage2_model's
|
||||
# stays disabled.
|
||||
cfg = gconfig.merge_cli_overrides(
|
||||
gconfig.DEFAULT_CONFIG,
|
||||
None,
|
||||
{
|
||||
"conditioning": {
|
||||
"particle": {"type": particle_conditioning},
|
||||
"material": {"type": material_conditioning},
|
||||
},
|
||||
"stage1_model": {"router": router_cfg},
|
||||
"stage2_model": {"router": {"enabled": False}, "k_max": K_MAX},
|
||||
},
|
||||
)
|
||||
run_setup_stage(
|
||||
Path(data),
|
||||
val_fraction=val_fraction,
|
||||
seed=seed,
|
||||
cfg=cfg,
|
||||
conditioning=conditioning,
|
||||
router_cfg=router_cfg,
|
||||
cache_setup=True,
|
||||
rebuild_setup_cache=rebuild,
|
||||
echo=echo,
|
||||
File diff suppressed because it is too large
Load Diff
@@ -102,6 +102,33 @@ def test_hist1d_overall_and_grouped():
|
||||
assert hg[11].sum() == 4
|
||||
|
||||
|
||||
def test_hist1d_clamps_extreme_values_and_drops_nan():
|
||||
# A rollout can emit a wildly out-of-range step_length (or an inf/NaN); the
|
||||
# fixed-edge binning must clamp rather than overflow the i32 bin cast.
|
||||
lf = pl.DataFrame(
|
||||
{"x": [5.0, 1.0725e10, float("inf"), -float("inf"), float("nan"), None]}
|
||||
).lazy()
|
||||
edges = np.linspace(0.0, 50.0, 6) # width 10
|
||||
h = R.hist1d(lf, pl.col("x"), edges)
|
||||
# 5 -> bin 0; 1e10 and +inf -> top bin; -inf -> bin 0; NaN/null dropped
|
||||
assert h[0].tolist() == [2, 0, 0, 0, 2]
|
||||
|
||||
|
||||
def test_profile_partial_clamps_extreme_values_and_drops_nan():
|
||||
lf = pl.DataFrame(
|
||||
{
|
||||
"event_id": [1, 1, 1, 1],
|
||||
"z": [5.0, 1.0725e10, float("nan"), 45.0],
|
||||
"w": [1.0, 2.0, 4.0, 8.0],
|
||||
}
|
||||
).lazy()
|
||||
edges = np.linspace(0.0, 50.0, 6)
|
||||
ev, mat = R.profile_partial(lf, pl.col("z"), edges, pl.col("w"))
|
||||
assert ev.tolist() == [1]
|
||||
# 1e10 clamps into the top bin alongside 45; the NaN row's weight is dropped
|
||||
assert mat[0].tolist() == [1.0, 0.0, 0.0, 0.0, 10.0]
|
||||
|
||||
|
||||
def test_physical_steps_drops_synthetic_rollout_rows_only():
|
||||
lf = _rollout_frame()
|
||||
phys = physical_steps(lf, Side.rollout).collect()
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import os
|
||||
import subprocess
|
||||
|
||||
from giant.tools import bump_dataset_version
|
||||
from scripts import bump_dataset_version
|
||||
|
||||
plan_bump_gen = bump_dataset_version.plan_bump_gen
|
||||
plan_bump_schema = bump_dataset_version.plan_bump_schema
|
||||
@@ -22,7 +22,9 @@ def test_git_user_name_returns_none_on_timeout(monkeypatch):
|
||||
|
||||
|
||||
def test_bump_gen_starts_at_gen1_when_none_exist(tmp_path):
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "first generation", None, "2026-01-01")
|
||||
dirs, log_line = plan_bump_gen(
|
||||
tmp_path, "steps", "first generation", None, "2026-01-01"
|
||||
)
|
||||
assert dirs == [
|
||||
tmp_path / "raw" / "steps" / "gen1",
|
||||
tmp_path / "processed" / "steps" / "gen1" / "schema1",
|
||||
@@ -54,7 +56,9 @@ def test_bump_gen_kinds_are_independent(tmp_path):
|
||||
|
||||
def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
dirs, log_line = plan_bump_schema(tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01")
|
||||
dirs, log_line = plan_bump_schema(
|
||||
tmp_path, "steps", "gen1", "added e_sec column", None, "2026-01-01"
|
||||
)
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema1"]
|
||||
assert "`gen1`/`schema1`" in log_line
|
||||
|
||||
@@ -62,14 +66,18 @@ def test_bump_schema_starts_at_schema1_for_a_fresh_gen(tmp_path):
|
||||
def test_bump_schema_increments_within_its_gen(tmp_path):
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema1").mkdir(parents=True)
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen1", "next schema", None, "2026-01-01")
|
||||
dirs, _ = plan_bump_schema(
|
||||
tmp_path, "steps", "gen1", "next schema", None, "2026-01-01"
|
||||
)
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen1" / "schema3"]
|
||||
|
||||
|
||||
def test_bump_schema_does_not_see_other_gens_schemas(tmp_path):
|
||||
(tmp_path / "processed" / "steps" / "gen1" / "schema5").mkdir(parents=True)
|
||||
(tmp_path / "raw" / "steps" / "gen2").mkdir(parents=True)
|
||||
dirs, _ = plan_bump_schema(tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01")
|
||||
dirs, _ = plan_bump_schema(
|
||||
tmp_path, "steps", "gen2", "fresh schema for gen2", None, "2026-01-01"
|
||||
)
|
||||
assert dirs == [tmp_path / "processed" / "steps" / "gen2" / "schema1"]
|
||||
|
||||
|
||||
@@ -83,7 +91,9 @@ def test_bump_schema_rejects_nonexistent_gen(tmp_path):
|
||||
|
||||
def test_bump_gen_to_specific_tag(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
dirs, log_line = plan_bump_gen(tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5")
|
||||
dirs, log_line = plan_bump_gen(
|
||||
tmp_path, "steps", "jump to gen5", None, "2026-01-01", target="gen5"
|
||||
)
|
||||
assert dirs[0] == tmp_path / "raw" / "steps" / "gen5"
|
||||
assert "`gen5`" in log_line
|
||||
|
||||
@@ -115,7 +125,9 @@ def test_bump_schema_to_specific_tag(tmp_path):
|
||||
def test_bump_schema_rejects_invalid_to_tag(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
try:
|
||||
plan_bump_schema(tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3")
|
||||
plan_bump_schema(
|
||||
tmp_path, "steps", "gen1", "bad tag", None, "2026-01-01", target="v3"
|
||||
)
|
||||
assert False, "expected SystemExit"
|
||||
except SystemExit:
|
||||
pass
|
||||
@@ -152,7 +164,15 @@ def _make_parquet(path):
|
||||
|
||||
|
||||
def test_update_manifest_bumps_to_specified_schema(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -174,7 +194,15 @@ def test_update_manifest_auto_detects_highest_schema(tmp_path):
|
||||
for schema in ("schema1", "schema2", "schema3"):
|
||||
d = tmp_path / "processed" / "steps" / "gen1" / schema / "pbwo4"
|
||||
d.mkdir(parents=True)
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema3"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
parquet.touch()
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -203,7 +231,15 @@ def test_update_manifest_reports_missing_targets(tmp_path):
|
||||
|
||||
|
||||
def test_update_manifest_skips_already_at_target(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -218,7 +254,15 @@ def test_update_manifest_skips_already_at_target(tmp_path):
|
||||
|
||||
|
||||
def test_update_manifest_preserves_comments_and_blanks(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -234,7 +278,15 @@ def test_update_manifest_preserves_comments_and_blanks(tmp_path):
|
||||
|
||||
|
||||
def test_update_manifest_bumps_gen(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema1" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen2"
|
||||
/ "schema1"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -251,7 +303,15 @@ def test_update_manifest_bumps_gen(tmp_path):
|
||||
|
||||
|
||||
def test_update_manifest_bumps_gen_and_schema(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen2" / "schema3" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen2"
|
||||
/ "schema3"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -268,7 +328,15 @@ def test_update_manifest_bumps_gen_and_schema(tmp_path):
|
||||
|
||||
|
||||
def test_apply_update_manifest_writes_file(tmp_path):
|
||||
parquet = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
parquet = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
_make_parquet(parquet)
|
||||
|
||||
manifest_dir = tmp_path / "pools" / "pbwo4"
|
||||
@@ -290,8 +358,24 @@ def test_apply_update_manifest_writes_file(tmp_path):
|
||||
|
||||
|
||||
def test_create_manifest_writes_relative_paths(tmp_path):
|
||||
pq1 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-000.parquet"
|
||||
pq2 = tmp_path / "processed" / "steps" / "gen1" / "schema2" / "pbwo4" / "shard-001.parquet"
|
||||
pq1 = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
pq2 = (
|
||||
tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema2"
|
||||
/ "pbwo4"
|
||||
/ "shard-001.parquet"
|
||||
)
|
||||
_make_parquet(pq1)
|
||||
_make_parquet(pq2)
|
||||
|
||||
@@ -335,7 +419,9 @@ def test_run_create_manifest_refuses_to_overwrite_existing_output(tmp_path):
|
||||
output.write_text("original contents\n")
|
||||
|
||||
try:
|
||||
bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output))
|
||||
bump_dataset_version.run_create_manifest(
|
||||
[str(pq)], execute=True, output=str(output)
|
||||
)
|
||||
assert False, "expected SystemExit"
|
||||
except SystemExit:
|
||||
pass
|
||||
@@ -349,7 +435,9 @@ def test_run_create_manifest_force_overwrites_existing_output(tmp_path):
|
||||
output.parent.mkdir(parents=True)
|
||||
output.write_text("original contents\n")
|
||||
|
||||
bump_dataset_version.run_create_manifest([str(pq)], execute=True, output=str(output), force=True)
|
||||
bump_dataset_version.run_create_manifest(
|
||||
[str(pq)], execute=True, output=str(output), force=True
|
||||
)
|
||||
assert output.read_text() != "original contents\n"
|
||||
|
||||
|
||||
|
||||
@@ -13,7 +13,9 @@ from tests.test_analysis_reduce import _reference_frame, _rollout_frame
|
||||
|
||||
def _build_ctx() -> Context:
|
||||
r, t = _rollout_frame(), _reference_frame()
|
||||
return build_context(r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000)
|
||||
return build_context(
|
||||
r, t, n_energy_bins=2, n_marginal_bins=10, top_k_pdg=3, sample_rows=1000
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
@@ -140,7 +142,10 @@ def test_chunked_matches_unchunked(ctx: Context, spec_id: str):
|
||||
|
||||
# 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)]
|
||||
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
|
||||
|
||||
@@ -1,280 +0,0 @@
|
||||
"""Tests for giant.checkpoint_io.load_for_inference (issues.md Issue 5) —
|
||||
the shared bootstrap `giant predict`/`giant rollout` use to go from a
|
||||
checkpoint path to ready-to-run models."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from giant import config as gconfig
|
||||
from giant.checkpoint_io import (
|
||||
CheckpointCompatibilityError,
|
||||
InferenceContext,
|
||||
conditioning_axes,
|
||||
load_for_inference,
|
||||
stage_cfg,
|
||||
)
|
||||
from giant.data.loader import TopNMap
|
||||
from giant.data.setup_cache import topnmap_to_json
|
||||
from giant.data.transforms import Normalizer
|
||||
from giant.model.network import build_models
|
||||
|
||||
PDG_MAP = {11: 0, 22: 1, -11: 2}
|
||||
MAT_MAP = {"G4_PbWO4": 0, "G4_AIR": 1}
|
||||
|
||||
|
||||
def _model_cfg(stage2_active: bool = True) -> dict:
|
||||
"""DEFAULT_CONFIG-derived, shrunk for speed — same pattern as
|
||||
tests/test_network.py::_minimal_model_config. Default `conditioning`
|
||||
(both axes "physical") needs no top-N vocab map, so this is a cheap,
|
||||
fully self-contained happy-path config."""
|
||||
cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
|
||||
cfg["conditioning"]["particle"]["emb_dim"] = 4
|
||||
cfg["conditioning"]["material"]["emb_dim"] = 4
|
||||
cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
|
||||
cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
|
||||
cfg["stage2_model"]["active"] = stage2_active
|
||||
return {
|
||||
"pdg_vocab": len(PDG_MAP),
|
||||
"mat_vocab": len(MAT_MAP),
|
||||
"conditioning": cfg["conditioning"],
|
||||
"stage1_model": cfg["stage1_model"],
|
||||
"stage2_model": cfg["stage2_model"],
|
||||
}
|
||||
|
||||
|
||||
def _norms() -> tuple[Normalizer, Normalizer, Normalizer]:
|
||||
rng = np.random.default_rng(0)
|
||||
cond = Normalizer().fit(rng.standard_normal((100, 15)).astype(np.float32))
|
||||
tgt = Normalizer().fit(rng.standard_normal((100, 9)).astype(np.float32))
|
||||
sec_phys = Normalizer().fit(rng.standard_normal((100, 2)).astype(np.float32))
|
||||
return cond, tgt, sec_phys
|
||||
|
||||
|
||||
def _write_checkpoint(tmp_path, model_cfg=None, ema: bool = False, **ckpt_overrides):
|
||||
cfg = model_cfg if model_cfg is not None else _model_cfg()
|
||||
built = build_models(cfg)
|
||||
stage1, stage2 = built["stage1"], built["stage2"]
|
||||
cond, tgt, sec_phys = _norms()
|
||||
ckpt: dict = {
|
||||
"model_config": cfg,
|
||||
"model": stage1.state_dict() if stage1 is not None else {},
|
||||
"sec_decoder": stage2.state_dict() if stage2 is not None else {},
|
||||
"pdg_map": PDG_MAP,
|
||||
"mat_map": MAT_MAP,
|
||||
"normalizer": {"cond": cond.to_dict(), "target": tgt.to_dict(), "sec_phys": sec_phys.to_dict()},
|
||||
"epoch": 3,
|
||||
"best_val_loss": 0.5,
|
||||
}
|
||||
if ema:
|
||||
ckpt["model_ema"] = stage1.state_dict() if stage1 is not None else {}
|
||||
ckpt["sec_decoder_ema"] = stage2.state_dict() if stage2 is not None else {}
|
||||
# DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
|
||||
# "onehot", and giant train's pipeline (gitea #29) now always writes a
|
||||
# sec_type_topn_map in that case — default one in here too, unless a
|
||||
# test explicitly overrides it, so fixtures represent a real, loadable
|
||||
# checkpoint by default rather than exercising the "missing" guard by
|
||||
# accident.
|
||||
particle_type_target = cfg.get("stage2_model", {}).get("particle_type", {}).get("target", "onehot")
|
||||
if particle_type_target == "onehot" and "sec_type_topn_map" not in ckpt_overrides:
|
||||
default_sec_type_topn = TopNMap(class_map=dict(zip(PDG_MAP, range(len(PDG_MAP)))), other_members={})
|
||||
ckpt["sec_type_topn_map"] = topnmap_to_json(default_sec_type_topn)
|
||||
ckpt.update(ckpt_overrides)
|
||||
path = tmp_path / "ckpt.pt"
|
||||
torch.save(ckpt, path)
|
||||
return path
|
||||
|
||||
|
||||
def _onehot_model_cfg() -> dict:
|
||||
cfg = _model_cfg()
|
||||
cfg["conditioning"]["particle"]["type"] = "onehot"
|
||||
return cfg
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Happy path
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_happy_path_returns_populated_context(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
|
||||
assert isinstance(ctx, InferenceContext)
|
||||
assert ctx.stage1 is not None and ctx.stage2 is not None
|
||||
assert not ctx.stage1.training
|
||||
assert not ctx.stage2.training
|
||||
assert next(ctx.stage1.parameters()).device == torch.device("cpu")
|
||||
assert ctx.pdg_map == PDG_MAP
|
||||
assert ctx.mat_map == MAT_MAP
|
||||
assert all(isinstance(k, int) for k in ctx.pdg_map)
|
||||
assert all(isinstance(k, str) for k in ctx.mat_map)
|
||||
assert ctx.particle_conditioning == "physical"
|
||||
assert ctx.material_conditioning == "physical"
|
||||
assert ctx.k_max == 3
|
||||
assert ctx.epoch == 3
|
||||
assert ctx.best_val_loss == 0.5
|
||||
assert ctx.model_config["stage1_model"]["hidden_dim"] == 8
|
||||
|
||||
|
||||
def test_happy_path_normalizer_values_round_trip(tmp_path):
|
||||
cond, tgt, sec_phys = _norms()
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
|
||||
assert ctx.cond_norm.mean is not None and cond.mean is not None
|
||||
assert ctx.tgt_norm.mean is not None and tgt.mean is not None
|
||||
assert ctx.sec_phys_norm.mean is not None and sec_phys.mean is not None
|
||||
np.testing.assert_allclose(ctx.cond_norm.mean, cond.mean)
|
||||
np.testing.assert_allclose(ctx.tgt_norm.mean, tgt.mean)
|
||||
np.testing.assert_allclose(ctx.sec_phys_norm.mean, sec_phys.mean)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Guards
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_missing_model_config_raises(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
ckpt = torch.load(checkpoint, weights_only=False)
|
||||
del ckpt["model_config"]
|
||||
torch.save(ckpt, checkpoint)
|
||||
|
||||
with pytest.raises(CheckpointCompatibilityError, match="no model_config"):
|
||||
load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
|
||||
|
||||
def test_missing_sec_decoder_raises(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
ckpt = torch.load(checkpoint, weights_only=False)
|
||||
del ckpt["sec_decoder"]
|
||||
torch.save(ckpt, checkpoint)
|
||||
|
||||
with pytest.raises(CheckpointCompatibilityError, match="no sec_decoder"):
|
||||
load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
|
||||
|
||||
def test_missing_sec_phys_normalizer_raises(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path)
|
||||
ckpt = torch.load(checkpoint, weights_only=False)
|
||||
del ckpt["normalizer"]["sec_phys"]
|
||||
torch.save(ckpt, checkpoint)
|
||||
|
||||
with pytest.raises(CheckpointCompatibilityError, match="no normalizer.sec_phys"):
|
||||
load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
|
||||
|
||||
def test_onehot_particle_conditioning_without_topn_map_raises(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path, model_cfg=_onehot_model_cfg())
|
||||
|
||||
with pytest.raises(CheckpointCompatibilityError, match="pdg_topn_map"):
|
||||
load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
|
||||
|
||||
def test_onehot_particle_conditioning_with_topn_map_succeeds(tmp_path):
|
||||
topn = TopNMap(class_map={11: 0, 22: 1}, other_members={})
|
||||
checkpoint = _write_checkpoint(
|
||||
tmp_path,
|
||||
model_cfg=_onehot_model_cfg(),
|
||||
pdg_topn_map=topnmap_to_json(topn),
|
||||
)
|
||||
ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
assert ctx.particle_conditioning == "onehot"
|
||||
assert ctx.pdg_topn_map is not None
|
||||
assert ctx.pdg_topn_map.class_map == {11: 0, 22: 1}
|
||||
|
||||
|
||||
def test_onehot_particle_type_target_without_sec_type_topn_map_raises(tmp_path):
|
||||
"""DEFAULT_CONFIG's stage2_model.particle_type.target="onehot" needs a
|
||||
sec_type_topn_map (gitea #29) — a checkpoint with neither key at all
|
||||
(not even the pre-#29 pdg_topn_map to fall back to) must fail loudly."""
|
||||
checkpoint = _write_checkpoint(tmp_path, sec_type_topn_map=None)
|
||||
ckpt = torch.load(checkpoint, weights_only=False)
|
||||
del ckpt["sec_type_topn_map"]
|
||||
torch.save(ckpt, checkpoint)
|
||||
|
||||
with pytest.raises(CheckpointCompatibilityError, match="sec_type_topn_map"):
|
||||
load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
|
||||
|
||||
def test_pre_gitea_29_checkpoint_falls_back_to_pdg_topn_map_for_sec_type(tmp_path):
|
||||
"""A checkpoint written before gitea #29 has no sec_type_topn_map key at
|
||||
all — conditioning and secondary-type onehot maps were always the same
|
||||
map, saved once under pdg_topn_map. load_for_inference must reproduce
|
||||
that exact pre-#29 behavior for such a checkpoint."""
|
||||
topn = TopNMap(class_map={11: 0, 22: 1, -11: 2}, other_members={})
|
||||
checkpoint = _write_checkpoint(
|
||||
tmp_path,
|
||||
model_cfg=_onehot_model_cfg(),
|
||||
pdg_topn_map=topnmap_to_json(topn),
|
||||
sec_type_topn_map=None,
|
||||
)
|
||||
ckpt = torch.load(checkpoint, weights_only=False)
|
||||
del ckpt["sec_type_topn_map"]
|
||||
torch.save(ckpt, checkpoint)
|
||||
|
||||
ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict")
|
||||
assert ctx.sec_type_topn_map is not None
|
||||
assert ctx.sec_type_topn_map.class_map == {11: 0, 22: 1, -11: 2}
|
||||
|
||||
|
||||
def test_ema_weights_requested_but_missing_raises(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path, ema=False)
|
||||
|
||||
with pytest.raises(CheckpointCompatibilityError, match="no EMA weights"):
|
||||
load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
|
||||
|
||||
|
||||
def test_ema_weights_requested_and_present_succeeds(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path, ema=True)
|
||||
ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", weights="ema")
|
||||
assert ctx.stage1 is not None and ctx.stage2 is not None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("command_name", ["predict", "rollout"])
|
||||
def test_inactive_stage_with_require_stage2_raises_with_command_name(tmp_path, command_name):
|
||||
checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
|
||||
|
||||
with pytest.raises(CheckpointCompatibilityError, match=f"{command_name} needs both"):
|
||||
load_for_inference(checkpoint, torch.device("cpu"), command_name)
|
||||
|
||||
|
||||
def test_inactive_stage_with_require_stage2_false_succeeds_with_stage2_none(tmp_path):
|
||||
checkpoint = _write_checkpoint(tmp_path, model_cfg=_model_cfg(stage2_active=False))
|
||||
|
||||
ctx = load_for_inference(checkpoint, torch.device("cpu"), "predict", require_stage2=False)
|
||||
assert ctx.stage1 is not None
|
||||
assert ctx.stage2 is None
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# conditioning_axes / stage_cfg
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_conditioning_axes_v02_flat_string_applies_to_both_axes():
|
||||
assert conditioning_axes({"conditioning": "embedding"}) == ("embedding", "embedding")
|
||||
|
||||
|
||||
def test_conditioning_axes_v03_nested_dict_independent_per_axis():
|
||||
model_cfg = {"conditioning": {"particle": {"type": "onehot"}, "material": {"type": "physical"}}}
|
||||
assert conditioning_axes(model_cfg) == ("onehot", "physical")
|
||||
|
||||
|
||||
def test_conditioning_axes_missing_key_uses_default():
|
||||
assert conditioning_axes({}, default="embedding") == ("embedding", "embedding")
|
||||
|
||||
|
||||
def test_stage_cfg_new_shape_returns_subdict():
|
||||
model_cfg = {"stage2_model": {"k_max": 7}}
|
||||
assert stage_cfg(model_cfg, "stage2") == {"k_max": 7}
|
||||
|
||||
|
||||
def test_stage_cfg_v02_flat_shape_returns_empty_dict():
|
||||
model_cfg = {"hidden_dim": 32, "n_blocks": 4}
|
||||
assert stage_cfg(model_cfg, "stage2") == {}
|
||||
@@ -37,14 +37,13 @@ def test_writes_config_with_overrides_applied(tmp_path: Path):
|
||||
with open(config_path, "rb") as f:
|
||||
cfg = tomllib.load(f)
|
||||
|
||||
assert cfg["stage1_model"]["generator"] == "ddpm"
|
||||
assert cfg["stage2_model"]["generator"] == "ddpm"
|
||||
assert cfg["train"]["mode"] == "ddpm"
|
||||
assert cfg["train"]["lr"] == 0.0005
|
||||
assert cfg["stage1_model"]["hidden_dim"] == 128
|
||||
assert cfg["stage1_model"]["n_res_blocks"] == 4
|
||||
assert cfg["model"]["hidden_dim"] == 128
|
||||
assert cfg["model"]["n_blocks"] == 4
|
||||
# untouched defaults still present
|
||||
assert cfg["train"]["epochs"] == 100
|
||||
assert "router" in cfg["stage1_model"]
|
||||
assert "router" in cfg["model"]
|
||||
|
||||
assert str(out_dir) in result.output
|
||||
assert "<data.parquet>" in result.output
|
||||
@@ -101,7 +100,9 @@ def test_force_guard_refuses_to_clobber_existing_checkpoints(tmp_path: Path):
|
||||
assert "already has last.pt" in result.output
|
||||
assert not (out_dir / "config.toml").exists()
|
||||
|
||||
result = runner.invoke(app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"])
|
||||
result = runner.invoke(
|
||||
app, ["new-run", "--out", str(out_dir), "--mode", "ddpm", "--force"]
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert (out_dir / "config.toml").exists()
|
||||
|
||||
|
||||
@@ -1,18 +1,13 @@
|
||||
import uuid
|
||||
|
||||
import torch
|
||||
import yaml
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from giant.cli import (
|
||||
_CEPH_PREDICTIONS,
|
||||
_resolve_prediction_output,
|
||||
_write_prediction_ref,
|
||||
app,
|
||||
)
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# _resolve_prediction_output
|
||||
@@ -121,7 +116,9 @@ def test_ref_yaml_includes_comment_when_provided(tmp_path):
|
||||
dataset = tmp_path / "full.manifest"
|
||||
pred_uuid = str(uuid.uuid4())
|
||||
|
||||
ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3")
|
||||
ref_path = _write_prediction_ref(
|
||||
checkpoint, pred_uuid, out, dataset, comment="baseline sweep run 3"
|
||||
)
|
||||
data = yaml.safe_load(ref_path.read_text())
|
||||
|
||||
assert data["comment"] == "baseline sweep run 3"
|
||||
@@ -136,7 +133,9 @@ def test_ref_timestamp_is_iso_format(tmp_path):
|
||||
checkpoint.touch()
|
||||
|
||||
pred_uuid = str(uuid.uuid4())
|
||||
ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
|
||||
ref_path = _write_prediction_ref(
|
||||
checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
|
||||
)
|
||||
data = yaml.safe_load(ref_path.read_text())
|
||||
|
||||
# Must parse without error and be timezone-aware (UTC).
|
||||
@@ -151,24 +150,9 @@ def test_ref_checkpoint_path_is_absolute(tmp_path):
|
||||
checkpoint.touch()
|
||||
|
||||
pred_uuid = str(uuid.uuid4())
|
||||
ref_path = _write_prediction_ref(checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d")
|
||||
ref_path = _write_prediction_ref(
|
||||
checkpoint, pred_uuid, tmp_path / "p.parquet", tmp_path / "d"
|
||||
)
|
||||
data = yaml.safe_load(ref_path.read_text())
|
||||
|
||||
assert data["checkpoint"].startswith("/")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bootstrap failure surfaces via the CLI (issues.md Issue 5 — confirms
|
||||
# CheckpointCompatibilityError -> typer.Exit(1) actually wires up end-to-end,
|
||||
# not just at the giant.checkpoint_io unit level).
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_predict_exits_1_on_checkpoint_missing_model_config(tmp_path):
|
||||
checkpoint = tmp_path / "bad.pt"
|
||||
torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
|
||||
|
||||
result = runner.invoke(app, ["predict", "dummy.parquet", "--checkpoint", str(checkpoint)])
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "checkpoint has no model_config" in result.output
|
||||
|
||||
@@ -1,33 +0,0 @@
|
||||
"""Thin CLI smoke coverage for `giant rollout` (issues.md Issue 5) — confirms
|
||||
the CheckpointCompatibilityError raised by giant.checkpoint_io.load_for_inference
|
||||
surfaces as a clean typer.Exit(1) with the expected message, end-to-end
|
||||
through the CLI, not just at the giant.checkpoint_io unit level."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import torch
|
||||
from typer.testing import CliRunner
|
||||
|
||||
from giant.cli import app
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
def test_rollout_exits_1_on_checkpoint_missing_model_config(tmp_path):
|
||||
checkpoint = tmp_path / "bad.pt"
|
||||
torch.save({"sec_decoder": {}, "normalizer": {"sec_phys": {}}}, checkpoint)
|
||||
|
||||
result = runner.invoke(
|
||||
app,
|
||||
[
|
||||
"rollout",
|
||||
"dummy.parquet",
|
||||
"--checkpoint",
|
||||
str(checkpoint),
|
||||
"--geometry",
|
||||
"dummy_geometry.pkl",
|
||||
],
|
||||
)
|
||||
|
||||
assert result.exit_code == 1
|
||||
assert "checkpoint has no model_config" in result.output
|
||||
@@ -1,177 +0,0 @@
|
||||
"""Tests for `giant train`'s stage-prefixed CLI flags: --stage1-*/--stage2-*
|
||||
must independently override each stage's config block, and must take precedence
|
||||
over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still
|
||||
apply the same value to both stages for backward compatibility."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
from typer.testing import CliRunner
|
||||
|
||||
import giant.cli as cli
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dict:
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
||||
captured["cfg"] = cfg
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
|
||||
result = runner.invoke(
|
||||
cli.app,
|
||||
["train", "dummy.parquet", "--out", str(tmp_path / "run")] + args,
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
return captured["cfg"]
|
||||
|
||||
|
||||
def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path):
|
||||
cfg = _invoke_and_capture_cfg(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
["--mode", "wgan", "--stage1-generator", "flow"],
|
||||
)
|
||||
assert cfg["stage1_model"]["generator"] == "flow"
|
||||
assert cfg["stage2_model"]["generator"] == "wgan"
|
||||
|
||||
|
||||
def test_stage2_only_knobs(monkeypatch, tmp_path):
|
||||
# --stage2-stage1-context is exercised separately at the overrides-dict
|
||||
# level (test_overrides_from_flags_stage2_only_knobs in test_config.py):
|
||||
# its only non-default value, "sampled", is rejected by validate_config
|
||||
# (issues.md Issue 1), so it can't appear in a full CLI invocation here.
|
||||
cfg = _invoke_and_capture_cfg(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
[
|
||||
"--stage2-decoder",
|
||||
"one_shot",
|
||||
"--stage2-k-max",
|
||||
"8",
|
||||
"--stage2-hidden-dim",
|
||||
"32",
|
||||
"--stage2-context-dim",
|
||||
"16",
|
||||
],
|
||||
)
|
||||
assert cfg["stage2_model"]["decoder"] == "one_shot"
|
||||
assert cfg["stage2_model"]["k_max"] == 8
|
||||
assert cfg["stage2_model"]["hidden_dim"] == 32
|
||||
assert cfg["stage2_model"]["context_dim"] == 16
|
||||
# untouched stage1 defaults
|
||||
assert cfg["stage1_model"]["hidden_dim"] == 256
|
||||
|
||||
|
||||
def test_stage1_hidden_dim_flag_overrides_legacy_hidden_dim_flag(monkeypatch, tmp_path):
|
||||
cfg = _invoke_and_capture_cfg(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
["--hidden-dim", "64", "--stage1-hidden-dim", "128"],
|
||||
)
|
||||
assert cfg["stage1_model"]["hidden_dim"] == 128
|
||||
|
||||
|
||||
def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path):
|
||||
cfg = _invoke_and_capture_cfg(
|
||||
monkeypatch,
|
||||
tmp_path,
|
||||
[
|
||||
"--mode",
|
||||
"wgan",
|
||||
"--n-critic",
|
||||
"5",
|
||||
"--stage1-n-critic",
|
||||
"3",
|
||||
"--stage2-gp-weight",
|
||||
"2.5",
|
||||
],
|
||||
)
|
||||
assert cfg["stage1_model"]["wgan"]["n_critic"] == 3
|
||||
assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0
|
||||
assert cfg["stage2_model"]["wgan"]["n_critic"] == 5
|
||||
assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5
|
||||
|
||||
|
||||
def test_batch_size_invalid_string_errors(monkeypatch, tmp_path):
|
||||
monkeypatch.setattr(cli, "run_train_job", lambda *a, **kw: None)
|
||||
result = runner.invoke(
|
||||
cli.app,
|
||||
["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "not-a-number"],
|
||||
)
|
||||
assert result.exit_code == 1
|
||||
assert "--batch-size must be an integer or 'auto'" in result.output
|
||||
|
||||
|
||||
def test_out_dir_resolution_prefers_explicit_out_over_resume(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
||||
captured["out_dir"] = out_dir
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
|
||||
resume_dir = tmp_path / "resumed_run"
|
||||
resume_dir.mkdir()
|
||||
(resume_dir / "last.pt").touch()
|
||||
explicit_out = tmp_path / "explicit_run"
|
||||
|
||||
result = runner.invoke(
|
||||
cli.app,
|
||||
["train", "dummy.parquet", "--out", str(explicit_out), "--resume", str(resume_dir / "last.pt")],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["out_dir"] == explicit_out
|
||||
|
||||
|
||||
def test_out_dir_resolution_falls_back_to_resume_parent(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
||||
captured["out_dir"] = out_dir
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
|
||||
resume_dir = tmp_path / "resumed_run"
|
||||
resume_dir.mkdir()
|
||||
(resume_dir / "last.pt").touch()
|
||||
|
||||
result = runner.invoke(cli.app, ["train", "dummy.parquet", "--resume", str(resume_dir / "last.pt")])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["out_dir"] == resume_dir
|
||||
|
||||
|
||||
def test_out_dir_resolution_defaults_when_neither_out_nor_resume_given(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, **kwargs):
|
||||
captured["out_dir"] = out_dir
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
monkeypatch.chdir(tmp_path)
|
||||
|
||||
result = runner.invoke(cli.app, ["train", "dummy.parquet"])
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["out_dir"] == Path("checkpoints") / cli.gconfig.default_out_dir_name(cli.gconfig.DEFAULT_CONFIG)
|
||||
|
||||
|
||||
def test_batch_size_auto_estimates_and_echoes(monkeypatch, tmp_path):
|
||||
captured: dict = {}
|
||||
|
||||
def _fake_run_train_job(*, data, cfg, out_dir, num_workers, **kwargs):
|
||||
captured["batch_size"] = cfg["train"]["batch_size"]
|
||||
|
||||
monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job)
|
||||
monkeypatch.setattr(cli.gconfig, "estimate_batch_size", lambda hidden_dim, n_blocks, device: 123)
|
||||
|
||||
result = runner.invoke(
|
||||
cli.app,
|
||||
["train", "dummy.parquet", "--out", str(tmp_path / "run"), "--batch-size", "auto"],
|
||||
)
|
||||
assert result.exit_code == 0, result.output
|
||||
assert captured["batch_size"] == 123
|
||||
assert "batch_size: 123 (auto-estimated from free GPU memory)" in result.output
|
||||
+21
-6
@@ -62,7 +62,9 @@ def _fake_venv(repo_dir: Path) -> None:
|
||||
giant.chmod(0o755)
|
||||
|
||||
|
||||
def _prep(rollout_yaml: Path, run_dir: str | Path | None = None, chunks: int = 1) -> Path:
|
||||
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,
|
||||
@@ -235,7 +237,9 @@ def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.M
|
||||
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)
|
||||
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
|
||||
@@ -245,7 +249,9 @@ 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)
|
||||
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] = {}
|
||||
@@ -262,7 +268,9 @@ def test_write_submit_rejects_n_chunks_mismatch_with_run_meta(tmp_path: Path):
|
||||
_job_walltimes instead of a clear error here."""
|
||||
run_dir = _prep(_write_inputs(tmp_path), chunks=2)
|
||||
_fake_venv(tmp_path)
|
||||
cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4)
|
||||
cfg = SubmitConfig(
|
||||
run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path, n_chunks=4
|
||||
)
|
||||
with pytest.raises(ValueError, match="n_chunks"):
|
||||
write_submit(cfg)
|
||||
|
||||
@@ -285,9 +293,16 @@ def test_write_submit_walltime_grows_with_chunk_rows(tmp_path: Path):
|
||||
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)
|
||||
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())}
|
||||
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
|
||||
|
||||
+240
-1039
File diff suppressed because it is too large
Load Diff
@@ -1,156 +0,0 @@
|
||||
"""Consumed-keys audit (issues.md Issue 5).
|
||||
|
||||
`validate_config_keys` (`giant/config.py`) only checks that a config key is
|
||||
*declared* — present somewhere in `DEFAULT_CONFIG`, which is generated from
|
||||
the frozen dataclasses. It says nothing about whether anything actually
|
||||
*reads* the value once parsed. Issues 1, 2 and 4 are three keys that slipped
|
||||
through exactly that gap: declared, round-tripped, silently ignored. This
|
||||
module walks every leaf path in `DEFAULT_CONFIG` and asserts each is either
|
||||
genuinely consumed by the model-building/training/rollout code, or explicitly
|
||||
recorded in `_KNOWN_UNUSED` with a reason.
|
||||
|
||||
"Consumed" is approximated by static analysis rather than true call-graph
|
||||
reachability: for each leaf path's field name, does it appear anywhere in a
|
||||
fixed whitelist of source files as a real attribute access, a dict-key-shaped
|
||||
string constant, or a function/constructor parameter name (the last of these
|
||||
because `Router` subclasses receive their config via `**kwargs` filtered by
|
||||
signature — see `giant.model.routers.build_router`)? Docstrings are excluded
|
||||
from the string-constant scan so prose mentioning a dotted config path in
|
||||
passing can't masquerade as a read of it. This whitelist-based approach is
|
||||
deliberately narrower than "anywhere in `giant/`": scanning the whole package
|
||||
produces false negatives from unrelated identifier collisions (e.g.
|
||||
`giant/analysis/router_gating.py`'s `_top1_shares(..., order: list, ...)`
|
||||
parameter would otherwise make `stage2_model.autoregressive.order` read as
|
||||
"consumed").
|
||||
"""
|
||||
|
||||
import ast
|
||||
from pathlib import Path
|
||||
|
||||
from giant.config import DEFAULT_CONFIG
|
||||
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[1]
|
||||
|
||||
# Files that legitimately consume model_config / training config at
|
||||
# build/train/rollout time. Not `giant/cli.py` (a CLI flag existing is not
|
||||
# consumption — that's precisely how Issue 1 slipped through), not
|
||||
# `giant/config.py` itself (declaring/parsing a field is not reading it), and
|
||||
# not `giant/model/_legacy.py` (the protected v0.2 migration surface, which
|
||||
# intentionally re-derives old flat keys under old names).
|
||||
_CONSUMER_ROOTS = ("giant/model", "giant/training")
|
||||
_CONSUMER_FILES = (
|
||||
"giant/sample.py",
|
||||
"giant/pipeline.py",
|
||||
"giant/rollout.py",
|
||||
"giant/checkpoint_io.py",
|
||||
"giant/particles.py",
|
||||
"giant/materials.py",
|
||||
)
|
||||
_EXCLUDED_FILES = ("giant/model/_legacy.py",)
|
||||
|
||||
# Leaf DEFAULT_CONFIG paths that are declared but not (yet) read anywhere in
|
||||
# the consumer whitelist above. Each entry must name the issue that tracks
|
||||
# it. If a key here starts showing up as consumed, the fix landed and this
|
||||
# entry is stale — see test_known_unused_allow_list_has_no_stale_entries.
|
||||
_KNOWN_UNUSED = {
|
||||
"stage2_model.stage1_context": (
|
||||
"issues.md Issue 1 — trainers.py hardcodes stage1_ctx to the "
|
||||
"ground-truth stage-1 output; 'sampled' is now rejected loudly by "
|
||||
"validate_config (not silently accepted), but the key still isn't "
|
||||
"read by any build/train consumer file since only 'truth' can pass "
|
||||
"validation — see Issue 16 for the real implementation"
|
||||
),
|
||||
"stage2_model.autoregressive.order": (
|
||||
"gitea #30 — validate_config now checks order is 'energy_desc', but "
|
||||
"nothing in the build/train/rollout consumer whitelist reads the "
|
||||
"value itself since it's still single-valued"
|
||||
),
|
||||
}
|
||||
|
||||
# "lambda" is a Python keyword, so the dataclasses expose the dict key
|
||||
# "lambda" as the field `lambda_weight` (giant/config.py:49-50).
|
||||
_FIELD_NAME_OVERRIDES = {"lambda": "lambda_weight"}
|
||||
|
||||
|
||||
def _leaf_paths(node: dict, prefix: str = "") -> list[str]:
|
||||
paths = []
|
||||
for key, value in node.items():
|
||||
if prefix == "" and key == "meta":
|
||||
continue
|
||||
path = f"{prefix}.{key}" if prefix else key
|
||||
if isinstance(value, dict):
|
||||
paths.extend(_leaf_paths(value, path))
|
||||
else:
|
||||
paths.append(path)
|
||||
return paths
|
||||
|
||||
|
||||
def _field_name(leaf_path: str) -> str:
|
||||
name = leaf_path.rsplit(".", 1)[-1]
|
||||
return _FIELD_NAME_OVERRIDES.get(name, name)
|
||||
|
||||
|
||||
def _is_docstring_expr(expr: ast.Expr) -> bool:
|
||||
return isinstance(expr.value, ast.Constant) and isinstance(expr.value.value, str)
|
||||
|
||||
|
||||
def _collect_names(source: str, filename: str) -> set[str]:
|
||||
tree = ast.parse(source, filename=filename)
|
||||
docstring_ids = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, (ast.Module, ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)):
|
||||
body = getattr(node, "body", [])
|
||||
if body and isinstance(body[0], ast.Expr) and _is_docstring_expr(body[0]):
|
||||
docstring_ids.add(id(body[0].value))
|
||||
|
||||
names: set[str] = set()
|
||||
for node in ast.walk(tree):
|
||||
if isinstance(node, ast.Attribute):
|
||||
names.add(node.attr)
|
||||
elif isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstring_ids:
|
||||
names.add(node.value)
|
||||
elif isinstance(node, ast.arg):
|
||||
names.add(node.arg)
|
||||
elif isinstance(node, ast.keyword) and node.arg is not None:
|
||||
names.add(node.arg)
|
||||
return names
|
||||
|
||||
|
||||
def _consumer_files() -> list[Path]:
|
||||
files: set[Path] = {_REPO_ROOT / f for f in _CONSUMER_FILES}
|
||||
for root in _CONSUMER_ROOTS:
|
||||
files |= set((_REPO_ROOT / root).rglob("*.py"))
|
||||
files -= {_REPO_ROOT / f for f in _EXCLUDED_FILES}
|
||||
return sorted(files)
|
||||
|
||||
|
||||
def _consumed_names() -> set[str]:
|
||||
names: set[str] = set()
|
||||
for path in _consumer_files():
|
||||
names |= _collect_names(path.read_text(), str(path))
|
||||
return names
|
||||
|
||||
|
||||
def test_every_config_key_is_consumed_or_allow_listed():
|
||||
consumed = _consumed_names()
|
||||
unconsumed = {p for p in _leaf_paths(DEFAULT_CONFIG) if _field_name(p) not in consumed}
|
||||
unexplained = unconsumed - _KNOWN_UNUSED.keys()
|
||||
assert not unexplained, (
|
||||
f"config key(s) {sorted(unexplained)} are declared in DEFAULT_CONFIG "
|
||||
"but not read anywhere in the build/train/rollout consumer files "
|
||||
f"({[str(f.relative_to(_REPO_ROOT)) for f in _consumer_files()]}) — "
|
||||
"either wire the key up, or add it to _KNOWN_UNUSED with a reason "
|
||||
"(see issues.md Issue 5)"
|
||||
)
|
||||
|
||||
|
||||
def test_known_unused_allow_list_has_no_stale_entries():
|
||||
consumed = _consumed_names()
|
||||
all_paths = set(_leaf_paths(DEFAULT_CONFIG))
|
||||
stale = {p for p in _KNOWN_UNUSED if p not in all_paths or _field_name(p) in consumed}
|
||||
assert not stale, (
|
||||
f"_KNOWN_UNUSED entry/entries {sorted(stale)} no longer belong on the "
|
||||
"allow-list — either the key was removed from DEFAULT_CONFIG, or it "
|
||||
"is now consumed (the underlying issue was fixed). Remove the stale "
|
||||
"entry/entries."
|
||||
)
|
||||
@@ -4,7 +4,7 @@ from pathlib import Path
|
||||
|
||||
import pytest
|
||||
|
||||
from giant.tools import create_root_files
|
||||
from scripts import create_root_files
|
||||
|
||||
parse_detector_spec = create_root_files.parse_detector_spec
|
||||
next_shard_index = create_root_files.next_shard_index
|
||||
@@ -16,7 +16,9 @@ SimJob = create_root_files.SimJob
|
||||
PlanError = create_root_files.PlanError
|
||||
|
||||
|
||||
def _write_fake_executable(path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0) -> Path:
|
||||
def _write_fake_executable(
|
||||
path: Path, *, output_count: int = 1, exit_code: int = 0, sleep: float = 0.0
|
||||
) -> Path:
|
||||
"""Stand-in for run_pbwo4/run_sampling: writes *output_count* .root files
|
||||
into its own cwd (so callers can verify each job gets an isolated workdir
|
||||
and that the workdir ends up holding *only* the .root output, matching
|
||||
@@ -87,13 +89,17 @@ def test_next_shard_index_continues_past_existing(tmp_path):
|
||||
|
||||
def test_plan_jobs_rejects_missing_gen(tmp_path):
|
||||
with pytest.raises(PlanError):
|
||||
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1")
|
||||
plan_jobs(
|
||||
["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="gen1"
|
||||
)
|
||||
|
||||
|
||||
def test_plan_jobs_rejects_malformed_gen(tmp_path):
|
||||
(tmp_path / "raw" / "steps" / "gen1").mkdir(parents=True)
|
||||
with pytest.raises(PlanError):
|
||||
plan_jobs(["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen")
|
||||
plan_jobs(
|
||||
["pbwo4"], num_files=2, dataset_root=tmp_path, kind="steps", gen="notagen"
|
||||
)
|
||||
|
||||
|
||||
def test_plan_jobs_continues_from_existing_shards(tmp_path):
|
||||
@@ -102,7 +108,9 @@ def test_plan_jobs_continues_from_existing_shards(tmp_path):
|
||||
(gen_dir / "pbwo4" / "shard-000.root").touch()
|
||||
(gen_dir / "pbwo4" / "shard-001.root").touch()
|
||||
|
||||
jobs = plan_jobs(["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1")
|
||||
jobs = plan_jobs(
|
||||
["pbwo4"], num_files=3, dataset_root=tmp_path, kind="steps", gen="gen1"
|
||||
)
|
||||
|
||||
assert [j.shard_index for j in jobs] == [2, 3, 4]
|
||||
assert all(j.detector == "pbwo4" and j.config is None for j in jobs)
|
||||
@@ -312,7 +320,9 @@ def test_run_all_caps_concurrency(tmp_path):
|
||||
assert {d.name for d in dests} == {f"shard-{i:03d}.root" for i in range(6)}
|
||||
|
||||
intervals = [json.loads(d.read_text()) for d in dests]
|
||||
events = sorted([(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals])
|
||||
events = sorted(
|
||||
[(p["start"], 1) for p in intervals] + [(p["end"], -1) for p in intervals]
|
||||
)
|
||||
concurrent = 0
|
||||
peak = 0
|
||||
for _, delta in events:
|
||||
|
||||
@@ -95,7 +95,7 @@ def _dummy_normalizer(width):
|
||||
|
||||
def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
|
||||
"""Two files that each restart event_id from 0 (one Geant4 job per file,
|
||||
see giant/tools/steps_to_parquet.py) must not have their same-numbered events
|
||||
see scripts/steps_to_parquet.py) must not have their same-numbered events
|
||||
collapsed together: every row from every file must show up in exactly one
|
||||
of train/val, and the number of distinct events must be the sum across
|
||||
files, not the union of raw ids."""
|
||||
@@ -126,8 +126,7 @@ def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path):
|
||||
target_normalizer=tgt_norm,
|
||||
batch_size=4,
|
||||
shuffle=False,
|
||||
particle_conditioning="embedding",
|
||||
material_conditioning="embedding",
|
||||
conditioning="embedding",
|
||||
)
|
||||
return sum(len(batch[0]) for batch in ds)
|
||||
|
||||
|
||||
+9
-7
@@ -3,15 +3,15 @@ from typer.testing import CliRunner
|
||||
from giant import cli as giant_cli
|
||||
from giant.config import Conditioning
|
||||
from giant.data import setup_cache
|
||||
from giant.tools import dwarf
|
||||
from giant.tools.dwarf import app
|
||||
from scripts import dwarf
|
||||
from scripts.dwarf import app
|
||||
from test_pipeline import _make_synthetic_steps
|
||||
|
||||
runner = CliRunner()
|
||||
|
||||
|
||||
def test_conditioning_enum_shared_across_both_clis():
|
||||
"""giant.cli and giant.tools.dwarf must use the one giant.config.Conditioning
|
||||
"""giant.cli and scripts.dwarf must use the one giant.config.Conditioning
|
||||
enum, not independently redefined copies that could silently drift apart
|
||||
on valid --conditioning values."""
|
||||
assert dwarf.Conditioning is Conditioning
|
||||
@@ -51,7 +51,9 @@ def test_convert_rejects_output_with_multiple_files(tmp_path):
|
||||
def test_convert_rejects_output_with_parallel_jobs(tmp_path):
|
||||
root_file = tmp_path / "shard.root"
|
||||
root_file.touch()
|
||||
result = runner.invoke(app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"])
|
||||
result = runner.invoke(
|
||||
app, ["convert", str(root_file), "--output", "out.parquet", "--jobs", "2"]
|
||||
)
|
||||
assert result.exit_code != 0
|
||||
assert "--output cannot be combined with --jobs > 1" in result.output
|
||||
|
||||
@@ -101,7 +103,7 @@ def test_warm_cache_writes_sidecar(tmp_path):
|
||||
assert loaded is not None
|
||||
assert loaded.vocab is not None
|
||||
assert loaded.event_index is not None
|
||||
assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
|
||||
assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
|
||||
|
||||
|
||||
def test_warm_cache_second_run_hits_cache(tmp_path):
|
||||
@@ -165,5 +167,5 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path):
|
||||
assert "fitting normalizer (streaming)" in result.output
|
||||
loaded = setup_cache.load(data, [data])
|
||||
assert loaded is not None
|
||||
assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
|
||||
assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers
|
||||
assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers
|
||||
assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers
|
||||
|
||||
+4
-15
@@ -1,23 +1,12 @@
|
||||
import torch
|
||||
from giant.constants import COND_DIM
|
||||
from giant.model.network import Stage1Model
|
||||
from giant.model.network import DenoisingMLP
|
||||
from giant.model.schedule import CosineSchedule, flow_matching_loss
|
||||
from giant.sample import sample_flow, sample_ddim
|
||||
|
||||
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
|
||||
|
||||
def _small_model():
|
||||
return Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
n_sec_head_k_max=15,
|
||||
)
|
||||
return DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
||||
|
||||
|
||||
def _batch(B=8):
|
||||
@@ -52,7 +41,7 @@ def test_sample_flow_shape():
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5)
|
||||
assert sample.shape == (B, 9)
|
||||
assert n_sec is not None and n_sec.shape == (B,)
|
||||
assert n_sec.shape == (B,)
|
||||
|
||||
|
||||
def test_ddpm_loss_nonneg():
|
||||
@@ -69,4 +58,4 @@ def test_sample_ddim_shape():
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5)
|
||||
assert sample.shape == (B, 9)
|
||||
assert n_sec is not None and n_sec.shape == (B,)
|
||||
assert n_sec.shape == (B,)
|
||||
|
||||
+6
-67
@@ -4,7 +4,6 @@ from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import numpy as np
|
||||
import pandas as pd
|
||||
import pytest
|
||||
|
||||
from giant import geometry as g
|
||||
@@ -13,70 +12,6 @@ from giant import geometry as g
|
||||
pytest.importorskip("sklearn")
|
||||
|
||||
|
||||
def _steps_frame(n=5, with_post=True):
|
||||
rng = np.random.default_rng(0)
|
||||
data = {
|
||||
"pre_x": rng.uniform(-10, 10, n),
|
||||
"pre_y": rng.uniform(-10, 10, n),
|
||||
"pre_z": rng.uniform(-10, 10, n),
|
||||
"material": ["G4_AIR"] * n,
|
||||
"layer_id": np.arange(n, dtype=np.int64),
|
||||
}
|
||||
if with_post:
|
||||
data["post_x"] = rng.uniform(-10, 10, n)
|
||||
data["post_y"] = rng.uniform(-10, 10, n)
|
||||
data["post_z"] = rng.uniform(-10, 10, n)
|
||||
return pd.DataFrame(data)
|
||||
|
||||
|
||||
def test_iter_point_batches_missing_columns_raises(tmp_path):
|
||||
path = tmp_path / "steps.parquet"
|
||||
pd.DataFrame({"pre_x": [0.0]}).to_parquet(path)
|
||||
with pytest.raises(ValueError, match="missing columns"):
|
||||
next(g._iter_point_batches(path))
|
||||
|
||||
|
||||
def test_iter_point_batches_without_post_columns_yields_pre_only(tmp_path):
|
||||
path = tmp_path / "steps.parquet"
|
||||
df = _steps_frame(n=5, with_post=False)
|
||||
df.to_parquet(path)
|
||||
|
||||
(pos, mat, lay) = next(g._iter_point_batches(path))
|
||||
|
||||
assert pos.shape == (5, 3)
|
||||
np.testing.assert_allclose(pos, df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
|
||||
assert list(mat) == ["G4_AIR"] * 5
|
||||
np.testing.assert_array_equal(lay, np.arange(5))
|
||||
|
||||
|
||||
def test_iter_point_batches_with_post_columns_doubles_and_concatenates_points(
|
||||
tmp_path,
|
||||
):
|
||||
path = tmp_path / "steps.parquet"
|
||||
df = _steps_frame(n=5, with_post=True)
|
||||
df.to_parquet(path)
|
||||
|
||||
(pos, mat, lay) = next(g._iter_point_batches(path))
|
||||
|
||||
# Every step contributes both its pre_pos and post_pos, sharing the
|
||||
# step's material/layer_id label — so batches double in length.
|
||||
assert pos.shape == (10, 3)
|
||||
np.testing.assert_allclose(pos[:5], df[["pre_x", "pre_y", "pre_z"]].to_numpy(dtype=np.float32))
|
||||
np.testing.assert_allclose(pos[5:], df[["post_x", "post_y", "post_z"]].to_numpy(dtype=np.float32))
|
||||
assert list(mat) == ["G4_AIR"] * 10
|
||||
np.testing.assert_array_equal(lay, np.concatenate([np.arange(5), np.arange(5)]))
|
||||
|
||||
|
||||
def test_iter_point_batches_respects_batch_size(tmp_path):
|
||||
path = tmp_path / "steps.parquet"
|
||||
df = _steps_frame(n=10, with_post=False)
|
||||
df.to_parquet(path, row_group_size=10)
|
||||
|
||||
batches = list(g._iter_point_batches(path, batch_size=4))
|
||||
|
||||
assert [len(pos) for pos, _, _ in batches] == [4, 4, 2]
|
||||
|
||||
|
||||
def _box_batch(n, rng):
|
||||
"""A labelled point cloud: inside a 100mm box -> PbWO4/0, else AIR/-1."""
|
||||
pos = rng.uniform(-200, 200, (n, 3)).astype(np.float32)
|
||||
@@ -202,7 +137,9 @@ def test_slab_classes_discovered():
|
||||
|
||||
def test_slab_query_labels_by_depth():
|
||||
orc = _build_slab()
|
||||
pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]) # layer 0, gap, layer 1
|
||||
pos = np.array(
|
||||
[[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0]]
|
||||
) # layer 0, gap, layer 1
|
||||
material, layer_id, escaped = orc.query(pos)
|
||||
assert list(material) == ["G4_PbWO4", "G4_AIR", "G4_W"]
|
||||
assert list(layer_id) == [0, -1, 1]
|
||||
@@ -231,7 +168,9 @@ def test_slab_save_load_roundtrip(tmp_path):
|
||||
orc.save(p)
|
||||
loaded = g.GeometryOracle.load(p)
|
||||
|
||||
pos = np.array([[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]])
|
||||
pos = np.array(
|
||||
[[0.0, 0.0, 50.0], [0.0, 0.0, 105.0], [0.0, 0.0, 150.0], [0.0, 0.0, 1e5]]
|
||||
)
|
||||
m0, l0, e0 = orc.query(pos)
|
||||
m1, l1, e1 = loaded.query(pos)
|
||||
assert (m0 == m1).all() and (l0 == l1).all() and (e0 == e1).all()
|
||||
|
||||
+6
-61
@@ -6,9 +6,7 @@ from giant.data.loader import (
|
||||
EVENT_ID_FILE_STRIDE,
|
||||
build_index_maps,
|
||||
build_index_maps_from_files,
|
||||
build_pdg_topn_map_from_files,
|
||||
build_process_map_from_files,
|
||||
build_topn_map_from_files,
|
||||
event_id_offset,
|
||||
find_parquet_files,
|
||||
iter_cond_chunks,
|
||||
@@ -180,63 +178,6 @@ def test_build_process_map_from_files_three_files_partial_overlap(tmp_path):
|
||||
assert proc_map["compt"] == 2
|
||||
|
||||
|
||||
# ── build_topn_map_from_files / build_pdg_topn_map_from_files ──────────────
|
||||
|
||||
|
||||
def test_build_topn_map_from_files_keeps_most_frequent(tmp_path):
|
||||
materials = ["G4_AIR"] * 5 + ["PbWO4"] * 3 + ["G4_Fe"] * 2 + ["G4_Pb"] * 1
|
||||
path = tmp_path / "a.parquet"
|
||||
pd.DataFrame({"material": materials}).to_parquet(path)
|
||||
|
||||
m = build_topn_map_from_files([path], "material", n_classes=3, cast=str)
|
||||
|
||||
assert m.class_map["G4_AIR"] == 0
|
||||
assert m.class_map["PbWO4"] == 1
|
||||
assert m.class_map["G4_Fe"] == 2 # "other" (n_classes - 1)
|
||||
assert m.class_map["G4_Pb"] == 2
|
||||
assert m.other_members == {"G4_Fe": 2, "G4_Pb": 1}
|
||||
|
||||
|
||||
def test_build_topn_map_from_files_fewer_values_than_n_classes(tmp_path):
|
||||
path = tmp_path / "a.parquet"
|
||||
pd.DataFrame({"material": ["G4_AIR", "PbWO4"]}).to_parquet(path)
|
||||
|
||||
m = build_topn_map_from_files([path], "material", n_classes=5, cast=str)
|
||||
|
||||
assert m.class_map == {"G4_AIR": 0, "PbWO4": 1}
|
||||
assert m.other_members == {}
|
||||
|
||||
|
||||
def test_build_pdg_topn_map_from_files_pools_primary_and_secondary_pdg(tmp_path):
|
||||
"""A species that's rare as a primary but common as a secondary must
|
||||
still rank by its pooled (primary + secondary) count, not just its
|
||||
primary-role count alone — the whole point of pooling both roles."""
|
||||
path = tmp_path / "a.parquet"
|
||||
# primary pdg: mostly 11 (electron), one lone 22 (photon)
|
||||
pdg = [11] * 5 + [22] * 1
|
||||
# secondaries: 22 (photon) appears often as a secondary despite being
|
||||
# rare as a primary above
|
||||
sec_pdg_list = [[22, 22]] * 5 + [[]] * 1
|
||||
pd.DataFrame({"pdg": pdg, "sec_pdg_list": sec_pdg_list}).to_parquet(path)
|
||||
|
||||
m = build_pdg_topn_map_from_files([path], n_classes=3)
|
||||
|
||||
# pooled: 11 -> 5, 22 -> 1 (primary) + 10 (secondary) = 11
|
||||
assert m.class_map[22] == 0
|
||||
assert m.class_map[11] == 1
|
||||
|
||||
|
||||
def test_build_pdg_topn_map_from_files_missing_sec_pdg_list_column(tmp_path):
|
||||
"""Files predating the parent->child join have no sec_pdg_list column —
|
||||
must not raise, just count the primary pdg column alone."""
|
||||
path = tmp_path / "a.parquet"
|
||||
pd.DataFrame({"pdg": [11, 11, 22]}).to_parquet(path)
|
||||
|
||||
m = build_pdg_topn_map_from_files([path], n_classes=3)
|
||||
|
||||
assert m.class_map == {11: 0, 22: 1}
|
||||
|
||||
|
||||
# ── build_index_maps (in-memory) ────────────────────────────────────────────
|
||||
|
||||
|
||||
@@ -315,7 +256,9 @@ def test_build_index_maps_from_files_ordering_independent_of_file_order(tmp_path
|
||||
|
||||
def test_build_index_maps_from_files_numeric_sort_for_nuclear_codes(tmp_path):
|
||||
path = tmp_path / "a.parquet"
|
||||
pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(path)
|
||||
pd.DataFrame({"pdg": [22, 1000060120, 11], "material": ["X", "X", "X"]}).to_parquet(
|
||||
path
|
||||
)
|
||||
|
||||
pdg_map, _ = build_index_maps_from_files([path])
|
||||
assert list(pdg_map.keys()) == [11, 22, 1000060120]
|
||||
@@ -396,7 +339,9 @@ def test_load_event_ids_applies_offset(tmp_path):
|
||||
path = tmp_path / "a.parquet"
|
||||
pd.DataFrame({"event_id": [0, 1, 2]}).to_parquet(path)
|
||||
offset = event_id_offset(1)
|
||||
np.testing.assert_array_equal(load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2])
|
||||
np.testing.assert_array_equal(
|
||||
load_event_ids(path, offset=offset), [offset, offset + 1, offset + 2]
|
||||
)
|
||||
|
||||
|
||||
def test_load_event_ids_raises_when_event_id_reaches_stride(tmp_path):
|
||||
|
||||
@@ -23,7 +23,11 @@ def test_get_material_properties_unfilled_entry_raises():
|
||||
|
||||
|
||||
def test_get_material_properties_returns_filled_entry_from_injected_table():
|
||||
table = {"G4_Pb": MaterialProperties(z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59)}
|
||||
table = {
|
||||
"G4_Pb": MaterialProperties(
|
||||
z_eff=82.0, a_eff=207.2, density=11.35, x0=0.5612, lambda_int=17.59
|
||||
)
|
||||
}
|
||||
props = get_material_properties("G4_Pb", table)
|
||||
assert props.z_eff == 82.0
|
||||
assert props.a_eff == 207.2
|
||||
|
||||
@@ -1,283 +0,0 @@
|
||||
"""Migration acceptance test for v0.3.0 step 2: "load a v0.2 checkpoint
|
||||
through migrate_config + the new build_models, and diff its outputs against
|
||||
v0.2 code on the same input batch — bit-identical, or the refactor has
|
||||
changed something it should not have."
|
||||
|
||||
No `/ceph` access on this machine (see CLAUDE.md's Compute environment
|
||||
section), so a real trained checkpoint can't be used here — a separate
|
||||
portal-machine follow-up with a real checkpoint is planned instead. This
|
||||
test is the synthetic stand-in: build a v0.2-shaped
|
||||
model from the frozen `tests/legacy/network_v02_snapshot.py` classes with
|
||||
fixed-seed random weights (playing the role of "a v0.2 checkpoint"), migrate
|
||||
its config and remap its state dict onto the new `build_models` output, and
|
||||
assert the two produce bit-identical output on the same random input batch.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
|
||||
from giant.model import network as net
|
||||
from tests.legacy import network_v02_snapshot as legacy
|
||||
|
||||
PDG_VOCAB = 12
|
||||
MAT_VOCAB = 4
|
||||
HIDDEN_DIM = 32
|
||||
N_BLOCKS = 2
|
||||
EMB_DIM = 8
|
||||
K = 6 # small k_max for a fast test
|
||||
BATCH = 5
|
||||
|
||||
|
||||
def _legacy_model_config(mode: str, conditioning: str) -> dict:
|
||||
return {
|
||||
"pdg_vocab": PDG_VOCAB,
|
||||
"mat_vocab": MAT_VOCAB,
|
||||
"hidden_dim": HIDDEN_DIM,
|
||||
"n_blocks": N_BLOCKS,
|
||||
"emb_dim": EMB_DIM,
|
||||
"dropout": 0.0,
|
||||
"k_max": K,
|
||||
"conditioning": conditioning,
|
||||
"router": {"enabled": False},
|
||||
"mode": mode,
|
||||
"noise_dim": 16,
|
||||
}
|
||||
|
||||
|
||||
def _random_batch(seed: int):
|
||||
g = torch.Generator().manual_seed(seed)
|
||||
cond_cont = torch.randn(BATCH, COND_DIM, generator=g)
|
||||
cond_cat = torch.randint(0, min(PDG_VOCAB, MAT_VOCAB), (BATCH, 2), generator=g)
|
||||
x1 = torch.randn(BATCH, X_DIM, generator=g)
|
||||
x2 = torch.randn(BATCH, K * SEC_SLOT_DIM, generator=g)
|
||||
t = torch.rand(BATCH, generator=g)
|
||||
return cond_cont, cond_cat, x1, x2, t
|
||||
|
||||
|
||||
def _assert_bit_identical(a: torch.Tensor, b: torch.Tensor, label: str) -> None:
|
||||
assert a.shape == b.shape, f"{label}: shape mismatch {a.shape} vs {b.shape}"
|
||||
assert torch.equal(a, b), f"{label}: outputs diverged, max abs diff = {(a - b).abs().max().item()}"
|
||||
|
||||
|
||||
def _run_migration_check(mode: str, conditioning: str) -> None:
|
||||
torch.manual_seed(0)
|
||||
legacy_cfg = _legacy_model_config(mode, conditioning)
|
||||
|
||||
if mode == "wgan":
|
||||
old_stage1 = legacy.WGANGenerator(
|
||||
pdg_vocab=PDG_VOCAB,
|
||||
mat_vocab=MAT_VOCAB,
|
||||
hidden_dim=HIDDEN_DIM,
|
||||
n_blocks=N_BLOCKS,
|
||||
emb_dim=EMB_DIM,
|
||||
noise_dim=16,
|
||||
dropout=0.0,
|
||||
k_max=K,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
old_stage2 = legacy.WGANSecondaryGenerator(
|
||||
pdg_vocab=PDG_VOCAB,
|
||||
mat_vocab=MAT_VOCAB,
|
||||
hidden_dim=HIDDEN_DIM,
|
||||
n_blocks=N_BLOCKS,
|
||||
emb_dim=EMB_DIM,
|
||||
sec_dim=K * SEC_SLOT_DIM,
|
||||
noise_dim=16,
|
||||
dropout=0.0,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
else:
|
||||
old_stage1 = legacy.DenoisingMLP(
|
||||
pdg_vocab=PDG_VOCAB,
|
||||
mat_vocab=MAT_VOCAB,
|
||||
hidden_dim=HIDDEN_DIM,
|
||||
n_blocks=N_BLOCKS,
|
||||
emb_dim=EMB_DIM,
|
||||
dropout=0.0,
|
||||
k_max=K,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
old_stage2 = legacy.SecondaryDecoder(
|
||||
pdg_vocab=PDG_VOCAB,
|
||||
mat_vocab=MAT_VOCAB,
|
||||
hidden_dim=HIDDEN_DIM,
|
||||
n_blocks=N_BLOCKS,
|
||||
emb_dim=EMB_DIM,
|
||||
sec_dim=K * SEC_SLOT_DIM,
|
||||
dropout=0.0,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
old_stage1.eval()
|
||||
old_stage2.eval()
|
||||
|
||||
cond_cont, cond_cat, x1, x2, t = _random_batch(seed=123)
|
||||
z1 = torch.randn(BATCH, 16, generator=torch.Generator().manual_seed(456))
|
||||
z2 = torch.randn(BATCH, 16, generator=torch.Generator().manual_seed(789))
|
||||
|
||||
with torch.no_grad():
|
||||
if mode == "wgan":
|
||||
old_out1 = old_stage1(z1, cond_cont, cond_cat)
|
||||
else:
|
||||
old_out1 = old_stage1(x1, t, cond_cont, cond_cat)
|
||||
old_n_sec = old_stage1.predict_n_sec(cond_cont, cond_cat)
|
||||
if mode == "wgan":
|
||||
old_out2 = old_stage2(z2, cond_cont, cond_cat, old_out1)
|
||||
else:
|
||||
old_out2 = old_stage2(x2, t, cond_cont, cond_cat, old_out1)
|
||||
|
||||
# --- migrate: config + state dict, through the new build_models ---
|
||||
new_models = net.build_models(legacy_cfg)
|
||||
new_stage1, new_stage2 = new_models["stage1"], new_models["stage2"]
|
||||
assert isinstance(new_stage1, net.Stage1Model)
|
||||
assert isinstance(new_stage2, net.Stage2OneShot)
|
||||
# n_sec.owner="stage1": n_sec lives on stage1, not stage2, for a
|
||||
# migrated v0.2 checkpoint.
|
||||
assert new_stage1.n_sec_head is not None
|
||||
assert new_stage2.n_sec_head is None
|
||||
|
||||
remapped1, remapped2 = net.migrate_legacy_state_dict(old_stage1.state_dict(), old_stage2.state_dict())
|
||||
missing1, unexpected1 = new_stage1.load_state_dict(remapped1, strict=True)
|
||||
missing2, unexpected2 = new_stage2.load_state_dict(remapped2, strict=True)
|
||||
assert not missing1 and not unexpected1
|
||||
assert not missing2 and not unexpected2
|
||||
new_stage1.eval()
|
||||
new_stage2.eval()
|
||||
|
||||
with torch.no_grad():
|
||||
if mode == "wgan":
|
||||
new_out1 = new_stage1(z1, cond_cont, cond_cat)
|
||||
else:
|
||||
new_out1 = new_stage1(x1, cond_cont, cond_cat, t=t)
|
||||
new_n_sec = new_stage1.predict_n_sec(cond_cont, cond_cat)
|
||||
if mode == "wgan":
|
||||
new_out2 = new_stage2(z2, cond_cont, cond_cat, new_out1)
|
||||
else:
|
||||
new_out2 = new_stage2(x2, cond_cont, cond_cat, new_out1, t=t)
|
||||
|
||||
_assert_bit_identical(old_out1, new_out1, f"stage1 output ({mode}, {conditioning})")
|
||||
_assert_bit_identical(old_n_sec, new_n_sec, f"n_sec logits ({mode}, {conditioning})")
|
||||
_assert_bit_identical(old_out2, new_out2, f"stage2 output ({mode}, {conditioning})")
|
||||
|
||||
|
||||
def test_migration_flow_embedding():
|
||||
_run_migration_check(mode="flow", conditioning="embedding")
|
||||
|
||||
|
||||
def test_migration_flow_physical():
|
||||
_run_migration_check(mode="flow", conditioning="physical")
|
||||
|
||||
|
||||
def test_migration_wgan_embedding():
|
||||
_run_migration_check(mode="wgan", conditioning="embedding")
|
||||
|
||||
|
||||
def test_migration_wgan_physical():
|
||||
_run_migration_check(mode="wgan", conditioning="physical")
|
||||
|
||||
|
||||
def test_migrate_legacy_model_config_shape():
|
||||
"""_migrate_legacy_model_config produces the nested shape build_models
|
||||
expects, with the n_sec.owner marker set so build_models routes the
|
||||
n_sec head back onto stage 1."""
|
||||
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
||||
migrated = net._migrate_legacy_model_config(legacy_cfg)
|
||||
assert migrated["pdg_vocab"] == PDG_VOCAB
|
||||
assert migrated["mat_vocab"] == MAT_VOCAB
|
||||
assert migrated["conditioning"]["particle"]["type"] == "physical"
|
||||
assert migrated["conditioning"]["particle"]["n_layers"] == 2
|
||||
assert migrated["conditioning"]["material"]["n_layers"] == 2
|
||||
assert migrated["stage1_model"]["hidden_dim"] == HIDDEN_DIM
|
||||
assert migrated["stage2_model"]["n_sec"]["owner"] == "stage1"
|
||||
assert migrated["stage2_model"]["decoder"] == "one_shot"
|
||||
|
||||
|
||||
def test_migrate_legacy_model_config_nonzero_expert_dims_raises():
|
||||
"""Regression: a v0.2 checkpoint's model_config carrying a non-default
|
||||
expert_hidden_dim/expert_n_blocks must fail loudly through this path too
|
||||
— not just giant.config.migrate_config's parallel TOML-load path.
|
||||
Silently dropping these keys (build_router's kwarg filtering) would
|
||||
resize the experts instead of refusing."""
|
||||
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
||||
legacy_cfg["router"] = {
|
||||
"enabled": True,
|
||||
"expert_hidden_dim": 128,
|
||||
"expert_n_blocks": 0,
|
||||
}
|
||||
try:
|
||||
net._migrate_legacy_model_config(legacy_cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "expert_hidden_dim" in str(e)
|
||||
|
||||
|
||||
def test_migrate_legacy_model_config_zero_expert_dims_dropped_silently():
|
||||
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
||||
legacy_cfg["router"] = {
|
||||
"enabled": True,
|
||||
"expert_hidden_dim": 0,
|
||||
"expert_n_blocks": 0,
|
||||
}
|
||||
migrated = net._migrate_legacy_model_config(legacy_cfg)
|
||||
assert "expert_hidden_dim" not in migrated["stage1_model"]["router"]
|
||||
assert "expert_n_blocks" not in migrated["stage1_model"]["router"]
|
||||
assert "expert_hidden_dim" not in migrated["stage2_model"]["router"]
|
||||
assert "expert_n_blocks" not in migrated["stage2_model"]["router"]
|
||||
|
||||
|
||||
def test_build_models_with_legacy_config_nonzero_expert_dims_raises():
|
||||
"""The same check must also fire through the actual caller,
|
||||
build_models, not just the internal helper directly."""
|
||||
legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical")
|
||||
legacy_cfg["router"] = {
|
||||
"enabled": True,
|
||||
"expert_hidden_dim": 128,
|
||||
"expert_n_blocks": 0,
|
||||
}
|
||||
try:
|
||||
net.build_models(legacy_cfg)
|
||||
assert False, "expected ValueError"
|
||||
except ValueError as e:
|
||||
assert "expert_hidden_dim" in str(e)
|
||||
|
||||
|
||||
def test_build_models_accepts_new_nested_shape_unchanged():
|
||||
"""A dict that already has a 'stage1_model' key (the new shape) is
|
||||
passed through build_models without going through the legacy migration
|
||||
path at all."""
|
||||
cfg = {
|
||||
"pdg_vocab": PDG_VOCAB,
|
||||
"mat_vocab": MAT_VOCAB,
|
||||
"conditioning": {
|
||||
"out_dim": 32,
|
||||
"particle": {"type": "physical", "emb_dim": EMB_DIM, "n_layers": 1},
|
||||
"material": {"type": "physical", "emb_dim": EMB_DIM, "n_layers": 1},
|
||||
},
|
||||
"stage1_model": {
|
||||
"active": True,
|
||||
"generator": "flow",
|
||||
"hidden_dim": HIDDEN_DIM,
|
||||
"n_res_blocks": N_BLOCKS,
|
||||
"dropout": 0.0,
|
||||
"flow": {"time_dim": 16},
|
||||
"router": {"enabled": False},
|
||||
},
|
||||
"stage2_model": {
|
||||
"active": True,
|
||||
"decoder": "one_shot",
|
||||
"generator": "flow",
|
||||
"hidden_dim": HIDDEN_DIM,
|
||||
"n_res_blocks": N_BLOCKS,
|
||||
"dropout": 0.0,
|
||||
"k_max": K,
|
||||
"context_dim": 16,
|
||||
"n_sec": {"mode": "head"},
|
||||
"flow": {"time_dim": 16},
|
||||
"router": {"enabled": False, "tie_to_stage1": False},
|
||||
},
|
||||
}
|
||||
models = net.build_models(cfg)
|
||||
assert isinstance(models["stage1"], net.Stage1Model)
|
||||
assert isinstance(models["stage2"], net.Stage2OneShot)
|
||||
# Fresh v0.3.0 config, n_sec.owner defaults to "stage2": n_sec lives on stage 2.
|
||||
assert models["stage1"].n_sec_head is None
|
||||
assert models["stage2"].n_sec_head is not None
|
||||
+8
-789
@@ -1,28 +1,6 @@
|
||||
import copy
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from giant import config as gconfig
|
||||
from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
|
||||
from giant.model.network import (
|
||||
AttentionHistory,
|
||||
ConditionEncoder,
|
||||
MarkovHistory,
|
||||
SinusoidalEmbedding,
|
||||
Stage1Model,
|
||||
Stage2Autoregressive,
|
||||
Stage2OneShot,
|
||||
build_critics,
|
||||
build_models,
|
||||
cat_col_layout,
|
||||
stage2_trunk_sec_dim,
|
||||
stage2_type_dim,
|
||||
)
|
||||
|
||||
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
ONEHOT_PARTICLE_CFG = {"type": "onehot", "emb_dim": 6, "n_layers": 1}
|
||||
ONEHOT_MATERIAL_CFG = {"type": "onehot", "emb_dim": 4, "n_layers": 1}
|
||||
from giant.constants import COND_DIM
|
||||
from giant.model.network import DenoisingMLP, SinusoidalEmbedding
|
||||
|
||||
|
||||
def test_sinusoidal_embedding_shape():
|
||||
@@ -37,15 +15,9 @@ def test_sinusoidal_embedding_batch_1():
|
||||
assert emb(t).shape == (1, 32)
|
||||
|
||||
|
||||
def test_stage1_model_output_shape():
|
||||
def test_denoising_mlp_output_shape():
|
||||
B = 8
|
||||
model = Stage1Model(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=3,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
n_sec_head_k_max=15,
|
||||
)
|
||||
model = DenoisingMLP(pdg_vocab=5, mat_vocab=3)
|
||||
x_t = torch.randn(B, 9)
|
||||
t = torch.rand(B)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
@@ -56,773 +28,20 @@ def test_stage1_model_output_shape():
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
out = model(x_t, cond_cont, cond_cat, t=t)
|
||||
out = model(x_t, t, cond_cont, cond_cat)
|
||||
assert out.shape == (B, 9)
|
||||
|
||||
|
||||
def test_stage1_model_gradients_flow():
|
||||
def test_denoising_mlp_gradients_flow():
|
||||
B = 4
|
||||
model = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
n_sec_head_k_max=15,
|
||||
)
|
||||
model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
||||
x_t = torch.randn(B, 9)
|
||||
t = torch.rand(B)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
# Both paths must be exercised to get gradients through all parameters.
|
||||
flow_loss = model(x_t, cond_cont, cond_cat, t=t).sum()
|
||||
flow_loss = model(x_t, t, cond_cont, cond_cat).sum()
|
||||
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
|
||||
(flow_loss + nsec_loss).backward()
|
||||
for name, p in model.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
|
||||
def test_stage1_model_no_n_sec_head_by_default():
|
||||
"""Fresh v0.3.0 construction (no n_sec_head_k_max) has no n_sec head —
|
||||
it moves to stage 2."""
|
||||
model = Stage1Model(pdg_vocab=3, mat_vocab=2, particle_cfg=PARTICLE_CFG, material_cfg=MATERIAL_CFG)
|
||||
assert model.n_sec_head is None
|
||||
|
||||
|
||||
# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim ---------------
|
||||
|
||||
|
||||
def test_cat_col_layout_neither_onehot():
|
||||
assert cat_col_layout("physical", "embedding") == (None, None)
|
||||
|
||||
|
||||
def test_cat_col_layout_particle_only():
|
||||
assert cat_col_layout("onehot", "physical") == (2, None)
|
||||
|
||||
|
||||
def test_cat_col_layout_material_only():
|
||||
assert cat_col_layout("physical", "onehot") == (None, 2)
|
||||
|
||||
|
||||
def test_cat_col_layout_both_onehot_particle_then_material():
|
||||
assert cat_col_layout("onehot", "onehot") == (2, 3)
|
||||
|
||||
|
||||
def test_stage2_type_dim_physical_is_particle_phys_dim():
|
||||
assert stage2_type_dim({"target": "physical"}, emb_dim=16) == PARTICLE_PHYS_DIM
|
||||
|
||||
|
||||
def test_stage2_type_dim_onehot_and_embedding_are_emb_dim():
|
||||
assert stage2_type_dim({"target": "onehot"}, emb_dim=16) == 16
|
||||
assert stage2_type_dim({"target": "embedding"}, emb_dim=16) == 16
|
||||
|
||||
|
||||
def test_stage2_trunk_sec_dim_physical_matches_v02_sec_dim():
|
||||
k_max = 15
|
||||
assert stage2_trunk_sec_dim({"target": "physical"}, "flow", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
|
||||
assert stage2_trunk_sec_dim({"target": "physical"}, "wgan", k_max, emb_dim=16) == k_max * SEC_SLOT_DIM
|
||||
|
||||
|
||||
def test_stage2_trunk_sec_dim_onehot_wgan_folds_type_in():
|
||||
k_max = 15
|
||||
assert stage2_trunk_sec_dim({"target": "onehot"}, "wgan", k_max, emb_dim=16) == k_max * (CONT_SLOT_DIM + 16)
|
||||
|
||||
|
||||
def test_stage2_trunk_sec_dim_onehot_flow_excludes_type():
|
||||
k_max = 15
|
||||
assert stage2_trunk_sec_dim({"target": "onehot"}, "flow", k_max, emb_dim=16) == k_max * CONT_SLOT_DIM
|
||||
|
||||
|
||||
# --- ConditionEncoder onehot mode -------------------------------------------
|
||||
|
||||
|
||||
def test_condition_encoder_onehot_forward_shape_and_gradients():
|
||||
B = 8
|
||||
particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
|
||||
material_emb_dim = int(ONEHOT_MATERIAL_CFG["emb_dim"])
|
||||
enc = ConditionEncoder(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=3,
|
||||
particle_cfg=ONEHOT_PARTICLE_CFG,
|
||||
material_cfg=ONEHOT_MATERIAL_CFG,
|
||||
out_dim=32,
|
||||
)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.stack(
|
||||
[
|
||||
torch.randint(0, 5, (B,)),
|
||||
torch.randint(0, 3, (B,)),
|
||||
torch.randint(0, particle_emb_dim, (B,)),
|
||||
torch.randint(0, material_emb_dim, (B,)),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
out = enc(cond_cont, cond_cat)
|
||||
assert out.shape == (B, 32)
|
||||
# onehot itself is unlearned, but the fusion MLP downstream still has
|
||||
# gradients — the encoder as a whole must still be trainable.
|
||||
out.sum().backward()
|
||||
assert enc.mlp[0].weight.grad is not None
|
||||
|
||||
|
||||
def test_condition_encoder_onehot_is_a_true_one_hot_vector():
|
||||
"""The onehot axis feeds a fixed, unlearned one-hot into the fusion MLP —
|
||||
verify the concatenated input segment really is one-hot, not e.g. an
|
||||
accidentally-learned embedding."""
|
||||
B = 4
|
||||
particle_emb_dim = int(ONEHOT_PARTICLE_CFG["emb_dim"])
|
||||
enc = ConditionEncoder(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=3,
|
||||
particle_cfg=ONEHOT_PARTICLE_CFG,
|
||||
material_cfg={"type": "physical", "emb_dim": 4, "n_layers": 1},
|
||||
out_dim=16,
|
||||
)
|
||||
cond_cont = torch.zeros(B, COND_DIM)
|
||||
idx = torch.tensor([0, 1, 2, 5])
|
||||
cond_cat = torch.stack(
|
||||
[
|
||||
torch.zeros(B, dtype=torch.long),
|
||||
torch.zeros(B, dtype=torch.long),
|
||||
idx.clamp(max=particle_emb_dim - 1),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
pdg_e = enc._particle_embed(cond_cont, cond_cat)
|
||||
assert pdg_e.shape == (B, particle_emb_dim)
|
||||
assert torch.all(pdg_e.sum(dim=-1) == 1.0)
|
||||
|
||||
|
||||
# --- Stage2OneShot particle_type architecture --------------------------------
|
||||
|
||||
|
||||
def _build_stage2(target: str, generator: str, emb_dim: int = 6) -> Stage2OneShot:
|
||||
particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
|
||||
if target != "physical":
|
||||
particle_cfg = dict(particle_cfg)
|
||||
if target == "embedding":
|
||||
particle_cfg["type"] = "embedding"
|
||||
k_max = 5
|
||||
sec_dim = stage2_trunk_sec_dim({"target": target}, generator, k_max, emb_dim)
|
||||
return Stage2OneShot(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=3,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
cond_out_dim=16,
|
||||
context_dim=8,
|
||||
sec_dim=sec_dim,
|
||||
generator=generator,
|
||||
k_max=k_max,
|
||||
particle_type_cfg={"target": target, "lambda": 1.0},
|
||||
)
|
||||
|
||||
|
||||
def test_stage2_oneshot_physical_has_no_type_head_regardless_of_generator():
|
||||
assert _build_stage2("physical", "flow").type_head is None
|
||||
assert _build_stage2("physical", "wgan").type_head is None
|
||||
|
||||
|
||||
def test_stage2_oneshot_onehot_flow_has_type_head():
|
||||
model = _build_stage2("onehot", "flow")
|
||||
assert model.type_head is not None
|
||||
|
||||
|
||||
def test_stage2_oneshot_onehot_wgan_has_no_type_head():
|
||||
"""Under wgan the type slice is folded into forward()'s own output and
|
||||
relaxed via ST-Gumbel by the trainer — no separate head needed."""
|
||||
model = _build_stage2("onehot", "wgan")
|
||||
assert model.type_head is None
|
||||
|
||||
|
||||
def test_stage2_oneshot_embedding_flow_has_type_head():
|
||||
model = _build_stage2("embedding", "flow")
|
||||
assert model.type_head is not None
|
||||
|
||||
|
||||
def test_stage2_oneshot_predict_type_shape():
|
||||
B, k_max, emb_dim = 4, 5, 6
|
||||
model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
out = model.predict_type(cond_cont, cond_cat, stage1_out)
|
||||
assert out.shape == (B, k_max, emb_dim)
|
||||
|
||||
|
||||
def test_stage2_oneshot_predict_type_raises_when_no_type_head():
|
||||
model = _build_stage2("physical", "flow")
|
||||
cond_cont = torch.randn(2, COND_DIM)
|
||||
cond_cat = torch.zeros(2, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(2, 9)
|
||||
try:
|
||||
model.predict_type(cond_cont, cond_cat, stage1_out)
|
||||
raise AssertionError("expected RuntimeError")
|
||||
except RuntimeError:
|
||||
pass
|
||||
|
||||
|
||||
def test_stage2_oneshot_forward_shape_onehot_wgan():
|
||||
B, k_max, emb_dim = 4, 5, 6
|
||||
model = _build_stage2("onehot", "wgan", emb_dim=emb_dim)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
z = torch.randn(B, model.noise_dim)
|
||||
out = model(z, cond_cont, cond_cat, stage1_out)
|
||||
assert out.shape == (B, k_max * (CONT_SLOT_DIM + emb_dim))
|
||||
|
||||
|
||||
def test_stage2_oneshot_forward_shape_onehot_flow_excludes_type():
|
||||
B, k_max, emb_dim = 4, 5, 6
|
||||
model = _build_stage2("onehot", "flow", emb_dim=emb_dim)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
x_t = torch.randn(B, k_max * CONT_SLOT_DIM)
|
||||
t = torch.rand(B)
|
||||
out = model(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
||||
assert out.shape == (B, k_max * CONT_SLOT_DIM)
|
||||
|
||||
|
||||
def test_stage2_oneshot_particle_type_n_classes_overrides_conditioning_emb_dim():
|
||||
"""gitea #29: stage2_model.particle_type.n_classes, not
|
||||
conditioning.particle.emb_dim, sizes the onehot type_head/type_dim when
|
||||
explicitly set — the two used to be silently the same number."""
|
||||
k_max = 5
|
||||
particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
|
||||
particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
|
||||
sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", k_max, 20)
|
||||
model = Stage2OneShot(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=3,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
cond_out_dim=16,
|
||||
context_dim=8,
|
||||
sec_dim=sec_dim,
|
||||
generator="flow",
|
||||
k_max=k_max,
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
)
|
||||
assert model.type_dim == 20 # not particle_cfg["emb_dim"] == 6
|
||||
assert model.type_head is not None
|
||||
assert model.type_head[-1].out_features == k_max * 20
|
||||
|
||||
|
||||
# --- MarkovHistory -----------------------------------------------------------
|
||||
|
||||
|
||||
def test_markov_history_shape():
|
||||
hist = MarkovHistory(in_dim=7, out_dim=12)
|
||||
B, K = 3, 5
|
||||
feat = torch.randn(B, K, 7)
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
out = hist(feat, has_prev)
|
||||
assert out.shape == (B, K, 12)
|
||||
|
||||
|
||||
def test_markov_history_uses_start_vector_when_no_prev():
|
||||
"""Slot 0's own raw feature must be ignored — a learned start vector is
|
||||
substituted there instead (a reasonable default, see
|
||||
Stage2Autoregressive's docstring)."""
|
||||
hist = MarkovHistory(in_dim=4, out_dim=6)
|
||||
B, K = 2, 3
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
feat_a = torch.randn(B, K, 4)
|
||||
feat_b = feat_a.clone()
|
||||
feat_b[:, 0] = torch.randn(B, 4) * 100
|
||||
out_a = hist(feat_a, has_prev)
|
||||
out_b = hist(feat_b, has_prev)
|
||||
assert torch.allclose(out_a[:, 0], out_b[:, 0])
|
||||
assert torch.allclose(out_a[:, 1:], out_b[:, 1:])
|
||||
|
||||
|
||||
# --- AttentionHistory (v0.3.0 step 7) ---------------------------------------
|
||||
|
||||
|
||||
def test_attention_history_shape():
|
||||
hist = AttentionHistory(in_dim=7, out_dim=12, n_heads=2, n_layers=2)
|
||||
B, K = 3, 5
|
||||
feat = torch.randn(B, K, 7)
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
out = hist(feat, has_prev)
|
||||
assert out.shape == (B, K, 12)
|
||||
|
||||
|
||||
def test_attention_history_uses_start_vector_when_no_prev():
|
||||
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=1)
|
||||
B, K = 2, 3
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
feat_a = torch.randn(B, K, 4)
|
||||
feat_b = feat_a.clone()
|
||||
feat_b[:, 0] = torch.randn(B, 4) * 100
|
||||
out_a = hist(feat_a, has_prev)
|
||||
out_b = hist(feat_b, has_prev)
|
||||
assert torch.allclose(out_a[:, 0], out_b[:, 0], atol=1e-5)
|
||||
|
||||
|
||||
def test_attention_history_is_causal():
|
||||
"""Position i's output must not depend on feat at positions > i — unlike
|
||||
MarkovHistory (which only ever looks at position i itself, already
|
||||
trivially "causal"), this is AttentionHistory's actual contribution:
|
||||
seeing the full prefix 0..i-1, never anything later."""
|
||||
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
|
||||
hist.eval()
|
||||
B, K = 2, 5
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
feat_a = torch.randn(B, K, 4)
|
||||
feat_b = feat_a.clone()
|
||||
feat_b[:, 3:] = torch.randn(B, K - 3, 4) * 100
|
||||
with torch.no_grad():
|
||||
out_a = hist(feat_a, has_prev)
|
||||
out_b = hist(feat_b, has_prev)
|
||||
assert torch.allclose(out_a[:, :3], out_b[:, :3], atol=1e-5)
|
||||
|
||||
|
||||
def test_attention_history_step_matches_forward():
|
||||
"""The incremental KV-cache path (`init_cache`/`step`,
|
||||
`giant/sample.py`'s AR loop) must reproduce `forward`'s parallel-pass
|
||||
output exactly, one position at a time."""
|
||||
hist = AttentionHistory(in_dim=4, out_dim=6, n_heads=2, n_layers=2)
|
||||
hist.eval()
|
||||
B, K = 3, 6
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
feat = torch.randn(B, K, 4)
|
||||
with torch.no_grad():
|
||||
expected = hist(feat, has_prev)
|
||||
|
||||
cache = hist.init_cache()
|
||||
outs = []
|
||||
for k in range(K):
|
||||
out_k, cache = hist.step(feat[:, k : k + 1], has_prev[:, k : k + 1], cache)
|
||||
outs.append(out_k)
|
||||
stepped = torch.cat(outs, dim=1)
|
||||
|
||||
assert torch.allclose(stepped, expected, atol=1e-5)
|
||||
|
||||
|
||||
# --- Stage2Autoregressive (v0.3.0 step 5) -----------------------------------
|
||||
|
||||
|
||||
def _build_stage2_ar(
|
||||
target: str,
|
||||
generator: str,
|
||||
emb_dim: int = 6,
|
||||
k_max: int = 5,
|
||||
history: str = "markov",
|
||||
) -> Stage2Autoregressive:
|
||||
particle_cfg = {"type": "physical", "emb_dim": emb_dim, "n_layers": 1}
|
||||
if target == "embedding":
|
||||
particle_cfg = dict(particle_cfg)
|
||||
particle_cfg["type"] = "embedding"
|
||||
return Stage2Autoregressive(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=3,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
cond_out_dim=16,
|
||||
context_dim=8,
|
||||
generator=generator,
|
||||
k_max=k_max,
|
||||
particle_type_cfg={"target": target, "lambda": 1.0},
|
||||
history=history,
|
||||
)
|
||||
|
||||
|
||||
def _ar_inputs(B: int, K: int, hist_dim: int):
|
||||
history_feat = torch.randn(B, K, hist_dim)
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
remaining_frac = torch.rand(B, K)
|
||||
slot_idx = torch.linspace(0, 1, K).unsqueeze(0).expand(B, -1)
|
||||
return history_feat, has_prev, remaining_frac, slot_idx
|
||||
|
||||
|
||||
def test_stage2_autoregressive_history_invalid_raises():
|
||||
with pytest.raises(ValueError):
|
||||
_build_stage2_ar("onehot", "wgan", history="bogus")
|
||||
|
||||
|
||||
def test_stage2_autoregressive_particle_type_n_classes_overrides_conditioning_emb_dim():
|
||||
"""gitea #29, Stage2Autoregressive side — see the Stage2OneShot version
|
||||
of this test for the full rationale."""
|
||||
particle_cfg = {"type": "physical", "emb_dim": 6, "n_layers": 1}
|
||||
particle_type_cfg = {"target": "onehot", "lambda": 1.0, "n_classes": 20}
|
||||
model = Stage2Autoregressive(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=3,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
cond_out_dim=16,
|
||||
context_dim=8,
|
||||
generator="flow",
|
||||
k_max=5,
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
)
|
||||
assert model.type_dim == 20 # not particle_cfg["emb_dim"] == 6
|
||||
assert model.type_head is not None
|
||||
assert model.type_head[-1].out_features == 20
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
||||
@pytest.mark.parametrize("generator", ["wgan", "flow"])
|
||||
@pytest.mark.parametrize("history", ["markov", "attention"])
|
||||
def test_stage2_autoregressive_forward_shape(target, generator, history):
|
||||
B, K, emb_dim = 4, 5, 6
|
||||
model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K, history=history)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
type_dim = stage2_type_dim({"target": target}, emb_dim)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
|
||||
token_dim = stage2_trunk_sec_dim({"target": target}, generator, 1, emb_dim)
|
||||
if generator == "wgan":
|
||||
x_t = torch.randn(B, K, model.noise_dim)
|
||||
t = None
|
||||
else:
|
||||
x_t = torch.randn(B, K, token_dim)
|
||||
t = torch.rand(B, K)
|
||||
out = model(
|
||||
x_t,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
t=t,
|
||||
)
|
||||
assert out.shape == (B, K, token_dim)
|
||||
|
||||
|
||||
def test_stage2_autoregressive_predict_n_sec_shape():
|
||||
B, k_max, emb_dim = 4, 5, 6
|
||||
model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=k_max)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
logits = model.predict_n_sec(cond_cont, cond_cat, stage1_out)
|
||||
assert logits.shape == (B, k_max + 1)
|
||||
|
||||
|
||||
def test_stage2_autoregressive_predict_type_shape():
|
||||
B, K, emb_dim = 4, 5, 6
|
||||
model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
|
||||
out = model.predict_type(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
)
|
||||
assert out.shape == (B, K, emb_dim)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target,generator", [("physical", "flow"), ("onehot", "wgan")])
|
||||
def test_stage2_autoregressive_predict_type_raises_when_no_type_head(target, generator):
|
||||
B, K, emb_dim = 2, 5, 6
|
||||
model = _build_stage2_ar(target, generator, emb_dim=emb_dim, k_max=K)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
type_dim = stage2_type_dim({"target": target}, emb_dim)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
|
||||
with pytest.raises(RuntimeError):
|
||||
model.predict_type(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
)
|
||||
|
||||
|
||||
def test_stage2_autoregressive_gradients_flow_wgan_onehot():
|
||||
B, K, emb_dim = 4, 5, 6
|
||||
model = _build_stage2_ar("onehot", "wgan", emb_dim=emb_dim, k_max=K)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
|
||||
z = torch.randn(B, K, model.noise_dim)
|
||||
gen_out = model(
|
||||
z,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
).sum()
|
||||
nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
||||
(gen_out + nsec_out).backward()
|
||||
for name, p in model.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
|
||||
def test_stage2_autoregressive_gradients_flow_onehot():
|
||||
B, K, emb_dim = 4, 5, 6
|
||||
model = _build_stage2_ar("onehot", "flow", emb_dim=emb_dim, k_max=K)
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
stage1_out = torch.randn(B, 9)
|
||||
type_dim = stage2_type_dim({"target": "onehot"}, emb_dim)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_inputs(B, K, CONT_SLOT_DIM + type_dim)
|
||||
token_dim = stage2_trunk_sec_dim({"target": "onehot"}, "flow", 1, emb_dim)
|
||||
x_t = torch.randn(B, K, token_dim)
|
||||
t = torch.rand(B, K)
|
||||
flow_out = model(
|
||||
x_t,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
t=t,
|
||||
).sum()
|
||||
nsec_out = model.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
||||
type_out = model.predict_type(
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
).sum()
|
||||
(flow_out + nsec_out + type_out).backward()
|
||||
for name, p in model.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
|
||||
def test_stage2_autoregressive_history_step_matches_parallel_history_encoder():
|
||||
"""`init_history_cache`/`history_step` (the incremental path
|
||||
`giant/sample.py`'s AR loop drives, one slot per call) must reproduce
|
||||
exactly what one parallel `self.history_encoder(history_feat, has_prev)`
|
||||
call over the whole shifted sequence would give at each position — the
|
||||
KV-cache correctness guarantee, exercised through `Stage2Autoregressive`
|
||||
itself rather than `AttentionHistory` in isolation
|
||||
(`test_attention_history_step_matches_forward` covers that lower layer)."""
|
||||
B, K, emb_dim = 3, 6, 6
|
||||
model = _build_stage2_ar("physical", "wgan", emb_dim=emb_dim, k_max=K, history="attention")
|
||||
model.eval()
|
||||
type_dim = stage2_type_dim({"target": "physical"}, emb_dim)
|
||||
hist_in_dim = CONT_SLOT_DIM + type_dim
|
||||
own_feat = torch.randn(B, K, hist_in_dim) # token i's own raw feature
|
||||
has_prev_full = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
history_feat = torch.cat([torch.zeros_like(own_feat[:, :1]), own_feat[:, :-1]], dim=1)
|
||||
|
||||
with torch.no_grad():
|
||||
expected = model.history_encoder(history_feat, has_prev_full)
|
||||
|
||||
cache = model.init_history_cache()
|
||||
outs = []
|
||||
prev = torch.zeros(B, 1, hist_in_dim)
|
||||
for k in range(K):
|
||||
has_prev_k = torch.full((B, 1), k >= 1, dtype=torch.bool)
|
||||
hist_k, cache = model.history_step(prev, has_prev_k, cache)
|
||||
outs.append(hist_k)
|
||||
prev = own_feat[:, k : k + 1]
|
||||
stepped = torch.cat(outs, dim=1)
|
||||
|
||||
assert torch.allclose(stepped, expected, atol=1e-5)
|
||||
|
||||
|
||||
def test_stage2_autoregressive_init_history_cache_is_none_for_markov():
|
||||
model = _build_stage2_ar("physical", "wgan", history="markov")
|
||||
assert model.init_history_cache() is None
|
||||
|
||||
|
||||
# ── build_models: conditioning.share_stages ─────────────────────────────────
|
||||
|
||||
|
||||
def _minimal_model_config(share_stages: bool) -> dict:
|
||||
cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG)
|
||||
cfg["conditioning"]["share_stages"] = share_stages
|
||||
cfg["conditioning"]["particle"]["emb_dim"] = 4
|
||||
cfg["conditioning"]["material"]["emb_dim"] = 4
|
||||
cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1})
|
||||
cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3})
|
||||
return {
|
||||
"pdg_vocab": 3,
|
||||
"mat_vocab": 2,
|
||||
"conditioning": cfg["conditioning"],
|
||||
"stage1_model": cfg["stage1_model"],
|
||||
"stage2_model": cfg["stage2_model"],
|
||||
}
|
||||
|
||||
|
||||
def test_build_models_share_stages_true_shares_condition_encoder_instance():
|
||||
built = build_models(_minimal_model_config(share_stages=True))
|
||||
stage1, stage2 = built["stage1"], built["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
assert stage1.cond_enc is stage2.cond_enc
|
||||
|
||||
|
||||
def test_build_models_share_stages_false_builds_independent_condition_encoders():
|
||||
built = build_models(_minimal_model_config(share_stages=False))
|
||||
stage1, stage2 = built["stage1"], built["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
assert stage1.cond_enc is not stage2.cond_enc
|
||||
|
||||
|
||||
def test_build_models_share_stages_true_shared_params_are_in_both_stage_parameter_lists():
|
||||
"""The shared encoder's parameters must actually appear in both stages'
|
||||
own `.parameters()` — that's what makes each stage's independent
|
||||
optimizer include (and update) them, which is the actual mechanism behind
|
||||
"shared weights, forced common representation", not just object identity
|
||||
on `.cond_enc`."""
|
||||
built = build_models(_minimal_model_config(share_stages=True))
|
||||
stage1, stage2 = built["stage1"], built["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
|
||||
shared_ids = {id(p) for p in stage1.cond_enc.parameters()}
|
||||
assert shared_ids
|
||||
assert shared_ids <= {id(p) for p in stage1.parameters()}
|
||||
assert shared_ids <= {id(p) for p in stage2.parameters()}
|
||||
|
||||
|
||||
def test_build_models_particle_type_n_classes_overrides_conditioning_emb_dim():
|
||||
"""gitea #29 end-to-end through build_models: setting
|
||||
stage2_model.particle_type.n_classes independently of
|
||||
conditioning.particle.emb_dim actually resizes the built stage2 model,
|
||||
not just the two lower-level unit tests above."""
|
||||
cfg = _minimal_model_config(share_stages=False) # conditioning.particle.emb_dim = 4
|
||||
cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 11}
|
||||
built = build_models(cfg)
|
||||
assert built["stage2"] is not None
|
||||
assert built["stage2"].type_dim == 11
|
||||
|
||||
|
||||
def test_build_critics_particle_type_n_classes_overrides_conditioning_emb_dim():
|
||||
cfg = _minimal_model_config(share_stages=False) # conditioning.particle.emb_dim = 4
|
||||
cfg["stage2_model"]["generator"] = "wgan"
|
||||
cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0, "n_classes": 4}
|
||||
default_n_classes_critic = build_critics(cfg)["stage2"]
|
||||
assert default_n_classes_critic is not None
|
||||
|
||||
cfg["stage2_model"]["particle_type"]["n_classes"] = 11
|
||||
wider_critic = build_critics(cfg)["stage2"]
|
||||
assert wider_critic is not None
|
||||
# k_max=3 slots, each CONT_SLOT_DIM + n_classes wide under wgan folding —
|
||||
# widening n_classes alone (emb_dim stays 4) must widen the critic input.
|
||||
assert wider_critic.input_proj.in_features > default_n_classes_critic.input_proj.in_features
|
||||
|
||||
|
||||
# ── build_models/build_critics: DEFAULT_CONFIG fallback drift (issues.md #1) ─
|
||||
|
||||
|
||||
def _partial_model_config() -> dict:
|
||||
"""A hand-built model_config that omits stage2_model.decoder and
|
||||
stage2_model.particle_type — deliberately not derived from
|
||||
DEFAULT_CONFIG, unlike _minimal_model_config above. Regression fixture
|
||||
for issues.md Issue 1: build_models/build_critics/StageSpec.from_config's
|
||||
own fallback defaults for these two keys must equal DEFAULT_CONFIG's
|
||||
("autoregressive" / "onehot"), not the old, now-wrong v0.2-shaped
|
||||
("one_shot" / "physical") literals that used to live in three separate
|
||||
.get(key, default) call sites."""
|
||||
return {
|
||||
"pdg_vocab": 3,
|
||||
"mat_vocab": 2,
|
||||
"conditioning": {
|
||||
"particle": {"type": "physical", "emb_dim": 4, "n_layers": 1},
|
||||
"material": {"type": "physical", "emb_dim": 4, "n_layers": 1},
|
||||
},
|
||||
"stage1_model": {"active": False},
|
||||
"stage2_model": {
|
||||
"generator": "wgan",
|
||||
"hidden_dim": 8,
|
||||
"n_res_blocks": 1,
|
||||
"k_max": 3,
|
||||
# decoder and particle_type deliberately omitted
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_build_models_omitted_decoder_and_particle_type_match_default_config():
|
||||
built = build_models(_partial_model_config())
|
||||
assert isinstance(built["stage2"], Stage2Autoregressive)
|
||||
assert built["stage2"].particle_type_cfg["target"] == "onehot"
|
||||
|
||||
|
||||
def test_build_critics_omitted_particle_type_matches_default_config():
|
||||
cfg = _partial_model_config()
|
||||
cfg["stage2_model"]["generator"] = "wgan"
|
||||
onehot_critic = build_critics(cfg)["stage2"]
|
||||
assert onehot_critic is not None
|
||||
onehot_in_dim = onehot_critic.input_proj.in_features
|
||||
|
||||
cfg["stage2_model"]["particle_type"] = {"target": "physical"}
|
||||
physical_critic = build_critics(cfg)["stage2"]
|
||||
assert physical_critic is not None
|
||||
physical_in_dim = physical_critic.input_proj.in_features
|
||||
|
||||
# onehot's per-slot type width is emb_dim classes vs. physical's fixed
|
||||
# (log-mass, charge) pair — different unless emb_dim happens to be 2, so
|
||||
# this also confirms the critic was actually built in onehot mode by
|
||||
# default, not silently falling back to physical.
|
||||
assert onehot_in_dim != physical_in_dim
|
||||
|
||||
|
||||
# ── build_critics: critic_hidden_dim/critic_n_res_blocks honoured (gitea #28) ─
|
||||
|
||||
|
||||
def test_build_critics_stage1_critic_hidden_dim_and_n_res_blocks_override_generator_size():
|
||||
cfg = _minimal_model_config(share_stages=False)
|
||||
cfg["stage1_model"]["generator"] = "wgan"
|
||||
cfg["stage1_model"]["hidden_dim"] = 8
|
||||
cfg["stage1_model"]["n_res_blocks"] = 1
|
||||
|
||||
inherited = build_critics(cfg)["stage1"]
|
||||
assert inherited is not None
|
||||
assert inherited.input_proj.out_features == 8
|
||||
assert len(inherited.blocks) == 1
|
||||
|
||||
cfg["stage1_model"]["wgan"]["critic_hidden_dim"] = 16
|
||||
cfg["stage1_model"]["wgan"]["critic_n_res_blocks"] = 3
|
||||
overridden = build_critics(cfg)["stage1"]
|
||||
assert overridden is not None
|
||||
assert overridden.input_proj.out_features == 16
|
||||
assert len(overridden.blocks) == 3
|
||||
|
||||
|
||||
def test_build_critics_stage2_critic_hidden_dim_and_n_res_blocks_override_generator_size():
|
||||
cfg = _minimal_model_config(share_stages=False)
|
||||
cfg["stage2_model"]["generator"] = "wgan"
|
||||
cfg["stage2_model"]["hidden_dim"] = 8
|
||||
cfg["stage2_model"]["n_res_blocks"] = 1
|
||||
|
||||
inherited = build_critics(cfg)["stage2"]
|
||||
assert inherited is not None
|
||||
assert inherited.input_proj.out_features == 8
|
||||
assert len(inherited.blocks) == 1
|
||||
|
||||
cfg["stage2_model"]["wgan"]["critic_hidden_dim"] = 16
|
||||
cfg["stage2_model"]["wgan"]["critic_n_res_blocks"] = 3
|
||||
overridden = build_critics(cfg)["stage2"]
|
||||
assert overridden is not None
|
||||
assert overridden.input_proj.out_features == 16
|
||||
assert len(overridden.blocks) == 3
|
||||
|
||||
+3
-101
@@ -1,11 +1,7 @@
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
from giant.data.loader import TopNMap
|
||||
from giant.particles import (
|
||||
decode_embedding_nearest,
|
||||
decode_topn_class,
|
||||
invert_dense_map,
|
||||
nearest_known_pdg,
|
||||
particle_mass_charge,
|
||||
particle_phys_array,
|
||||
@@ -49,7 +45,9 @@ def test_ground_state_nucleus_resolved_via_particle_package():
|
||||
"""He-4 (Z=2, A=4) is a common nuclide in `particle`'s ground-state table."""
|
||||
mass, charge = particle_mass_charge(1000020040)
|
||||
assert charge == pytest.approx(2.0)
|
||||
assert mass == pytest.approx(4 * 931.494, rel=0.05) # near A*amu, binding-energy-corrected
|
||||
assert mass == pytest.approx(
|
||||
4 * 931.494, rel=0.05
|
||||
) # near A*amu, binding-energy-corrected
|
||||
|
||||
|
||||
def test_nuclear_isomer_falls_back_to_z_a_decode():
|
||||
@@ -128,99 +126,3 @@ def test_nearest_known_pdg_shape():
|
||||
)
|
||||
assert result.shape == (n,)
|
||||
assert set(result.tolist()) <= set(candidates)
|
||||
|
||||
|
||||
# ── invert_dense_map ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_invert_dense_map_round_trips():
|
||||
pdg_map = {22: 0, 11: 1, -11: 2, 2212: 3}
|
||||
inv = invert_dense_map(pdg_map)
|
||||
for pdg, idx in pdg_map.items():
|
||||
assert inv[idx] == pdg
|
||||
|
||||
|
||||
# ── decode_topn_class ────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _topn_fixture():
|
||||
# n_classes=4: photon/electron/positron get their own class (0,1,2),
|
||||
# everything else (proton, neutron) falls into "other" (class 3).
|
||||
class_map = {22: 0, 11: 1, -11: 2, 2212: 3, 2112: 3}
|
||||
other_members = {2212: 7, 2112: 3}
|
||||
return TopNMap(class_map=class_map, other_members=other_members), 4
|
||||
|
||||
|
||||
def test_decode_topn_class_known_classes_are_exact():
|
||||
topn_map, n_classes = _topn_fixture()
|
||||
out = decode_topn_class(np.array([0, 1, 2]), topn_map, n_classes)
|
||||
np.testing.assert_array_equal(out, [22, 11, -11])
|
||||
|
||||
|
||||
def test_decode_topn_class_other_modal_picks_most_frequent():
|
||||
topn_map, n_classes = _topn_fixture()
|
||||
out = decode_topn_class(np.array([3, 3]), topn_map, n_classes, other_policy="modal")
|
||||
assert (out == 2212).all() # count 7 > 3
|
||||
|
||||
|
||||
def test_decode_topn_class_other_drop_returns_zero_sentinel():
|
||||
topn_map, n_classes = _topn_fixture()
|
||||
out = decode_topn_class(np.array([3]), topn_map, n_classes, other_policy="drop")
|
||||
assert out[0] == 0
|
||||
|
||||
|
||||
def test_decode_topn_class_other_sample_stays_within_members():
|
||||
topn_map, n_classes = _topn_fixture()
|
||||
rng = np.random.default_rng(0)
|
||||
out = decode_topn_class(np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng)
|
||||
assert set(out.tolist()) <= {2212, 2112}
|
||||
|
||||
|
||||
def test_decode_topn_class_unknown_other_policy_raises():
|
||||
topn_map, n_classes = _topn_fixture()
|
||||
with pytest.raises(ValueError):
|
||||
decode_topn_class(np.array([3]), topn_map, n_classes, other_policy="bogus")
|
||||
|
||||
|
||||
def test_decode_topn_class_empty_other_members_raises():
|
||||
class_map = {22: 0, 11: 1}
|
||||
topn_map = TopNMap(class_map=class_map, other_members={})
|
||||
with pytest.raises(ValueError):
|
||||
decode_topn_class(np.array([1]), topn_map, 2, other_policy="sample")
|
||||
|
||||
|
||||
def test_decode_topn_class_preserves_shape():
|
||||
topn_map, n_classes = _topn_fixture()
|
||||
idx = np.array([[0, 1], [2, 3]])
|
||||
out = decode_topn_class(idx, topn_map, n_classes, other_policy="modal")
|
||||
assert out.shape == (2, 2)
|
||||
|
||||
|
||||
# ── decode_embedding_nearest ─────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_decode_embedding_nearest_exact_row_recovers_pdg():
|
||||
emb_weight = np.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]])
|
||||
idx_to_pdg = {0: 22, 1: 11, 2: 2212}
|
||||
vectors = np.array([[0.0, 1.0], [-1.0, -1.0]]) # exact rows 1, 2
|
||||
pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
|
||||
np.testing.assert_array_equal(pdg, [11, 2212])
|
||||
np.testing.assert_allclose(dist, [0.0, 0.0], atol=1e-8)
|
||||
|
||||
|
||||
def test_decode_embedding_nearest_off_manifold_snaps_to_closest_row():
|
||||
emb_weight = np.array([[1.0, 0.0], [0.0, 1.0]])
|
||||
idx_to_pdg = {0: 22, 1: 11}
|
||||
vectors = np.array([[0.9, 0.2]]) # closer to row 0
|
||||
pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
|
||||
assert pdg[0] == 22
|
||||
assert dist[0] > 0.0
|
||||
|
||||
|
||||
def test_decode_embedding_nearest_preserves_leading_shape():
|
||||
emb_weight = np.array([[1.0, 0.0], [0.0, 1.0]])
|
||||
idx_to_pdg = {0: 22, 1: 11}
|
||||
vectors = np.random.default_rng(0).standard_normal((3, 4, 2))
|
||||
pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
|
||||
assert pdg.shape == (3, 4)
|
||||
assert dist.shape == (3, 4)
|
||||
|
||||
+114
-150
@@ -4,64 +4,44 @@ import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from giant.constants import (
|
||||
COND_DIM,
|
||||
CONT_SLOT_DIM,
|
||||
K_MAX,
|
||||
PARTICLE_PHYS_DIM,
|
||||
SEC_DIM,
|
||||
X_DIM,
|
||||
)
|
||||
from giant.model.network import Stage1Model, Stage2Autoregressive, Stage2OneShot
|
||||
from giant.model.schedule import (
|
||||
flow_matching_loss_secondary,
|
||||
flow_matching_loss_secondary_ar,
|
||||
)
|
||||
from giant.constants import COND_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, X_DIM
|
||||
from giant.model.network import DenoisingMLP, SecondaryDecoder
|
||||
from giant.model.schedule import flow_matching_loss_secondary
|
||||
from giant.sample import sample_secondaries
|
||||
|
||||
|
||||
# ── helpers ──────────────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def _particle_material_cfg(conditioning: str) -> tuple[dict, dict]:
|
||||
cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
|
||||
return dict(cfg), dict(cfg)
|
||||
|
||||
|
||||
def _stage1(pdg=3, mat=2, conditioning="embedding"):
|
||||
particle_cfg, material_cfg = _particle_material_cfg(conditioning)
|
||||
return Stage1Model(
|
||||
return DenoisingMLP(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
n_sec_head_k_max=K_MAX,
|
||||
n_blocks=2,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
|
||||
|
||||
def _sec_decoder(pdg=3, mat=2, conditioning="embedding"):
|
||||
particle_cfg, material_cfg = _particle_material_cfg(conditioning)
|
||||
return Stage2OneShot(
|
||||
return SecondaryDecoder(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator="flow",
|
||||
time_dim=16,
|
||||
n_blocks=2,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
|
||||
|
||||
def _cond(B=8, pdg=3, mat=2):
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
|
||||
cond_cat = torch.stack(
|
||||
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
|
||||
)
|
||||
return cond_cont, cond_cat
|
||||
|
||||
|
||||
# ── Stage1Model Phase-2 additions ────────────────────────────────────────────
|
||||
# ── DenoisingMLP Phase-2 additions ───────────────────────────────────────────
|
||||
|
||||
|
||||
def test_predict_n_sec_shape():
|
||||
@@ -102,7 +82,7 @@ def test_condition_encoder_embedding_mode_has_embedding_tables():
|
||||
assert not hasattr(model.cond_enc, "particle_mlp")
|
||||
|
||||
|
||||
# ── Stage2OneShot ─────────────────────────────────────────────────────────────
|
||||
# ── SecondaryDecoder ──────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("conditioning", ["embedding", "physical"])
|
||||
@@ -113,7 +93,7 @@ def test_sec_decoder_output_shape(conditioning):
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
||||
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert out.shape == (B, SEC_DIM)
|
||||
|
||||
|
||||
@@ -124,7 +104,7 @@ def test_sec_decoder_no_nan():
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
||||
out = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
|
||||
@@ -135,9 +115,7 @@ def test_sec_decoder_gradients():
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
flow_out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t).sum()
|
||||
nsec_out = decoder.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
||||
(flow_out + nsec_out).backward()
|
||||
decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
|
||||
for name, p in decoder.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
@@ -152,7 +130,9 @@ def test_flow_matching_loss_secondary_scalar():
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
|
||||
loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
|
||||
loss = flow_matching_loss_secondary(
|
||||
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
|
||||
)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() >= 0.0
|
||||
|
||||
@@ -165,7 +145,9 @@ def test_flow_matching_loss_secondary_mask_zeros_padding():
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
sec_mask = torch.zeros(B, K_MAX, dtype=torch.bool)
|
||||
loss = flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask)
|
||||
loss = flow_matching_loss_secondary(
|
||||
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
|
||||
)
|
||||
assert loss.item() == pytest.approx(0.0, abs=1e-6)
|
||||
|
||||
|
||||
@@ -176,102 +158,8 @@ def test_flow_matching_loss_secondary_has_grad():
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
sec_mask = torch.ones(B, K_MAX, dtype=torch.bool)
|
||||
flow_matching_loss_secondary(decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask).backward()
|
||||
assert any(p.grad is not None for p in decoder.parameters())
|
||||
|
||||
|
||||
# ── masked flow matching loss — autoregressive (v0.3.0 step 5) ─────────────
|
||||
|
||||
|
||||
def _sec_decoder_ar(pdg=3, mat=2, k_max=K_MAX):
|
||||
particle_cfg, material_cfg = _particle_material_cfg("embedding")
|
||||
return Stage2Autoregressive(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator="flow",
|
||||
time_dim=16,
|
||||
k_max=k_max,
|
||||
)
|
||||
|
||||
|
||||
def _ar_history_inputs(B, K, hist_dim):
|
||||
history_feat = torch.randn(B, K, hist_dim)
|
||||
has_prev = (torch.arange(K) >= 1).unsqueeze(0).expand(B, -1)
|
||||
remaining_frac = torch.rand(B, K)
|
||||
slot_idx = torch.linspace(0, 1, K).unsqueeze(0).expand(B, -1)
|
||||
return history_feat, has_prev, remaining_frac, slot_idx
|
||||
|
||||
|
||||
def test_flow_matching_loss_secondary_ar_scalar():
|
||||
B, K, pdg, mat = 8, K_MAX, 3, 2
|
||||
decoder = _sec_decoder_ar(pdg, mat, k_max=K)
|
||||
x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
|
||||
sec_mask = torch.ones(B, K, dtype=torch.bool)
|
||||
loss = flow_matching_loss_secondary_ar(
|
||||
decoder,
|
||||
x1,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
sec_mask,
|
||||
)
|
||||
assert loss.shape == ()
|
||||
assert loss.item() >= 0.0
|
||||
|
||||
|
||||
def test_flow_matching_loss_secondary_ar_mask_zeros_padding():
|
||||
B, K, pdg, mat = 4, K_MAX, 3, 2
|
||||
decoder = _sec_decoder_ar(pdg, mat, k_max=K)
|
||||
x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
|
||||
sec_mask = torch.zeros(B, K, dtype=torch.bool)
|
||||
loss = flow_matching_loss_secondary_ar(
|
||||
decoder,
|
||||
x1,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
sec_mask,
|
||||
)
|
||||
assert loss.item() == pytest.approx(0.0, abs=1e-6)
|
||||
|
||||
|
||||
def test_flow_matching_loss_secondary_ar_has_grad():
|
||||
B, K, pdg, mat = 4, K_MAX, 3, 2
|
||||
decoder = _sec_decoder_ar(pdg, mat, k_max=K)
|
||||
x1 = torch.randn(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
history_feat, has_prev, remaining_frac, slot_idx = _ar_history_inputs(B, K, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
|
||||
sec_mask = torch.ones(B, K, dtype=torch.bool)
|
||||
flow_matching_loss_secondary_ar(
|
||||
decoder,
|
||||
x1,
|
||||
cond_cont,
|
||||
cond_cat,
|
||||
stage1_out,
|
||||
history_feat,
|
||||
has_prev,
|
||||
remaining_frac,
|
||||
slot_idx,
|
||||
sec_mask,
|
||||
flow_matching_loss_secondary(
|
||||
decoder, x1, cond_cont, cond_cat, stage1_out, sec_mask
|
||||
).backward()
|
||||
assert any(p.grad is not None for p in decoder.parameters())
|
||||
|
||||
@@ -285,7 +173,9 @@ def test_sample_secondaries_shapes():
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
|
||||
sec_cont, sec_phys, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3)
|
||||
sec_cont, sec_phys, sec_valid = sample_secondaries(
|
||||
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=3
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_phys.shape == (B, K_MAX, PARTICLE_PHYS_DIM)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
@@ -298,7 +188,9 @@ def test_sample_secondaries_valid_mask_matches_n_sec():
|
||||
cond_cont, cond_cat = _cond(B, pdg, mat)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.tensor([0, 1, 3, K_MAX])
|
||||
_, _, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
|
||||
_, _, sec_valid = sample_secondaries(
|
||||
decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2
|
||||
)
|
||||
for i, n in enumerate(n_sec_pred.tolist()):
|
||||
assert sec_valid[i, :n].all()
|
||||
assert not sec_valid[i, n:].any()
|
||||
@@ -332,7 +224,9 @@ def test_encode_secondaries_energy_conservation():
|
||||
pre_dir = rng.standard_normal((N, 3)).astype(np.float32)
|
||||
pre_dir /= np.linalg.norm(pre_dir, axis=1, keepdims=True)
|
||||
|
||||
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
|
||||
sec_cont = encode_secondaries(
|
||||
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
|
||||
)
|
||||
assert sec_cont.shape == (N, K_MAX, 6)
|
||||
assert np.isfinite(sec_cont).all()
|
||||
|
||||
@@ -379,7 +273,9 @@ def test_encode_secondaries_stick_logits_match_naive_reference():
|
||||
logit = np.clip(logit, -_STICK_LOGIT_CLIP, _STICK_LOGIT_CLIP)
|
||||
expected[row, i] = logit
|
||||
|
||||
np.testing.assert_allclose(stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4)
|
||||
np.testing.assert_allclose(
|
||||
stick_logits[sec_valid], expected[sec_valid], rtol=1e-4, atol=1e-4
|
||||
)
|
||||
|
||||
|
||||
def test_encode_secondaries_direction_encoding():
|
||||
@@ -491,7 +387,9 @@ def test_encode_secondaries_physical_columns_match_ground_truth_pdg():
|
||||
sec_valid[0, 0] = True
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
|
||||
sec_cont = encode_secondaries(
|
||||
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
|
||||
)
|
||||
mass, charge = particle_mass_charge(11)
|
||||
assert sec_cont[0, 0, 4] == pytest.approx(log_transform(np.array([mass]))[0])
|
||||
assert sec_cont[0, 0, 5] == pytest.approx(charge)
|
||||
@@ -525,7 +423,9 @@ def test_decode_secondaries_valid_slots_sum_to_e_sec():
|
||||
e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, e_sec, pre_dir
|
||||
)
|
||||
|
||||
valid_sum = (sec_E * sec_valid).sum(axis=1)
|
||||
has_secondaries = n_sec > 0
|
||||
@@ -548,7 +448,9 @@ def test_decode_secondaries_zero_n_sec_has_zero_energy():
|
||||
e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, e_sec, pre_dir
|
||||
)
|
||||
|
||||
assert not sec_valid.any()
|
||||
np.testing.assert_allclose(sec_E, 0.0)
|
||||
@@ -568,7 +470,9 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
|
||||
e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
|
||||
pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)
|
||||
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
|
||||
sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, e_sec, pre_dir
|
||||
)
|
||||
|
||||
for i, k in enumerate(n_sec):
|
||||
if k == 0:
|
||||
@@ -592,8 +496,12 @@ def test_decode_secondaries_rescale_preserves_relative_shares():
|
||||
n_sec = np.array([4])
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_E_small, _, _, _, sec_valid = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
|
||||
sec_E_large, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir)
|
||||
sec_E_small, _, _, _, sec_valid = decode_secondaries(
|
||||
sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir
|
||||
)
|
||||
sec_E_large, _, _, _, _ = decode_secondaries(
|
||||
sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir
|
||||
)
|
||||
|
||||
ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
|
||||
ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0]
|
||||
@@ -615,14 +523,70 @@ def test_decode_secondaries_mass_charge_round_trip_with_normalizer():
|
||||
sec_valid[0, 0] = True
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = encode_secondaries(sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list)
|
||||
sec_cont = encode_secondaries(
|
||||
sec_E_list, sec_dir_list, sec_valid, e_sec, pre_dir, sec_pdg_list=sec_pdg_list
|
||||
)
|
||||
norm = Normalizer()
|
||||
norm.mean = np.array([-2.0, 0.5], dtype=np.float32)
|
||||
norm.std = np.array([3.0, 1.5], dtype=np.float32)
|
||||
sec_cont_normed = sec_cont.copy()
|
||||
sec_cont_normed[:, :, 4:6] = norm.transform(sec_cont[:, :, 4:6].reshape(-1, 2)).reshape(N, K_MAX, 2)
|
||||
sec_cont_normed[:, :, 4:6] = norm.transform(
|
||||
sec_cont[:, :, 4:6].reshape(-1, 2)
|
||||
).reshape(N, K_MAX, 2)
|
||||
|
||||
n_sec = np.array([1])
|
||||
_, _, sec_mass, sec_charge, _ = decode_secondaries(sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm)
|
||||
_, _, sec_mass, sec_charge, _ = decode_secondaries(
|
||||
sec_cont_normed, n_sec, e_sec, pre_dir, sec_phys_normalizer=norm
|
||||
)
|
||||
assert sec_mass[0, 0] == pytest.approx(938.27208943, abs=1e-2)
|
||||
assert sec_charge[0, 0] == pytest.approx(1.0, abs=1e-4)
|
||||
|
||||
|
||||
def test_decode_secondaries_extreme_negative_log_mass_stays_nonnegative():
|
||||
from giant.data.transforms import decode_secondaries, log_transform
|
||||
|
||||
N = 1
|
||||
e_sec = np.array([5.0], dtype=np.float32)
|
||||
n_sec = np.array([1])
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
|
||||
sec_cont[0, 0, 0] = 10.0 # stick logit -> ~all of e_sec
|
||||
sec_cont[0, 0, 1:4] = [0, 0, 1]
|
||||
sec_cont[0, 0, 4] = -50.0 # raw model prediction: extremely negative log_mass
|
||||
sec_cont[0, 0, 5] = 1.0
|
||||
|
||||
_, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
|
||||
|
||||
# A raw model prediction isn't itself the output of log_transform, so
|
||||
# naively applying inv_log_transform can undershoot zero (see
|
||||
# decode_secondaries) — which then crashes the next log_transform call
|
||||
# once this mass is fed back in as conditioning during rollout. The
|
||||
# float32 residual from clipping can land a hair below zero, but must
|
||||
# stay well above -eps so log_transform(mass) stays finite.
|
||||
assert sec_mass[0, 0] > -1e-8
|
||||
log_transform(sec_mass[0, 0])
|
||||
|
||||
|
||||
def test_decode_secondaries_extreme_positive_log_mass_stays_finite():
|
||||
from giant.data.transforms import decode_secondaries, log_transform
|
||||
|
||||
N = 1
|
||||
e_sec = np.array([5.0], dtype=np.float32)
|
||||
n_sec = np.array([1])
|
||||
pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)
|
||||
|
||||
sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32)
|
||||
sec_cont[0, 0, 0] = 10.0 # stick logit -> ~all of e_sec
|
||||
sec_cont[0, 0, 1:4] = [0, 0, 1]
|
||||
sec_cont[0, 0, 4] = 200.0 # raw model prediction: extremely positive log_mass
|
||||
sec_cont[0, 0, 5] = 1.0
|
||||
|
||||
_, _, sec_mass, _, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
|
||||
|
||||
# Mirror image of the extreme-negative case above: exp(log_mass)
|
||||
# overflows float32 to inf for an unclipped raw prediction this large,
|
||||
# which then crashes the next log_transform call the same way a
|
||||
# negative mass would.
|
||||
assert np.isfinite(sec_mass[0, 0])
|
||||
log_transform(sec_mass[0, 0])
|
||||
|
||||
+17
-231
@@ -6,10 +6,8 @@ import pytest
|
||||
import torch
|
||||
|
||||
from giant import config as gconfig
|
||||
from giant.constants import COND_DIM
|
||||
from giant.data import setup_cache
|
||||
from giant.data.transforms import Normalizer
|
||||
from giant.pipeline import _seed_energy_router, run_train_job
|
||||
from giant.pipeline import run_train_job
|
||||
|
||||
|
||||
def _unit(v):
|
||||
@@ -48,7 +46,9 @@ def _make_synthetic_steps(path, n_events=20, seed=0):
|
||||
pre_dir = np.array([0.0, 0.0, 1.0])
|
||||
post_dir = _unit(rng.normal(size=3))
|
||||
post_pos = pre_pos + step_length * pre_dir
|
||||
sec_energies = list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
|
||||
sec_energies = (
|
||||
list(rng.dirichlet(np.ones(n_sec)) * e_sec) if n_sec > 0 else []
|
||||
)
|
||||
sec_pdgs = [pdgs[(row_idx + j) % 2] for j in range(n_sec)]
|
||||
sec_dirs = [_unit(rng.normal(size=3)) for _ in range(n_sec)]
|
||||
rows.append(
|
||||
@@ -99,19 +99,10 @@ def _tiny_cfg(**train_overrides):
|
||||
"warmup_epochs": 0,
|
||||
"validate_every": 0,
|
||||
"max_val_batches": 1,
|
||||
"wandb": False,
|
||||
}
|
||||
)
|
||||
cfg["train"].update(train_overrides)
|
||||
cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0})
|
||||
cfg["stage2_model"].update(
|
||||
# decoder="autoregressive" is DEFAULT_CONFIG's default (v0.3.0 step 5)
|
||||
# and left as-is here on purpose, so this pipeline-level fixture
|
||||
# exercises the real default end-to-end against actual data.
|
||||
{"hidden_dim": 8, "n_res_blocks": 1, "dropout": 0.0}
|
||||
)
|
||||
cfg["conditioning"]["particle"]["emb_dim"] = 4
|
||||
cfg["conditioning"]["material"]["emb_dim"] = 4
|
||||
cfg["model"].update({"hidden_dim": 8, "n_blocks": 1, "emb_dim": 4, "dropout": 0.0})
|
||||
return cfg
|
||||
|
||||
|
||||
@@ -152,106 +143,17 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch):
|
||||
assert "normalizer: cache hit" in joined
|
||||
|
||||
|
||||
def test_run_train_job_builds_caches_and_persists_sec_type_topn_map(tmp_path, data):
|
||||
"""DEFAULT_CONFIG's stage2_model.particle_type.target defaults to
|
||||
"onehot" while conditioning.particle.type stays "physical" — a plain
|
||||
_tiny_cfg() run must build the secondary-type-only pdg top-N map (gitea
|
||||
#29: no longer shared with any conditioning-side onehot map), cache it in
|
||||
the setup-cache sidecar, and persist it into the checkpoint's
|
||||
sec_type_topn_map key, with no extra config needed. pdg_topn_map
|
||||
(conditioning-only) stays unbuilt since conditioning.particle.type is
|
||||
"physical" here."""
|
||||
echo1 = _run(data, tmp_path / "out1")
|
||||
assert any("building pdg top-N map" in m for m in echo1)
|
||||
|
||||
loaded = setup_cache.load(data, [data])
|
||||
assert loaded is not None
|
||||
# stage2_model.particle_type.n_classes = 0 -> conditioning.particle.emb_dim = 4
|
||||
key = setup_cache.topn_key("pdg", 4)
|
||||
assert key in loaded.topn_maps
|
||||
assert set(loaded.topn_maps[key].class_map.keys()) >= {11, 22}
|
||||
|
||||
ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
|
||||
assert ckpt.get("pdg_topn_map") is None
|
||||
assert "sec_type_topn_map" in ckpt
|
||||
assert set(ckpt["sec_type_topn_map"]["class_map"].keys()) >= {"11", "22"}
|
||||
|
||||
echo2 = _run(data, tmp_path / "out2")
|
||||
assert any("pdg top-N map: cache hit" in m for m in echo2)
|
||||
|
||||
|
||||
def test_run_train_job_independent_cond_and_sec_type_topn_maps(tmp_path, data):
|
||||
"""conditioning.particle.type="onehot" and
|
||||
stage2_model.particle_type.target="onehot" with different class counts
|
||||
(gitea #29's fix: stage2_model.particle_type.n_classes decouples the two)
|
||||
build two distinct top-N maps, cached under their own (axis, n_classes)
|
||||
key and persisted under two distinct checkpoint keys — no longer forced
|
||||
to share conditioning.particle.emb_dim."""
|
||||
cfg = _tiny_cfg()
|
||||
cfg["conditioning"]["particle"]["type"] = "onehot" # emb_dim = 4, from _tiny_cfg
|
||||
cfg["stage2_model"]["particle_type"]["n_classes"] = 3
|
||||
echo = _run(data, tmp_path / "out", cfg=cfg)
|
||||
assert any("mapped to 4 classes" in m for m in echo)
|
||||
assert any("mapped to 3 classes" in m for m in echo)
|
||||
|
||||
loaded = setup_cache.load(data, [data])
|
||||
assert loaded is not None
|
||||
cond_key = setup_cache.topn_key("pdg", 4)
|
||||
type_key = setup_cache.topn_key("pdg", 3)
|
||||
assert cond_key in loaded.topn_maps
|
||||
assert type_key in loaded.topn_maps
|
||||
|
||||
ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
|
||||
assert ckpt.get("pdg_topn_map") is not None
|
||||
assert ckpt.get("sec_type_topn_map") is not None
|
||||
|
||||
|
||||
def test_run_train_job_builds_caches_and_persists_material_topn_map(tmp_path, data):
|
||||
"""conditioning.material.type="onehot" is an independent axis from the
|
||||
pdg one above, with its own build/cache-hit branch in run_setup_stage —
|
||||
exercise both here the same way the pdg test above does."""
|
||||
cfg = _tiny_cfg()
|
||||
cfg["conditioning"]["material"]["type"] = "onehot"
|
||||
echo1 = _run(data, tmp_path / "out1", cfg=cfg)
|
||||
assert any("building material top-N map" in m for m in echo1)
|
||||
|
||||
loaded = setup_cache.load(data, [data])
|
||||
assert loaded is not None
|
||||
key = setup_cache.topn_key("material", 4) # conditioning.material.emb_dim = 4
|
||||
assert key in loaded.topn_maps
|
||||
assert set(loaded.topn_maps[key].class_map.keys()) >= {"G4_AIR", "G4_Fe"}
|
||||
|
||||
ckpt = torch.load(tmp_path / "out1" / "last.pt", weights_only=False)
|
||||
assert "mat_topn_map" in ckpt
|
||||
assert set(ckpt["mat_topn_map"]["class_map"].keys()) >= {"G4_AIR", "G4_Fe"}
|
||||
|
||||
echo2 = _run(data, tmp_path / "out2", cfg=cfg)
|
||||
assert any("material top-N map: cache hit" in m for m in echo2)
|
||||
|
||||
|
||||
def test_run_train_job_no_topn_map_for_physical_target(tmp_path, data):
|
||||
cfg = _tiny_cfg()
|
||||
cfg["stage2_model"]["particle_type"] = {"target": "physical", "lambda": 1.0}
|
||||
echo = _run(data, tmp_path / "out", cfg=cfg)
|
||||
assert not any("top-N map" in m for m in echo)
|
||||
|
||||
loaded = setup_cache.load(data, [data])
|
||||
assert loaded is not None
|
||||
assert loaded.topn_maps == {}
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
|
||||
def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(tmp_path, data, monkeypatch):
|
||||
# num_workers>0 makes DataLoader actually fork worker subprocesses
|
||||
# (unlike every other test here, which runs with num_workers=0) — pytest
|
||||
# itself is multi-threaded, hence Python's fork-safety warning below.
|
||||
def test_run_train_job_warns_when_num_workers_exceeds_shared_quota(
|
||||
tmp_path, data, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2
|
||||
echo = _run(data, tmp_path / "out", num_workers=3)
|
||||
assert any("num-workers=3" in m and "exceeds" in m for m in echo)
|
||||
|
||||
|
||||
@pytest.mark.filterwarnings("ignore::DeprecationWarning:multiprocessing.popen_fork")
|
||||
def test_run_train_job_no_warning_when_num_workers_within_shared_quota(tmp_path, data, monkeypatch):
|
||||
def test_run_train_job_no_warning_when_num_workers_within_shared_quota(
|
||||
tmp_path, data, monkeypatch
|
||||
):
|
||||
monkeypatch.setattr("giant.pipeline.os.cpu_count", lambda: 8) # quota = 2
|
||||
echo = _run(data, tmp_path / "out", num_workers=2)
|
||||
assert not any("exceeds" in m for m in echo)
|
||||
@@ -291,70 +193,6 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat
|
||||
assert "fitting normalizer (streaming)" in joined
|
||||
|
||||
|
||||
def test_run_train_job_custom_k_max_end_to_end(tmp_path, data):
|
||||
"""Regression: stage2_model.k_max other
|
||||
than the K_MAX module constant's default (15) must not produce a shape
|
||||
mismatch between the data pipeline (loader.py/transforms.py padding) and
|
||||
the model (network.py's trunks, sized from this same config value)."""
|
||||
cfg = _tiny_cfg()
|
||||
cfg["stage2_model"]["k_max"] = 3
|
||||
_run(data, tmp_path / "out", cfg=cfg)
|
||||
ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
|
||||
assert ckpt["model_config"]["stage2_model"]["k_max"] == 3
|
||||
|
||||
|
||||
def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data):
|
||||
"""Regression: conditioning.particle.type
|
||||
and conditioning.material.type are configured independently and may mix
|
||||
freely — e.g. particle "embedding" with
|
||||
material "physical" — end-to-end through the real data pipeline, not
|
||||
just accepted by validate_config."""
|
||||
cfg = _tiny_cfg()
|
||||
cfg["conditioning"]["particle"]["type"] = "embedding"
|
||||
cfg["conditioning"]["material"]["type"] = "physical"
|
||||
_run(data, tmp_path / "out", cfg=cfg)
|
||||
ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
|
||||
cond_cfg = ckpt["model_config"]["conditioning"]
|
||||
assert cond_cfg["particle"]["type"] == "embedding"
|
||||
assert cond_cfg["material"]["type"] == "physical"
|
||||
|
||||
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
|
||||
# Particle block ([COND_DIM_BASE:COND_DIM_BASE+PARTICLE_PHYS_DIM]) stays
|
||||
# unfitted (mean=0/std=1) since "embedding" never computes real values
|
||||
# for it; the material block is fit for real under "physical".
|
||||
from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM
|
||||
|
||||
assert cond_norm.mean is not None and cond_norm.std is not None
|
||||
np.testing.assert_allclose(cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0)
|
||||
np.testing.assert_allclose(cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0)
|
||||
material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :]
|
||||
assert np.all(material_std > 0) and not np.allclose(material_std, 1.0)
|
||||
|
||||
|
||||
def test_run_train_job_share_stages_end_to_end(tmp_path, data):
|
||||
"""Regression: conditioning.share_stages
|
||||
= true must actually train (not raise NotImplementedError), and the
|
||||
resulting checkpoint's two stages must reload into a single shared
|
||||
ConditionEncoder instance rather than two independent ones."""
|
||||
from giant.model.network import build_models
|
||||
|
||||
cfg = _tiny_cfg()
|
||||
cfg["conditioning"]["share_stages"] = True
|
||||
_run(data, tmp_path / "out", cfg=cfg)
|
||||
ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False)
|
||||
assert ckpt["model_config"]["conditioning"]["share_stages"] is True
|
||||
|
||||
built = build_models(ckpt["model_config"])
|
||||
stage1, stage2 = built["stage1"], built["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
assert stage1.cond_enc is stage2.cond_enc
|
||||
|
||||
stage1.load_state_dict(ckpt["model"])
|
||||
stage2.load_state_dict(ckpt["sec_decoder"])
|
||||
for p1, p2 in zip(stage1.cond_enc.parameters(), stage2.cond_enc.parameters()):
|
||||
assert torch.equal(p1, p2)
|
||||
|
||||
|
||||
def test_run_train_job_matches_uncached_output(tmp_path, data):
|
||||
_run(data, tmp_path / "uncached", cache_setup=False)
|
||||
_run(data, tmp_path / "cached1", cache_setup=True)
|
||||
@@ -364,63 +202,11 @@ def test_run_train_job_matches_uncached_output(tmp_path, data):
|
||||
cached = torch.load(tmp_path / "cached2" / "last.pt", weights_only=False)
|
||||
|
||||
for key in ("cond", "target", "sec_phys"):
|
||||
np.testing.assert_allclose(uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"])
|
||||
np.testing.assert_allclose(uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"])
|
||||
np.testing.assert_allclose(
|
||||
uncached["normalizer"][key]["mean"], cached["normalizer"][key]["mean"]
|
||||
)
|
||||
np.testing.assert_allclose(
|
||||
uncached["normalizer"][key]["std"], cached["normalizer"][key]["std"]
|
||||
)
|
||||
assert uncached["pdg_map"] == cached["pdg_map"]
|
||||
assert uncached["mat_map"] == cached["mat_map"]
|
||||
|
||||
|
||||
def _fitted_cond_norm(seed=0):
|
||||
rng = np.random.default_rng(seed)
|
||||
return Normalizer().fit(rng.normal(size=(64, COND_DIM)).astype(np.float32))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"router_cfg",
|
||||
[
|
||||
{"enabled": False, "type": "energy", "n_experts": 4},
|
||||
{"enabled": True, "type": "pdg", "n_experts": 4},
|
||||
],
|
||||
)
|
||||
def test_seed_energy_router_noop_when_not_an_enabled_energy_router(router_cfg):
|
||||
cond_norm = _fitted_cond_norm()
|
||||
echoed = []
|
||||
_seed_energy_router(router_cfg, cond_norm, np.array([1.0, 2.0]), 3, echoed.append)
|
||||
assert "centers_init" not in router_cfg
|
||||
assert echoed == []
|
||||
|
||||
|
||||
def test_seed_energy_router_falls_back_to_default_and_warns_when_no_samples():
|
||||
router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
|
||||
cond_norm = _fitted_cond_norm()
|
||||
echoed = []
|
||||
_seed_energy_router(router_cfg, cond_norm, np.empty(0), energy_idx=3, echo=echoed.append)
|
||||
assert "centers_init" not in router_cfg
|
||||
assert len(echoed) == 1
|
||||
assert "falls back to default centers" in echoed[0]
|
||||
|
||||
|
||||
def test_seed_energy_router_seeds_centers_from_data_quantiles():
|
||||
router_cfg = {"enabled": True, "type": "energy", "n_experts": 4}
|
||||
cond_norm = _fitted_cond_norm()
|
||||
energy_idx = 3
|
||||
# A grid of "raw" quantile values as setup_cache.energy_quantiles_from_sample
|
||||
# would produce them: monotonically increasing, in the same (log-energy)
|
||||
# units as the conditioning column being normalized against.
|
||||
energy_quantiles = np.linspace(1.0, 10.0, 33).astype(np.float32)
|
||||
echoed = []
|
||||
_seed_energy_router(router_cfg, cond_norm, energy_quantiles, energy_idx, echoed.append)
|
||||
|
||||
assert "centers_init" in router_cfg
|
||||
centers = np.asarray(router_cfg["centers_init"], dtype=np.float32)
|
||||
assert centers.shape == (router_cfg["n_experts"],)
|
||||
|
||||
levels = np.linspace(0.0, 1.0, router_cfg["n_experts"])
|
||||
raw_centers = setup_cache.energy_quantile_at(energy_quantiles, levels)
|
||||
expected = (raw_centers - cond_norm.mean[energy_idx]) / cond_norm.std[energy_idx]
|
||||
np.testing.assert_allclose(centers, expected, rtol=1e-5)
|
||||
# Quantile levels are increasing, and the normalizer's std is positive, so
|
||||
# the seeded centers must preserve that order rather than e.g. reversing it.
|
||||
assert np.all(np.diff(centers) > 0)
|
||||
assert len(echoed) == 1
|
||||
assert "seeded EnergyRouter centers" in echoed[0]
|
||||
|
||||
@@ -4,12 +4,10 @@ from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
|
||||
pytest.importorskip("plotstyle")
|
||||
|
||||
from giant.analysis import render as render_mod # noqa: E402
|
||||
from giant.analysis.reduced import Reduced # noqa: E402
|
||||
|
||||
|
||||
@@ -21,152 +19,6 @@ def _try_render(reduced: list[Reduced], out: Path):
|
||||
return render_all(out / "reduced", out / "plots")
|
||||
|
||||
|
||||
def test_render_router_diagnostics_and_edge_cases(tmp_path: Path):
|
||||
reduced = [
|
||||
Reduced(
|
||||
"rg",
|
||||
"router",
|
||||
"router_gating",
|
||||
"Router gating",
|
||||
"pre-step energy [MeV]",
|
||||
{
|
||||
"n_experts": 2,
|
||||
"log_x": True,
|
||||
"router_type": "energy",
|
||||
"rollout": {
|
||||
"centers": [1.0, 10.0, 100.0],
|
||||
"means": [[0.6, 0.4], [0.5, 0.5], [0.4, 0.6]],
|
||||
},
|
||||
"reference": {
|
||||
"centers": [1.0, 10.0, 100.0],
|
||||
"means": [[0.55, 0.45], [0.5, 0.5], [0.45, 0.55]],
|
||||
},
|
||||
},
|
||||
),
|
||||
Reduced(
|
||||
"rs",
|
||||
"router",
|
||||
"router_share",
|
||||
"Router share",
|
||||
"species",
|
||||
{
|
||||
"categories": ["e-", "gamma"],
|
||||
"n_experts": 2,
|
||||
"router_type": "energy",
|
||||
"rollout": {"e-": [0.7, 0.3], "gamma": [0.2, 0.8]},
|
||||
"reference": {"e-": [0.6, 0.4], "gamma": [0.3, 0.7]},
|
||||
},
|
||||
),
|
||||
Reduced(
|
||||
"ru",
|
||||
"router",
|
||||
"unavailable",
|
||||
"Router unavailable",
|
||||
"x",
|
||||
{"note": "router diagnostics unavailable: no router in this run"},
|
||||
),
|
||||
Reduced(
|
||||
"g4",
|
||||
"marginals",
|
||||
"grouped_hist",
|
||||
"Grouped (4)",
|
||||
"x",
|
||||
{
|
||||
"edges": [0, 1, 2],
|
||||
"groups": {lbl: {"rollout": [1, 2], "reference": [2, 1]} for lbl in ("a", "b", "c", "d")},
|
||||
"log_y": True,
|
||||
},
|
||||
),
|
||||
Reduced(
|
||||
"sl",
|
||||
"species",
|
||||
"single_hist",
|
||||
"Single (log-x)",
|
||||
"x",
|
||||
{"edges": [1, 10, 100], "rollout": [5, 1], "log_x": True, "log_y": True},
|
||||
),
|
||||
]
|
||||
try:
|
||||
pdfs = _try_render(reduced, tmp_path)
|
||||
except RuntimeError as e: # LaTeX missing at render time
|
||||
pytest.skip(f"LaTeX rendering unavailable: {e}")
|
||||
assert len(pdfs) == len(reduced)
|
||||
assert all(p.exists() for p in pdfs)
|
||||
|
||||
|
||||
def test_render_all_run_gallery_invokes_subprocess(tmp_path: Path, monkeypatch):
|
||||
# render_mod.subprocess *is* the stdlib subprocess module, so a blanket
|
||||
# patch of .run would also swallow the real subprocess.run calls
|
||||
# matplotlib's texmanager makes to compile LaTeX during savefig — only
|
||||
# intercept the "gallery generate" call itself and pass everything else
|
||||
# (LaTeX included) through to the real subprocess.run.
|
||||
calls = []
|
||||
real_run = render_mod.subprocess.run
|
||||
|
||||
def fake_run(*a, **k):
|
||||
if a and a[0] and a[0][0] == "gallery":
|
||||
calls.append((a, k))
|
||||
return None
|
||||
return real_run(*a, **k)
|
||||
|
||||
monkeypatch.setattr(render_mod.subprocess, "run", fake_run)
|
||||
reduced = [
|
||||
Reduced(
|
||||
"s",
|
||||
"species",
|
||||
"single_hist",
|
||||
"Single",
|
||||
"x",
|
||||
{"edges": [0, 1, 2], "rollout": [5, 1]},
|
||||
)
|
||||
]
|
||||
for r in reduced:
|
||||
r.save(tmp_path / "reduced" / f"{r.id}.json")
|
||||
try:
|
||||
render_mod.render_all(tmp_path / "reduced", tmp_path / "plots", run_gallery=True)
|
||||
except RuntimeError as e:
|
||||
pytest.skip(f"LaTeX rendering unavailable: {e}")
|
||||
|
||||
assert len(calls) == 1
|
||||
args, kwargs = calls[0]
|
||||
assert args[0] == ["gallery", "generate", "--source", str(tmp_path / "plots")]
|
||||
assert kwargs == {"check": True}
|
||||
|
||||
|
||||
def test_render_run_glues_condor_run_meta_into_render_all(tmp_path: Path, monkeypatch):
|
||||
from giant.analysis import condor as condor_mod
|
||||
|
||||
run_dir = tmp_path / "run"
|
||||
(run_dir / "reduced").mkdir(parents=True)
|
||||
|
||||
merge_calls = []
|
||||
monkeypatch.setattr(condor_mod, "merge_all", lambda rd: merge_calls.append(Path(rd)))
|
||||
meta = condor_mod.RunMeta(
|
||||
rollout="rollout.parquet",
|
||||
reference="reference.parquet",
|
||||
run_dir=str(run_dir),
|
||||
title="my-run",
|
||||
plot_meta={"checkpoint": "ckpt/best.pt"},
|
||||
)
|
||||
monkeypatch.setattr(condor_mod.RunMeta, "load", classmethod(lambda cls, p: meta))
|
||||
|
||||
Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1], "rollout": [1]}).save(
|
||||
run_dir / "reduced" / "s.json"
|
||||
)
|
||||
|
||||
try:
|
||||
pdfs = render_mod.render_run(run_dir)
|
||||
except RuntimeError as e:
|
||||
pytest.skip(f"LaTeX rendering unavailable: {e}")
|
||||
|
||||
assert merge_calls == [run_dir]
|
||||
assert len(pdfs) == 1
|
||||
plot_meta = (run_dir / "plots" / "species" / "s.yaml").read_text()
|
||||
assert "checkpoint" in plot_meta
|
||||
root_meta = (run_dir / "plots" / "metadata.yaml").read_text()
|
||||
assert "my-run" in root_meta
|
||||
|
||||
|
||||
def test_render_one_of_each_kind(tmp_path: Path):
|
||||
reduced = [
|
||||
Reduced(
|
||||
@@ -238,133 +90,3 @@ def test_render_one_of_each_kind(tmp_path: Path):
|
||||
assert len(pdfs) == len(reduced)
|
||||
assert all(p.exists() for p in pdfs)
|
||||
assert (tmp_path / "plots" / "metadata.yaml").exists()
|
||||
|
||||
|
||||
# ── pure-function helpers: no matplotlib figure needed ──────────────────
|
||||
|
||||
|
||||
def test_density_zero_total_returns_counts_unchanged():
|
||||
counts = np.array([0.0, 0.0, 0.0])
|
||||
out = render_mod._density(counts, np.array([0.0, 1.0, 2.0, 3.0]))
|
||||
np.testing.assert_array_equal(out, counts)
|
||||
|
||||
|
||||
def test_density_normalizes_by_total_and_bin_width():
|
||||
counts = [1, 3]
|
||||
edges = np.array([0.0, 2.0, 4.0]) # bin width 2
|
||||
out = render_mod._density(counts, edges)
|
||||
np.testing.assert_allclose(out, np.array([1, 3]) / (4 * 2))
|
||||
|
||||
|
||||
def test_router_summary_disabled_is_off():
|
||||
assert render_mod._router_summary({"enabled": False, "type": "energy"}) == "off"
|
||||
assert render_mod._router_summary({}) == "off"
|
||||
|
||||
|
||||
def test_router_summary_enabled_formats_type_and_n_experts():
|
||||
cfg = {"enabled": True, "type": "energy", "n_experts": 8}
|
||||
assert render_mod._router_summary(cfg) == "energy×8"
|
||||
|
||||
|
||||
def test_figure_params_v2_basics_and_router_and_epoch():
|
||||
mc = {
|
||||
"stage1_model": {
|
||||
"hidden_dim": 256,
|
||||
"n_res_blocks": 4,
|
||||
"generator": "flow",
|
||||
"router": {"enabled": True, "type": "energy", "n_experts": 4},
|
||||
},
|
||||
"conditioning": {"particle": {"type": "physical"}},
|
||||
}
|
||||
run_meta = {"training_epoch": 12, "best_val_loss": 0.123456, "steps": 10}
|
||||
params = render_mod._figure_params(run_meta | {"model_config": mc})
|
||||
assert params == {
|
||||
"hidden_dim": 256,
|
||||
"n_res_blocks": 4,
|
||||
"mode": "flow",
|
||||
"conditioning": "physical",
|
||||
"router": "energy×4",
|
||||
"epoch": 12,
|
||||
"best_val_loss": 0.1235,
|
||||
"steps": 10,
|
||||
}
|
||||
|
||||
|
||||
def test_figure_params_v2_wgan_reports_noise_dim_not_steps():
|
||||
mc = {
|
||||
"stage1_model": {
|
||||
"generator": "wgan",
|
||||
"wgan": {"noise_dim": 32},
|
||||
},
|
||||
}
|
||||
run_meta = {"model_config": mc, "steps": 10}
|
||||
params = render_mod._figure_params(run_meta)
|
||||
assert params["mode"] == "wgan"
|
||||
assert params["noise_dim"] == 32
|
||||
assert "steps" not in params
|
||||
|
||||
|
||||
def test_figure_params_v2_reports_mode_s2_only_when_it_differs():
|
||||
same = {
|
||||
"stage1_model": {"generator": "flow"},
|
||||
"stage2_model": {"generator": "flow"},
|
||||
}
|
||||
assert "mode_s2" not in render_mod._figure_params({"model_config": same})
|
||||
|
||||
mixed = {
|
||||
"stage1_model": {"generator": "flow"},
|
||||
"stage2_model": {"generator": "wgan"},
|
||||
}
|
||||
params = render_mod._figure_params({"model_config": mixed})
|
||||
assert params["mode_s2"] == "wgan"
|
||||
|
||||
|
||||
def test_figure_params_old_shape_basics():
|
||||
run_meta = {
|
||||
"model_config": {
|
||||
"hidden_dim": 128,
|
||||
"n_blocks": 3,
|
||||
"mode": "ddpm",
|
||||
"conditioning": "embedding",
|
||||
"router": {"enabled": False},
|
||||
},
|
||||
"training_epoch": 5,
|
||||
"best_val_loss": 0.5,
|
||||
"steps": 20,
|
||||
}
|
||||
params = render_mod._figure_params(run_meta)
|
||||
assert params == {
|
||||
"hidden_dim": 128,
|
||||
"n_blocks": 3,
|
||||
"mode": "ddpm",
|
||||
"conditioning": "embedding",
|
||||
"router": "off",
|
||||
"epoch": 5,
|
||||
"best_val_loss": 0.5,
|
||||
"steps": 20,
|
||||
}
|
||||
|
||||
|
||||
def test_figure_params_old_shape_wgan_reports_noise_dim_not_steps():
|
||||
run_meta = {
|
||||
"model_config": {"mode": "wgan", "noise_dim": 16},
|
||||
"steps": 20,
|
||||
}
|
||||
params = render_mod._figure_params(run_meta)
|
||||
assert params["noise_dim"] == 16
|
||||
assert "steps" not in params
|
||||
|
||||
|
||||
def test_plot_metadata_includes_note_and_run_meta_parameters():
|
||||
r = Reduced("u", "router", "unavailable", "Unavailable", "x", {"note": "no router data"})
|
||||
meta = render_mod._plot_metadata(r, {"title": "run-1", "checkpoint": "ckpt.pt"})
|
||||
assert meta["note"] == "no router data"
|
||||
assert meta["parameters"] == {"checkpoint": "ckpt.pt"}
|
||||
assert "title" not in meta["parameters"]
|
||||
|
||||
|
||||
def test_plot_metadata_omits_parameters_when_run_meta_empty():
|
||||
r = Reduced("s", "species", "single_hist", "Single", "x", {"edges": [0, 1]})
|
||||
meta = render_mod._plot_metadata(r, {})
|
||||
assert "parameters" not in meta
|
||||
assert "note" not in meta
|
||||
|
||||
+15
-451
@@ -8,16 +8,10 @@ import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX
|
||||
from giant.data.loader import TopNMap
|
||||
from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG
|
||||
from giant.data.transforms import Normalizer
|
||||
from giant.model.network import (
|
||||
Stage1Model,
|
||||
Stage2Autoregressive,
|
||||
Stage2OneShot,
|
||||
stage2_trunk_sec_dim,
|
||||
)
|
||||
from giant.rollout import L1DistCollector, make_seed_frontier, rollout
|
||||
from giant.model.network import DenoisingMLP, SecondaryDecoder
|
||||
from giant.rollout import make_seed_frontier, rollout
|
||||
|
||||
pytest.importorskip("sklearn")
|
||||
from giant import geometry as g # noqa: E402
|
||||
@@ -27,26 +21,11 @@ MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1}
|
||||
|
||||
|
||||
def _models(conditioning="embedding"):
|
||||
particle_cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
|
||||
material_cfg = {"type": conditioning, "emb_dim": 16, "n_layers": 1}
|
||||
s1 = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
n_sec_head_k_max=K_MAX,
|
||||
s1 = DenoisingMLP(
|
||||
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, conditioning=conditioning
|
||||
)
|
||||
s2 = Stage2OneShot(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator="flow",
|
||||
time_dim=16,
|
||||
s2 = SecondaryDecoder(
|
||||
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, conditioning=conditioning
|
||||
)
|
||||
return s1.eval(), s2.eval()
|
||||
|
||||
@@ -111,8 +90,7 @@ def _run(
|
||||
batch_size=128,
|
||||
max_tracks_per_event=max_tracks_per_event,
|
||||
escape_threshold=escape_threshold,
|
||||
particle_conditioning=conditioning,
|
||||
material_conditioning=conditioning,
|
||||
conditioning=conditioning,
|
||||
)
|
||||
|
||||
|
||||
@@ -121,8 +99,12 @@ def fake_material_props(monkeypatch):
|
||||
import giant.materials as gm
|
||||
|
||||
fake = {
|
||||
"G4_AIR": gm.MaterialProperties(z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5),
|
||||
"G4_PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7),
|
||||
"G4_AIR": gm.MaterialProperties(
|
||||
z_eff=7.3, a_eff=14.4, density=1.2e-3, x0=3.0e4, lambda_int=7.0e5
|
||||
),
|
||||
"G4_PbWO4": gm.MaterialProperties(
|
||||
z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
|
||||
),
|
||||
}
|
||||
monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
|
||||
return fake
|
||||
@@ -169,7 +151,7 @@ def test_seed_frontier_embedding_mode_skips_unresolvable_pdg_lookup():
|
||||
frontier construction — mass/charge are simply zero-filled, unused."""
|
||||
seeds = _seeds(3)
|
||||
seeds["pdg"] = np.full(3, 999999999, dtype=np.int64)
|
||||
fr, _counts = make_seed_frontier(**seeds, particle_conditioning="embedding")
|
||||
fr, _counts = make_seed_frontier(**seeds, conditioning="embedding")
|
||||
np.testing.assert_array_equal(fr["mass"], 0.0)
|
||||
np.testing.assert_array_equal(fr["charge"], 0.0)
|
||||
|
||||
@@ -343,421 +325,3 @@ def test_on_chunk_never_buffers_full_records():
|
||||
assert rec.termination_reason_counts == {"natural_end": 1}
|
||||
with pytest.raises(AssertionError):
|
||||
rec.to_dict()
|
||||
|
||||
|
||||
# ── v0.3.0 step 6: per-stage generators, AR decoder, particle_type.target ───
|
||||
|
||||
# emb_dim=3: "other" (class idx 2) is shared by -11 and 13 (muon), matching
|
||||
# the real shape build_pdg_topn_map_from_files produces — see
|
||||
# decode_topn_class's docstring.
|
||||
PDG_TOPN_MAP = TopNMap(
|
||||
class_map={22: 0, 11: 1, -11: 2, 13: 2},
|
||||
other_members={-11: 5, 13: 1},
|
||||
)
|
||||
|
||||
|
||||
def _models_v3(
|
||||
conditioning="physical",
|
||||
decoder="one_shot",
|
||||
target="physical",
|
||||
generator1="flow",
|
||||
generator2="flow",
|
||||
k_max=6,
|
||||
emb_dim=4,
|
||||
stage2_has_n_sec_head=True,
|
||||
):
|
||||
particle_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
|
||||
material_cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
|
||||
# A fresh v0.3.0 Stage1Model — no n_sec_head_k_max, unlike _models() above
|
||||
# (n_sec ownership moves to stage 2 by default).
|
||||
s1 = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator=generator1,
|
||||
noise_dim=8,
|
||||
)
|
||||
particle_type_cfg = {"target": target}
|
||||
# Explicit kwargs rather than a shared **common dict: a dict() call whose
|
||||
# values have heterogeneous types (str/int/dict/bool) widens under static
|
||||
# analysis to dict[str, <big union>], which then makes every constructor
|
||||
# keyword not itself part of that union (router, cond_enc, ...) look like
|
||||
# a type mismatch to `ty` even though every actual value passed is fine.
|
||||
if decoder == "one_shot":
|
||||
sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator2, k_max, emb_dim)
|
||||
s2 = Stage2OneShot(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator=generator2,
|
||||
time_dim=16,
|
||||
noise_dim=8,
|
||||
k_max=k_max,
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
build_n_sec_head=stage2_has_n_sec_head,
|
||||
sec_dim=sec_dim,
|
||||
)
|
||||
else:
|
||||
s2 = Stage2Autoregressive(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator=generator2,
|
||||
time_dim=16,
|
||||
noise_dim=8,
|
||||
k_max=k_max,
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
build_n_sec_head=stage2_has_n_sec_head,
|
||||
)
|
||||
return s1.eval(), s2.eval()
|
||||
|
||||
|
||||
def _run_v3(
|
||||
s1,
|
||||
s2,
|
||||
escape_threshold=1e9,
|
||||
energy_cutoff=1.0,
|
||||
max_steps=15,
|
||||
max_tracks_per_event=100,
|
||||
seeds=None,
|
||||
conditioning="physical",
|
||||
sec_type_topn_map=None,
|
||||
other_policy="sample",
|
||||
seed=0,
|
||||
stage1_ddpm_steps=1000,
|
||||
l1_dist_collector=None,
|
||||
):
|
||||
torch.manual_seed(0)
|
||||
np.random.seed(0)
|
||||
cond, tgt, sec_phys = _norms()
|
||||
return rollout(
|
||||
s1,
|
||||
s2,
|
||||
_oracle(),
|
||||
seeds or _seeds(),
|
||||
cond,
|
||||
tgt,
|
||||
sec_phys,
|
||||
PDG_MAP,
|
||||
MAT_MAP,
|
||||
energy_cutoff=energy_cutoff,
|
||||
max_steps=max_steps,
|
||||
steps=3,
|
||||
batch_size=128,
|
||||
max_tracks_per_event=max_tracks_per_event,
|
||||
escape_threshold=escape_threshold,
|
||||
particle_conditioning=conditioning,
|
||||
material_conditioning=conditioning,
|
||||
sec_type_topn_map=sec_type_topn_map,
|
||||
other_policy=other_policy,
|
||||
seed=seed,
|
||||
stage1_ddpm_steps=stage1_ddpm_steps,
|
||||
l1_dist_collector=l1_dist_collector,
|
||||
)
|
||||
|
||||
|
||||
def test_rollout_stage2_owns_n_sec_when_stage1_has_no_head(fake_material_props):
|
||||
"""A fresh v0.3.0 Stage1Model has no n_sec_head — n_sec must come from
|
||||
Stage2's own head instead, and the run must still complete and
|
||||
conserve energy."""
|
||||
s1, s2 = _models_v3()
|
||||
rec = _run_v3(s1, s2)
|
||||
assert len(rec["event_id"]) > 0
|
||||
seeds = _seeds()
|
||||
for i, ev in enumerate(seeds["event_id"]):
|
||||
m = rec["event_id"] == ev
|
||||
dep = rec["edep"][m].sum()
|
||||
leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
|
||||
assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
|
||||
|
||||
|
||||
def test_resolve_n_sec_raises_when_neither_stage_owns_head(fake_material_props):
|
||||
"""Neither stage owning n_sec_head only happens for a
|
||||
stage2_model.n_sec.mode other than "head" — not a valid rollout-capable
|
||||
checkpoint, and must fail with a clear error rather than crash deep
|
||||
inside predict_n_sec."""
|
||||
s1, s2 = _models_v3(stage2_has_n_sec_head=False)
|
||||
with pytest.raises(RuntimeError, match="n_sec_head"):
|
||||
_run_v3(s1, s2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
|
||||
@pytest.mark.parametrize("generator2", ["flow", "wgan"])
|
||||
def test_rollout_physical_target_decoder_generator_matrix(fake_material_props, decoder, generator2):
|
||||
"""Every (decoder, stage2 generator) combination under
|
||||
particle_type.target="physical" must run to completion and conserve
|
||||
energy."""
|
||||
s1, s2 = _models_v3(decoder=decoder, generator2=generator2)
|
||||
rec = _run_v3(s1, s2)
|
||||
assert len(rec["event_id"]) > 0
|
||||
seeds = _seeds()
|
||||
for i, ev in enumerate(seeds["event_id"]):
|
||||
m = rec["event_id"] == ev
|
||||
dep = rec["edep"][m].sum()
|
||||
leak = rec["pre_E"][m & (rec["termination_reason"] == TERM_ESCAPED)].sum()
|
||||
assert dep + leak == pytest.approx(seeds["pre_E"][i], rel=1e-4)
|
||||
|
||||
|
||||
def test_sample_stage1_dispatches_ddpm_by_generator_kind():
|
||||
"""stage1_model.generator="ddpm" must be dispatched to sample_ddpm
|
||||
(previously _step_chunk silently fell through to the flow ODE sampler
|
||||
regardless of the checkpoint's actual generator — see giant.sample.sample_stage1).
|
||||
A short T avoids the reverse-diffusion numerical blowup an untrained,
|
||||
random-weight network produces over many steps; that instability is a
|
||||
property of sampling from an untrained net, not of the dispatch logic
|
||||
under test here, so a full oracle-driven rollout isn't needed."""
|
||||
from giant.sample import sample_stage1
|
||||
|
||||
s1, _ = _models_v3(generator1="ddpm")
|
||||
cond_cont = torch.randn(6, 15)
|
||||
cond_cat = torch.zeros(6, 2, dtype=torch.long)
|
||||
sample, n_sec = sample_stage1(s1, cond_cont, cond_cat, steps=10, ddpm_steps=5)
|
||||
assert sample.shape == (6, 9)
|
||||
assert n_sec is None # fresh v0.3.0 Stage1Model owns no n_sec_head
|
||||
|
||||
|
||||
@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
|
||||
def test_rollout_onehot_target_end_to_end(fake_material_props, decoder):
|
||||
"""particle_type.target="onehot" resolves a concrete PDG via
|
||||
decode_topn_class (argmax + other_policy), and that PDG's real physics
|
||||
(giant.particles.particle_phys_array) become the secondary's identity —
|
||||
unlike "physical", not just a reporting label."""
|
||||
s1, s2 = _models_v3(decoder=decoder, target="onehot", emb_dim=3)
|
||||
rec = _run_v3(s1, s2, sec_type_topn_map=PDG_TOPN_MAP, other_policy="modal")
|
||||
assert len(rec["event_id"]) > 0
|
||||
# Every spawned secondary's nominal pdg must be one decode_topn_class can
|
||||
# actually produce (the topn map's known classes + its "other" members).
|
||||
possible = set(PDG_TOPN_MAP.class_map.keys()) | set(PDG_TOPN_MAP.other_members.keys())
|
||||
secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
|
||||
assert secondary_pdgs <= possible
|
||||
|
||||
|
||||
def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props):
|
||||
s1, s2 = _models_v3(target="onehot", emb_dim=3)
|
||||
with pytest.raises(RuntimeError, match="sec_type_topn_map"):
|
||||
_run_v3(s1, s2, sec_type_topn_map=None)
|
||||
|
||||
|
||||
# --- conditioning.{particle,material}.type = "onehot" — a separate axis from
|
||||
# stage2_model.particle_type.target above: this is what feeds cond_cat's
|
||||
# extra top-N columns for ConditionEncoder's own "onehot" mode, not the
|
||||
# secondary-species decode. ---------------------------------------------
|
||||
|
||||
COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={})
|
||||
COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={})
|
||||
|
||||
|
||||
def _onehot_conditioning_models():
|
||||
particle_cfg = {"type": "onehot", "emb_dim": len(PDG_MAP), "n_layers": 1}
|
||||
material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
|
||||
s1 = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
)
|
||||
s2 = Stage2OneShot(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
sec_dim=stage2_trunk_sec_dim({"target": "physical"}, "flow", K_MAX, 3),
|
||||
generator="flow",
|
||||
time_dim=16,
|
||||
)
|
||||
return s1.eval(), s2.eval()
|
||||
|
||||
|
||||
def _run_onehot_conditioning(pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP):
|
||||
s1, s2 = _onehot_conditioning_models()
|
||||
cond, tgt, sec_phys = _norms()
|
||||
return rollout(
|
||||
s1,
|
||||
s2,
|
||||
_oracle(),
|
||||
_seeds(),
|
||||
cond,
|
||||
tgt,
|
||||
sec_phys,
|
||||
PDG_MAP,
|
||||
MAT_MAP,
|
||||
energy_cutoff=1.0,
|
||||
max_steps=30,
|
||||
steps=4,
|
||||
batch_size=128,
|
||||
max_tracks_per_event=300,
|
||||
escape_threshold=1e9,
|
||||
particle_conditioning="onehot",
|
||||
material_conditioning="onehot",
|
||||
pdg_topn_map=pdg_topn_map,
|
||||
mat_topn_map=mat_topn_map,
|
||||
)
|
||||
|
||||
|
||||
def test_rollout_conditioning_onehot_end_to_end(fake_material_props):
|
||||
rec = _run_onehot_conditioning()
|
||||
assert len(rec["event_id"]) > 0
|
||||
assert set(rec["event_id"].tolist()) == set(range(6))
|
||||
|
||||
|
||||
def test_rollout_conditioning_onehot_particle_missing_topn_map_raises(
|
||||
fake_material_props,
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="pdg_topn_map"):
|
||||
_run_onehot_conditioning(pdg_topn_map=None)
|
||||
|
||||
|
||||
def test_rollout_conditioning_onehot_material_missing_topn_map_raises(
|
||||
fake_material_props,
|
||||
):
|
||||
with pytest.raises(RuntimeError, match="mat_topn_map"):
|
||||
_run_onehot_conditioning(mat_topn_map=None)
|
||||
|
||||
|
||||
SEC_TYPE_TOPN_MAP_DIFFERENT_N = TopNMap(class_map={22: 0, 11: 1, -11: 2, 13: 3}, other_members={2112: 3, 2212: 1})
|
||||
|
||||
|
||||
def _run_conditioning_and_type_onehot_different_n_classes():
|
||||
"""Both conditioning.particle.type="onehot" and
|
||||
stage2_model.particle_type.target="onehot" active at once, with
|
||||
stage2_model.particle_type.n_classes deliberately different from
|
||||
conditioning.particle.emb_dim (gitea #29)."""
|
||||
cond_emb_dim = len(PDG_MAP) # 3
|
||||
type_n_classes = 5 # deliberately different from cond_emb_dim
|
||||
particle_cfg = {"type": "onehot", "emb_dim": cond_emb_dim, "n_layers": 1}
|
||||
material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1}
|
||||
particle_type_cfg = {"target": "onehot", "n_classes": type_n_classes}
|
||||
|
||||
s1 = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
).eval()
|
||||
sec_dim = stage2_trunk_sec_dim(particle_type_cfg, "flow", K_MAX, type_n_classes)
|
||||
s2 = Stage2OneShot(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
sec_dim=sec_dim,
|
||||
generator="flow",
|
||||
time_dim=16,
|
||||
k_max=K_MAX,
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
).eval()
|
||||
# Sanity: the model's own type_dim followed n_classes, not cond_emb_dim.
|
||||
assert s2.type_dim == type_n_classes
|
||||
|
||||
cond, tgt, sec_phys = _norms()
|
||||
return rollout(
|
||||
s1,
|
||||
s2,
|
||||
_oracle(),
|
||||
_seeds(),
|
||||
cond,
|
||||
tgt,
|
||||
sec_phys,
|
||||
PDG_MAP,
|
||||
MAT_MAP,
|
||||
energy_cutoff=1.0,
|
||||
max_steps=15,
|
||||
steps=3,
|
||||
batch_size=128,
|
||||
max_tracks_per_event=100,
|
||||
escape_threshold=1e9,
|
||||
particle_conditioning="onehot",
|
||||
material_conditioning="onehot",
|
||||
pdg_topn_map=COND_PDG_TOPN_MAP,
|
||||
mat_topn_map=COND_MAT_TOPN_MAP,
|
||||
sec_type_topn_map=SEC_TYPE_TOPN_MAP_DIFFERENT_N,
|
||||
other_policy="modal",
|
||||
)
|
||||
|
||||
|
||||
def test_rollout_conditioning_and_type_onehot_with_different_n_classes(fake_material_props):
|
||||
"""gitea #29 end-to-end: conditioning.particle.type="onehot" and
|
||||
stage2_model.particle_type.target="onehot" now use independently sized
|
||||
top-N maps (stage2_model.particle_type.n_classes != conditioning.particle
|
||||
.emb_dim), and rollout must decode secondaries using the type-side map,
|
||||
not silently reuse the conditioning-side one (the pre-#29 bug)."""
|
||||
rec = _run_conditioning_and_type_onehot_different_n_classes()
|
||||
assert len(rec["event_id"]) > 0
|
||||
possible = set(SEC_TYPE_TOPN_MAP_DIFFERENT_N.class_map.keys()) | set(
|
||||
SEC_TYPE_TOPN_MAP_DIFFERENT_N.other_members.keys()
|
||||
)
|
||||
secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
|
||||
assert secondary_pdgs <= possible
|
||||
|
||||
|
||||
@pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"])
|
||||
def test_rollout_embedding_target_end_to_end(decoder):
|
||||
"""particle_type.target="embedding" L1-snaps to the nearest row of the
|
||||
conditioning's own particle embedding table, so every resolved PDG must
|
||||
be a real member of the dense training vocab (pdg_map) — unlike
|
||||
"onehot", there is no "other" bucket to fall outside of."""
|
||||
s1, s2 = _models_v3(conditioning="embedding", decoder=decoder, target="embedding", emb_dim=4)
|
||||
rec = _run_v3(s1, s2, conditioning="embedding")
|
||||
assert len(rec["event_id"]) > 0
|
||||
secondary_pdgs = set(rec["pdg"][rec["generation"] > 0].tolist())
|
||||
assert secondary_pdgs <= set(PDG_MAP.keys())
|
||||
|
||||
|
||||
def test_l1_dist_collector_populated_only_for_embedding_target():
|
||||
"""The L1-distance diagnostic only makes sense under
|
||||
particle_type.target="embedding" — a physical-target run must leave the
|
||||
collector empty rather than silently accumulating garbage."""
|
||||
s1, s2 = _models_v3(target="physical")
|
||||
collector = L1DistCollector()
|
||||
_run_v3(s1, s2, l1_dist_collector=collector)
|
||||
assert collector.n == 0
|
||||
assert collector.summary() is None
|
||||
|
||||
|
||||
def test_l1_dist_collector_accumulates_for_embedding_target():
|
||||
s1, s2 = _models_v3(conditioning="embedding", target="embedding", emb_dim=4)
|
||||
collector = L1DistCollector()
|
||||
rec = _run_v3(s1, s2, conditioning="embedding", l1_dist_collector=collector)
|
||||
n_secondaries = int((rec["generation"] > 0).sum())
|
||||
assert n_secondaries > 0 # sanity: the tiny model does spawn secondaries
|
||||
summary = collector.summary()
|
||||
assert summary is not None
|
||||
assert summary["n"] == collector.n > 0
|
||||
assert summary["min"] <= summary["mean"] <= summary["max"]
|
||||
assert summary["std"] >= 0.0
|
||||
assert len(summary["hist_edges"]) == len(summary["hist_counts"]) + 1
|
||||
assert sum(summary["hist_counts"]) <= summary["n"] # some may fall outside [lo, hi)
|
||||
|
||||
|
||||
def test_l1_dist_collector_add_ignores_invalid_slots():
|
||||
collector = L1DistCollector()
|
||||
dist = np.array([[1.0, 5.0, 9.0]])
|
||||
valid = np.array([[True, False, True]])
|
||||
collector.add(dist, valid)
|
||||
assert collector.n == 2
|
||||
assert collector.minimum == 1.0
|
||||
assert collector.maximum == 9.0
|
||||
|
||||
|
||||
def test_l1_dist_collector_add_empty_is_noop():
|
||||
collector = L1DistCollector()
|
||||
collector.add(np.zeros((0, 3)), np.zeros((0, 3), dtype=bool))
|
||||
assert collector.n == 0
|
||||
assert collector.summary() is None
|
||||
|
||||
+202
-241
@@ -6,55 +6,47 @@ import torch
|
||||
from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
|
||||
from giant.model.network import (
|
||||
ComposedRouter,
|
||||
DenoisingMLP,
|
||||
EnergyRouter,
|
||||
MonolithicTrunk,
|
||||
PdgRouter,
|
||||
ProcessRouter,
|
||||
ROUTER_REGISTRY,
|
||||
RoutedTrunk,
|
||||
Stage1Model,
|
||||
Stage2OneShot,
|
||||
RoutedDenoisingMLP,
|
||||
RoutedSecondaryDecoder,
|
||||
SecondaryDecoder,
|
||||
build_composed_router,
|
||||
build_models,
|
||||
build_router,
|
||||
)
|
||||
|
||||
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
|
||||
|
||||
def _cond(B=8, pdg=3, mat=2):
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
|
||||
cond_cat = torch.stack(
|
||||
[torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1
|
||||
)
|
||||
return cond_cont, cond_cat
|
||||
|
||||
|
||||
def _routed_stage1(n_experts=4, pdg=3, mat=2, **router_kwargs):
|
||||
router = build_router("energy", n_experts, **router_kwargs)
|
||||
return Stage1Model(
|
||||
return RoutedDenoisingMLP(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=2,
|
||||
router=router,
|
||||
n_sec_head_k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
)
|
||||
|
||||
|
||||
def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs):
|
||||
router = build_router("energy", n_experts, **router_kwargs)
|
||||
return Stage2OneShot(
|
||||
return RoutedSecondaryDecoder(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=2,
|
||||
generator="flow",
|
||||
time_dim=16,
|
||||
router=router,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
)
|
||||
|
||||
|
||||
@@ -76,7 +68,9 @@ def test_energy_router_gate_partition_of_unity():
|
||||
def test_energy_router_top1_matches_gate_argmax():
|
||||
router = EnergyRouter(n_experts=4)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
|
||||
assert torch.equal(
|
||||
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
||||
)
|
||||
|
||||
|
||||
def test_energy_router_hardens_as_temperature_shrinks():
|
||||
@@ -98,7 +92,7 @@ def test_energy_router_balance_loss_is_nonnegative_scalar():
|
||||
|
||||
|
||||
def test_build_router_ignores_unrecognized_kwargs():
|
||||
# lambda_balance is a router config key but not an EnergyRouter kwarg
|
||||
# lambda_balance is a model_config.router key but not an EnergyRouter kwarg
|
||||
router = build_router("energy", 4, temperature=0.3, lambda_balance=0.5)
|
||||
assert isinstance(router, EnergyRouter)
|
||||
assert router.temperature == 0.3
|
||||
@@ -124,8 +118,12 @@ def test_energy_router_centers_init_wrong_length_raises():
|
||||
|
||||
|
||||
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)
|
||||
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)
|
||||
|
||||
@@ -152,7 +150,9 @@ def test_energy_router_learn_width_matches_fixed_temperature_at_init():
|
||||
per-expert width must reproduce the fixed-temperature gate exactly."""
|
||||
centers_init = [-1.0, 0.0, 0.5, 1.5]
|
||||
fixed = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init)
|
||||
learned = EnergyRouter(n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True)
|
||||
learned = EnergyRouter(
|
||||
n_experts=4, temperature=0.3, centers_init=centers_init, learn_width=True
|
||||
)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
torch.testing.assert_close(
|
||||
learned.gate(cond_cont, cond_cat),
|
||||
@@ -185,15 +185,21 @@ def test_energy_router_learn_width_and_temperature_mutually_exclusive_raises():
|
||||
EnergyRouter(n_experts=4, learn_width=True, learn_temperature=True)
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for learn_width and learn_temperature both set")
|
||||
raise AssertionError(
|
||||
"expected ValueError for learn_width and learn_temperature both set"
|
||||
)
|
||||
|
||||
|
||||
def test_energy_router_width_ratio_bounds_must_bracket_one_raises():
|
||||
try:
|
||||
EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0)
|
||||
EnergyRouter(
|
||||
n_experts=4, learn_width=True, width_min_ratio=1.0, width_max_ratio=2.0
|
||||
)
|
||||
except ValueError:
|
||||
return
|
||||
raise AssertionError("expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0")
|
||||
raise AssertionError(
|
||||
"expected ValueError for width_min_ratio/width_max_ratio not bracketing 1.0"
|
||||
)
|
||||
|
||||
|
||||
def test_energy_router_effective_width_stays_within_bounds():
|
||||
@@ -230,7 +236,9 @@ def test_energy_router_learn_width_hardens_when_pushed_to_floor():
|
||||
"""Pushing every expert's width toward the (tiny) floor should harden the
|
||||
gate to a one-hot at the nearest center, generalizing the fixed-
|
||||
temperature->0 hardening test to the per-expert path."""
|
||||
router = EnergyRouter(n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0)
|
||||
router = EnergyRouter(
|
||||
n_experts=4, learn_width=True, width_min_ratio=1e-4, width_max_ratio=10.0
|
||||
)
|
||||
with torch.no_grad():
|
||||
router.raw_width.fill_(-1e6)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
@@ -246,7 +254,9 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
|
||||
expert's own gate share, without needing to touch any other expert's
|
||||
width — the "each expert learns its own coverage independently" property
|
||||
this feature is meant to add."""
|
||||
router = EnergyRouter(n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0])
|
||||
router = EnergyRouter(
|
||||
n_experts=2, temperature=1.0, learn_width=True, centers_init=[0.0, 10.0]
|
||||
)
|
||||
cond_cont, cond_cat = _cond(4)
|
||||
cond_cont[:, 3] = 3.0 # fixed energy, unequal distance to each center
|
||||
|
||||
@@ -260,7 +270,9 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
|
||||
|
||||
|
||||
def test_build_router_threads_learn_width_kwargs_through():
|
||||
router = build_router("energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0)
|
||||
router = build_router(
|
||||
"energy", 4, learn_width=True, width_min_ratio=0.2, width_max_ratio=8.0
|
||||
)
|
||||
assert isinstance(router, EnergyRouter)
|
||||
assert router.learn_width is True
|
||||
assert isinstance(router.raw_width, torch.nn.Parameter)
|
||||
@@ -363,18 +375,18 @@ def test_build_router_from_cfg_sets_gumbel_for_composed_router():
|
||||
assert router.gumbel is True
|
||||
|
||||
|
||||
def test_routed_stage1_forward_runs_with_gumbel_enabled():
|
||||
def test_routed_denoising_mlp_forward_runs_with_gumbel_enabled():
|
||||
"""End-to-end forward through _route_forward's train branch with
|
||||
straight-through Gumbel-softmax combine weights enabled."""
|
||||
B = 8
|
||||
model = _routed_stage1(n_experts=3)
|
||||
model.trunk.router.gumbel = True
|
||||
model.trunk.router.gumbel_tau = 0.5
|
||||
model.router.gumbel = True
|
||||
model.router.gumbel_tau = 0.5
|
||||
model.train()
|
||||
x_t = torch.randn(B, X_DIM)
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
out = model(x_t, cond_cont, cond_cat, t=t)
|
||||
out = model(x_t, t, cond_cont, cond_cat)
|
||||
assert out.shape == (B, X_DIM)
|
||||
assert torch.isfinite(out).all()
|
||||
|
||||
@@ -397,7 +409,9 @@ def test_pdg_router_gate_partition_of_unity():
|
||||
def test_pdg_router_top1_matches_gate_argmax():
|
||||
router = PdgRouter(n_experts=4, pdg_vocab=3)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
|
||||
assert torch.equal(
|
||||
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
||||
)
|
||||
|
||||
|
||||
def test_pdg_router_hardens_as_temperature_shrinks():
|
||||
@@ -449,89 +463,59 @@ def test_build_router_pdg_type_uses_pdg_vocab():
|
||||
assert router.pdg_emb.num_embeddings == 5
|
||||
|
||||
|
||||
def _nested_cfg(
|
||||
pdg_vocab,
|
||||
mat_vocab,
|
||||
stage1_router=None,
|
||||
stage2_router=None,
|
||||
particle_type="physical",
|
||||
material_type="physical",
|
||||
**overrides,
|
||||
):
|
||||
"""Minimal new-shape (v0.3.0) model_config for build_models, with
|
||||
optional router sub-blocks. `overrides` deep-patches stage1_model."""
|
||||
stage1_model = {
|
||||
"active": True,
|
||||
"generator": "flow",
|
||||
"hidden_dim": 16,
|
||||
"n_res_blocks": 2,
|
||||
"dropout": 0.0,
|
||||
"flow": {"time_dim": 16},
|
||||
"router": stage1_router or {"enabled": False},
|
||||
}
|
||||
stage1_model.update(overrides)
|
||||
return {
|
||||
"pdg_vocab": pdg_vocab,
|
||||
"mat_vocab": mat_vocab,
|
||||
"conditioning": {
|
||||
"out_dim": 32,
|
||||
"particle": {"type": particle_type, "emb_dim": 8, "n_layers": 1},
|
||||
"material": {"type": material_type, "emb_dim": 8, "n_layers": 1},
|
||||
},
|
||||
"stage1_model": stage1_model,
|
||||
"stage2_model": {
|
||||
"active": True,
|
||||
"decoder": "one_shot",
|
||||
"generator": "flow",
|
||||
"hidden_dim": 16,
|
||||
"n_res_blocks": 2,
|
||||
"dropout": 0.0,
|
||||
"k_max": K_MAX,
|
||||
"context_dim": 16,
|
||||
"n_sec": {"mode": "head"},
|
||||
"flow": {"time_dim": 16},
|
||||
"router": stage2_router or {"enabled": False, "tie_to_stage1": False},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def test_build_models_routed_with_pdg_router():
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
particle_type="embedding",
|
||||
material_type="embedding",
|
||||
stage1_router={"enabled": True, "type": "pdg", "n_experts": 3},
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "pdg",
|
||||
"n_experts": 3,
|
||||
},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1 = models["stage1"]
|
||||
assert isinstance(stage1, Stage1Model)
|
||||
assert isinstance(stage1.trunk, RoutedTrunk)
|
||||
assert isinstance(stage1.trunk.router, PdgRouter)
|
||||
assert len(stage1.trunk.experts) == 3
|
||||
assert stage1.trunk.router.pdg_emb.num_embeddings == 4
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(stage1.router, PdgRouter)
|
||||
assert len(stage1.experts) == 3
|
||||
assert stage1.router.pdg_emb.num_embeddings == 4
|
||||
|
||||
|
||||
def test_build_models_rejects_pdg_router_with_physical_conditioning():
|
||||
"""conditioning.particle.type="physical" is meant to generalize beyond the
|
||||
training PDG vocab; PdgRouter always uses a training-vocab nn.Embedding
|
||||
regardless of conditioning, so the combination must raise rather than
|
||||
silently building a model that can't actually generalize the way it
|
||||
claims to."""
|
||||
cfg = _nested_cfg(
|
||||
"""conditioning="physical" is meant to generalize beyond the training PDG
|
||||
vocab; PdgRouter always uses a training-vocab nn.Embedding regardless of
|
||||
conditioning, so the combination must raise rather than silently building
|
||||
a model that can't actually generalize the way it claims to."""
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
stage1_router={"enabled": True, "type": "pdg", "n_experts": 3},
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
conditioning="physical",
|
||||
router={"enabled": True, "type": "pdg", "n_experts": 3},
|
||||
)
|
||||
with pytest.raises(ValueError, match="physical"):
|
||||
build_models(cfg)
|
||||
build_models(model_config)
|
||||
|
||||
|
||||
def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditioning():
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
stage1_router={
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
conditioning="physical",
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
@@ -541,7 +525,7 @@ def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditi
|
||||
},
|
||||
)
|
||||
with pytest.raises(ValueError, match="physical"):
|
||||
build_models(cfg)
|
||||
build_models(model_config)
|
||||
|
||||
|
||||
# ── ProcessRouter ────────────────────────────────────────────────────────────
|
||||
@@ -562,7 +546,9 @@ def test_process_router_gate_partition_of_unity():
|
||||
def test_process_router_top1_matches_gate_argmax():
|
||||
router = ProcessRouter(n_experts=4, pdg_vocab=3, mat_vocab=2)
|
||||
cond_cont, cond_cat = _cond(16)
|
||||
assert torch.equal(router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1))
|
||||
assert torch.equal(
|
||||
router.top1(cond_cont, cond_cat), router.gate(cond_cont, cond_cat).argmax(-1)
|
||||
)
|
||||
|
||||
|
||||
def test_process_router_balance_loss_is_nonnegative_scalar():
|
||||
@@ -611,38 +597,43 @@ def test_build_router_process_type_uses_pdg_mat_vocab():
|
||||
|
||||
|
||||
def test_build_models_routed_with_process_router():
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
particle_type="embedding",
|
||||
material_type="embedding",
|
||||
stage1_router={
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "process",
|
||||
"n_experts": 3,
|
||||
"lambda_proc": 1.0,
|
||||
},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1 = models["stage1"]
|
||||
assert isinstance(stage1, Stage1Model)
|
||||
assert isinstance(stage1.trunk, RoutedTrunk)
|
||||
assert isinstance(stage1.trunk.router, ProcessRouter)
|
||||
assert len(stage1.trunk.experts) == 3
|
||||
assert stage1.trunk.router.pdg_emb.num_embeddings == 4
|
||||
assert stage1.trunk.router.mat_emb.num_embeddings == 2
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(stage1.router, ProcessRouter)
|
||||
assert len(stage1.experts) == 3
|
||||
assert stage1.router.pdg_emb.num_embeddings == 4
|
||||
assert stage1.router.mat_emb.num_embeddings == 2
|
||||
|
||||
|
||||
# ── ComposedRouter ───────────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_composed_router_n_experts_is_product():
|
||||
router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
assert router.n_experts == 12
|
||||
|
||||
|
||||
def test_composed_router_gate_partition_of_unity():
|
||||
router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
g = router.gate(cond_cont, cond_cat)
|
||||
assert g.shape == (16, 12)
|
||||
@@ -679,7 +670,9 @@ def test_composed_router_top1_factors_into_per_axis_argmax():
|
||||
|
||||
|
||||
def test_composed_router_supports_different_expert_counts_per_axis():
|
||||
router = ComposedRouter([EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)])
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=5), PdgRouter(n_experts=2, pdg_vocab=5)]
|
||||
)
|
||||
assert router.n_experts == 10
|
||||
cond_cont, cond_cat = _cond(8, pdg=5)
|
||||
assert router.gate(cond_cont, cond_cat).shape == (8, 10)
|
||||
@@ -687,7 +680,9 @@ def test_composed_router_supports_different_expert_counts_per_axis():
|
||||
|
||||
def test_composed_router_classify_loss_sums_sub_router_losses():
|
||||
"""energy/pdg both default to zero, so the composed loss should too."""
|
||||
router = ComposedRouter([EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)])
|
||||
router = ComposedRouter(
|
||||
[EnergyRouter(n_experts=4), PdgRouter(n_experts=3, pdg_vocab=5)]
|
||||
)
|
||||
cond_cont, cond_cat = _cond(16, pdg=5)
|
||||
labels = torch.randint(0, 4, (16,))
|
||||
loss = router.classify_loss(cond_cont, cond_cat, labels)
|
||||
@@ -782,12 +777,15 @@ def test_build_composed_router_resolves_per_axis_specs():
|
||||
|
||||
|
||||
def test_build_models_routed_with_composed_router():
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=5,
|
||||
mat_vocab=2,
|
||||
particle_type="embedding",
|
||||
material_type="embedding",
|
||||
stage1_router={
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
@@ -795,58 +793,29 @@ def test_build_models_routed_with_composed_router():
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 3,
|
||||
},
|
||||
stage2_router={
|
||||
"enabled": True,
|
||||
"tie_to_stage1": False,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
"axis0_n_experts": 4,
|
||||
"axis1_type": "pdg",
|
||||
"axis1_n_experts": 3,
|
||||
},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1, stage2 = models["stage1"], models["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
assert isinstance(stage1.trunk, RoutedTrunk)
|
||||
assert isinstance(stage1.trunk.router, ComposedRouter)
|
||||
assert len(stage1.trunk.experts) == 12
|
||||
assert len(stage2.trunk.experts) == 12
|
||||
# stage1 and stage2 must not share router weights when tie_to_stage1 is
|
||||
# false (same convention as v0.2's two-independent-routers behaviour).
|
||||
assert stage1.trunk.router is not stage2.trunk.router
|
||||
|
||||
|
||||
def test_build_models_routed_stage2_ties_to_stage1_router():
|
||||
cfg = _nested_cfg(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
stage1_router={"enabled": True, "type": "energy", "n_experts": 3},
|
||||
stage2_router={
|
||||
"enabled": True,
|
||||
"tie_to_stage1": True,
|
||||
"type": "energy",
|
||||
"n_experts": 3,
|
||||
},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1, stage2 = models["stage1"], models["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
assert stage1.trunk.router is stage2.trunk.router
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(stage1.router, ComposedRouter)
|
||||
assert len(stage1.experts) == 12
|
||||
assert len(sec_decoder.experts) == 12
|
||||
# stage1 and sec_decoder must not share router weights (same convention
|
||||
# as the single-axis routers built by build_models).
|
||||
assert stage1.router is not sec_decoder.router
|
||||
|
||||
|
||||
def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
|
||||
from giant.sample import sample_flow, sample_secondaries
|
||||
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
# PdgRouter (axis1) always builds its own training-vocab embedding,
|
||||
# incompatible with conditioning.particle.type="physical" (the
|
||||
# _nested_cfg default) — see _check_router_conditioning_compat.
|
||||
particle_type="embedding",
|
||||
material_type="embedding",
|
||||
stage1_router={
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=8,
|
||||
expert_n_blocks=1,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "composed",
|
||||
"axis0_type": "energy",
|
||||
@@ -855,31 +824,24 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
|
||||
"axis1_n_experts": 2,
|
||||
},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1, stage2 = models["stage1"], models["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
B = 5
|
||||
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
|
||||
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
|
||||
assert stage1_norm.shape == (B, X_DIM)
|
||||
# A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) —
|
||||
# sample_flow returns n_sec_pred=None here, and n_sec must be
|
||||
# asked of stage2 instead, using the just-sampled stage1_norm as context.
|
||||
assert n_sec_pred is None
|
||||
n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
|
||||
assert n_sec_pred.shape == (B,)
|
||||
|
||||
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
||||
stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
||||
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
|
||||
|
||||
# ── Stage1Model with a routed trunk ─────────────────────────────────────────
|
||||
# ── RoutedDenoisingMLP ───────────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_routed_stage1_output_shape_train_and_eval():
|
||||
def test_routed_denoising_mlp_output_shape_train_and_eval():
|
||||
B = 8
|
||||
model = _routed_stage1()
|
||||
x_t = torch.randn(B, X_DIM)
|
||||
@@ -887,16 +849,16 @@ def test_routed_stage1_output_shape_train_and_eval():
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
|
||||
model.train()
|
||||
out_train = model(x_t, cond_cont, cond_cat, t=t)
|
||||
out_train = model(x_t, t, cond_cont, cond_cat)
|
||||
assert out_train.shape == (B, X_DIM)
|
||||
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
out_eval = model(x_t, cond_cont, cond_cat, t=t)
|
||||
out_eval = model(x_t, t, cond_cont, cond_cat)
|
||||
assert out_eval.shape == (B, X_DIM)
|
||||
|
||||
|
||||
def test_routed_stage1_gradients_flow_in_train_mode():
|
||||
def test_routed_denoising_mlp_gradients_flow_in_train_mode():
|
||||
"""Soft mixture in train mode should touch every expert's parameters."""
|
||||
B = 8
|
||||
model = _routed_stage1(n_experts=3)
|
||||
@@ -904,14 +866,14 @@ def test_routed_stage1_gradients_flow_in_train_mode():
|
||||
t = torch.rand(B)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
model.train()
|
||||
flow_loss = model(x_t, cond_cont, cond_cat, t=t).sum()
|
||||
flow_loss = model(x_t, t, cond_cont, cond_cat).sum()
|
||||
nsec_loss = model.predict_n_sec(cond_cont, cond_cat).sum()
|
||||
(flow_loss + nsec_loss).backward()
|
||||
for name, p in model.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
|
||||
def test_routed_stage1_eval_dispatch_matches_manual_grouping():
|
||||
def test_routed_denoising_mlp_eval_dispatch_matches_manual_grouping():
|
||||
"""Eval-mode grouped top-1 dispatch must equal running each row through
|
||||
its assigned expert individually (batch order shouldn't matter)."""
|
||||
B = 12
|
||||
@@ -922,20 +884,20 @@ def test_routed_stage1_eval_dispatch_matches_manual_grouping():
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
|
||||
with torch.no_grad():
|
||||
batched = model(x_t, cond_cont, cond_cat, t=t)
|
||||
batched = model(x_t, t, cond_cont, cond_cat)
|
||||
|
||||
t_emb = model.time_emb(t)
|
||||
c_emb = model.cond_enc(cond_cont, cond_cat)
|
||||
cond = torch.cat([t_emb, c_emb], dim=-1)
|
||||
idx = model.trunk.router.top1(cond_cont, cond_cat)
|
||||
idx = model.router.top1(cond_cont, cond_cat)
|
||||
manual = torch.zeros_like(x_t)
|
||||
for i in range(B):
|
||||
manual[i] = model.trunk.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
|
||||
manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0]
|
||||
|
||||
torch.testing.assert_close(batched, manual, atol=1e-5, rtol=1e-4)
|
||||
|
||||
|
||||
def test_routed_stage1_predict_n_sec_shape():
|
||||
def test_routed_denoising_mlp_predict_n_sec_shape():
|
||||
B = 6
|
||||
model = _routed_stage1()
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
@@ -943,15 +905,15 @@ def test_routed_stage1_predict_n_sec_shape():
|
||||
assert logits.shape == (B, K_MAX + 1)
|
||||
|
||||
|
||||
def test_routed_stage1_has_no_pdg_embedding_weight_method():
|
||||
def test_routed_denoising_mlp_has_no_pdg_embedding_weight_method():
|
||||
model = _routed_stage1(pdg=5, mat=2)
|
||||
assert not hasattr(model, "pdg_embedding_weight")
|
||||
|
||||
|
||||
# ── Stage2OneShot with a routed trunk ────────────────────────────────────────
|
||||
# ── RoutedSecondaryDecoder ───────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_routed_stage2_output_shape_train_and_eval():
|
||||
def test_routed_secondary_decoder_output_shape_train_and_eval():
|
||||
B = 8
|
||||
decoder = _routed_sec_decoder()
|
||||
x_t = torch.randn(B, SEC_DIM)
|
||||
@@ -960,16 +922,16 @@ def test_routed_stage2_output_shape_train_and_eval():
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
|
||||
decoder.train()
|
||||
out_train = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
||||
out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert out_train.shape == (B, SEC_DIM)
|
||||
|
||||
decoder.eval()
|
||||
with torch.no_grad():
|
||||
out_eval = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t)
|
||||
out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out)
|
||||
assert out_eval.shape == (B, SEC_DIM)
|
||||
|
||||
|
||||
def test_routed_stage2_gradients_flow():
|
||||
def test_routed_secondary_decoder_gradients_flow():
|
||||
B = 4
|
||||
decoder = _routed_sec_decoder(n_experts=3)
|
||||
x_t = torch.randn(B, SEC_DIM)
|
||||
@@ -977,9 +939,7 @@ def test_routed_stage2_gradients_flow():
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
decoder.train()
|
||||
flow_loss = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t).sum()
|
||||
nsec_loss = decoder.predict_n_sec(cond_cont, cond_cat, stage1_out).sum()
|
||||
(flow_loss + nsec_loss).backward()
|
||||
decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward()
|
||||
for name, p in decoder.named_parameters():
|
||||
assert p.grad is not None, f"no grad for {name}"
|
||||
|
||||
@@ -988,33 +948,46 @@ def test_routed_stage2_gradients_flow():
|
||||
|
||||
|
||||
def test_build_models_monolith_when_router_absent():
|
||||
cfg = _nested_cfg(pdg_vocab=4, mat_vocab=2)
|
||||
models = build_models(cfg)
|
||||
stage1, stage2 = models["stage1"], models["stage2"]
|
||||
assert isinstance(stage1, Stage1Model)
|
||||
assert isinstance(stage2, Stage2OneShot)
|
||||
assert isinstance(stage1.trunk, MonolithicTrunk)
|
||||
assert isinstance(stage2.trunk, MonolithicTrunk)
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
hidden_dim=32,
|
||||
n_blocks=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, DenoisingMLP)
|
||||
assert isinstance(sec_decoder, SecondaryDecoder)
|
||||
|
||||
|
||||
def test_build_models_monolith_when_router_disabled():
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
stage1_router={"enabled": False, "type": "energy", "n_experts": 4},
|
||||
hidden_dim=32,
|
||||
n_blocks=2,
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
router={"enabled": False, "type": "energy", "n_experts": 4},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1, stage2 = models["stage1"], models["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
assert isinstance(stage1.trunk, MonolithicTrunk)
|
||||
assert isinstance(stage2.trunk, MonolithicTrunk)
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, DenoisingMLP)
|
||||
assert isinstance(sec_decoder, SecondaryDecoder)
|
||||
|
||||
|
||||
def test_build_models_routed_when_enabled():
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=4,
|
||||
mat_vocab=2,
|
||||
stage1_router={
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=16,
|
||||
expert_n_blocks=2,
|
||||
router={
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 4,
|
||||
@@ -1022,49 +995,37 @@ def test_build_models_routed_when_enabled():
|
||||
"learn_centers": True,
|
||||
"lambda_balance": 0.0,
|
||||
},
|
||||
stage2_router={
|
||||
"enabled": True,
|
||||
"tie_to_stage1": False,
|
||||
"type": "energy",
|
||||
"n_experts": 4,
|
||||
"temperature": 0.5,
|
||||
"learn_centers": True,
|
||||
},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1, stage2 = models["stage1"], models["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
assert isinstance(stage1.trunk, RoutedTrunk)
|
||||
assert isinstance(stage2.trunk, RoutedTrunk)
|
||||
assert len(stage1.trunk.experts) == 4
|
||||
assert len(stage2.trunk.experts) == 4
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
assert isinstance(stage1, RoutedDenoisingMLP)
|
||||
assert isinstance(sec_decoder, RoutedSecondaryDecoder)
|
||||
assert len(stage1.experts) == 4
|
||||
assert len(sec_decoder.experts) == 4
|
||||
|
||||
|
||||
def test_build_models_routed_pair_is_drop_in_for_sample_flow():
|
||||
"""Exercise the exact calling convention giant/sample.py uses."""
|
||||
from giant.sample import sample_flow, sample_secondaries
|
||||
|
||||
cfg = _nested_cfg(
|
||||
model_config = dict(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
stage1_router={"enabled": True, "type": "energy", "n_experts": 2},
|
||||
emb_dim=16,
|
||||
dropout=0.1,
|
||||
k_max=K_MAX,
|
||||
expert_hidden_dim=8,
|
||||
expert_n_blocks=1,
|
||||
router={"enabled": True, "type": "energy", "n_experts": 2},
|
||||
)
|
||||
models = build_models(cfg)
|
||||
stage1, stage2 = models["stage1"], models["stage2"]
|
||||
assert stage1 is not None and stage2 is not None
|
||||
stage1, sec_decoder = build_models(model_config)
|
||||
B = 5
|
||||
cond_cont, cond_cat = _cond(B, pdg=3, mat=2)
|
||||
stage1_norm, n_sec_pred = sample_flow(stage1, cond_cont, cond_cat, steps=2)
|
||||
assert stage1_norm.shape == (B, X_DIM)
|
||||
# A fresh v0.3.0 Stage1Model has no n_sec_head (it moves to stage 2) —
|
||||
# sample_flow returns n_sec_pred=None here, and n_sec must be
|
||||
# asked of stage2 instead, using the just-sampled stage1_norm as context.
|
||||
assert n_sec_pred is None
|
||||
n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
|
||||
assert n_sec_pred.shape == (B,)
|
||||
|
||||
sec_cont, sec_type_emb, sec_valid = sample_secondaries(
|
||||
stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
||||
sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
|
||||
@@ -36,8 +36,7 @@ def _model_cfg() -> dict:
|
||||
|
||||
def _write_checkpoint(tmp_path) -> str:
|
||||
cfg = _model_cfg()
|
||||
stage1 = build_models(cfg)["stage1"]
|
||||
assert stage1 is not None
|
||||
stage1, _ = build_models(cfg)
|
||||
norm = Normalizer()
|
||||
norm.mean = np.zeros(15, dtype=np.float32)
|
||||
norm.std = np.ones(15, dtype=np.float32)
|
||||
|
||||
@@ -1,224 +0,0 @@
|
||||
"""Tests for giant/sample.py's v0.3.0 stage-model sampling — the AR loop
|
||||
(`sample_secondaries_ar`) and non-"physical" `particle_type.target` coverage
|
||||
for the one-shot samplers."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from giant.constants import COND_DIM, CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, X_DIM
|
||||
from giant.model.network import (
|
||||
Stage1Model,
|
||||
Stage2Autoregressive,
|
||||
Stage2OneShot,
|
||||
stage2_trunk_sec_dim,
|
||||
)
|
||||
from giant.sample import (
|
||||
sample_flow,
|
||||
sample_secondaries,
|
||||
sample_secondaries_ar,
|
||||
sample_secondaries_wgan,
|
||||
sample_wgan,
|
||||
)
|
||||
|
||||
_PHYS_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
|
||||
|
||||
def _particle_material_cfg(conditioning: str, emb_dim: int) -> tuple[dict, dict]:
|
||||
cfg = {"type": conditioning, "emb_dim": emb_dim, "n_layers": 1}
|
||||
return dict(cfg), dict(cfg)
|
||||
|
||||
|
||||
def _cond(B: int, pdg: int = 3, mat: int = 2) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.stack([torch.randint(0, pdg, (B,)), torch.randint(0, mat, (B,))], dim=1)
|
||||
return cond_cont, cond_cat
|
||||
|
||||
|
||||
def _conditioning_for(target: str) -> str:
|
||||
# target="embedding" regresses against the conditioning's own embedding
|
||||
# table — only meaningful when the
|
||||
# conditioning axis is itself "embedding".
|
||||
return "embedding" if target == "embedding" else "physical"
|
||||
|
||||
|
||||
def _stage2_oneshot(target: str, generator: str, emb_dim: int = 6, pdg: int = 3, mat: int = 2) -> Stage2OneShot:
|
||||
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
|
||||
particle_type_cfg = {"target": target}
|
||||
# build_models (giant/model/network.py) computes sec_dim this same way
|
||||
# before constructing Stage2OneShot — its own default (SEC_DIM, the
|
||||
# "physical" width) is only correct for target="physical".
|
||||
sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator, K_MAX, emb_dim)
|
||||
return Stage2OneShot(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator=generator,
|
||||
time_dim=16,
|
||||
noise_dim=8,
|
||||
sec_dim=sec_dim,
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
).eval()
|
||||
|
||||
|
||||
def _stage2_ar(
|
||||
target: str,
|
||||
generator: str,
|
||||
emb_dim: int = 6,
|
||||
pdg: int = 3,
|
||||
mat: int = 2,
|
||||
k_max: int = 5,
|
||||
history: str = "markov",
|
||||
) -> Stage2Autoregressive:
|
||||
particle_cfg, material_cfg = _particle_material_cfg(_conditioning_for(target), emb_dim)
|
||||
return Stage2Autoregressive(
|
||||
pdg_vocab=pdg,
|
||||
mat_vocab=mat,
|
||||
particle_cfg=particle_cfg,
|
||||
material_cfg=material_cfg,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator=generator,
|
||||
time_dim=16,
|
||||
noise_dim=8,
|
||||
k_max=k_max,
|
||||
particle_type_cfg={"target": target},
|
||||
history=history,
|
||||
attn_n_heads=2,
|
||||
attn_n_layers=1,
|
||||
).eval()
|
||||
|
||||
|
||||
def _expected_type_dim(target: str, emb_dim: int) -> int:
|
||||
return PARTICLE_PHYS_DIM if target == "physical" else emb_dim
|
||||
|
||||
|
||||
# ── Stage-1 n_sec ownership ──────────────────────────────────────────────────
|
||||
|
||||
|
||||
def test_sample_flow_returns_none_n_sec_when_stage1_owns_no_head():
|
||||
model = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=_PHYS_CFG,
|
||||
material_cfg=_PHYS_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
)
|
||||
cond_cont, cond_cat = _cond(4)
|
||||
sample, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
|
||||
assert sample.shape == (4, X_DIM)
|
||||
assert n_sec is None
|
||||
|
||||
|
||||
def test_sample_wgan_returns_none_n_sec_when_stage1_owns_no_head():
|
||||
model = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=_PHYS_CFG,
|
||||
material_cfg=_PHYS_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
generator="wgan",
|
||||
noise_dim=8,
|
||||
)
|
||||
cond_cont, cond_cat = _cond(4)
|
||||
sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
|
||||
assert sample.shape == (4, X_DIM)
|
||||
assert n_sec is None
|
||||
|
||||
|
||||
def test_sample_flow_returns_n_sec_for_legacy_stage1():
|
||||
model = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=_PHYS_CFG,
|
||||
material_cfg=_PHYS_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
n_sec_head_k_max=K_MAX,
|
||||
)
|
||||
cond_cont, cond_cat = _cond(5)
|
||||
_, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2)
|
||||
assert n_sec is not None and n_sec.shape == (5,)
|
||||
|
||||
|
||||
# ── Stage2OneShot: non-"physical" particle_type.target ──────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
||||
def test_sample_secondaries_flow_shapes_by_target(target):
|
||||
B, emb_dim = 5, 6
|
||||
decoder = _stage2_oneshot(target, "flow", emb_dim=emb_dim)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
|
||||
sec_cont, sec_type, sec_valid = sample_secondaries(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
|
||||
assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
|
||||
assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
assert torch.isfinite(sec_cont).all()
|
||||
assert torch.isfinite(sec_type).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
||||
def test_sample_secondaries_wgan_shapes_by_target(target):
|
||||
B, emb_dim = 5, 6
|
||||
decoder = _stage2_oneshot(target, "wgan", emb_dim=emb_dim)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.randint(0, K_MAX + 1, (B,))
|
||||
sec_cont, sec_type, sec_valid = sample_secondaries_wgan(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
||||
assert sec_cont.shape == (B, K_MAX, CONT_SLOT_DIM)
|
||||
assert sec_type.shape == (B, K_MAX, _expected_type_dim(target, emb_dim))
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
|
||||
|
||||
# ── Stage2Autoregressive ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.parametrize("history", ["markov", "attention"])
|
||||
@pytest.mark.parametrize("generator", ["flow", "wgan"])
|
||||
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
||||
def test_sample_secondaries_ar_shapes(target, generator, history):
|
||||
B, k_max, emb_dim = 4, 5, 6
|
||||
decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max, history=history)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.randint(0, k_max + 1, (B,))
|
||||
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
|
||||
assert sec_cont.shape == (B, k_max, CONT_SLOT_DIM)
|
||||
assert sec_type.shape == (B, k_max, _expected_type_dim(target, emb_dim))
|
||||
assert sec_valid.shape == (B, k_max)
|
||||
assert torch.isfinite(sec_cont).all()
|
||||
assert torch.isfinite(sec_type).all()
|
||||
|
||||
|
||||
@pytest.mark.parametrize("generator", ["flow", "wgan"])
|
||||
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
||||
def test_sample_secondaries_ar_valid_mask_matches_n_sec(target, generator):
|
||||
B, k_max, emb_dim = 3, 5, 6
|
||||
decoder = _stage2_ar(target, generator, emb_dim=emb_dim, k_max=k_max)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.tensor([0, 2, k_max])
|
||||
_, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
|
||||
for i, n in enumerate(n_sec_pred.tolist()):
|
||||
assert sec_valid[i, :n].all()
|
||||
assert not sec_valid[i, n:].any()
|
||||
|
||||
|
||||
def test_sample_secondaries_ar_first_slot_has_no_history():
|
||||
"""Slot 0 always has has_prev=False internally — nothing to assert on
|
||||
the public API directly, but a k_max=1 run should not crash on the
|
||||
"previous token" path at all (has_prev never true)."""
|
||||
B, emb_dim = 3, 6
|
||||
decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=1)
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.tensor([0, 1, 1])
|
||||
sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
|
||||
assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
|
||||
assert sec_valid.tolist() == [[False], [True], [True]]
|
||||
+13
-63
@@ -7,14 +7,15 @@ import pandas as pd
|
||||
import pytest
|
||||
|
||||
from giant.data import setup_cache
|
||||
from giant.data.loader import TopNMap
|
||||
from giant.data.setup_cache import NormalizerEntry, SetupCache
|
||||
from giant.data.transforms import Normalizer
|
||||
|
||||
|
||||
def _touch_parquet(path, n=1):
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
pd.DataFrame({"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}).to_parquet(path)
|
||||
pd.DataFrame(
|
||||
{"pdg": [11] * n, "material": ["G4_AIR"] * n, "process": ["eIoni"] * n}
|
||||
).to_parquet(path)
|
||||
return path
|
||||
|
||||
|
||||
@@ -27,7 +28,9 @@ def _normalizer(width=3):
|
||||
|
||||
def _entry(n_train_steps=100, sample=None):
|
||||
sample = np.array([1.0, 2.0, 3.0], dtype=np.float32) if sample is None else sample
|
||||
return NormalizerEntry(_normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample)
|
||||
return NormalizerEntry(
|
||||
_normalizer(), _normalizer(), _normalizer(2), n_train_steps, sample
|
||||
)
|
||||
|
||||
|
||||
# ── sidecar_path ─────────────────────────────────────────────────────────
|
||||
@@ -35,7 +38,9 @@ def _entry(n_train_steps=100, sample=None):
|
||||
|
||||
def test_sidecar_path_single_file(tmp_path):
|
||||
f = tmp_path / "shard.parquet"
|
||||
assert setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
|
||||
assert (
|
||||
setup_cache.sidecar_path(f) == tmp_path / "shard.parquet.giant_train_cache.json"
|
||||
)
|
||||
|
||||
|
||||
def test_sidecar_path_directory(tmp_path):
|
||||
@@ -45,7 +50,10 @@ def test_sidecar_path_directory(tmp_path):
|
||||
|
||||
def test_sidecar_path_manifest(tmp_path):
|
||||
m = tmp_path / "pools" / "full.manifest"
|
||||
assert setup_cache.sidecar_path(m) == tmp_path / "pools" / "full.manifest.giant_train_cache.json"
|
||||
assert (
|
||||
setup_cache.sidecar_path(m)
|
||||
== tmp_path / "pools" / "full.manifest.giant_train_cache.json"
|
||||
)
|
||||
|
||||
|
||||
# ── fingerprint_files ────────────────────────────────────────────────────
|
||||
@@ -93,37 +101,6 @@ def test_save_load_round_trip(tmp_path):
|
||||
np.testing.assert_allclose(entry.energy_quantiles, [1.0, 2.0, 3.0])
|
||||
|
||||
|
||||
def test_save_load_round_trip_topn_maps(tmp_path):
|
||||
data = _touch_parquet(tmp_path / "shard.parquet")
|
||||
files = [data]
|
||||
|
||||
cache = SetupCache.empty(files)
|
||||
cache.topn_maps[setup_cache.topn_key("pdg", 3)] = TopNMap(
|
||||
class_map={22: 0, 11: 1, 2212: 2}, other_members={2212: 5}
|
||||
)
|
||||
cache.topn_maps[setup_cache.topn_key("material", 2)] = TopNMap(
|
||||
class_map={"G4_AIR": 0, "PbWO4": 1}, other_members={}
|
||||
)
|
||||
|
||||
setup_cache.save(data, files, cache)
|
||||
loaded = setup_cache.load(data, files)
|
||||
|
||||
assert loaded is not None
|
||||
pdg_m = loaded.topn_maps[setup_cache.topn_key("pdg", 3)]
|
||||
assert pdg_m.class_map == {22: 0, 11: 1, 2212: 2}
|
||||
assert pdg_m.other_members == {2212: 5}
|
||||
# key type is int (matches pdg_map's own key type), not str
|
||||
assert all(isinstance(k, int) for k in pdg_m.class_map)
|
||||
|
||||
mat_m = loaded.topn_maps[setup_cache.topn_key("material", 2)]
|
||||
assert mat_m.class_map == {"G4_AIR": 0, "PbWO4": 1}
|
||||
|
||||
|
||||
def test_topn_key_unknown_axis_raises():
|
||||
with pytest.raises(ValueError, match="unknown top-N map axis"):
|
||||
setup_cache.topn_key("process", 4)
|
||||
|
||||
|
||||
def test_load_missing_sidecar_returns_none(tmp_path):
|
||||
data = _touch_parquet(tmp_path / "shard.parquet")
|
||||
assert setup_cache.load(data, [data]) is None
|
||||
@@ -172,25 +149,6 @@ def test_load_invalidates_on_file_content_change(tmp_path):
|
||||
assert setup_cache.load(data, files) is None
|
||||
|
||||
|
||||
def test_load_returns_none_on_malformed_cache_body(tmp_path):
|
||||
"""format_version/dims/fingerprint all check out, but the cache body
|
||||
itself doesn't match SetupCache.from_json's expected shape (e.g. hand-
|
||||
edited or written by a version that changed a nested key) — a clean
|
||||
miss, not a crash."""
|
||||
data = _touch_parquet(tmp_path / "shard.parquet")
|
||||
files = [data]
|
||||
setup_cache.save(data, files, SetupCache.empty(files))
|
||||
|
||||
path = setup_cache.sidecar_path(data)
|
||||
raw = json.loads(path.read_text())
|
||||
raw["vocab"] = {"pdg_map": {"11": 0}} # missing required "mat_map" key
|
||||
path.write_text(json.dumps(raw))
|
||||
|
||||
echoed = []
|
||||
assert setup_cache.load(data, files, echo=echoed.append) is None
|
||||
assert any("malformed" in m for m in echoed)
|
||||
|
||||
|
||||
def test_load_soft_warns_on_git_hash_mismatch_but_still_hits(tmp_path, capsys):
|
||||
data = _touch_parquet(tmp_path / "shard.parquet")
|
||||
files = [data]
|
||||
@@ -367,11 +325,3 @@ def test_compute_event_index_from_files_single_file_unaffected(tmp_path):
|
||||
|
||||
np.testing.assert_array_equal(unique_ids, [5, 7])
|
||||
np.testing.assert_array_equal(counts, [2, 1])
|
||||
|
||||
|
||||
def test_compute_event_index_from_files_empty_file_list():
|
||||
unique_ids, counts = setup_cache.compute_event_index_from_files([])
|
||||
assert unique_ids.size == 0
|
||||
assert counts.size == 0
|
||||
assert unique_ids.dtype == np.int64
|
||||
assert counts.dtype == np.int64
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
import polars as pl
|
||||
|
||||
from giant.tools import steps_to_parquet
|
||||
from scripts import steps_to_parquet
|
||||
|
||||
|
||||
def _frame() -> pl.DataFrame:
|
||||
@@ -24,14 +24,18 @@ def _frame() -> pl.DataFrame:
|
||||
def test_e_sec_sums_child_first_step_energy():
|
||||
out, n_orphaned = steps_to_parquet._add_secondary_attributes(_frame())
|
||||
assert n_orphaned == 0
|
||||
e_sec = dict(zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"]))
|
||||
e_sec = dict(
|
||||
zip(zip(out["track_id"], out["step_no"], out["event_id"]), out["e_sec"])
|
||||
)
|
||||
assert e_sec[(1, 0, 0)] == 15.0 # one child, first-step pre_E 15
|
||||
assert e_sec[(1, 0, 1)] == 50.0 # two children, 20 + 30
|
||||
|
||||
|
||||
def test_e_sec_zero_when_no_children():
|
||||
out, _ = steps_to_parquet._add_secondary_attributes(_frame())
|
||||
childless = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1))
|
||||
childless = out.filter(
|
||||
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 1)
|
||||
)
|
||||
assert childless["e_sec"].item() == 0.0
|
||||
|
||||
|
||||
@@ -65,7 +69,9 @@ def test_orphaned_child_track_is_dropped_not_nulled():
|
||||
}
|
||||
)
|
||||
out, n_orphaned = steps_to_parquet._add_secondary_attributes(df)
|
||||
row = out.filter((pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0))
|
||||
row = out.filter(
|
||||
(pl.col("event_id") == 0) & (pl.col("track_id") == 1) & (pl.col("step_no") == 0)
|
||||
)
|
||||
|
||||
assert n_orphaned == 1
|
||||
assert row["child_track_ids"].to_list() == [[2]]
|
||||
|
||||
@@ -2,7 +2,7 @@ import json
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
from giant.tools import steps_to_parquet_parallel
|
||||
from scripts import steps_to_parquet_parallel
|
||||
|
||||
run_parallel = steps_to_parquet_parallel.run_parallel
|
||||
resolve_destination = steps_to_parquet_parallel.resolve_destination
|
||||
@@ -124,13 +124,31 @@ def _make_dataset(tmp_path: Path, schemas: list[str] | None = None) -> Path:
|
||||
def test_resolve_destination_uses_latest_schema(tmp_path):
|
||||
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3", "schema2"])
|
||||
dest = resolve_destination(root_file, tmp_path, schema_override=None)
|
||||
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema3" / "pbwo4" / "shard-000.parquet"
|
||||
assert (
|
||||
dest
|
||||
== tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema3"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_destination_schema_override_wins(tmp_path):
|
||||
root_file = _make_dataset(tmp_path, schemas=["schema1", "schema3"])
|
||||
dest = resolve_destination(root_file, tmp_path, schema_override="schema9")
|
||||
assert dest == tmp_path / "processed" / "steps" / "gen1" / "schema9" / "pbwo4" / "shard-000.parquet"
|
||||
assert (
|
||||
dest
|
||||
== tmp_path
|
||||
/ "processed"
|
||||
/ "steps"
|
||||
/ "gen1"
|
||||
/ "schema9"
|
||||
/ "pbwo4"
|
||||
/ "shard-000.parquet"
|
||||
)
|
||||
|
||||
|
||||
def test_resolve_destination_errors_without_any_schema(tmp_path):
|
||||
|
||||
+68
-720
@@ -1,49 +1,6 @@
|
||||
"""Tests for giant/training/."""
|
||||
"""Tests for giant/train.py helpers."""
|
||||
|
||||
import copy
|
||||
import csv
|
||||
import math
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from giant.constants import (
|
||||
COND_DIM,
|
||||
CONT_SLOT_DIM,
|
||||
K_MAX,
|
||||
PARTICLE_PHYS_DIM,
|
||||
SEC_SLOT_DIM,
|
||||
X_DIM,
|
||||
)
|
||||
from giant.data.dataset import StepBatch
|
||||
from giant.model.network import build_critics, build_models
|
||||
from giant.training import (
|
||||
FlowDDPMStageTrainer,
|
||||
StageSpec,
|
||||
WGANStageTrainer,
|
||||
build_stage_trainers,
|
||||
train,
|
||||
)
|
||||
from giant.training.metrics import _wandb_run_config
|
||||
from giant.training.stage2_inputs import (
|
||||
_ar_has_prev,
|
||||
_assemble_stage2_ar_inputs,
|
||||
_assemble_stage2_ar_target,
|
||||
_assemble_stage2_real,
|
||||
_gumbel_tau,
|
||||
_relax_onehot_type_slice,
|
||||
_remaining_energy_fraction,
|
||||
_shift_prev,
|
||||
_stage2_tf_prob,
|
||||
_stick_fraction,
|
||||
_type_repr,
|
||||
)
|
||||
|
||||
PDG_VOCAB = 6
|
||||
MAT_VOCAB = 3
|
||||
from giant.train import _gumbel_tau, _wandb_run_config
|
||||
|
||||
|
||||
def test_gumbel_tau_at_step_zero_is_start():
|
||||
@@ -69,681 +26,72 @@ def test_gumbel_tau_handles_zero_total_steps():
|
||||
assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9
|
||||
|
||||
|
||||
def test_wandb_run_config_includes_full_cfg_and_param_counts():
|
||||
cfg = {
|
||||
"train": {"lr": 3e-4},
|
||||
"conditioning": {"out_dim": 128},
|
||||
"stage1_model": {"generator": "flow"},
|
||||
"stage2_model": {"generator": "wgan"},
|
||||
}
|
||||
wcfg = _wandb_run_config(cfg, model_config={"pdg_vocab": 3}, param_counts={"stage1": 100})
|
||||
assert wcfg["train"] == {"lr": 3e-4}
|
||||
assert wcfg["stage1_model"] == {"generator": "flow"}
|
||||
assert wcfg["stage2_model"] == {"generator": "wgan"}
|
||||
assert wcfg["model_config"] == {"pdg_vocab": 3}
|
||||
assert wcfg["param_counts"] == {"stage1": 100}
|
||||
def _base_wandb_kwargs(**overrides):
|
||||
kwargs = dict(
|
||||
mode="flow",
|
||||
epochs=30,
|
||||
lr=3e-4,
|
||||
warmup_epochs=3,
|
||||
weight_decay=0.01,
|
||||
ema_decay=0.9999,
|
||||
lambda_nsec=0.1,
|
||||
lambda_s2=1.0,
|
||||
lambda_balance=0.035,
|
||||
lambda_proc=0.0,
|
||||
lambda_entropy=0.0,
|
||||
gumbel_tau_start=1.0,
|
||||
gumbel_tau_end=0.1,
|
||||
n_critic=5,
|
||||
gp_weight=10.0,
|
||||
model_config={"router": {"enabled": False}},
|
||||
stage1_params=100,
|
||||
sec_decoder_params=50,
|
||||
critic_params=0,
|
||||
sec_critic_params=0,
|
||||
total_params=150,
|
||||
)
|
||||
kwargs.update(overrides)
|
||||
return kwargs
|
||||
|
||||
|
||||
def test_wandb_run_config_omits_router_knobs_when_router_disabled():
|
||||
cfg = _wandb_run_config(**_base_wandb_kwargs())
|
||||
for key in (
|
||||
"lambda_balance",
|
||||
"lambda_proc",
|
||||
"lambda_entropy",
|
||||
"gumbel_tau_start",
|
||||
"gumbel_tau_end",
|
||||
):
|
||||
assert key not in cfg
|
||||
# still present, nested, regardless of router state
|
||||
assert cfg["model"] == {"router": {"enabled": False}}
|
||||
|
||||
|
||||
def test_wandb_run_config_includes_router_knobs_when_router_enabled():
|
||||
cfg = _wandb_run_config(
|
||||
**_base_wandb_kwargs(model_config={"router": {"enabled": True}})
|
||||
)
|
||||
assert cfg["lambda_balance"] == 0.035
|
||||
assert cfg["lambda_proc"] == 0.0
|
||||
assert cfg["lambda_entropy"] == 0.0
|
||||
assert cfg["gumbel_tau_start"] == 1.0
|
||||
assert cfg["gumbel_tau_end"] == 0.1
|
||||
|
||||
|
||||
def test_wandb_run_config_omits_wgan_knobs_when_mode_is_not_wgan():
|
||||
cfg = _wandb_run_config(**_base_wandb_kwargs(mode="flow"))
|
||||
assert "n_critic" not in cfg
|
||||
assert "gp_weight" not in cfg
|
||||
|
||||
|
||||
def test_wandb_run_config_includes_wgan_knobs_when_mode_is_wgan():
|
||||
cfg = _wandb_run_config(**_base_wandb_kwargs(mode="wgan"))
|
||||
assert cfg["n_critic"] == 5
|
||||
assert cfg["gp_weight"] == 10.0
|
||||
|
||||
|
||||
def test_wandb_run_config_handles_missing_model_config():
|
||||
cfg = {"train": {}, "conditioning": {}, "stage1_model": {}, "stage2_model": {}}
|
||||
wcfg = _wandb_run_config(cfg, model_config=None, param_counts={})
|
||||
assert wcfg["model_config"] == {}
|
||||
|
||||
|
||||
# --- AR helper functions (v0.3.0 step 5) ---------
|
||||
|
||||
|
||||
def test_stick_fraction_matches_sigmoid_of_logit():
|
||||
sec_cont = torch.zeros(2, 3, SEC_SLOT_DIM)
|
||||
sec_cont[..., 0] = torch.tensor([[0.0, 2.0, -2.0], [1.0, -1.0, 0.0]])
|
||||
frac = _stick_fraction(sec_cont)
|
||||
assert torch.allclose(frac, torch.sigmoid(sec_cont[..., 0]))
|
||||
|
||||
|
||||
def test_remaining_energy_fraction_hand_computed():
|
||||
fraction = torch.tensor([[0.5, 0.5, 1.0]])
|
||||
remaining = _remaining_energy_fraction(fraction)
|
||||
assert torch.allclose(remaining, torch.tensor([[1.0, 0.5, 0.25]]))
|
||||
|
||||
|
||||
def test_shift_prev_shifts_and_zero_pads_slot0():
|
||||
x = torch.arange(2 * 4 * 3).reshape(2, 4, 3).float()
|
||||
shifted = _shift_prev(x)
|
||||
assert torch.all(shifted[:, 0] == 0)
|
||||
assert torch.equal(shifted[:, 1:], x[:, :-1])
|
||||
|
||||
|
||||
def test_ar_has_prev_false_only_at_slot_zero():
|
||||
has_prev = _ar_has_prev(5, torch.device("cpu"))
|
||||
assert has_prev.shape == (1, 5)
|
||||
assert has_prev.tolist() == [[False, True, True, True, True]]
|
||||
|
||||
|
||||
# --- _stage2_tf_prob (v0.3.0 step 7) -----------
|
||||
|
||||
|
||||
def test_stage2_tf_prob_always_is_constant_one():
|
||||
assert _stage2_tf_prob("always", 1.0, 0.0, 0, 10) == 1.0
|
||||
assert _stage2_tf_prob("always", 1.0, 0.0, 9, 10) == 1.0
|
||||
|
||||
|
||||
def test_stage2_tf_prob_never_is_constant_zero():
|
||||
assert _stage2_tf_prob("never", 1.0, 1.0, 0, 10) == 0.0
|
||||
assert _stage2_tf_prob("never", 1.0, 1.0, 9, 10) == 0.0
|
||||
|
||||
|
||||
def test_stage2_tf_prob_scheduled_interpolates_linearly():
|
||||
assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 11) == 1.0
|
||||
assert abs(_stage2_tf_prob("scheduled", 1.0, 0.0, 5, 11) - 0.5) < 1e-9
|
||||
assert _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11) == 0.0
|
||||
|
||||
|
||||
def test_stage2_tf_prob_scheduled_clamps_beyond_total_epochs():
|
||||
end = _stage2_tf_prob("scheduled", 1.0, 0.0, 10, 11)
|
||||
beyond = _stage2_tf_prob("scheduled", 1.0, 0.0, 50, 11)
|
||||
assert beyond == end
|
||||
|
||||
|
||||
def test_stage2_tf_prob_scheduled_handles_single_epoch():
|
||||
# total_epochs=1 is guarded to a denominator of 1 internally (like
|
||||
# _gumbel_tau's total_steps=0 guard) — epoch=0 gives zero progress.
|
||||
assert _stage2_tf_prob("scheduled", 1.0, 0.0, 0, 1) == 1.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"])
|
||||
def test_type_repr_shapes_and_values(target):
|
||||
B, K, emb_dim = 3, 4, 6
|
||||
sec_cont = torch.randn(B, K, SEC_SLOT_DIM)
|
||||
sec_type_idx = torch.randint(0, emb_dim, (B, K))
|
||||
cond_enc = torch.nn.Module()
|
||||
if target == "embedding":
|
||||
cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
|
||||
repr_ = _type_repr(sec_type_idx, sec_cont, {"target": target}, cond_enc, emb_dim)
|
||||
expected_width = PARTICLE_PHYS_DIM if target == "physical" else emb_dim
|
||||
assert repr_.shape == (B, K, expected_width)
|
||||
if target == "physical":
|
||||
assert torch.equal(repr_, sec_cont[..., CONT_SLOT_DIM : CONT_SLOT_DIM + PARTICLE_PHYS_DIM])
|
||||
if target == "onehot":
|
||||
assert torch.all(repr_.sum(-1) == 1.0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"target,generator",
|
||||
[
|
||||
("physical", "flow"),
|
||||
("physical", "wgan"),
|
||||
("onehot", "flow"),
|
||||
("onehot", "wgan"),
|
||||
("embedding", "flow"),
|
||||
("embedding", "wgan"),
|
||||
],
|
||||
)
|
||||
def test_assemble_stage2_ar_target_matches_assemble_stage2_real_flattened(target, generator):
|
||||
"""Regression test tying the refactor together: _assemble_stage2_real is
|
||||
now defined as _assemble_stage2_ar_target(...).flatten(1)."""
|
||||
B, emb_dim = 4, 6
|
||||
sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
|
||||
sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
|
||||
cond_enc = torch.nn.Module()
|
||||
if target == "embedding":
|
||||
cond_enc.pdg_emb = torch.nn.Embedding(emb_dim, emb_dim)
|
||||
particle_type_cfg = {"target": target}
|
||||
flat = _assemble_stage2_real(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
|
||||
unflat = _assemble_stage2_ar_target(sec_cont, sec_type_idx, particle_type_cfg, generator, cond_enc, emb_dim)
|
||||
assert torch.equal(unflat.flatten(1), flat)
|
||||
|
||||
|
||||
def test_assemble_stage2_ar_inputs_shapes_and_history_feat_width():
|
||||
B, emb_dim = 3, 6
|
||||
sec_cont = torch.randn(B, K_MAX, SEC_SLOT_DIM)
|
||||
sec_type_idx = torch.randint(0, emb_dim, (B, K_MAX))
|
||||
cond_enc = torch.nn.Module()
|
||||
out = _assemble_stage2_ar_inputs(sec_cont, sec_type_idx, {"target": "physical"}, cond_enc, emb_dim)
|
||||
assert out["history_feat"].shape == (B, K_MAX, CONT_SLOT_DIM + PARTICLE_PHYS_DIM)
|
||||
assert out["has_prev"].shape == (B, K_MAX)
|
||||
assert out["remaining_frac"].shape == (B, K_MAX)
|
||||
assert out["slot_idx"].shape == (B, K_MAX)
|
||||
assert torch.all(out["slot_idx"][:, 0] == 0.0)
|
||||
assert torch.all(out["slot_idx"][:, -1] == 1.0)
|
||||
|
||||
|
||||
def test_relax_onehot_type_slice_grad_probe_populates_both_norms():
|
||||
B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
|
||||
x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
|
||||
grad_probe: dict[str, float] = {}
|
||||
out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5, grad_probe=grad_probe)
|
||||
out.sum().backward()
|
||||
assert grad_probe["cont"] >= 0.0
|
||||
assert grad_probe["type"] >= 0.0
|
||||
|
||||
|
||||
def test_relax_onehot_type_slice_grad_probe_none_is_backward_compatible():
|
||||
B, k_max, cont_dim, type_dim = 4, K_MAX, CONT_SLOT_DIM, 6
|
||||
x_flat = torch.randn(B, k_max * (cont_dim + type_dim), requires_grad=True)
|
||||
out = _relax_onehot_type_slice(x_flat, k_max, cont_dim, type_dim, tau=0.5)
|
||||
out.sum().backward()
|
||||
assert x_flat.grad is not None
|
||||
|
||||
|
||||
# --- end-to-end train() integration tests -----------------------------------
|
||||
|
||||
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
|
||||
|
||||
def _base_cfg():
|
||||
return {
|
||||
"conditioning": {
|
||||
"out_dim": 32,
|
||||
"share_stages": False,
|
||||
"particle": dict(PARTICLE_CFG),
|
||||
"material": dict(MATERIAL_CFG),
|
||||
},
|
||||
"stage1_model": {
|
||||
"active": True,
|
||||
"generator": "flow",
|
||||
"hidden_dim": 24,
|
||||
"n_res_blocks": 2,
|
||||
"dropout": 0.0,
|
||||
"lambda": 1.0,
|
||||
"flow": {"time_dim": 16},
|
||||
"ddpm": {"time_dim": 16, "n_steps": 50},
|
||||
"wgan": {
|
||||
"noise_dim": 16,
|
||||
"n_critic": 2,
|
||||
"gp_weight": 10.0,
|
||||
"critic_lr": 0.0,
|
||||
},
|
||||
"router": {"enabled": False},
|
||||
},
|
||||
"stage2_model": {
|
||||
"active": True,
|
||||
"decoder": "one_shot",
|
||||
"generator": "wgan",
|
||||
"hidden_dim": 24,
|
||||
"n_res_blocks": 2,
|
||||
"dropout": 0.0,
|
||||
"lambda": 1.0,
|
||||
"k_max": K_MAX,
|
||||
"context_dim": 16,
|
||||
"n_sec": {"mode": "head", "lambda": 0.1},
|
||||
# Explicit, not relying on the fallback default (which is
|
||||
# "onehot", matching DEFAULT_CONFIG — see issues.md Issue 1):
|
||||
# the "physical"-labelled cases below (and this fixture's own
|
||||
# comment history) intend this as the base "physical" case,
|
||||
# with "*_onehot"/"*_embedding" cases opting in explicitly.
|
||||
"particle_type": {"target": "physical", "lambda": 1.0},
|
||||
"flow": {"time_dim": 16},
|
||||
"ddpm": {"time_dim": 16, "n_steps": 50},
|
||||
"wgan": {
|
||||
"noise_dim": 16,
|
||||
"n_critic": 2,
|
||||
"gp_weight": 10.0,
|
||||
"critic_lr": 0.0,
|
||||
},
|
||||
"router": {"enabled": False, "tie_to_stage1": False},
|
||||
},
|
||||
"train": {
|
||||
"epochs": 2,
|
||||
"batch_size": 8,
|
||||
"lr": 3e-4,
|
||||
"weight_decay": 0.01,
|
||||
"ema_decay": 0.999,
|
||||
"warmup_epochs": 0,
|
||||
"val_fraction": 0.1,
|
||||
"max_val_batches": 0,
|
||||
"num_workers": 0,
|
||||
"seed": 0,
|
||||
"validate_every": 0,
|
||||
"validate_steps": 2,
|
||||
"wandb": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _fake_batches(n_batches, batch_size, seed=0):
|
||||
g = torch.Generator().manual_seed(seed)
|
||||
batches = []
|
||||
for _ in range(n_batches):
|
||||
cond_cont = torch.randn(batch_size, COND_DIM, generator=g)
|
||||
cond_cat = torch.stack(
|
||||
[
|
||||
torch.randint(0, PDG_VOCAB, (batch_size,), generator=g),
|
||||
torch.randint(0, MAT_VOCAB, (batch_size,), generator=g),
|
||||
],
|
||||
dim=1,
|
||||
)
|
||||
x1 = torch.randn(batch_size, X_DIM, generator=g)
|
||||
n_sec = torch.randint(0, K_MAX, (batch_size,), generator=g)
|
||||
sec_cont = torch.randn(batch_size, K_MAX, SEC_SLOT_DIM, generator=g)
|
||||
proc_idx = torch.zeros(batch_size, dtype=torch.long)
|
||||
sec_type_idx = torch.zeros(batch_size, K_MAX, dtype=torch.long)
|
||||
batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
|
||||
return batches
|
||||
|
||||
|
||||
def _model_config(cfg):
|
||||
return {
|
||||
"pdg_vocab": PDG_VOCAB,
|
||||
"mat_vocab": MAT_VOCAB,
|
||||
"conditioning": cfg["conditioning"],
|
||||
"stage1_model": cfg["stage1_model"],
|
||||
"stage2_model": cfg["stage2_model"],
|
||||
}
|
||||
|
||||
|
||||
def _run_train(cfg, out_dir, resume_path=None):
|
||||
model_config = _model_config(cfg)
|
||||
models = build_models(model_config)
|
||||
critics = build_critics(model_config)
|
||||
train_loader = _fake_batches(4, cfg["train"]["batch_size"])
|
||||
val_loader = _fake_batches(2, cfg["train"]["batch_size"], seed=1)
|
||||
train(
|
||||
cfg=cfg,
|
||||
models=models,
|
||||
critics=critics,
|
||||
train_loader=train_loader,
|
||||
val_loader=val_loader,
|
||||
device=torch.device("cpu"),
|
||||
out_dir=out_dir,
|
||||
normalizer_dict={"cond": {}, "target": {}, "sec_phys": {}},
|
||||
pdg_map={"22": 0},
|
||||
mat_map={"G4_AIR": 0},
|
||||
proc_map=None,
|
||||
model_config=model_config,
|
||||
total_train_batches=4,
|
||||
resume_path=resume_path,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"label,mutate",
|
||||
[
|
||||
("both_flow", lambda cfg: None),
|
||||
("both_wgan", lambda cfg: cfg["stage1_model"].__setitem__("generator", "wgan")),
|
||||
(
|
||||
"mixed_stage1_flow_stage2_wgan",
|
||||
lambda cfg: None, # already the default
|
||||
),
|
||||
(
|
||||
"mixed_stage1_wgan_stage2_flow",
|
||||
lambda cfg: (
|
||||
cfg["stage1_model"].__setitem__("generator", "wgan"),
|
||||
cfg["stage2_model"].__setitem__("generator", "flow"),
|
||||
),
|
||||
),
|
||||
("stage1_only", lambda cfg: cfg["stage2_model"].__setitem__("active", False)),
|
||||
("stage2_only", lambda cfg: cfg["stage1_model"].__setitem__("active", False)),
|
||||
(
|
||||
"both_ddpm_stage1_flow_stage2",
|
||||
lambda cfg: (
|
||||
cfg["stage1_model"].__setitem__("generator", "ddpm"),
|
||||
cfg["stage2_model"].__setitem__("generator", "flow"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"routed_stage1_energy_gumbel",
|
||||
lambda cfg: cfg["stage1_model"].__setitem__(
|
||||
"router",
|
||||
{
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 3,
|
||||
"temperature": 0.5,
|
||||
"learn_centers": True,
|
||||
"lambda_balance": 0.1,
|
||||
"lambda_entropy": 0.01,
|
||||
"gumbel": True,
|
||||
"gumbel_tau_start": 1.0,
|
||||
"gumbel_tau_end": 0.1,
|
||||
},
|
||||
),
|
||||
),
|
||||
(
|
||||
"stage2_onehot_target_wgan",
|
||||
lambda cfg: cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
|
||||
),
|
||||
(
|
||||
"stage2_onehot_target_flow",
|
||||
lambda cfg: (
|
||||
cfg["stage2_model"].__setitem__("generator", "flow"),
|
||||
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
|
||||
),
|
||||
),
|
||||
(
|
||||
"stage2_embedding_target_wgan",
|
||||
lambda cfg: (
|
||||
cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
|
||||
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
|
||||
cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
|
||||
),
|
||||
),
|
||||
(
|
||||
"stage2_embedding_target_flow",
|
||||
lambda cfg: (
|
||||
cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
|
||||
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
|
||||
cfg["stage2_model"].__setitem__("generator", "flow"),
|
||||
cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
|
||||
),
|
||||
),
|
||||
(
|
||||
"ar_wgan_onehot",
|
||||
lambda cfg: (
|
||||
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
|
||||
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
|
||||
),
|
||||
),
|
||||
(
|
||||
"ar_wgan_physical",
|
||||
lambda cfg: cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
|
||||
),
|
||||
(
|
||||
"ar_flow_onehot",
|
||||
lambda cfg: (
|
||||
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
|
||||
cfg["stage2_model"].__setitem__("generator", "flow"),
|
||||
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
|
||||
),
|
||||
),
|
||||
(
|
||||
"ar_flow_embedding",
|
||||
lambda cfg: (
|
||||
cfg["conditioning"]["particle"].__setitem__("type", "embedding"),
|
||||
cfg["conditioning"]["material"].__setitem__("type", "embedding"),
|
||||
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
|
||||
cfg["stage2_model"].__setitem__("generator", "flow"),
|
||||
cfg["stage2_model"].__setitem__("particle_type", {"target": "embedding", "lambda": 1.0}),
|
||||
),
|
||||
),
|
||||
(
|
||||
"ar_stage2_only",
|
||||
lambda cfg: (
|
||||
cfg["stage1_model"].__setitem__("active", False),
|
||||
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
|
||||
),
|
||||
),
|
||||
(
|
||||
"ar_mixed_stage1_wgan_stage2_flow_onehot",
|
||||
lambda cfg: (
|
||||
cfg["stage1_model"].__setitem__("generator", "wgan"),
|
||||
cfg["stage2_model"].__setitem__("generator", "flow"),
|
||||
cfg["stage2_model"].__setitem__("decoder", "autoregressive"),
|
||||
cfg["stage2_model"].__setitem__("particle_type", {"target": "onehot", "lambda": 1.0}),
|
||||
),
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_train_end_to_end(label, mutate):
|
||||
cfg = _base_cfg()
|
||||
mutate(cfg)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
out_dir = Path(tmp) / "run"
|
||||
_run_train(cfg, out_dir)
|
||||
assert (out_dir / "last.pt").exists()
|
||||
assert (out_dir / "metrics.csv").exists()
|
||||
ckpt = torch.load(out_dir / "last.pt", weights_only=False)
|
||||
if cfg["stage1_model"]["active"]:
|
||||
assert "model" in ckpt
|
||||
else:
|
||||
assert "model" not in ckpt
|
||||
if cfg["stage2_model"]["active"]:
|
||||
assert "sec_decoder" in ckpt
|
||||
else:
|
||||
assert "sec_decoder" not in ckpt
|
||||
|
||||
|
||||
def test_train_resume_continues_from_checkpoint():
|
||||
cfg = _base_cfg()
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
out_dir = Path(tmp) / "run"
|
||||
_run_train(cfg, out_dir)
|
||||
ckpt_before = torch.load(out_dir / "last.pt", weights_only=False)
|
||||
assert ckpt_before["epoch"] == 2
|
||||
|
||||
cfg2 = copy.deepcopy(cfg)
|
||||
cfg2["train"]["epochs"] = 3
|
||||
_run_train(cfg2, out_dir, resume_path=out_dir / "last.pt")
|
||||
ckpt_after = torch.load(out_dir / "last.pt", weights_only=False)
|
||||
assert ckpt_after["epoch"] == 3
|
||||
assert ckpt_after["global_step"] > ckpt_before["global_step"]
|
||||
|
||||
|
||||
def test_train_raises_when_no_active_stage():
|
||||
cfg = _base_cfg()
|
||||
cfg["stage1_model"]["active"] = False
|
||||
cfg["stage2_model"]["active"] = False
|
||||
model_config = _model_config(cfg)
|
||||
models = build_models(model_config)
|
||||
critics = build_critics(model_config)
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
with pytest.raises(ValueError, match="no active stage"):
|
||||
train(
|
||||
cfg=cfg,
|
||||
models=models,
|
||||
critics=critics,
|
||||
train_loader=_fake_batches(1, 8),
|
||||
val_loader=_fake_batches(1, 8),
|
||||
device=torch.device("cpu"),
|
||||
out_dir=Path(tmp) / "run",
|
||||
model_config=model_config,
|
||||
total_train_batches=1,
|
||||
)
|
||||
|
||||
|
||||
def test_metrics_csv_columns_are_stage_prefixed():
|
||||
cfg = _base_cfg()
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
out_dir = Path(tmp) / "run"
|
||||
_run_train(cfg, out_dir)
|
||||
header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
|
||||
assert "stage1/train/loss" in header
|
||||
assert "stage2/train/d_loss" in header
|
||||
assert "val/loss" in header
|
||||
assert "epoch" in header
|
||||
|
||||
|
||||
def test_wgan_stage_trainer_skips_generator_step_when_no_grad_this_batch():
|
||||
"""Regression test: on a non-generator-step batch, if this stage's model
|
||||
has no n_sec_head (n_sec defaults to stage 2), g_loss is a
|
||||
graph-less zero — .backward() must not be called on it."""
|
||||
cfg = _base_cfg()
|
||||
cfg["stage1_model"]["generator"] = "wgan"
|
||||
model_config = _model_config(cfg)
|
||||
models = build_models(model_config)
|
||||
critics = build_critics(model_config)
|
||||
assert models["stage1"] is not None and critics["stage1"] is not None
|
||||
spec = StageSpec(
|
||||
name="stage1",
|
||||
is_stage2=False,
|
||||
generator="wgan",
|
||||
n_critic=1000, # never a generator step in this test
|
||||
ema_decay=0.0,
|
||||
steps_per_epoch=4,
|
||||
)
|
||||
trainer = WGANStageTrainer(spec, models["stage1"], critics["stage1"], torch.device("cpu"))
|
||||
assert trainer.model.n_sec_head is None
|
||||
batch = _fake_batches(1, 8)[0]
|
||||
stats = trainer.step(batch, torch.device("cpu"), global_step=1)
|
||||
assert stats["did_g_step"] is False
|
||||
|
||||
|
||||
def test_flow_stage_trainer_ddpm_not_implemented_for_stage2():
|
||||
spec = StageSpec(name="stage2", is_stage2=True, generator="ddpm", ddpm_n_steps=50, ema_decay=0.0)
|
||||
with pytest.raises(NotImplementedError):
|
||||
FlowDDPMStageTrainer(spec, torch.nn.Linear(1, 1), torch.device("cpu"))
|
||||
|
||||
|
||||
def test_stage_spec_from_config_omitted_decoder_and_particle_type_match_default_config():
|
||||
"""Regression for issues.md Issue 1: StageSpec.from_config's own fallback
|
||||
defaults for stage2_model.decoder/particle_type must equal
|
||||
DEFAULT_CONFIG's ("autoregressive" / "onehot"), not the old, now-wrong
|
||||
("one_shot" / "physical") literals a .get(key, default) call used to
|
||||
supply when a hand-built cfg omitted these keys."""
|
||||
cfg = _base_cfg()
|
||||
del cfg["stage2_model"]["decoder"]
|
||||
del cfg["stage2_model"]["particle_type"]
|
||||
spec = StageSpec.from_config(cfg, "stage2", is_stage2=True, steps_per_epoch=1)
|
||||
assert spec.decoder == "autoregressive"
|
||||
assert spec.particle_type.target == "onehot"
|
||||
|
||||
|
||||
def _routed_stage1_trainer(lambda_balance, lambda_proc, lambda_entropy):
|
||||
cfg = _base_cfg()
|
||||
cfg["stage1_model"]["router"] = {
|
||||
"enabled": True,
|
||||
"type": "energy",
|
||||
"n_experts": 3,
|
||||
"temperature": 0.5,
|
||||
"learn_centers": True,
|
||||
"lambda_balance": lambda_balance,
|
||||
"lambda_proc": lambda_proc,
|
||||
"lambda_entropy": lambda_entropy,
|
||||
}
|
||||
model_config = _model_config(cfg)
|
||||
models = build_models(model_config)
|
||||
critics = build_critics(model_config)
|
||||
trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
|
||||
return trainers["stage1"]
|
||||
|
||||
|
||||
def test_router_aux_losses_skipped_when_lambda_zero_but_run_when_positive():
|
||||
"""Gitea #31: FlowDDPMStageTrainer._compute must not call
|
||||
router.balance_loss/classify_loss/entropy_loss when the corresponding
|
||||
lambda is 0 (the default) -- those calls do their own router.gate(...)
|
||||
forward pass that is wasted once the term is masked out of the total
|
||||
loss anyway. Checked both ways: zero lambdas must skip all three calls,
|
||||
positive lambdas must still make them (the guard must not accidentally
|
||||
suppress the real path)."""
|
||||
batch = _fake_batches(1, 4)[0]
|
||||
device = torch.device("cpu")
|
||||
|
||||
trainer_zero = _routed_stage1_trainer(0.0, 0.0, 0.0)
|
||||
router_zero = trainer_zero.router
|
||||
router_zero.balance_loss = MagicMock(wraps=router_zero.balance_loss)
|
||||
router_zero.classify_loss = MagicMock(wraps=router_zero.classify_loss)
|
||||
router_zero.entropy_loss = MagicMock(wraps=router_zero.entropy_loss)
|
||||
stats_zero = trainer_zero.step(batch, device, global_step=1)
|
||||
assert router_zero.balance_loss.call_count == 0
|
||||
assert router_zero.classify_loss.call_count == 0
|
||||
assert router_zero.entropy_loss.call_count == 0
|
||||
assert stats_zero["loss_balance"] == 0.0
|
||||
assert stats_zero["loss_proc"] == 0.0
|
||||
assert stats_zero["loss_entropy"] == 0.0
|
||||
|
||||
trainer_pos = _routed_stage1_trainer(0.1, 0.1, 0.01)
|
||||
router_pos = trainer_pos.router
|
||||
router_pos.balance_loss = MagicMock(wraps=router_pos.balance_loss)
|
||||
router_pos.classify_loss = MagicMock(wraps=router_pos.classify_loss)
|
||||
router_pos.entropy_loss = MagicMock(wraps=router_pos.entropy_loss)
|
||||
trainer_pos.step(batch, device, global_step=1)
|
||||
assert router_pos.balance_loss.call_count == 1
|
||||
assert router_pos.classify_loss.call_count == 1
|
||||
assert router_pos.entropy_loss.call_count == 1
|
||||
|
||||
|
||||
# --- AR trainer wiring (v0.3.0 step 5) --------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("teacher_forcing", ["always", "scheduled", "never"])
|
||||
@pytest.mark.parametrize("history", ["markov", "attention"])
|
||||
@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
|
||||
def test_build_stage_trainers_ar_scheduled_and_attention_step_runs(teacher_forcing, history, stage2_generator):
|
||||
"""v0.3.0 step 7: history='attention' and teacher_forcing in
|
||||
{'scheduled', 'never'} must actually train — a stage-2 AR trainer.step()
|
||||
must run and produce a finite loss, for every {history} x
|
||||
{teacher_forcing} x {generator} combination."""
|
||||
cfg = _base_cfg()
|
||||
cfg["stage2_model"]["decoder"] = "autoregressive"
|
||||
cfg["stage2_model"]["generator"] = stage2_generator
|
||||
cfg["stage2_model"]["autoregressive"] = {
|
||||
"history": history,
|
||||
"teacher_forcing": teacher_forcing,
|
||||
"tf_p_start": 1.0,
|
||||
"tf_p_end": 0.0,
|
||||
"attn_n_heads": 2,
|
||||
"attn_n_layers": 1,
|
||||
}
|
||||
model_config = _model_config(cfg)
|
||||
models = build_models(model_config)
|
||||
critics = build_critics(model_config)
|
||||
trainers = build_stage_trainers(cfg, models, critics, torch.device("cpu"), total_train_batches=4)
|
||||
trainer = trainers["stage2"]
|
||||
batch = _fake_batches(1, 4)[0]
|
||||
stats = trainer.step(batch, torch.device("cpu"), global_step=1)
|
||||
loss_key = "g_loss" if stage2_generator == "wgan" else "loss"
|
||||
assert math.isfinite(stats[loss_key])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("stage2_generator", ["wgan", "flow"])
|
||||
def test_train_end_to_end_ar_attention_history_scheduled_teacher_forcing(
|
||||
stage2_generator,
|
||||
):
|
||||
"""Full `train()` run (not just one `trainer.step()` call) with
|
||||
history='attention' AND teacher_forcing='scheduled' together — the
|
||||
combination v0.3.0 step 7 exists to land — must complete and write a
|
||||
checkpoint + metrics.csv with finite losses throughout."""
|
||||
cfg = _base_cfg()
|
||||
cfg["stage2_model"]["decoder"] = "autoregressive"
|
||||
cfg["stage2_model"]["generator"] = stage2_generator
|
||||
cfg["stage2_model"]["autoregressive"] = {
|
||||
"history": "attention",
|
||||
"teacher_forcing": "scheduled",
|
||||
"tf_p_start": 1.0,
|
||||
"tf_p_end": 0.0,
|
||||
"attn_n_heads": 2,
|
||||
"attn_n_layers": 1,
|
||||
}
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
out_dir = Path(tmp) / "run"
|
||||
_run_train(cfg, out_dir)
|
||||
assert (out_dir / "last.pt").exists()
|
||||
with open(out_dir / "metrics.csv", newline="") as f:
|
||||
rows = list(csv.DictReader(f))
|
||||
assert len(rows) == cfg["train"]["epochs"]
|
||||
loss_col = "stage2/train/g_loss" if stage2_generator == "wgan" else "stage2/train/loss"
|
||||
assert all(math.isfinite(float(r[loss_col])) for r in rows)
|
||||
|
||||
|
||||
def test_ar_wgan_onehot_grad_norm_instrumentation_populates_metrics():
|
||||
"""Differentiability validation-obligation instrumentation: the
|
||||
trunk-gradient-norm-by-slice columns must appear and actually fire for
|
||||
generator='wgan' + particle_type.target='onehot' under decoder=
|
||||
'autoregressive' (added at v0.3.0 step 5 to accrue evidence during the
|
||||
architecture comparison)."""
|
||||
cfg = _base_cfg()
|
||||
cfg["stage2_model"]["decoder"] = "autoregressive"
|
||||
cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0}
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
out_dir = Path(tmp) / "run"
|
||||
_run_train(cfg, out_dir)
|
||||
with open(out_dir / "metrics.csv", newline="") as f:
|
||||
rows = list(csv.DictReader(f))
|
||||
assert "stage2/train/grad_norm_type_slice" in rows[0]
|
||||
assert "stage2/train/grad_norm_cont_slice" in rows[0]
|
||||
assert any(float(r["stage2/train/grad_norm_type_slice"]) > 0 for r in rows)
|
||||
assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows)
|
||||
|
||||
|
||||
def test_wgan_onehot_one_shot_also_gets_grad_norm_instrumentation():
|
||||
"""The instrumentation is decoder-agnostic — one_shot + wgan + onehot
|
||||
must populate the same columns."""
|
||||
cfg = _base_cfg()
|
||||
cfg["stage2_model"]["particle_type"] = {"target": "onehot", "lambda": 1.0}
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
out_dir = Path(tmp) / "run"
|
||||
_run_train(cfg, out_dir)
|
||||
with open(out_dir / "metrics.csv", newline="") as f:
|
||||
rows = list(csv.DictReader(f))
|
||||
assert any(float(r["stage2/train/grad_norm_type_slice"]) > 0 for r in rows)
|
||||
assert any(float(r["stage2/train/grad_norm_cont_slice"]) > 0 for r in rows)
|
||||
|
||||
|
||||
def test_wgan_physical_omits_grad_norm_slice_columns():
|
||||
cfg = _base_cfg() # _base_cfg's stage2_model.particle_type.target is "physical"
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
out_dir = Path(tmp) / "run"
|
||||
_run_train(cfg, out_dir)
|
||||
header = (out_dir / "metrics.csv").read_text().splitlines()[0].split(",")
|
||||
assert "stage2/train/grad_norm_type_slice" not in header
|
||||
assert "stage2/train/grad_norm_cont_slice" not in header
|
||||
cfg = _wandb_run_config(**_base_wandb_kwargs(model_config=None))
|
||||
assert cfg["model"] == {}
|
||||
assert "lambda_balance" not in cfg
|
||||
|
||||
+49
-63
@@ -162,7 +162,9 @@ def test_local_frame_rotation_normalizes_non_unit_pre_dir():
|
||||
post_dir = rng.standard_normal((N, 3)).astype(np.float32)
|
||||
post_dir /= np.linalg.norm(post_dir, axis=1, keepdims=True)
|
||||
|
||||
pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(np.float32)
|
||||
pre_dir_scaled = pre_dir_unit * rng.uniform(0.9, 1.1, size=(N, 1)).astype(
|
||||
np.float32
|
||||
)
|
||||
expected = local_frame_rotation(pre_dir_unit, post_dir)
|
||||
result = local_frame_rotation(pre_dir_scaled, post_dir)
|
||||
np.testing.assert_allclose(result, expected, atol=1e-4)
|
||||
@@ -198,10 +200,14 @@ def test_reconstruct_post_pos_straight_line():
|
||||
step_length = rng.uniform(0.1, 5.0, size=N).astype(np.float32)
|
||||
post_pos = pre_pos + step_length[:, None] * pre_dir
|
||||
|
||||
travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
|
||||
travel_dir_local = local_frame_rotation(
|
||||
pre_dir, travel_direction(pre_pos, post_pos)
|
||||
)
|
||||
np.testing.assert_allclose(travel_dir_local, np.tile([0, 0, 1], (N, 1)), atol=1e-4)
|
||||
|
||||
reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
|
||||
reconstructed = reconstruct_post_pos(
|
||||
pre_pos, pre_dir, step_length, travel_dir_local
|
||||
)
|
||||
np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
|
||||
|
||||
|
||||
@@ -215,8 +221,12 @@ def test_reconstruct_post_pos_general_roundtrip():
|
||||
post_pos = pre_pos + rng.standard_normal((N, 3)).astype(np.float32)
|
||||
step_length = np.linalg.norm(post_pos - pre_pos, axis=1).astype(np.float32)
|
||||
|
||||
travel_dir_local = local_frame_rotation(pre_dir, travel_direction(pre_pos, post_pos))
|
||||
reconstructed = reconstruct_post_pos(pre_pos, pre_dir, step_length, travel_dir_local)
|
||||
travel_dir_local = local_frame_rotation(
|
||||
pre_dir, travel_direction(pre_pos, post_pos)
|
||||
)
|
||||
reconstructed = reconstruct_post_pos(
|
||||
pre_pos, pre_dir, step_length, travel_dir_local
|
||||
)
|
||||
np.testing.assert_allclose(reconstructed, post_pos, atol=1e-4)
|
||||
|
||||
|
||||
@@ -309,7 +319,7 @@ def test_build_features_clamps_n_sec_label_to_k_max():
|
||||
pdg_map = {11: 0}
|
||||
mat_map = {"PbWO4": 0}
|
||||
|
||||
_, _, _, n_sec, _, _, _, _, _ = build_features(data, pdg_map, mat_map)
|
||||
_, _, _, n_sec, _, _, _, _ = build_features(data, pdg_map, mat_map)
|
||||
|
||||
assert n_sec.max() <= K_MAX
|
||||
np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])
|
||||
@@ -346,20 +356,24 @@ def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
|
||||
|
||||
|
||||
def test_build_features_proc_idx_zero_without_proc_map():
|
||||
data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
|
||||
data = _minimal_step_data(
|
||||
3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
|
||||
)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
*_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map)
|
||||
*_, proc_idx, _, _ = build_features(data, pdg_map, mat_map)
|
||||
|
||||
np.testing.assert_array_equal(proc_idx, [0, 0, 0])
|
||||
|
||||
|
||||
def test_build_features_proc_idx_looks_up_proc_map():
|
||||
data = _minimal_step_data(3, process=np.array(["compt", "phot", "eIoni"], dtype=object))
|
||||
data = _minimal_step_data(
|
||||
3, process=np.array(["compt", "phot", "eIoni"], dtype=object)
|
||||
)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
proc_map = {"compt": 0, "phot": 1, "eIoni": 2}
|
||||
|
||||
*_, proc_idx, _, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map)
|
||||
*_, proc_idx, _, _ = build_features(data, pdg_map, mat_map, proc_map=proc_map)
|
||||
|
||||
np.testing.assert_array_equal(proc_idx, [0, 1, 2])
|
||||
|
||||
@@ -382,7 +396,9 @@ def test_build_features_require_secondaries_ok_when_no_secondaries():
|
||||
data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
_, _, _, _, sec_cont, *_ = build_features(data, pdg_map, mat_map, require_secondaries=True)
|
||||
_, _, _, _, sec_cont, *_ = build_features(
|
||||
data, pdg_map, mat_map, require_secondaries=True
|
||||
)
|
||||
|
||||
assert not sec_cont.any()
|
||||
|
||||
@@ -397,7 +413,11 @@ def fake_material_props(monkeypatch):
|
||||
in by the user (see giant.materials.MaterialPropertiesNotFilledError)."""
|
||||
import giant.materials as gm
|
||||
|
||||
fake = {"PbWO4": gm.MaterialProperties(z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7)}
|
||||
fake = {
|
||||
"PbWO4": gm.MaterialProperties(
|
||||
z_eff=75.6, a_eff=205.3, density=8.28, x0=0.89, lambda_int=20.7
|
||||
)
|
||||
}
|
||||
monkeypatch.setattr(gm, "MATERIAL_PROPERTIES", fake)
|
||||
return fake
|
||||
|
||||
@@ -406,13 +426,7 @@ def test_build_features_embedding_mode_zero_fills_physical_columns():
|
||||
data = _minimal_step_data(3)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
cond_cont, *_ = build_features(
|
||||
data,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
particle_conditioning="embedding",
|
||||
material_conditioning="embedding",
|
||||
)
|
||||
cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="embedding")
|
||||
|
||||
assert cond_cont.shape[1] == COND_DIM
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0)
|
||||
@@ -424,18 +438,14 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props):
|
||||
data = _minimal_step_data(3)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
cond_cont, *_ = build_features(
|
||||
data,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
particle_conditioning="physical",
|
||||
material_conditioning="physical",
|
||||
)
|
||||
cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="physical")
|
||||
|
||||
assert cond_cont.shape[1] == COND_DIM
|
||||
mass, charge = particle_mass_charge(11)
|
||||
expected_log_mass = log_transform(np.array([mass]))[0]
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5)
|
||||
np.testing.assert_allclose(
|
||||
cond_cont[:, COND_DIM_BASE], expected_log_mass, atol=1e-5
|
||||
)
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], charge)
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 2], 75.6) # z_eff
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 3], 205.3) # a_eff
|
||||
@@ -451,13 +461,7 @@ def test_build_features_physical_mode_unfilled_material_raises():
|
||||
pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0}
|
||||
|
||||
with pytest.raises(MaterialPropertiesNotFilledError):
|
||||
build_features(
|
||||
data,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
particle_conditioning="physical",
|
||||
material_conditioning="physical",
|
||||
)
|
||||
build_features(data, pdg_map, mat_map, conditioning="physical")
|
||||
|
||||
|
||||
def test_build_cond_features_mass_charge_override(fake_material_props):
|
||||
@@ -470,15 +474,11 @@ def test_build_cond_features_mass_charge_override(fake_material_props):
|
||||
data["charge"] = np.array([2.0, -2.0], dtype=np.float32)
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
cond_cont, _ = build_cond_features(
|
||||
data,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
particle_conditioning="physical",
|
||||
material_conditioning="physical",
|
||||
)
|
||||
cond_cont, _ = build_cond_features(data, pdg_map, mat_map, conditioning="physical")
|
||||
|
||||
np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])))
|
||||
np.testing.assert_allclose(
|
||||
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])
|
||||
|
||||
|
||||
@@ -494,12 +494,7 @@ def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode():
|
||||
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,
|
||||
particle_conditioning="embedding",
|
||||
material_conditioning="embedding",
|
||||
data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding"
|
||||
)
|
||||
|
||||
assert cond_cont.shape[-1] == COND_DIM
|
||||
@@ -526,8 +521,7 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode(
|
||||
pdg_map,
|
||||
mat_map,
|
||||
cond_normalizer=legacy_norm,
|
||||
particle_conditioning="physical",
|
||||
material_conditioning="physical",
|
||||
conditioning="physical",
|
||||
)
|
||||
|
||||
|
||||
@@ -616,23 +610,13 @@ def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_materi
|
||||
}
|
||||
|
||||
cond_cont, cond_cat = build_cond_features(
|
||||
data,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
particle_conditioning="physical",
|
||||
material_conditioning="physical",
|
||||
data, pdg_map, mat_map, conditioning="physical"
|
||||
)
|
||||
assert cond_cont.shape[-1] == COND_DIM
|
||||
np.testing.assert_array_equal(cond_cat, [[0, 0]]) # dummy indices, no raise
|
||||
|
||||
with pytest.raises(KeyError):
|
||||
build_cond_features(
|
||||
data,
|
||||
pdg_map,
|
||||
mat_map,
|
||||
particle_conditioning="embedding",
|
||||
material_conditioning="embedding",
|
||||
)
|
||||
build_cond_features(data, pdg_map, mat_map, conditioning="embedding")
|
||||
|
||||
|
||||
# ── _WelfordAccumulator ──────────────────────────────────────────────────────
|
||||
@@ -690,7 +674,9 @@ def test_welford_accumulator_matches_naive_running_mean_reference():
|
||||
naive_M2 = np.zeros(F)
|
||||
naive_n = 0
|
||||
for chunk in chunks:
|
||||
naive_mean, naive_M2, naive_n = naive_update(naive_mean, naive_M2, naive_n, chunk)
|
||||
naive_mean, naive_M2, naive_n = naive_update(
|
||||
naive_mean, naive_M2, naive_n, chunk
|
||||
)
|
||||
|
||||
acc = _WelfordAccumulator(F)
|
||||
for chunk in chunks:
|
||||
|
||||
@@ -1,43 +0,0 @@
|
||||
"""Tests for the secondary-type embedding-distance diagnostic
|
||||
(giant.analysis.type_embedding_distance)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from giant.analysis.type_embedding_distance import compute_type_embedding_l1_distance
|
||||
|
||||
|
||||
def _summary(n=100):
|
||||
return {
|
||||
"n": n,
|
||||
"mean": 1.23,
|
||||
"std": 0.45,
|
||||
"min": 0.01,
|
||||
"max": 9.87,
|
||||
"hist_edges": [0.0, 1.0, 2.0, 3.0],
|
||||
"hist_counts": [30, 40, 30],
|
||||
}
|
||||
|
||||
|
||||
def test_none_is_unavailable():
|
||||
r = compute_type_embedding_l1_distance(None)
|
||||
assert r.kind == "unavailable"
|
||||
assert r.id == "type_embedding_l1_distance"
|
||||
assert r.payload["note"]
|
||||
|
||||
|
||||
def test_summary_produces_single_hist():
|
||||
r = compute_type_embedding_l1_distance(_summary())
|
||||
assert r.kind == "single_hist"
|
||||
assert r.id == "type_embedding_l1_distance"
|
||||
assert r.payload["edges"] == [0.0, 1.0, 2.0, 3.0]
|
||||
assert r.payload["rollout"] == [30, 40, 30]
|
||||
assert r.payload["log_x"] is True
|
||||
assert r.payload["log_y"] is True
|
||||
assert "n=100" in r.payload["note"]
|
||||
|
||||
|
||||
def test_single_hist_payload_shape_matches_render_contract():
|
||||
"""_render_single (giant.analysis.render) requires len(rollout) ==
|
||||
len(edges) - 1."""
|
||||
r = compute_type_embedding_l1_distance(_summary())
|
||||
assert len(r.payload["rollout"]) == len(r.payload["edges"]) - 1
|
||||
+23
-88
@@ -1,63 +1,30 @@
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
|
||||
from giant.data.dataset import StepBatch
|
||||
from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim
|
||||
from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM
|
||||
from giant.model.network import DenoisingMLP, SecondaryDecoder
|
||||
from giant.validate import validate_marginals
|
||||
|
||||
_PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
_MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
_K_MAX = 5
|
||||
|
||||
|
||||
def _tiny_models(particle_type_cfg: dict | None = None):
|
||||
"""A fresh v0.3.0 pair: Stage1Model owns no n_sec_head, so n_sec always
|
||||
comes from Stage2OneShot."""
|
||||
s1 = Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=_PARTICLE_CFG,
|
||||
material_cfg=_MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
)
|
||||
target = (particle_type_cfg or {}).get("target", "physical")
|
||||
sec_dim = stage2_trunk_sec_dim(
|
||||
particle_type_cfg or {"target": "physical"},
|
||||
"flow",
|
||||
_K_MAX,
|
||||
int(_PARTICLE_CFG["emb_dim"]),
|
||||
)
|
||||
s2 = Stage2OneShot(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=_PARTICLE_CFG,
|
||||
material_cfg=_MATERIAL_CFG,
|
||||
hidden_dim=16,
|
||||
n_res_blocks=1,
|
||||
generator="flow",
|
||||
time_dim=16,
|
||||
k_max=_K_MAX,
|
||||
sec_dim=sec_dim,
|
||||
particle_type_cfg=particle_type_cfg,
|
||||
)
|
||||
assert s2.particle_type_cfg.get("target", "physical") == target
|
||||
def _tiny_models():
|
||||
s1 = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=16, n_blocks=1)
|
||||
s2 = SecondaryDecoder(pdg_vocab=3, mat_vocab=2, hidden_dim=16, n_blocks=1)
|
||||
return s1.eval(), s2.eval()
|
||||
|
||||
|
||||
def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8):
|
||||
"""A val_loader matching StreamingStepsDataset's StepBatch shape."""
|
||||
def _zero_secondaries_loader(B=4, n_batches=2):
|
||||
"""A val_loader whose every batch has n_sec=0 (real side) — matches the
|
||||
(cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) tuple shape
|
||||
StreamingStepsDataset yields."""
|
||||
batches = []
|
||||
for _ in range(n_batches):
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
cond_cat = torch.zeros(B, 2, dtype=torch.long)
|
||||
x1 = torch.randn(B, X_DIM)
|
||||
n_sec = torch.full((B,), n_sec_value, dtype=torch.long)
|
||||
sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM)
|
||||
n_sec = torch.zeros(B, dtype=torch.long)
|
||||
sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM)
|
||||
proc_idx = torch.zeros(B, dtype=torch.long)
|
||||
sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long)
|
||||
batches.append(StepBatch(cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx))
|
||||
batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx))
|
||||
return batches
|
||||
|
||||
|
||||
@@ -65,54 +32,22 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch
|
||||
"""If n_sec_pred collapses to 0 across the whole validated set (realistic
|
||||
during early/unstable training), phys_kl must degrade to NaN instead of
|
||||
crashing on the empty-array .min()/.max() reduction inside
|
||||
_histogram_kl."""
|
||||
_histogram_kl -- a regression the old species/bincount code this
|
||||
replaced explicitly guarded against."""
|
||||
s1, s2 = _tiny_models()
|
||||
loader = _loader(n_sec_value=0)
|
||||
loader = _zero_secondaries_loader()
|
||||
|
||||
def _fake_resolve_n_sec(stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred):
|
||||
return torch.zeros(cond_cont.size(0), dtype=torch.long)
|
||||
# Force the Stage-1 n_sec head's prediction to 0 for every sample too, so
|
||||
# the generated side's valid-slot mask is also empty (real side is
|
||||
# already all n_sec=0 by construction of the fake loader above).
|
||||
def _fake_sample_flow(model, cond_cont, cond_cat, **kw):
|
||||
B = cond_cont.size(0)
|
||||
return torch.randn(B, X_DIM), torch.zeros(B, dtype=torch.long)
|
||||
|
||||
monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec)
|
||||
monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow)
|
||||
|
||||
result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2, steps=2)
|
||||
result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2)
|
||||
|
||||
assert np.asarray(result["phys_real"]).shape == (0, 2)
|
||||
assert np.asarray(result["phys_generated"]).shape == (0, 2)
|
||||
assert np.isnan(np.asarray(result["phys_kl"])).all()
|
||||
|
||||
|
||||
def test_validate_marginals_physical_target_shapes():
|
||||
s1, s2 = _tiny_models()
|
||||
loader = _loader(n_sec_value=2)
|
||||
|
||||
result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
|
||||
|
||||
assert np.asarray(result["real"]).shape == (4, X_DIM)
|
||||
assert np.asarray(result["generated"]).shape == (4, X_DIM)
|
||||
assert np.asarray(result["kl_divergence"]).shape == (X_DIM,)
|
||||
assert "phys_real" in result and "phys_generated" in result and "phys_kl" in result
|
||||
assert "type_class_real" not in result
|
||||
|
||||
|
||||
def test_validate_marginals_onehot_type_class_marginal():
|
||||
particle_type_cfg = {"target": "onehot"}
|
||||
s1, s2 = _tiny_models(particle_type_cfg)
|
||||
loader = _loader(n_sec_value=2, n_classes=s2.type_dim)
|
||||
|
||||
result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2)
|
||||
|
||||
assert "phys_real" not in result
|
||||
# Real/generated valid-slot counts need not agree (real: ground-truth
|
||||
# n_sec=2 always; generated: the untrained n_sec_head's own prediction).
|
||||
assert np.asarray(result["type_class_real"]).ndim == 1
|
||||
assert np.asarray(result["type_class_gen"]).ndim == 1
|
||||
assert np.asarray(result["type_class_real"]).shape[0] > 0
|
||||
|
||||
|
||||
def test_validate_marginals_without_sec_decoder_returns_stage1_only():
|
||||
s1, _ = _tiny_models()
|
||||
loader = _loader(n_sec_value=0)
|
||||
|
||||
result = validate_marginals(s1, loader, n_batches=1, steps=2)
|
||||
|
||||
assert set(result) == {"real", "generated", "kl_divergence"}
|
||||
|
||||
+22
-47
@@ -1,13 +1,15 @@
|
||||
import torch
|
||||
|
||||
from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM
|
||||
from giant.model.network import CriticModel, Stage1Model, Stage2OneShot
|
||||
from giant.model.network import (
|
||||
Critic,
|
||||
SecondaryCritic,
|
||||
WGANGenerator,
|
||||
WGANSecondaryGenerator,
|
||||
)
|
||||
from giant.model.wgan import critic_loss, generator_loss, gradient_penalty
|
||||
from giant.sample import sample_secondaries_wgan, sample_wgan
|
||||
|
||||
PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1}
|
||||
|
||||
|
||||
def _cond(B=8):
|
||||
cond_cont = torch.randn(B, COND_DIM)
|
||||
@@ -16,56 +18,23 @@ def _cond(B=8):
|
||||
|
||||
|
||||
def _small_generator():
|
||||
return Stage1Model(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator="wgan",
|
||||
noise_dim=8,
|
||||
n_sec_head_k_max=K_MAX,
|
||||
return WGANGenerator(
|
||||
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8
|
||||
)
|
||||
|
||||
|
||||
def _small_critic():
|
||||
return CriticModel(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
in_dim=X_DIM,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
stage="stage1",
|
||||
)
|
||||
return Critic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
||||
|
||||
|
||||
def _small_sec_generator():
|
||||
return Stage2OneShot(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
generator="wgan",
|
||||
noise_dim=8,
|
||||
return WGANSecondaryGenerator(
|
||||
pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8
|
||||
)
|
||||
|
||||
|
||||
def _small_sec_critic():
|
||||
return CriticModel(
|
||||
pdg_vocab=3,
|
||||
mat_vocab=2,
|
||||
particle_cfg=PARTICLE_CFG,
|
||||
material_cfg=MATERIAL_CFG,
|
||||
in_dim=SEC_DIM,
|
||||
hidden_dim=32,
|
||||
n_res_blocks=2,
|
||||
stage="stage2",
|
||||
)
|
||||
return SecondaryCritic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2)
|
||||
|
||||
|
||||
def _mask(B, n_sec):
|
||||
@@ -120,7 +89,7 @@ def test_sample_wgan_shape():
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
sample, n_sec = sample_wgan(model, cond_cont, cond_cat)
|
||||
assert sample.shape == (B, X_DIM)
|
||||
assert n_sec is not None and n_sec.shape == (B,)
|
||||
assert n_sec.shape == (B,)
|
||||
|
||||
|
||||
# --- Stage-2 generator/critic ---
|
||||
@@ -152,7 +121,9 @@ def test_sample_secondaries_wgan_shape():
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
stage1_out = torch.randn(B, X_DIM)
|
||||
n_sec_pred = torch.randint(0, K_MAX, (B,))
|
||||
sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(model, cond_cont, cond_cat, stage1_out, n_sec_pred)
|
||||
sec_cont, sec_phys, sec_valid = sample_secondaries_wgan(
|
||||
model, cond_cont, cond_cat, stage1_out, n_sec_pred
|
||||
)
|
||||
assert sec_cont.shape == (B, K_MAX, 4)
|
||||
assert sec_phys.shape == (B, K_MAX, 2)
|
||||
assert sec_valid.shape == (B, K_MAX)
|
||||
@@ -181,7 +152,9 @@ def test_gradient_penalty_masked():
|
||||
mask = _mask(B, n_sec)
|
||||
real = torch.randn(B, SEC_DIM) * mask
|
||||
fake = torch.randn(B, SEC_DIM) * mask
|
||||
gp = gradient_penalty(lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask)
|
||||
gp = gradient_penalty(
|
||||
lambda x: sec_critic(x, cond_cont, cond_cat, stage1_out), real, fake, mask=mask
|
||||
)
|
||||
assert gp.item() >= 0.0
|
||||
|
||||
|
||||
@@ -191,7 +164,9 @@ def test_critic_loss_scalar_and_grad():
|
||||
cond_cont, cond_cat = _cond(B)
|
||||
real = torch.randn(B, X_DIM)
|
||||
fake = torch.randn(B, X_DIM)
|
||||
loss = critic_loss(lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0)
|
||||
loss = critic_loss(
|
||||
lambda x: critic(x, cond_cont, cond_cat), real, fake.detach(), gp_weight=10.0
|
||||
)
|
||||
assert loss.shape == ()
|
||||
loss.backward()
|
||||
assert any(p.grad is not None for p in critic.parameters())
|
||||
|
||||
@@ -357,105 +357,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/ae/8c/469afb6465b853afff216f9528ffda78a915ff880ed58813ba4faf4ba0b6/contourpy-1.3.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b7448cb5a725bb1e35ce88771b86fba35ef418952474492cf7c764059933ff8b", size = 203831, upload-time = "2025-07-26T12:02:51.449Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "coverage"
|
||||
version = "7.15.4"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/be/c3/4f2195f512fb172aa425a8803a874b2baa9ba7f80ff7b6080998761fc701/coverage-7.15.4.tar.gz", hash = "sha256:0548198fff07ccf4faf469520bce1c2eceb1ce3e62891921138dec10907f9d00", size = 936952, upload-time = "2026-08-06T13:50:24.442Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/1d/48/bc8d4ba7b37551a767bd863f15b3f80182b271c2f55975356f5f7dbe94c2/coverage-7.15.4-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:d4fedd1f7f428f9fe83b1ead5e7cc87a43427be31aadafbac3ac0636dc7abb22", size = 222543, upload-time = "2026-08-06T13:47:37.562Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/20/dd/88d6f83f1fffc974a3691a34a97951c5b12df7512a6782c5963883cbc058/coverage-7.15.4-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:37e2f0cdf58e2e1fed4e4d5a8f8786ae2f7eb80b478016876667dc4a01d60a97", size = 222905, upload-time = "2026-08-06T13:47:38.927Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/bd/5c/54ee0d4748585bb0acab9891cd8d92f2d3593165b4e59fc9de113bfb3140/coverage-7.15.4-cp312-cp312-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:fb55d0e70bb15f2e81477613627286581414693d74ac7963c93a790dd453ca9d", size = 254407, upload-time = "2026-08-06T13:47:40.488Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/8c/3f/f0642a372f494bd0d7dad3b497083b910194a5f1c88be2c94fef707c3b59/coverage-7.15.4-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:899b9da30f3c6c336566e3707495bb23e8302d39d862f01fa78c48b99b9437e2", size = 257145, upload-time = "2026-08-06T13:47:41.931Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/71/17/8b46d0ed68251016002ec972c8fc0119961a765d0984cafb8bf317c43758/coverage-7.15.4-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d15715e8c46552827e5e4f30a35575a2dbcad14454cf3284c54483946bd16931", size = 258257, upload-time = "2026-08-06T13:47:43.527Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/30/b8/8498a0e72d0adbe15477dd07463d2b3bb2c9f6a4815e8589e50939e2c3ae/coverage-7.15.4-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:002a438859f7b430bc99afeaf01a6d187dad1d0dc907b64cdeffc632a5db8fd8", size = 260517, upload-time = "2026-08-06T13:47:45.121Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/41/e1/7dce19c3bdb1e3dd63e769508216500edad81bd5f69a26d724e32aceaf78/coverage-7.15.4-cp312-cp312-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:e4193a04b518f7968f3099755f5509ee7cccc6dc2b92a6b14841934d22e222c9", size = 254785, upload-time = "2026-08-06T13:47:46.541Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/dd/b1/e1494703c675a2561723cd9b89f45c9168782c31280c611b1f767851e57c/coverage-7.15.4-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:e98dcc55d572b38e69d117da7e8e8efb8500f1f5eaf81ecd460a63220790b839", size = 256176, upload-time = "2026-08-06T13:47:48.155Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/73/76/a5629d270fb638a43a4b10466f51e2f49d532c1aa4da2913cbbb150bbe0a/coverage-7.15.4-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:af6c538498ce66c10d3fd541c2a8d5b03da5850355add34e6cba564210cb9e72", size = 254321, upload-time = "2026-08-06T13:47:49.757Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ff/4f/9c44447218435d5766b911534f9d798144a5560f85e9a54ebe5f3f5d19f9/coverage-7.15.4-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:1d10025d96ea89fc2f73714dbc4cbd433fe012c1ac9e23f895d7728b238b6e52", size = 258390, upload-time = "2026-08-06T13:47:51.248Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/de/36/c1e127616fb3fa18a9ff71e76c417f2fd7424332a4870015ac224ef4c039/coverage-7.15.4-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:d802e1947603162ded419bff83ac7489820355d2b856dfb09206574e3a37ac0c", size = 253894, upload-time = "2026-08-06T13:47:52.816Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e9/b9/fdb92c8ae7a8bb9b850cc253b7b3b9c8526f68130002048b5671cd510d09/coverage-7.15.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c2de40895718f91951b86712b4c5b694acaf9a0a49be13874896f599a1eed3f4", size = 255763, upload-time = "2026-08-06T13:47:54.296Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6f/c0/a7d51b2587c7bdb76e71b0896d2565bf7d60436b5122fc83e511adb1f7cd/coverage-7.15.4-cp312-cp312-win32.whl", hash = "sha256:5c3431b2161279b7db5c2a1aa58ae02e5cb8c3c42d93a5094be3f5537bd5b11b", size = 224597, upload-time = "2026-08-06T13:47:56.074Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/49/b9/5c5f80cc55f5acaaca6dee677626bfcec8c87204a7809b438b08e84f4571/coverage-7.15.4-cp312-cp312-win_amd64.whl", hash = "sha256:6befeab5fb2b51c958ca4ac6c5d141a1e8240f4f76e46350f1911963deda49cd", size = 225135, upload-time = "2026-08-06T13:47:57.52Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/47/e4/2a4561f89ff6bf7c925c287d0f2cce8bdf139c3a33735c87e3203401cf94/coverage-7.15.4-cp312-cp312-win_arm64.whl", hash = "sha256:67bc345491ab55b837277d76f5775d057e8c7f1ac44d890d8c2c82adde258c6f", size = 224515, upload-time = "2026-08-06T13:47:58.977Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f1/84/651a9310859673aaa3b3203f1aa1641ca60fcf2494683e1c9474c7172780/coverage-7.15.4-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:c705b28feb2775dc82a25f1d473a370bc37ff93f5177f4e29ce2425f560f6921", size = 222565, upload-time = "2026-08-06T13:48:00.796Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/82/f9/4dcf700137e8af550670f4d74d1b63828ce93e1e2b05e5f10710eb2ea987/coverage-7.15.4-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:3ff205ab5e3ecc670f6a4dd19d9cbf12ede53dd41cfc1e15716ec961ea6d314e", size = 222936, upload-time = "2026-08-06T13:48:02.391Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/07/4a/612ff1e780b3fbfd637486f542f84adc5503873d8b5d279dec1ffeef9414/coverage-7.15.4-cp313-cp313-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5172326e861a38b48b48befca15e0f477a26b283337a33a739c8fed229934e36", size = 253926, upload-time = "2026-08-06T13:48:04.382Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b0/04/d1cff1c2ead4708a6a79c01d3736b6a25bd38a36678398f72a8dd33dfad9/coverage-7.15.4-cp313-cp313-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:12b59c90084e3234fb11184886bf4a40f4f16a8c8f867be2e087b81f8e8868d4", size = 256523, upload-time = "2026-08-06T13:48:05.996Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b9/80/d34e13fb4b293cbdb9665838cf5522077b8ad14ef947550631a4bced36a5/coverage-7.15.4-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:349062d66f00b40fa2c1c222438bad25fabf755631b5d82937fe985c8008615c", size = 257759, upload-time = "2026-08-06T13:48:08.036Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0f/e7/2c5fe7636fdb0732fe0f09f308a5b066864078b7fc61f6678e8478554f2e/coverage-7.15.4-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:4256ced708e598e05209bc1a8ab4074e04a51dba4c62fb45926a229af675ace7", size = 259890, upload-time = "2026-08-06T13:48:09.834Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/92/28/9689f0858dfff59c2ea688938ab9fa2925631235df67126a42b6c5c70ae1/coverage-7.15.4-cp313-cp313-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:d80f974b20782d9612c8b4c9beeca867074c7cf4079d1419843fa25a26428b25", size = 254121, upload-time = "2026-08-06T13:48:11.459Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f9/e2/785077c230c157243eb5aa9a26c3be260ecd02001bead54a3cada3df8e03/coverage-7.15.4-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:2e179f19bfe1d31f8eeeaa12990194d761c4f62f0759661000bca6cd8729f40b", size = 255891, upload-time = "2026-08-06T13:48:13.209Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/d4/90/e20371b17b40f912f21305c2db2f30efa3de306f7320fc916804872c85a4/coverage-7.15.4-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:8bc16bb47b7679670eceff71d78bfb7d6e5b143f6c2cd117487ec7c75e0d4b78", size = 253859, upload-time = "2026-08-06T13:48:14.736Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/05/49/25371987ee459a5f67c0427fb75c74f9358e65f2c71fe75bf41c1b6c5fcb/coverage-7.15.4-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:1cd685005cd2c4200adfc14cf39a603b9320efab3f18a8f7f156d20c9cc3345f", size = 258011, upload-time = "2026-08-06T13:48:16.464Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/30/6e/32e67467f6154bf4f1c4f63b05acc5097cba4237d45bbeeea446b52e8ac1/coverage-7.15.4-cp313-cp313-musllinux_1_2_riscv64.whl", hash = "sha256:337399ad2c93b3acd2a937627dae8b3e86b66707cd3d3e856347999aadf1ef8d", size = 253676, upload-time = "2026-08-06T13:48:18.493Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/03/c1/8b24192e89286399765155251f99ee9f070a9d637109018ac23d99b99f6f/coverage-7.15.4-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:96e257121228ec5cd2bb919276e94ac11074471bc37d68dbae0e8308cce15fff", size = 255453, upload-time = "2026-08-06T13:48:20.057Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/16/6f/8b41ebdf67c87854e17c035336a90f1cfbad0c14c2a584301be6ff148718/coverage-7.15.4-cp313-cp313-win32.whl", hash = "sha256:c65a9e0dfc6143491879da4e13b5e30f8be192055de508d737fb14601edbd22c", size = 224605, upload-time = "2026-08-06T13:48:21.655Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/e0/e2/2946c7f0b42b152ecb21ff1bdad72e3d301e790c0c487e4a86e8c9f69347/coverage-7.15.4-cp313-cp313-win_amd64.whl", hash = "sha256:2ff8f5e9b8f7a94f0c11c45631eee103dbcb7d63274edd12c56efe1be690b3b4", size = 225148, upload-time = "2026-08-06T13:48:23.376Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9e/83/3f4a69957f48ae7a0aba76c34743f88963d607b19e03f3f8e66f91cae0f9/coverage-7.15.4-cp313-cp313-win_arm64.whl", hash = "sha256:6e0a8a5083b096487d6cfced94cdd514d8f5db6f113610fb36c0620edb1028cf", size = 224536, upload-time = "2026-08-06T13:48:25.117Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ea/ac/748cf29eeb2d6be34a3176ce26a4f49e38085ee08e8935f05f6f26ed7e0f/coverage-7.15.4-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:770e9325ab5ea6d56f77e59b29ecfe0ac20b57a82a601876f90494a4dda0386f", size = 222608, upload-time = "2026-08-06T13:48:26.806Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0b/02/1abbf5c984677b0aa439cdacaccbf38d248939d8ef8fe1cc7a50d73edb77/coverage-7.15.4-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d12b33a3a50a1676b7784dc8d00a0c6d66a9f2add4b85a041c19b6a7e53ef23c", size = 222940, upload-time = "2026-08-06T13:48:28.432Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/eb/e1/ff8f9f53d9fcf586125b55d0b1f04ec1c14955fee41e83d5814bee141bb5/coverage-7.15.4-cp314-cp314-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:5669c8378ebde86f5def7a25d29586631b58acc27ffde04399f678f3dfc6e082", size = 253985, upload-time = "2026-08-06T13:48:29.995Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/26/595759762e514e81be1d7d01ed03444303bcd152226a6529998d253f9201/coverage-7.15.4-cp314-cp314-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:ff97a14362eef486483ed44042ca2027ea257df6ff768e62358ee0c9776925ac", size = 256492, upload-time = "2026-08-06T13:48:31.634Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/24/68/b79aabac54d482be23b5fcdd4f4662bff24a78edc4ee29201726929936d5/coverage-7.15.4-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5a325e815318638aed1655d9c06e6d7c2d3d46c09231ce988070428a8762d734", size = 257837, upload-time = "2026-08-06T13:48:33.186Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/09/0f/bf7f297885a5bf6fd71e5782404e0ff059ca09e8711ceb3a08544abde45a/coverage-7.15.4-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:474223409d88eb20d2d6a0d37ea60e8647a65a90cc008dc1f0410af5f64f1e0d", size = 260152, upload-time = "2026-08-06T13:48:34.75Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/fd/f1/296744e854ff8368542343457414380465e9ceefb9192342feb9d3bc461d/coverage-7.15.4-cp314-cp314-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:7f2f62ae3cd189dd2e13aece758c57b3eecbd27be070dbd4cbd10936049e5dbf", size = 253978, upload-time = "2026-08-06T13:48:36.434Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/55/b0/bbdb2e9057493e66220a2e149ca2d301ba0e3a58a83bd6b90de9826d16f3/coverage-7.15.4-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:39ece820e29e0a2ba34b3ecb3be83c27e997eed8926f2ba6fe7ce7a0bda5843b", size = 255846, upload-time = "2026-08-06T13:48:38.317Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/96/e4/38015b2b6d21258713bd17e76b59d033b191efb5703589cffd037dfbca20/coverage-7.15.4-cp314-cp314-musllinux_1_2_i686.whl", hash = "sha256:f21b56dcace11dfe013014201f577dcd592b2a9b72182d930361b47cf6f73f25", size = 253808, upload-time = "2026-08-06T13:48:39.993Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0b/64/0d515c1e60ee6fbfd1a0e79c07cd87d388a233b7adc37758735677203808/coverage-7.15.4-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:93a3a0b662abcc10c73a47cbc72cd60f63618d6989fb2d1286e50eacd974f303", size = 258081, upload-time = "2026-08-06T13:48:41.971Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/91/71/04d9e7a3642146c6351338aef4ef85ab11dbbb54744c13245caba1aad1c0/coverage-7.15.4-cp314-cp314-musllinux_1_2_riscv64.whl", hash = "sha256:141fae2cabf5569b782c10afc4c850ce10f618c13f8db54765cba99cc839da1f", size = 253624, upload-time = "2026-08-06T13:48:43.731Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b4/a7/6c28b74c81ebff66987b0e2522ba5cffa3e90b0c33cb6a2eb264d4ee8cf1/coverage-7.15.4-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:81294c7e6ab30c5f74c0353b11b2fd6320e72d9bee6ac73b357caa8b916323a5", size = 255280, upload-time = "2026-08-06T13:48:45.58Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/52/af/bc19996a7014b98d7bbb0f0939453c67074af65784a3aa16a789a07381fa/coverage-7.15.4-cp314-cp314-win32.whl", hash = "sha256:7bbd7d6418e0dab31a206af5203bd43ae36edb8e7fba1940b055d3e9249290d7", size = 224768, upload-time = "2026-08-06T13:48:47.525Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ee/90/219484e476d6e101ba0a444852579e05f5b75c37c611a42ed1190f73ef62/coverage-7.15.4-cp314-cp314-win_amd64.whl", hash = "sha256:f0204ed122758782970526057093f448051a39db9d810d4e344bb87a3546f425", size = 225259, upload-time = "2026-08-06T13:48:49.513Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b7/66/fa77daf4e383e5f776dac62c2409b6af81910ae6fe326bd5170dba74cc63/coverage-7.15.4-cp314-cp314-win_arm64.whl", hash = "sha256:9e71e7bc71c686a123347ae47a0de33a175e797a85bb57b791492adf4eec8ed8", size = 224684, upload-time = "2026-08-06T13:48:51.235Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/58/5b/f03bf0ce362bbf3f785fa5219620d00778d4ac6fc9e407734828e9c672f6/coverage-7.15.4-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:7c922735321eef3f87c280a3d39afff6b646723a2880b862cda4ac7a093b8aa8", size = 223338, upload-time = "2026-08-06T13:48:52.896Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/0f/76/e77d0ae22501831cc9f92193e8a957a5caa1dd177f90a6d1d9b106242d92/coverage-7.15.4-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:f41c17c4668a655ce96d090d8d5ffdc24ef64b5a02f9753884d08483e8a4a41a", size = 223609, upload-time = "2026-08-06T13:48:54.688Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/82/1a/b1f089da8d38ac612fa2dd6dc7f4a1a7657d12f3e261d2996edd3a838d0b/coverage-7.15.4-cp314-cp314t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:46822e9b6ff1c6a72b518c162c44a8f45a61a1d609c51084bf5b16c023c5037b", size = 264970, upload-time = "2026-08-06T13:48:56.403Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/bf/31/e66d98d6e9c7fcc88470f1e234eaf6b1950dc0dfbf797f7282c1c861da24/coverage-7.15.4-cp314-cp314t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:3d6f4955b73b5445271379a59e3792b0d978f42d4a01e0cf7a67d9c33a3bb0a5", size = 267088, upload-time = "2026-08-06T13:48:58.41Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/59/a1/ae94eb2c541add426378408379f233591e069040b1e2cdb33df9498a0682/coverage-7.15.4-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3fc9e047706fb4a9abb54f719d3aa643e80e5bb3818182c40aee01ac0f0247ba", size = 269508, upload-time = "2026-08-06T13:49:00.42Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/9c/c7/88a10694a1c6a213569766aba9f25847b28155d4ac731b13226db216356d/coverage-7.15.4-cp314-cp314t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:05e491d4f3165d62d4f5c8fd48dfeabf2ae8f42cbbd484319af33ea851b78982", size = 270629, upload-time = "2026-08-06T13:49:02.234Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b3/34/d8b8232e5e55169933b59aabcef2fedfa4b9d8897361bb80fcbda146505f/coverage-7.15.4-cp314-cp314t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:226c66e80ec0598d3b9b4874123df167ccca342aca8714f77cac6829688ee09c", size = 264043, upload-time = "2026-08-06T13:49:04.102Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/7e/35/58b009dbf8c471c7224716478b9fed4a7e1af15320e1ed41660978504663/coverage-7.15.4-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:ac41cc14bebda0dbfb0628036b7f75706935c95bcc07fefe9a0f93614aa60a57", size = 266963, upload-time = "2026-08-06T13:49:05.821Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/62/aa/57fbda1b42c892968273c56b6ee9dc0f1310850859230a507bc7873b1f65/coverage-7.15.4-cp314-cp314t-musllinux_1_2_i686.whl", hash = "sha256:8af623e5cd92080acddd02b38f2f406a2c3a0893c38950b211890361448fbf26", size = 264569, upload-time = "2026-08-06T13:49:07.706Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/98/8a/360e6e7f24d477b7e889703af0afa878d15b6d4d8d2a822b2835c169a879/coverage-7.15.4-cp314-cp314t-musllinux_1_2_ppc64le.whl", hash = "sha256:07545711d4f0f32852a18f18ad11f76f0109909d09e78b9008b4cfc67e829429", size = 268299, upload-time = "2026-08-06T13:49:09.587Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4e/89/6f701261aee21b6b5fa8f7872229406dc917e125069448292223bf213606/coverage-7.15.4-cp314-cp314t-musllinux_1_2_riscv64.whl", hash = "sha256:a0865421cfdc53654b342d515e5a233187590882d20b95752150e53f65460017", size = 263413, upload-time = "2026-08-06T13:49:11.604Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/3f/0f/6f04036edc260ed425af83e834f627fad48941ce97b50bfe6edd8b6fa623/coverage-7.15.4-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:460115e32ee40566476db5048f9bec1e842c127ad8e6f8be745aad3ac9cbc839", size = 265725, upload-time = "2026-08-06T13:49:13.38Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c4/ce/d19b5d4d5c49a7bfb925fd74310fee7d28bc99520ac3367ccbc54e662518/coverage-7.15.4-cp314-cp314t-win32.whl", hash = "sha256:cbde877ef9dd7baf272b9bfef2b8a25edd45d9170fc326951dd20eb480335e85", size = 225079, upload-time = "2026-08-06T13:49:15.265Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/26/bb/7aa1b3b173faee0679037ca950bbbe1247273656697994d8d13f80f8d4b4/coverage-7.15.4-cp314-cp314t-win_amd64.whl", hash = "sha256:3da9e92d1c551fd7563833e9ade686efb0c4b7363ab7681a94283958c950bf5e", size = 225911, upload-time = "2026-08-06T13:49:17.279Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/81/1c/4ea9e47426d80038d9222db3c4534cb6021a74b237d3ff97ffd33b6600dd/coverage-7.15.4-cp314-cp314t-win_arm64.whl", hash = "sha256:3a54f5a0d85050c73a38f6793090ee83974531e67fe5e57a1da9bee11398aa5e", size = 225219, upload-time = "2026-08-06T13:49:19.293Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/2b/c4/dc5d2ac8f9142e7ec7de66e7bf0591db29d78955a040bd915870d9c0e657/coverage-7.15.4-cp315-cp315-macosx_10_15_x86_64.whl", hash = "sha256:2c9872e4d9dc5d3cf616bf4b382f5a00359305a5be666a3dd0b5cdb4e49597f9", size = 222604, upload-time = "2026-08-06T13:49:21.279Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/70/39/33e63df81fe2ee100897451841c821467635923e58e37c6bd4b46dd8106c/coverage-7.15.4-cp315-cp315-macosx_11_0_arm64.whl", hash = "sha256:e101dbb4b9b72f0cddd8cdc8c9c5b47f456766f5e0ac82dbfb75e5c55409b78a", size = 222944, upload-time = "2026-08-06T13:49:23.187Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/99/1f/ef3ffb5557febc75a0d97aa459d0266d7d741110265121cc6d8539343d44/coverage-7.15.4-cp315-cp315-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:7d1abebdb047729e852b9c77a00497dfbeb11eb3a117e037d7dbc3ac8e5f5c54", size = 254050, upload-time = "2026-08-06T13:49:25.008Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6f/f5/1f0f6f77698c3601ca0ae7431e34b24c62ca2f06fecb23b73ed1f651d2be/coverage-7.15.4-cp315-cp315-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:d28a4a899354d0ea6214cc59b4fa19eefbce1b9ff1688ab579acf49e894bd3fb", size = 256967, upload-time = "2026-08-06T13:49:26.896Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/03/7a/2ed9bed79925f4367c83c77f66a89e5ca7229c288d2d19ad5f36d1ca0070/coverage-7.15.4-cp315-cp315-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ffb3c2aacea411cc7e1d27712490c11108e2de1d39019ae32915493a59a8b9ed", size = 258587, upload-time = "2026-08-06T13:49:28.692Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/45/8c/fa34044f71b7cc4ecb6da9c2408770959b0591fa9b5fb6fb6bca38f94298/coverage-7.15.4-cp315-cp315-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:a9447978a92f405d301123cfd39ff49895490efb769a758fe2734c7f631bf8ce", size = 260785, upload-time = "2026-08-06T13:49:30.472Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4f/54/d5727ce36b4524a7394ab9f5f1df378e1f23affcdab01037dc8655185cc7/coverage-7.15.4-cp315-cp315-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:050467a7983b8e2fe7dd41a78bb30c3e7f8c0b8cafda14b1c46f8b5e3cf2dd3c", size = 254545, upload-time = "2026-08-06T13:49:32.271Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/dc/e6/6e3783e576719590194bdffb6dd6d85490801785b7c331e35a245d8cb8b5/coverage-7.15.4-cp315-cp315-musllinux_1_2_aarch64.whl", hash = "sha256:d003b7a5708ddad5c206c79607a6b92abb6fc13c57d99d8a4468cc03a2941ced", size = 256682, upload-time = "2026-08-06T13:49:34.089Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/dc/f2/bacdbde18b69ed2de424fcf64d9fb0a4913753d4f0eca8bae9daad69f4bd/coverage-7.15.4-cp315-cp315-musllinux_1_2_i686.whl", hash = "sha256:c38efe30fd74e5c19e9433f11fb1f5dc9c6522770971b7c6145bbaa413dc8800", size = 254560, upload-time = "2026-08-06T13:49:36.052Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/6c/a3/1fb927196e3477c1b48831169ab58ba08f451ba87ae311ff1de68b26a616/coverage-7.15.4-cp315-cp315-musllinux_1_2_ppc64le.whl", hash = "sha256:1f4f826d70f772ab8b0c052329580d7fe8b8abd191e4ce0c8f81aec6614665d3", size = 258792, upload-time = "2026-08-06T13:49:38.01Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/41/58/30d4c149c69053de0edfe325614c1d28d508f62b1783e0e4a234d2e49136/coverage-7.15.4-cp315-cp315-musllinux_1_2_riscv64.whl", hash = "sha256:4a4bf917c9953f57c957be31c1cd504e3bd2f34d4a352b9d391a3025336f6768", size = 253968, upload-time = "2026-08-06T13:49:39.934Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/89/e4/77f639371b918aad30dda4051f95404b43578f7f2e2f87ba73e02ed1ff37/coverage-7.15.4-cp315-cp315-musllinux_1_2_x86_64.whl", hash = "sha256:1c9bf40ebef178a45192c75c4964760bb261b0e6ad725da5fc4c93f674f19753", size = 255893, upload-time = "2026-08-06T13:49:41.825Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/5c/62/13be29b3ddab35f14c87967a4820a05106d2a3eccb4fa4ff550bf30b75e0/coverage-7.15.4-cp315-cp315-win32.whl", hash = "sha256:43619d04c3671792d2c4706ae8bf45e265dc87bbd4078189ef8b847ea1e74be2", size = 224768, upload-time = "2026-08-06T13:49:44.08Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a1/70/af0c6be0f964af6954f6b74bc109b0dbca02824696d2520fb17fe1ab06e3/coverage-7.15.4-cp315-cp315-win_amd64.whl", hash = "sha256:be619439dbcd31a2eab10b32de9fff62c26ed4bab69dc32b8363fdaaa0882809", size = 225242, upload-time = "2026-08-06T13:49:45.899Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/4f/2d/f3bd3aab899fc9efc18b53133ee68f5f98574ef480649b23e12962226387/coverage-7.15.4-cp315-cp315-win_arm64.whl", hash = "sha256:def597967dafc2e8d97c9097ea453c464e0bb8ed38f193a43070f10dc623bb6d", size = 224674, upload-time = "2026-08-06T13:49:48.322Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/f5/ca/f69251cd63eabc6438321aea22148754cce758a26bde07dd490e3fe7cfc5/coverage-7.15.4-cp315-cp315t-macosx_10_15_x86_64.whl", hash = "sha256:c7dbc748ac8a1e3e59a2b28bea47675e6e778081dbbf081bde0d75def2fcbe1d", size = 223333, upload-time = "2026-08-06T13:49:50.293Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/a7/a7/037b53b2885b0d8447064432491a4d5a1014cd9f97a594d53acd0c04541a/coverage-7.15.4-cp315-cp315t-macosx_11_0_arm64.whl", hash = "sha256:2413074a5ecbb61a01a7888fc72db0ca324d13588c5b38bc0dd8564cdcdfea26", size = 223630, upload-time = "2026-08-06T13:49:52.637Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/80/4f/152b8a4779ae90da11bb24f7467df8a59f0be48a5c52acb856325ca48289/coverage-7.15.4-cp315-cp315t-manylinux1_i686.manylinux_2_28_i686.manylinux_2_5_i686.whl", hash = "sha256:4e6f6f632b7b2f714bf7a1346e8f97b650ee71f3c298aaad42a2ab60f0f07645", size = 264489, upload-time = "2026-08-06T13:49:54.52Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/10/2d/84b4b9e0e1dd6528a51920ff7031f35b789382e467a28ec6a5a578cb8812/coverage-7.15.4-cp315-cp315t-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:8df457da2249d3c75ca2e5e835d59c725abfe92d27fdff6cd99eed85b51d5e9a", size = 267567, upload-time = "2026-08-06T13:49:56.721Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/53/fc/ba01cc25299f9f8a2c8b02d3b28c53f3543d9fbfbe4e74fa2760b48f163e/coverage-7.15.4-cp315-cp315t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:050f66a08805acb5b8a23c6d4a517b1ecf82c08e81ed0e4bd727df065e5c6624", size = 270123, upload-time = "2026-08-06T13:49:58.736Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/cf/d0/db2647cbf40b14f8c308f94ff7bf89c06d564e59f396906edf50086ec788/coverage-7.15.4-cp315-cp315t-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1587fb771d1ccceef708fdde1e5af8c7ed24b486b61d13a321acb7d8145390aa", size = 271107, upload-time = "2026-08-06T13:50:00.811Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/70/ff/4d2d17924552c458bb4f77dd631f0e3bc92fbbdf2d2d916cd4b33bbfd5b1/coverage-7.15.4-cp315-cp315t-manylinux_2_31_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:8b4f1c3a69ca580f3fbd6b2046915f536d7f586874f25c1bb23add2a3c88d50f", size = 264955, upload-time = "2026-08-06T13:50:03.023Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/ee/de/dc010c7a3691f396d93bbc26bfcafa1c2a3a351cd520470f15faf5795bd5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_aarch64.whl", hash = "sha256:ffb58d7eff5b7f6ecc6fa21d6288ab7f968a212cb67d682c269c09b9eba3b66f", size = 267949, upload-time = "2026-08-06T13:50:05.557Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/78/ea/dc96a11375e83c045c2f7c61fb6918277cfe9401db7c0f7b1d111a84b2e5/coverage-7.15.4-cp315-cp315t-musllinux_1_2_i686.whl", hash = "sha256:d9df165544774574ee004b953023d1bebada1894a80b1052a43d798b0f676e67", size = 264421, upload-time = "2026-08-06T13:50:07.612Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c8/86/b77131a0f9503ce461cd577076147d7a9040f0c5dda772686f729e2cc9cb/coverage-7.15.4-cp315-cp315t-musllinux_1_2_ppc64le.whl", hash = "sha256:f9de0a24a4079b53e523b5c5e2c5945ec251ab486652659955187cf255a259bc", size = 269121, upload-time = "2026-08-06T13:50:09.58Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/24/24/944bc35007862955e7ebf05754e645419dcf5d7526c52735cfa2715e8ebf/coverage-7.15.4-cp315-cp315t-musllinux_1_2_riscv64.whl", hash = "sha256:150089274bdc9f940628552cb92844e0223c987f1902ab8efe9f45a2ec758d88", size = 264565, upload-time = "2026-08-06T13:50:11.722Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c7/cc/a3bb9f93e7e740659163e2ea584f8196ddcd2c456a5dbe15f6c50105fec1/coverage-7.15.4-cp315-cp315t-musllinux_1_2_x86_64.whl", hash = "sha256:a58a94fed5da6997d258e8f7668c1e195fbd04a691d781b7558f1e468f9e68bc", size = 266522, upload-time = "2026-08-06T13:50:13.786Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/49/dd/e0e40f3560d878d888c580698ff5ad1179f5e1c3ac949684ef66b41a3817/coverage-7.15.4-cp315-cp315t-win32.whl", hash = "sha256:ebd5a6d8466ff30836572f3ba2cae8a5e8f85029b1c6d5e2ed338dc472a5166a", size = 225068, upload-time = "2026-08-06T13:50:15.825Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c6/7e/37732ea80eebc30e976e4cdab15c190bc42d96959a42e38ddf6f8c60468f/coverage-7.15.4-cp315-cp315t-win_amd64.whl", hash = "sha256:288bde2a2d7ab6b6c2d7252fcde8b524387f2d970bdba9658fc6f8bbcaef0f9b", size = 225895, upload-time = "2026-08-06T13:50:17.928Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/c6/08/1e00f7923eaaba45fb3d51dd794125fc766304b1df264f3a9c6557bfb30e/coverage-7.15.4-cp315-cp315t-win_arm64.whl", hash = "sha256:68be5e1de60ff13c9095bbec0e5a7fa45b33b101752215b91345ea1f61c4a278", size = 225213, upload-time = "2026-08-06T13:50:19.981Z" },
|
||||
{ url = "https://files.pythonhosted.org/packages/b4/d9/e70c286c979378f061d8266e279b686ab0b0b688e1fe0af864684f23a77d/coverage-7.15.4-py3-none-any.whl", hash = "sha256:964730a1e9de9c0cf11be6a1a3c79ce419c34882842abd256086ba4698705e84", size = 214332, upload-time = "2026-08-06T13:50:22.192Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "cramjam"
|
||||
version = "2.11.0"
|
||||
@@ -633,7 +534,7 @@ wheels = [
|
||||
|
||||
[[package]]
|
||||
name = "giant"
|
||||
version = "0.3.0"
|
||||
version = "0.1.0"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "numpy" },
|
||||
@@ -671,7 +572,6 @@ dev = [
|
||||
{ name = "plotstyle" },
|
||||
{ name = "polars" },
|
||||
{ name = "pytest" },
|
||||
{ name = "pytest-cov" },
|
||||
{ name = "ruff" },
|
||||
{ name = "scikit-learn" },
|
||||
{ name = "ty" },
|
||||
@@ -699,7 +599,6 @@ requires-dist = [
|
||||
{ name = "polars", marker = "extra == 'convert'", specifier = ">=1.0,<2" },
|
||||
{ name = "pyarrow", specifier = ">=16,<25" },
|
||||
{ name = "pytest", marker = "extra == 'dev'", specifier = ">=8,<10" },
|
||||
{ name = "pytest-cov", marker = "extra == 'dev'", specifier = ">=5,<8" },
|
||||
{ name = "pyyaml", specifier = ">=6,<7" },
|
||||
{ name = "ruff", marker = "extra == 'dev'", specifier = ">=0.15,<1" },
|
||||
{ name = "scikit-learn", marker = "extra == 'geometry'", specifier = ">=1.4,<2" },
|
||||
@@ -1814,20 +1713,6 @@ wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/8b/5a/ba30a81239b909821b3153e303e7def45178bf353da4f72380e6c5e8793b/pytest-9.1.0-py3-none-any.whl", hash = "sha256:8ebb0e7888bdf2bdfc602ec51f8f62d50200af37356c74e503c79a94f5c81f32", size = 386453, upload-time = "2026-06-13T18:52:44.045Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pytest-cov"
|
||||
version = "7.1.0"
|
||||
source = { registry = "https://pypi.org/simple" }
|
||||
dependencies = [
|
||||
{ name = "coverage" },
|
||||
{ name = "pluggy" },
|
||||
{ name = "pytest" },
|
||||
]
|
||||
sdist = { url = "https://files.pythonhosted.org/packages/b1/51/a849f96e117386044471c8ec2bd6cfebacda285da9525c9106aeb28da671/pytest_cov-7.1.0.tar.gz", hash = "sha256:30674f2b5f6351aa09702a9c8c364f6a01c27aae0c1366ae8016160d1efc56b2", size = 55592, upload-time = "2026-03-21T20:11:16.284Z" }
|
||||
wheels = [
|
||||
{ url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" },
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "python-dateutil"
|
||||
version = "2.9.0.post0"
|
||||
|
||||
Reference in New Issue
Block a user