diff --git a/.gitea/workflows/ci.yml b/.gitea/workflows/ci.yml index e69fc65..478bbb9 100644 --- a/.gitea/workflows/ci.yml +++ b/.gitea/workflows/ci.yml @@ -82,7 +82,11 @@ 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 + - run: uv run pytest --cov --cov-report=term-missing --cov-report=xml + - uses: actions/upload-artifact@v3 + with: + name: coverage-report + path: coverage.xml sync-version-on-tag: name: Sync project version with tag diff --git a/.gitignore b/.gitignore index 95aeaf8..fdad9a8 100644 --- a/.gitignore +++ b/.gitignore @@ -20,3 +20,8 @@ checkpoints/ # giant analyze run directories (shared.json, reduced/, plots/, condor logs) /analysis_runs/ + +# Coverage artifacts +.coverage +coverage.xml +htmlcov/ diff --git a/CLAUDE.md b/CLAUDE.md index ccab334..7c89dab 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -91,4 +91,6 @@ 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. + **Condor-submitted GPU training/rollout (in progress, `condor-gpu-train-rollout` branch, not yet merged):** moves `giant train`/`giant rollout` off the shared portal GPU dev machines (see Compute environment) onto remote-GPU HTCondor submission on TOpAS/NEMO2 (`giant/condor.py`). Partway between "needs major features" and feature-complete — not ready to merge yet. diff --git a/README.md b/README.md index 02c395f..c2244f5 100644 --- a/README.md +++ b/README.md @@ -1,20 +1,29 @@ # giant -**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. +**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate. -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. +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. -Training is driven entirely from parquet files of the miniCaloSim steps tree. No Geant4 runtime dependency. +## 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. ## Architecture -A **two-stage model**, both stages checkpointed together, with a choice of generative mode per stage (`--mode`): +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). -- **`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`):** +**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning: | Index | Variable | Encoding | |-------|----------|----------| @@ -23,36 +32,29 @@ A **two-stage model**, both stages checkpointed together, with a choice of gener | 3–5 | `post_dir` in local frame | unit vector | | 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector | -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. +- 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)`. -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. +**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`): -**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. +- `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` -**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: +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). -- **`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). +**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. -`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. +**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. ## Data -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. +- 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. ## Project structure @@ -64,15 +66,20 @@ giant/ │ │ ├── transforms.py # log transforms, local-frame rotation, energy simplex, secondary encode/decode │ │ └── dataset.py # StepsDataset / StreamingStepsDataset (PyTorch) │ ├── model/ -│ │ ├── network.py # ConditionEncoder, DenoisingMLP, SecondaryDecoder, Router/MoE, WGAN generator/critic +│ │ ├── network.py # ConditionEncoder, Stage1Model, Stage2OneShot/Stage2Autoregressive, Router/MoE, CriticModel │ │ ├── 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; nearest-known-PDG lookup +│ ├── particles.py # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode │ ├── materials.py # material name → (Z_eff, A_eff, density, X0, λ_int) │ ├── config.py # default hyperparameters, TOML config merging, device autodetect │ ├── pipeline.py # builds datasets/normalizers and kicks off a training run (with setup-stage caching) -│ ├── train.py # two-stage training loop, checkpointing, graceful shutdown, W&B logging +│ ├── 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) │ ├── 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 @@ -106,39 +113,48 @@ 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 extras selecting the torch build (pinned to 2.3.x); plain `uv sync` installs no torch at all. See `CLAUDE.md` for details. +`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details. -## Training, prediction, and rollout +## Training, prediction, rollout ```bash -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 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 predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt -# Full-shower rollout needs a geometry oracle (position → material/layer_id): -dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl +dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl # position → material/layer_id giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl ``` -`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. +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 `//` 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. ## Validation and analysis -`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: +- `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): ```bash giant analyze submit rollout.yaml --accounting-group cms # prep + one HTCondor job per plot (compute only) giant analyze render --gallery # local: styled PDFs + HTML gallery (needs LaTeX) ``` -`` 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. +`` 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. ## Development diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index fc665e3..458965b 100644 --- a/giant/analysis/catalog.py +++ b/giant/analysis/catalog.py @@ -58,6 +58,7 @@ 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 @@ -71,6 +72,11 @@ 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( @@ -80,6 +86,7 @@ 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. @@ -103,6 +110,7 @@ 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, ) @@ -158,9 +166,7 @@ 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)) @@ -232,9 +238,7 @@ 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: @@ -257,9 +261,7 @@ 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 @@ -309,9 +311,7 @@ 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,9 +480,7 @@ 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", @@ -513,18 +511,8 @@ 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()} @@ -560,9 +548,7 @@ 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", @@ -609,11 +595,9 @@ 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) @@ -651,13 +635,14 @@ _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) +) # --------------------------------------------------------------------------- @@ -678,9 +663,7 @@ 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: @@ -688,12 +671,8 @@ 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), ) ) @@ -701,9 +680,7 @@ 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, @@ -721,9 +698,7 @@ 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, @@ -735,9 +710,7 @@ 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, @@ -762,9 +735,7 @@ 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, @@ -831,6 +802,13 @@ 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 diff --git a/giant/analysis/condor.py b/giant/analysis/condor.py index 0312ca4..f512621 100644 --- a/giant/analysis/condor.py +++ b/giant/analysis/condor.py @@ -83,6 +83,12 @@ _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", ) @@ -155,9 +161,7 @@ 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) — @@ -247,6 +251,7 @@ 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. @@ -261,11 +266,15 @@ 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 " - f"n_chunks={effective_n} (chunkable={spec.chunkable})" + f"{spec_id}: chunk_index={chunk_index} out of range for n_chunks={effective_n} (chunkable={spec.chunkable})" ) bundle = Bundle.open( - rollout, reference, ctx, checkpoint=checkpoint, chunk=(chunk_index, effective_n) + rollout, + reference, + ctx, + checkpoint=checkpoint, + chunk=(chunk_index, effective_n), + type_embedding_l1_dist=type_embedding_l1_dist, ) partial = Partial( id=spec_id, @@ -291,6 +300,7 @@ 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"), ) @@ -314,10 +324,7 @@ 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( @@ -362,11 +369,7 @@ 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" @@ -386,9 +389,7 @@ 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``; @@ -464,9 +465,7 @@ 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) diff --git a/giant/analysis/context.py b/giant/analysis/context.py index f8cd84b..090ead5 100644 --- a/giant/analysis/context.py +++ b/giant/analysis/context.py @@ -74,9 +74,7 @@ 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)) @@ -104,31 +102,15 @@ 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) @@ -145,8 +127,7 @@ 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. diff --git a/giant/analysis/grouping.py b/giant/analysis/grouping.py index e5ec095..bc42005 100644 --- a/giant/analysis/grouping.py +++ b/giant/analysis/grouping.py @@ -68,9 +68,7 @@ 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: @@ -88,9 +86,7 @@ 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 diff --git a/giant/analysis/reduce.py b/giant/analysis/reduce.py index 15089f6..b050563 100644 --- a/giant/analysis/reduce.py +++ b/giant/analysis/reduce.py @@ -142,9 +142,7 @@ 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 ) @@ -254,10 +252,7 @@ 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") ) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index 3941981..f5f57ea 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -41,11 +41,47 @@ def _overlay(ax, edges: np.ndarray, series: dict[str, list], log_y: bool) -> Non ax.set_yscale("log") -def _router_summary(model_config: dict) -> str: - r = model_config.get("router") or {} - if not r.get("enabled"): +def _router_summary(router_cfg: dict) -> str: + if not router_cfg.get("enabled"): return "off" - return f"{r.get('type', '?')}×{r.get('n_experts', '?')}" + 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 def _figure_params(run_meta: dict) -> dict: @@ -59,8 +95,14 @@ 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: @@ -71,7 +113,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) + params["router"] = _router_summary(mc.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: @@ -97,11 +139,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") @@ -143,9 +185,7 @@ 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") @@ -157,9 +197,7 @@ 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") @@ -171,9 +209,7 @@ 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, {}) @@ -182,9 +218,7 @@ 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") @@ -303,9 +337,7 @@ 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 @@ -326,9 +358,7 @@ 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: @@ -356,6 +386,4 @@ 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) diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py index 6dbc3f8..2a6a70f 100644 --- a/giant/analysis/router_gating.py +++ b/giant/analysis/router_gating.py @@ -61,7 +61,8 @@ class _RouterHandle: pdg_map: dict[int, int] mat_map: dict[str, int] cond_normalizer: "Normalizer" - conditioning: str + particle_conditioning: str + material_conditioning: str router_type: str @@ -69,32 +70,43 @@ 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 {} - router_cfg = model_cfg.get("router") + # 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") + ) if not router_cfg or not router_cfg.get("enabled"): return None - stage1, _ = build_models(model_cfg) + built = build_models(model_cfg) + stage1 = built["stage1"] + if stage1 is None: + return None 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=stage1.router, + router=router, pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()}, mat_map={str(k): v for k, v in ckpt["mat_map"].items()}, cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]), - conditioning=model_cfg.get("conditioning", "embedding"), + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, 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) @@ -102,9 +114,7 @@ def _subsample( 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 @@ -128,13 +138,9 @@ def _gate_for_df( 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(), @@ -144,12 +150,11 @@ def _gate_for_df( handle.pdg_map, handle.mat_map, cond_normalizer=handle.cond_normalizer, - conditioning=handle.conditioning, + particle_conditioning=handle.particle_conditioning, + material_conditioning=handle.material_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 @@ -172,9 +177,7 @@ 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 @@ -194,10 +197,7 @@ def _top1_shares( 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,9 +233,7 @@ 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", @@ -271,9 +269,7 @@ 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)} @@ -318,9 +314,7 @@ 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 = [], {} diff --git a/giant/analysis/runtime_estimate.py b/giant/analysis/runtime_estimate.py index f5279f7..7f6cc37 100644 --- a/giant/analysis/runtime_estimate.py +++ b/giant/analysis/runtime_estimate.py @@ -54,9 +54,7 @@ _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 diff --git a/giant/analysis/sources.py b/giant/analysis/sources.py index 2248c90..abedc74 100644 --- a/giant/analysis/sources.py +++ b/giant/analysis/sources.py @@ -93,10 +93,7 @@ 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: @@ -121,11 +118,7 @@ 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)) @@ -137,9 +130,7 @@ 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: diff --git a/giant/analysis/type_embedding_distance.py b/giant/analysis/type_embedding_distance.py new file mode 100644 index 0000000..8304a2a --- /dev/null +++ b/giant/analysis/type_embedding_distance.py @@ -0,0 +1,70 @@ +"""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" + ), + }, + ) diff --git a/giant/checkpoint_io.py b/giant/checkpoint_io.py new file mode 100644 index 0000000..4ee92cc --- /dev/null +++ b/giant/checkpoint_io.py @@ -0,0 +1,208 @@ +"""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 + + +@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 + 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" + ) + 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, + 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"), + ) diff --git a/giant/cli.py b/giant/cli.py index 070a00c..ea254bb 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -34,24 +34,20 @@ from giant.data.loader import ( from giant.data.transforms import ( build_features, build_cond_features, - decode_secondaries, energy_simplex_decode, inv_local_frame_rotation, inv_log_transform, reconstruct_post_pos, - Normalizer, ) +from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference from giant.geometry import GeometryOracle -from giant.model.network import build_models -from giant.particles import nearest_known_pdg from giant.pipeline import run_train_job -from giant.rollout import rollout as run_rollout -from giant.sample import ( - sample_flow, - sample_secondaries, - sample_wgan, - sample_secondaries_wgan, +from giant.rollout import ( + L1DistCollector, + decode_secondary_identity, + rollout as run_rollout, ) +from giant.sample import resolve_n_sec, sample_stage1, sample_stage2 app = typer.Typer(no_args_is_help=True) @@ -66,35 +62,45 @@ def _router_total_experts(router_cfg: dict) -> int: """ if router_cfg.get("type") == "composed": axis_counts = { - m.group(1): int(v) - for k, v in router_cfg.items() - if (m := re.match(r"^axis(\d+)_n_experts$", k)) + m.group(1): int(v) for k, v in router_cfg.items() if (m := re.match(r"^axis(\d+)_n_experts$", k)) } return math.prod(axis_counts.values()) if axis_counts else 1 return int(router_cfg.get("n_experts", 1)) -def _batch_size_estimate_dims(model_cfg: dict, training: bool) -> tuple[int, int]: +def _batch_size_estimate_dims(model_cfg: dict, training: bool, stage: str = "stage1") -> tuple[int, int]: """Pick the (hidden_dim, n_blocks) that dominate per-call activation memory. - Routed models spend their FLOPs in the (smaller) expert trunks, not the - monolith's hidden_dim/n_blocks, so estimate_batch_size needs the expert - dims instead when routing is enabled. Training runs the full soft mixture - (every expert on the whole batch), so its activation memory scales with - the expert count; inference does top-1 dispatch (each row hits one - expert), so the batch just partitions across experts and one expert's - dims already bound it. estimate_batch_size scales memory linearly with - hidden_dim * n_blocks, so the training multiplier folds into n_blocks. + `model_cfg` is either the new nested shape (has a `f"{stage}_model"` key + — the merged training `cfg`, or a checkpoint's new-format `model_config`) + or a v0.2 checkpoint's flat `model_config`. v0.3.0 dropped per-expert + sizing (giant.model.network's routed trunks always inherit the stage's + own hidden_dim/n_res_blocks — no more `resolve_expert_dims`), so the new + shape needs no special-casing there; the legacy flat shape may still + carry a v0.2 `expert_hidden_dim`/`expert_n_blocks` override, honoured + only when that checkpoint's router was actually enabled. + + Routed models spend their FLOPs in the (smaller) expert trunks. Training + runs the full soft mixture (every expert on the whole batch), so its + activation memory scales with the expert count; inference does top-1 + dispatch (each row hits one expert), so the batch just partitions across + experts and one expert's dims already bound it. estimate_batch_size + scales memory linearly with hidden_dim * n_blocks, so the training + multiplier folds into n_blocks. """ - router_cfg = model_cfg.get("router") - if router_cfg and router_cfg.get("enabled"): - hidden_dim, n_blocks = gconfig.resolve_expert_dims( - router_cfg, model_cfg["hidden_dim"], model_cfg["n_blocks"] - ) - if training: - n_blocks *= _router_total_experts(router_cfg) - return hidden_dim, n_blocks - return model_cfg["hidden_dim"], model_cfg["n_blocks"] + if f"{stage}_model" in model_cfg: + stage_cfg = model_cfg[f"{stage}_model"] + hidden_dim, n_blocks = stage_cfg["hidden_dim"], stage_cfg["n_res_blocks"] + router_cfg = stage_cfg.get("router") + else: + hidden_dim, n_blocks = model_cfg["hidden_dim"], model_cfg["n_blocks"] + router_cfg = model_cfg.get("router") + if router_cfg and router_cfg.get("enabled"): + hidden_dim = model_cfg.get("expert_hidden_dim") or hidden_dim + n_blocks = model_cfg.get("expert_n_blocks") or n_blocks + if router_cfg and router_cfg.get("enabled") and training: + n_blocks = n_blocks * _router_total_experts(router_cfg) + return hidden_dim, n_blocks def _coerce_scalar(value: str) -> object: @@ -211,6 +217,16 @@ class Mode(str, Enum): wgan = "wgan" +class Decoder(str, Enum): + one_shot = "one_shot" + autoregressive = "autoregressive" + + +class Stage1Context(str, Enum): + truth = "truth" + sampled = "sampled" + + # Conditioning itself lives in giant.config (imported below as gconfig) — # shared with scripts/dwarf.py's Typer commands so the two CLIs can't # silently drift apart on the option's valid values. @@ -227,44 +243,12 @@ class Weights(str, Enum): ema = "ema" -def _load_model_weights( - model: torch.nn.Module, - sec_decoder: torch.nn.Module, - ckpt: dict, - weights: "Weights", - checkpoint_path: Path, -) -> None: - """Load either the raw or EMA state dicts from a training checkpoint. - - EMA weights (giant.train's shadow copy, see --ema-decay) only exist in - checkpoints written after that feature landed, so `ema` fails loudly - rather than silently falling back to raw weights a caller didn't ask for. - """ - if weights == 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: - typer.echo( - f"error: {checkpoint_path} has no EMA weights (trained before " - "--ema-decay, or with --ema-decay 0) — use --weights raw", - err=True, - ) - raise typer.Exit(1) - model.load_state_dict(ckpt[model_key]) - sec_decoder.load_state_dict(ckpt[sec_key]) - - @app.command() def train( - data: Annotated[ - Path, typer.Argument(help="Parquet file or directory of parquet files") - ], + data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")], config: Annotated[ Optional[Path], - typer.Option( - "--config", "-c", help="TOML config file (overridden by explicit flags)" - ), + typer.Option("--config", "-c", help="TOML config file (overridden by explicit flags)"), ] = None, mode: Annotated[ Optional[Mode], @@ -276,8 +260,7 @@ def train( typer.Option( "--batch-size", "-b", - help="Integer, or 'auto' to estimate from free GPU memory " - "(cuda devices only)", + help="Integer, or 'auto' to estimate from free GPU memory (cuda devices only)", ), ] = None, lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None, @@ -293,16 +276,91 @@ def train( "alongside the raw weights in checkpoints (0 disables; default: 0.9999)", ), ] = None, - warmup_epochs: Annotated[ - Optional[int], typer.Option("--warmup-epochs", "-w") - ] = None, + warmup_epochs: Annotated[Optional[int], typer.Option("--warmup-epochs", "-w")] = None, hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None, n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None, emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None, dropout: Annotated[ Optional[float], + typer.Option("--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"), + ] = None, + stage1_generator: Annotated[ + Optional[Mode], typer.Option( - "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)" + "--stage1-generator", + help="Stage 1's generative objective — overrides --mode for stage 1 only", + ), + ] = None, + stage1_hidden_dim: Annotated[ + Optional[int], + typer.Option( + "--stage1-hidden-dim", + help="Overrides --hidden-dim for stage 1 only (same effect today; " + "--hidden-dim is kept as a shorthand since stage 1 was the only " + "target before stage2_model got its own flags)", + ), + ] = None, + stage1_n_res_blocks: Annotated[ + Optional[int], + typer.Option("--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"), + ] = None, + stage1_dropout: Annotated[ + Optional[float], + typer.Option("--stage1-dropout", help="Overrides --dropout for stage 1 only"), + ] = None, + stage2_generator: Annotated[ + Optional[Mode], + typer.Option( + "--stage2-generator", + help="Stage 2's generative objective — overrides --mode for stage 2 " + "only, e.g. combine with --stage1-generator flow for a mixed " + "flow/wgan run", + ), + ] = None, + stage2_hidden_dim: Annotated[ + Optional[int], + typer.Option("--stage2-hidden-dim", help="Stage 2 trunk width"), + ] = None, + stage2_n_res_blocks: Annotated[ + Optional[int], + typer.Option("--stage2-n-res-blocks", help="Stage 2 trunk depth"), + ] = None, + stage2_dropout: Annotated[ + Optional[float], + typer.Option("--stage2-dropout", help="Dropout inside stage 2's ResBlocks"), + ] = None, + stage2_decoder: Annotated[ + Optional[Decoder], + typer.Option( + "--stage2-decoder", + help="one_shot: predict all k_max secondary slots at once (v0.2 " + "behaviour). autoregressive: emit one secondary at a time in " + "descending-energy order (default)", + ), + ] = None, + stage2_k_max: Annotated[ + Optional[int], + typer.Option( + "--stage2-k-max", + help="Maximum secondary slots (fixed width under one_shot, a " + "generation-loop safety cap under autoregressive; default: 15)", + ), + ] = None, + stage2_context_dim: Annotated[ + Optional[int], + typer.Option( + "--stage2-context-dim", + help="Width of the projected stage-1 outcome fed into stage 2's conditioning (default: 64)", + ), + ] = None, + stage2_stage1_context: Annotated[ + Optional[Stage1Context], + typer.Option( + "--stage2-stage1-context", + help="What stage 2 conditions on during training: 'truth' (the " + "ground-truth stage-1 target, detached — default) or 'sampled' " + "(stage 1's own sampled output, closing the train/inference gap " + "at the cost of an extra sampling pass per batch)", ), ] = None, conditioning: Annotated[ @@ -324,13 +382,9 @@ def train( ] = None, router_type: Annotated[ Optional[str], - typer.Option( - "--router-type", help="Router implementation name (see ROUTER_REGISTRY)" - ), - ] = None, - n_experts: Annotated[ - Optional[int], typer.Option("--n-experts", help="Number of routed experts") + typer.Option("--router-type", help="Router implementation name (see ROUTER_REGISTRY)"), ] = None, + n_experts: Annotated[Optional[int], typer.Option("--n-experts", help="Number of routed experts")] = None, router_axis: Annotated[ Optional[list[str]], typer.Option( @@ -345,37 +399,67 @@ def train( Optional[int], typer.Option( "--n-critic", - help="WGAN-GP (--mode wgan only): critic updates per generator " - "update (default: 5)", + help="WGAN-GP (--mode wgan only): critic updates per generator update (default: 5)", ), ] = None, gp_weight: Annotated[ Optional[float], typer.Option( "--gp-weight", - help="WGAN-GP (--mode wgan only): gradient-penalty coefficient " - "(default: 10.0)", + help="WGAN-GP (--mode wgan only): gradient-penalty coefficient (default: 10.0)", ), ] = None, noise_dim: Annotated[ Optional[int], typer.Option( "--noise-dim", - help="WGAN (--mode wgan only): generator input noise-vector " - "width (default: 64)", + help="WGAN (--mode wgan only): generator input noise-vector width (default: 64)", ), ] = None, critic_lr: Annotated[ Optional[float], typer.Option( "--critic-lr", - help="WGAN-GP (--mode wgan only): critic learning rate " - "(default: same as --lr)", + help="WGAN-GP (--mode wgan only): critic learning rate (default: same as --lr)", ), ] = None, - val_fraction: Annotated[ - Optional[float], typer.Option("--val-fraction", "-f") + stage1_n_critic: Annotated[ + Optional[int], + typer.Option("--stage1-n-critic", help="Overrides --n-critic for stage 1 only"), ] = None, + stage1_gp_weight: Annotated[ + Optional[float], + typer.Option("--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"), + ] = None, + stage1_noise_dim: Annotated[ + Optional[int], + typer.Option("--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"), + ] = None, + stage1_critic_lr: Annotated[ + Optional[float], + typer.Option("--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"), + ] = None, + stage2_n_critic: Annotated[ + Optional[int], + typer.Option("--stage2-n-critic", help="Overrides --n-critic for stage 2 only"), + ] = None, + stage2_gp_weight: Annotated[ + Optional[float], + typer.Option("--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"), + ] = None, + stage2_noise_dim: Annotated[ + Optional[int], + typer.Option( + "--stage2-noise-dim", + help="Overrides --noise-dim for stage 2 only; under " + "--stage2-decoder autoregressive a fresh draw is made per token", + ), + ] = None, + stage2_critic_lr: Annotated[ + Optional[float], + typer.Option("--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"), + ] = None, + val_fraction: Annotated[Optional[float], typer.Option("--val-fraction", "-f")] = None, seed: Annotated[ Optional[int], typer.Option("--seed", "-s", help="Random seed for reproducibility"), @@ -401,15 +485,12 @@ def train( Optional[int], typer.Option( "--max-val-batches", - help="Cap the per-epoch val-loss pass to N batches (0 = full " - "val set every epoch; default: 200)", + help="Cap the per-epoch val-loss pass to N batches (0 = full val set every epoch; default: 200)", ), ] = None, shuffle_buffer: Annotated[ int, - typer.Option( - "--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling" - ), + typer.Option("--shuffle-buffer", "-B", help="Rows held in RAM per worker for shuffling"), ] = 65536, cache_setup: Annotated[ bool, @@ -453,8 +534,7 @@ def train( Optional[bool], typer.Option( "--wandb/--no-wandb", - help="Log per-epoch training metrics to Weights & Biases " - "(requires `uv sync --extra wandb`)", + help="Log per-epoch training metrics to Weights & Biases (requires `uv sync --extra wandb`)", ), ] = None, wandb_project: Annotated[ @@ -485,72 +565,77 @@ def train( batch_size_value = int(batch_size) except ValueError: typer.echo( - f"error: --batch-size must be an integer or 'auto', " - f"got {batch_size!r}", + f"error: --batch-size must be an integer or 'auto', got {batch_size!r}", err=True, ) raise typer.Exit(1) - cli_train = { - k: v - for k, v in { - "mode": mode.value if mode is not None else None, - "epochs": epochs, - "batch_size": batch_size_value, - "lr": lr, - "weight_decay": weight_decay, - "ema_decay": ema_decay, - "warmup_epochs": warmup_epochs, - "val_fraction": val_fraction, - "num_workers": num_workers, - "seed": seed, - "validate_every": validate_every, - "validate_steps": validate_steps, - "max_val_batches": max_val_batches, - "n_critic": n_critic, - "gp_weight": gp_weight, - "critic_lr": critic_lr, - "wandb": wandb, - "wandb_project": wandb_project, - "wandb_run_name": wandb_run_name, - "wandb_log_every": wandb_log_every, - }.items() - if v is not None - } - cli_model: dict[str, object] = { - k: v - for k, v in { - "hidden_dim": hidden_dim, - "n_blocks": n_blocks, - "emb_dim": emb_dim, - "dropout": dropout, - "conditioning": conditioning.value if conditioning is not None else None, - "noise_dim": noise_dim, - }.items() - if v is not None - } cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) - if cli_router: - cli_model["router"] = cli_router - cfg = gconfig.merge_cli_overrides( - gconfig.DEFAULT_CONFIG, config, cli_train, cli_model - ) - t, m = cfg["train"], cfg["model"] + flag_values: dict[str, object] = { + "epochs": epochs, + "batch_size": batch_size_value, + "lr": lr, + "weight_decay": weight_decay, + "ema_decay": ema_decay, + "warmup_epochs": warmup_epochs, + "val_fraction": val_fraction, + "num_workers": num_workers, + "seed": seed, + "validate_every": validate_every, + "validate_steps": validate_steps, + "max_val_batches": max_val_batches, + "wandb": wandb, + "wandb_project": wandb_project, + "wandb_run_name": wandb_run_name, + "wandb_log_every": wandb_log_every, + "hidden_dim": hidden_dim, + "n_blocks": n_blocks, + "dropout": dropout, + "stage1_hidden_dim": stage1_hidden_dim, + "stage1_n_res_blocks": stage1_n_res_blocks, + "stage1_dropout": stage1_dropout, + "stage2_hidden_dim": stage2_hidden_dim, + "stage2_n_res_blocks": stage2_n_res_blocks, + "stage2_dropout": stage2_dropout, + "stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None, + "stage2_k_max": stage2_k_max, + "stage2_context_dim": stage2_context_dim, + "stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None, + "mode": mode.value if mode is not None else None, + "stage1_generator": stage1_generator.value if stage1_generator is not None else None, + "stage2_generator": stage2_generator.value if stage2_generator is not None else None, + "conditioning": conditioning.value if conditioning is not None else None, + "emb_dim": emb_dim, + "router_config": cli_router or None, + "n_critic": n_critic, + "gp_weight": gp_weight, + "noise_dim": noise_dim, + "critic_lr": critic_lr, + "stage1_n_critic": stage1_n_critic, + "stage1_gp_weight": stage1_gp_weight, + "stage1_noise_dim": stage1_noise_dim, + "stage1_critic_lr": stage1_critic_lr, + "stage2_n_critic": stage2_n_critic, + "stage2_gp_weight": stage2_gp_weight, + "stage2_noise_dim": stage2_noise_dim, + "stage2_critic_lr": stage2_critic_lr, + } + overrides = gconfig.overrides_from_flags(flag_values) + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides) + gconfig.validate_config(cfg) + t = cfg["train"] _device = torch.device(device) if device else gconfig.auto_device() if batch_size_auto: - est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(m, training=True) + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(cfg, training=True, stage="stage1") try: - t["batch_size"] = gconfig.estimate_batch_size( - est_hidden_dim, est_n_blocks, _device - ) + t["batch_size"] = gconfig.estimate_batch_size(est_hidden_dim, est_n_blocks, _device) except ValueError as exc: typer.echo(f"error: {exc}", err=True) raise typer.Exit(1) - typer.echo( - f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)" - ) + typer.echo(f"batch_size: {t['batch_size']} (auto-estimated from free GPU memory)") if out is not None: out_dir = out @@ -567,7 +652,7 @@ def train( # Name only encodes what's non-default (see default_out_dir_name), so # two runs with identical hyperparams in the same to-the-minute # timestamp would otherwise collide on this name — which also - # doubles as the W&B run id (giant.train) — hence the suffix loop in + # doubles as the W&B run id (giant.training) — hence the suffix loop in # resolve_default_out_dir. out_dir = gconfig.resolve_default_out_dir(cfg) @@ -606,9 +691,19 @@ def new_run( n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None, emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None, dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None, - conditioning: Annotated[ - Optional[Conditioning], typer.Option("--conditioning") - ] = None, + stage1_generator: Annotated[Optional[Mode], typer.Option("--stage1-generator")] = None, + stage1_hidden_dim: Annotated[Optional[int], typer.Option("--stage1-hidden-dim")] = None, + stage1_n_res_blocks: Annotated[Optional[int], typer.Option("--stage1-n-res-blocks")] = None, + stage1_dropout: Annotated[Optional[float], typer.Option("--stage1-dropout")] = None, + stage2_generator: Annotated[Optional[Mode], typer.Option("--stage2-generator")] = None, + stage2_hidden_dim: Annotated[Optional[int], typer.Option("--stage2-hidden-dim")] = None, + stage2_n_res_blocks: Annotated[Optional[int], typer.Option("--stage2-n-res-blocks")] = None, + stage2_dropout: Annotated[Optional[float], typer.Option("--stage2-dropout")] = None, + stage2_decoder: Annotated[Optional[Decoder], typer.Option("--stage2-decoder")] = None, + stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None, + stage2_context_dim: Annotated[Optional[int], typer.Option("--stage2-context-dim")] = None, + stage2_stage1_context: Annotated[Optional[Stage1Context], typer.Option("--stage2-stage1-context")] = None, + conditioning: Annotated[Optional[Conditioning], typer.Option("--conditioning")] = None, router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None, router_type: Annotated[Optional[str], typer.Option("--router-type")] = None, n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None, @@ -619,16 +714,13 @@ def new_run( ] = None, comment: Annotated[ Optional[str], - typer.Option( - "--comment", help="Free-text note recorded in config.toml's meta section" - ), + typer.Option("--comment", help="Free-text note recorded in config.toml's meta section"), ] = None, data: Annotated[ Optional[Path], typer.Option( "--data", - help="Dataset path to fill in the printed next-step command " - "(not stored in the config)", + help="Dataset path to fill in the printed next-step command (not stored in the config)", ), ] = None, force: Annotated[ @@ -640,9 +732,7 @@ def new_run( ] = False, dry_run: Annotated[ bool, - typer.Option( - "--dry-run", help="Print the resolved config without writing anything" - ), + typer.Option("--dry-run", help="Print the resolved config without writing anything"), ] = False, ) -> None: """Scaffold a new training run: resolve hyperparams to a config.toml and lay out its run dir. @@ -655,34 +745,35 @@ def new_run( (with the full dataset-derived meta section), so this scaffold's meta section is just a placeholder recording what was asked for and when. """ - cli_train = { - k: v - for k, v in { - "mode": mode.value if mode is not None else None, - "epochs": epochs, - "batch_size": batch_size, - "lr": lr, - }.items() - if v is not None - } - cli_model: dict[str, object] = { - k: v - for k, v in { - "hidden_dim": hidden_dim, - "n_blocks": n_blocks, - "emb_dim": emb_dim, - "dropout": dropout, - "conditioning": conditioning.value if conditioning is not None else None, - }.items() - if v is not None - } cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) - if cli_router: - cli_model["router"] = cli_router + flag_values: dict[str, object] = { + "epochs": epochs, + "batch_size": batch_size, + "lr": lr, + "hidden_dim": hidden_dim, + "n_blocks": n_blocks, + "dropout": dropout, + "stage1_hidden_dim": stage1_hidden_dim, + "stage1_n_res_blocks": stage1_n_res_blocks, + "stage1_dropout": stage1_dropout, + "stage2_hidden_dim": stage2_hidden_dim, + "stage2_n_res_blocks": stage2_n_res_blocks, + "stage2_dropout": stage2_dropout, + "stage2_decoder": stage2_decoder.value if stage2_decoder is not None else None, + "stage2_k_max": stage2_k_max, + "stage2_context_dim": stage2_context_dim, + "stage2_stage1_context": stage2_stage1_context.value if stage2_stage1_context is not None else None, + "mode": mode.value if mode is not None else None, + "stage1_generator": stage1_generator.value if stage1_generator is not None else None, + "stage2_generator": stage2_generator.value if stage2_generator is not None else None, + "conditioning": conditioning.value if conditioning is not None else None, + "emb_dim": emb_dim, + "router_config": cli_router or None, + } + overrides = gconfig.overrides_from_flags(flag_values) - cfg = gconfig.merge_cli_overrides( - gconfig.DEFAULT_CONFIG, config, cli_train, cli_model - ) + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides) + gconfig.validate_config(cfg) run_dir = (out or gconfig.resolve_default_out_dir(cfg)).resolve() if not force: @@ -699,7 +790,7 @@ def new_run( if dry_run: typer.echo("dry-run: not writing anything. Resolved config:") - for section in ("train", "model"): + for section in ("train", "conditioning", "stage1_model", "stage2_model"): typer.echo(f"[{section}]") for k, v in cfg[section].items(): if k == "router": @@ -708,6 +799,11 @@ def new_run( return meta = { + # Tags the written config.toml as v0.3-shaped so a later + # `migrate_config` load (e.g. `giant train --config ...`) treats it + # as already-migrated instead of misreading it as v0.2 and dropping + # its stage1_model/stage2_model/conditioning content. + "config_version": gconfig.CONFIG_VERSION, "git_hash": gconfig.git_hash(), "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"), "created_by": "giant new-run", @@ -728,9 +824,7 @@ def new_run( @app.command() def predict( - data: Annotated[ - Path, typer.Argument(help="Parquet file or directory of parquet files") - ], + data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")], checkpoint: Annotated[ Path, typer.Option( @@ -763,8 +857,7 @@ def predict( typer.Option( "--batch-size", "-b", - help="Inference batch size, or 'auto' to estimate from free GPU " - "memory (cuda devices only)", + help="Inference batch size, or 'auto' to estimate from free GPU memory (cuda devices only)", ), ] = "4096", steps: Annotated[ @@ -816,35 +909,25 @@ def predict( typer.echo(f"device: {_device}") # --- Load checkpoint --- - ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) - if "model_config" not in ckpt: - typer.echo( - "error: checkpoint has no model_config — retrain with the current code", - err=True, - ) + try: + ctx = load_for_inference(checkpoint, _device, "predict", weights=weights.value) + except CheckpointCompatibilityError as exc: + typer.echo(f"error: {exc}", err=True) raise typer.Exit(1) + typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})") - if "sec_decoder" not in ckpt: - typer.echo( - "error: checkpoint has no sec_decoder — retrain with the current code", - err=True, - ) - raise typer.Exit(1) - - if "sec_phys" not in ckpt.get("normalizer", {}): - typer.echo( - "error: checkpoint has no normalizer.sec_phys — retrain with the " - "current code", - err=True, - ) - raise typer.Exit(1) - - model_cfg = ckpt["model_config"] + assert ctx.stage1 is not None and ctx.stage2 is not None # require_stage2=True (default) guarantees this + model, sec_decoder = ctx.stage1, ctx.stage2 + cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm + pdg_map, mat_map = ctx.pdg_map, ctx.mat_map + pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map + particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning + other_policy = ctx.other_policy + stage1_ddpm_steps = ctx.stage1_ddpm_steps + stage2_k_max = ctx.k_max if batch_size_auto: - est_hidden_dim, est_n_blocks = _batch_size_estimate_dims( - model_cfg, training=False - ) + est_hidden_dim, est_n_blocks = _batch_size_estimate_dims(ctx.model_config, training=False) try: batch_size_value = gconfig.estimate_batch_size( est_hidden_dim, @@ -855,26 +938,10 @@ def predict( except ValueError as exc: typer.echo(f"error: {exc}", err=True) raise typer.Exit(1) - typer.echo( - f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)" - ) + typer.echo(f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)") assert batch_size_value is not None bs = batch_size_value - conditioning = model_cfg.get("conditioning", "embedding") - 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"]) - - model, sec_decoder = build_models(model_cfg) - _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) - model.to(_device).eval() - sec_decoder.to(_device).eval() - - typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})") - gconfig.warn_if_checkpoint_config_mismatch(checkpoint) # --- Output path --- out, dataset_path, pred_uuid = _resolve_prediction_output(data, out) @@ -890,45 +957,60 @@ def predict( skipped = 0 unknown_pdg_counts: Counter[int] = Counter() total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files) - chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks - def _concat( - a: dict[str, np.ndarray], b: dict[str, np.ndarray] - ) -> dict[str, np.ndarray]: + def chunk_iter(path: Path, offset: int): + if coord == Coord.local: + return iter_file_chunks(path, offset=offset, k_max=stage2_k_max) + return iter_cond_chunks(path, offset=offset) + + # load_for_inference already guarantees pdg_topn_map/mat_topn_map are not + # None whenever the matching conditioning axis is "onehot" — the extra + # `is not None` conjuncts below are redundant at runtime, just narrowing + # for the type checker. + cond_pdg_topn = pdg_topn_map.class_map if pdg_topn_map is not None and particle_conditioning == "onehot" else None + cond_mat_topn = mat_topn_map.class_map if mat_topn_map is not None and material_conditioning == "onehot" else None + + def _concat(a: dict[str, np.ndarray], b: dict[str, np.ndarray]) -> dict[str, np.ndarray]: return {k: np.concatenate([a[k], b[k]], axis=0) for k in a} def _process(piece: dict[str, np.ndarray]) -> None: nonlocal writer, total if coord == Coord.local: - cond_cont, cond_cat, target_raw, _, _, _, _, _ = build_features( - piece, pdg_map, mat_map, conditioning=conditioning + cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features( + piece, + pdg_map, + mat_map, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, + pdg_topn_map=cond_pdg_topn, + mat_topn_map=cond_mat_topn, + k_max=stage2_k_max, ) cond_cont = cond_norm.transform(cond_cont) else: cond_cont, cond_cat = build_cond_features( - piece, pdg_map, mat_map, cond_norm, conditioning=conditioning + piece, + pdg_map, + mat_map, + cond_norm, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, + pdg_topn_map=cond_pdg_topn, + mat_topn_map=cond_mat_topn, ) cc = torch.from_numpy(cond_cont).float().to(_device) ck = torch.from_numpy(cond_cat).long().to(_device) - if model_cfg.get("mode") == "wgan": - stage1_norm, n_sec_pred = sample_wgan(model, cc, ck) - else: - stage1_norm, n_sec_pred = sample_flow(model, cc, ck, steps=steps) + stage1_norm, n_sec_pred = sample_stage1(model, cc, ck, steps=steps, ddpm_steps=stage1_ddpm_steps) if coord == Coord.global_: - if model_cfg.get("mode") == "wgan": - sec_cont, sec_phys, _sec_valid_pred = sample_secondaries_wgan( - sec_decoder, cc, ck, stage1_norm, n_sec_pred - ) - else: - sec_cont, sec_phys, _sec_valid_pred = sample_secondaries( - sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps=steps - ) - sec_full_np = torch.cat([sec_cont, sec_phys], dim=-1).cpu().numpy() + # A fresh v0.3.0 Stage1Model owns no n_sec_head — + # sample_stage1 returns n_sec_pred=None then, so ask stage 2. + n_sec_pred = resolve_n_sec(model, sec_decoder, cc, ck, stage1_norm, n_sec_pred) + sec_cont, sec_type, _sec_valid_pred = sample_stage2(sec_decoder, cc, ck, stage1_norm, n_sec_pred, steps) + n_sec_pred_np = n_sec_pred.cpu().numpy() - n_sec_pred_np = n_sec_pred.cpu().numpy() pred = stage1_norm.cpu().numpy() # normalised # Inverse-normalise → local frame, log-scaled scalars @@ -949,14 +1031,8 @@ def predict( "material": piece["material"], "layer_id": piece["layer_id"], "n_sec": piece["n_sec"], - **{ - f"pred_{name}": raw[:, j] - for j, name in enumerate(LOCAL_TARGET_NAMES) - }, - **{ - f"true_{name}": target_raw[:, j] - for j, name in enumerate(LOCAL_TARGET_NAMES) - }, + **{f"pred_{name}": raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)}, + **{f"true_{name}": target_raw[:, j] for j, name in enumerate(LOCAL_TARGET_NAMES)}, } ) else: @@ -966,9 +1042,7 @@ def predict( # (hence delta_e == edep + e_sec) holds by construction. e_sec_pred # doubles as the stick-breaking energy budget for the Stage-2 decode # below, since the model has no other source for it at inference. - edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode( - raw[:, 1:3], piece["pre_E"] - ) + edep, e_sec_pred, _post_E, delta_e = energy_simplex_decode(raw[:, 1:3], piece["pre_E"]) # Normalise predicted direction then rotate back to world frame post_dir_local = raw[:, 3:6].copy() @@ -981,36 +1055,31 @@ def predict( travel_dir_local = raw[:, 6:9].copy() norms = np.linalg.norm(travel_dir_local, axis=1, keepdims=True) travel_dir_local /= np.where(norms < 1e-8, 1.0, norms) - post_pos_world = reconstruct_post_pos( - piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local - ) + post_pos_world = reconstruct_post_pos(piece["pre_pos"], piece["pre_dir"], step_length, travel_dir_local) - sec_E, sec_dir_world, sec_mass, sec_charge, _sec_valid = decode_secondaries( - sec_full_np, + # particle_type.target="physical": sec_pdg_code is a reporting- + # only nearest-known-PDG label (never fed back into the model — + # "no snapping at inference"). "onehot"/"embedding": PDG + # resolution IS the secondary's identity — see + # decode_secondary_identity's docstring. + sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, _l1_dist = decode_secondary_identity( + sec_decoder, + sec_cont, + sec_type, n_sec_pred_np, e_sec_pred, piece["pre_dir"], - sec_phys_normalizer=sec_phys_norm, + sec_phys_norm, + pdg_map, + pdg_topn_map, + other_policy, + None, ) - # Reporting-only nearest-known-PDG label (never fed back into the - # model) for the sec_pdg_list output column — see - # giant/particles.py and the "no snapping at inference" design. - sec_pdg_code = nearest_known_pdg( - sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys() - ).reshape(sec_mass.shape) - sec_pdg_list = [ - sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np) - ] + sec_pdg_list = [sec_pdg_code[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)] sec_E_list = [sec_E[i, :n].tolist() for i, n in enumerate(n_sec_pred_np)] - sec_dx_list = [ - sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np) - ] - sec_dy_list = [ - sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np) - ] - sec_dz_list = [ - sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np) - ] + sec_dx_list = [sec_dir_world[i, :n, 0].tolist() for i, n in enumerate(n_sec_pred_np)] + sec_dy_list = [sec_dir_world[i, :n, 1].tolist() for i, n in enumerate(n_sec_pred_np)] + sec_dz_list = [sec_dir_world[i, :n, 2].tolist() for i, n in enumerate(n_sec_pred_np)] table = pa.table( { @@ -1092,9 +1161,7 @@ def predict( typer.echo(f"reference: {ref_path}") if skipped: - codes = ", ".join( - f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items()) - ) + codes = ", ".join(f"{pdg} ({count})" for pdg, count in sorted(unknown_pdg_counts.items())) typer.echo( f"warning: skipped {skipped:,} row(s) with unknown PDG code(s): {codes}", err=True, @@ -1142,9 +1209,7 @@ def _seed_from_data(files: list[Path], n_events: int | None) -> dict[str, np.nda @app.command() def rollout( - data: Annotated[ - Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)") - ], + data: Annotated[Path, typer.Argument(help="Parquet file/dir to seed showers from (real events)")], checkpoint: Annotated[ Path, typer.Option("--checkpoint", "-c", help="Checkpoint .pt (best.pt/last.pt)"), @@ -1164,9 +1229,7 @@ def rollout( help="Stop a track when its energy drops below this [MeV]", ), ] = 0.1, - max_steps: Annotated[ - int, typer.Option("--max-steps", help="Max steps per individual track") - ] = 1000, + max_steps: Annotated[int, typer.Option("--max-steps", help="Max steps per individual track")] = 1000, steps: Annotated[ int, typer.Option( @@ -1184,9 +1247,7 @@ def rollout( "requires a checkpoint trained with EMA enabled.", ), ] = Weights.raw, - batch_size: Annotated[ - int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward") - ] = 4096, + batch_size: Annotated[int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")] = 4096, max_tracks_per_event: Annotated[ Optional[int], typer.Option( @@ -1201,15 +1262,9 @@ def rollout( help="Override the oracle's NN-distance escape threshold [mm]", ), ] = None, - n_events: Annotated[ - Optional[int], typer.Option("--n-events", help="Cap number of seed events") - ] = None, - device: Annotated[ - Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)") - ] = None, - out: Annotated[ - Optional[Path], typer.Option("--out", "-o", help="Output steps parquet") - ] = None, + n_events: Annotated[Optional[int], typer.Option("--n-events", help="Cap number of seed events")] = None, + device: Annotated[Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")] = None, + out: Annotated[Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")] = None, seed: Annotated[ Optional[int], typer.Option("--seed", help="Torch/numpy seed for reproducibility"), @@ -1223,45 +1278,27 @@ def rollout( _device = torch.device(device) if device else gconfig.auto_device() typer.echo(f"device: {_device}") - ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) - for key in ("model_config", "sec_decoder"): - if key not in ckpt: - typer.echo( - f"error: checkpoint has no {key} — retrain with the current code", - err=True, - ) - raise typer.Exit(1) - - if "sec_phys" not in ckpt.get("normalizer", {}): - typer.echo( - "error: checkpoint has no normalizer.sec_phys — retrain with the " - "current code", - err=True, - ) + try: + ctx = load_for_inference(checkpoint, _device, "rollout", weights=weights.value) + except CheckpointCompatibilityError as exc: + typer.echo(f"error: {exc}", err=True) raise typer.Exit(1) - - gconfig.warn_if_checkpoint_config_mismatch(checkpoint) - training_cfg = gconfig.load_checkpoint_config(checkpoint) - - model_cfg = ckpt["model_config"] - conditioning = model_cfg.get("conditioning", "embedding") - pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()} - 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"]) - - model, sec_decoder = build_models(model_cfg) - _load_model_weights(model, sec_decoder, ckpt, weights, checkpoint) - model.to(_device).eval() - sec_decoder.to(_device).eval() typer.echo(f"loaded checkpoint: {checkpoint} (weights: {weights.value})") + training_cfg = gconfig.load_checkpoint_config(checkpoint) + + assert ctx.stage1 is not None and ctx.stage2 is not None # require_stage2=True (default) guarantees this + model, sec_decoder = ctx.stage1, ctx.stage2 + cond_norm, tgt_norm, sec_phys_norm = ctx.cond_norm, ctx.tgt_norm, ctx.sec_phys_norm + pdg_map, mat_map = ctx.pdg_map, ctx.mat_map + pdg_topn_map, mat_topn_map = ctx.pdg_topn_map, ctx.mat_topn_map + particle_conditioning, material_conditioning = ctx.particle_conditioning, ctx.material_conditioning + other_policy = ctx.other_policy + stage1_ddpm_steps, stage2_ddpm_steps = ctx.stage1_ddpm_steps, ctx.stage2_ddpm_steps + model_cfg = ctx.model_config + oracle = GeometryOracle.load(geometry) - typer.echo( - f"loaded geometry oracle: {geometry} " - f"(escape_threshold={oracle.escape_threshold:.3f})" - ) + typer.echo(f"loaded geometry oracle: {geometry} (escape_threshold={oracle.escape_threshold:.3f})") files = find_parquet_files(data) seeds = _seed_from_data(files, n_events) @@ -1289,6 +1326,10 @@ def rollout( writer = pq.ParquetWriter(out, table.schema) writer.write_table(table) + # Only meaningful under particle_type.target="embedding" — a + # no-op collector otherwise, cheaper than branching the call itself. + l1_dist_collector = L1DistCollector() + summary = run_rollout( model, sec_decoder, @@ -1307,12 +1348,21 @@ def rollout( max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, on_chunk=_write_chunk, - conditioning=conditioning, - mode=model_cfg.get("mode", "flow"), + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, + pdg_topn_map=pdg_topn_map, + mat_topn_map=mat_topn_map, + other_policy=other_policy, + seed=seed, + stage1_ddpm_steps=stage1_ddpm_steps, + stage2_ddpm_steps=stage2_ddpm_steps, + l1_dist_collector=l1_dist_collector, ) if writer is not None: writer.close() + l1_summary = l1_dist_collector.summary() + ref_path = _write_prediction_ref(checkpoint, pred_uuid, out, dataset_path) ref = yaml.safe_load(ref_path.read_text()) ref.update( @@ -1332,13 +1382,18 @@ def rollout( "rollout_seed": seed, "n_rows": summary["n_rows"], "termination_reason_counts": summary["termination_reason_counts"], + # Diagnostic — only present under + # stage2_model.particle_type.target="embedding"; omitted (not + # written as null) otherwise, so giant.analysis can tell "not + # applicable to this checkpoint" apart from "collector empty". + **({"type_embedding_l1_dist": l1_summary} if l1_summary is not None else {}), # Full architecture spec baked into the checkpoint — includes the # entire router sub-dict, not just a hand-picked subset, so any # model knob (router type/n_experts, noise_dim, vocab sizes, ...) # is available downstream without touching this command again. "model_config": dict(model_cfg), - "training_epoch": ckpt.get("epoch"), - "best_val_loss": ckpt.get("best_val_loss"), + "training_epoch": ctx.epoch, + "best_val_loss": ctx.best_val_loss, # [train]/[meta] from the sibling config.toml (giant.config.save_config) # — empty dicts if the checkpoint has no config.toml next to it. "training_config": dict(training_cfg.get("train", {})), @@ -1363,9 +1418,7 @@ app.add_typer(analyze_app, name="analyze") def analyze_prep( rollout_yaml: Annotated[ Path, - typer.Argument( - help="giant rollout YAML sidecar (names the rollout + reference files)" - ), + typer.Argument(help="giant rollout YAML sidecar (names the rollout + reference files)"), ], run_dir: Annotated[ Optional[Path], @@ -1380,9 +1433,7 @@ def analyze_prep( top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6, chunks: Annotated[ int, - typer.Option( - "--chunks", help="Split each plot's data into this many event_id chunks" - ), + typer.Option("--chunks", help="Split each plot's data into this many event_id chunks"), ] = 1, ) -> None: """Read the rollout YAML → shared.json + run_meta.json in the run directory.""" @@ -1402,15 +1453,9 @@ def analyze_prep( @analyze_app.command("compute-one") def analyze_compute_one( - id: Annotated[ - str, typer.Option("--id", help="Catalog plot id (see `analyze list`)") - ], - run_dir: Annotated[ - Path, typer.Option("--run-dir", help="Run directory from `analyze prep`") - ], - chunk: Annotated[ - int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)") - ] = 0, + id: Annotated[str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")], + run_dir: Annotated[Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")], + chunk: Annotated[int, typer.Option("--chunk", help="Chunk index (see `analyze prep --chunks`)")] = 0, ) -> None: """Run one (plot, chunk)'s streaming reduction (this is what each condor job runs).""" from giant.analysis import compute_one @@ -1421,12 +1466,8 @@ def analyze_compute_one( @analyze_app.command("merge-one") def analyze_merge_one( - id: Annotated[ - str, typer.Option("--id", help="Catalog plot id (see `analyze list`)") - ], - run_dir: Annotated[ - Path, typer.Option("--run-dir", help="Run directory from `analyze prep`") - ], + id: Annotated[str, typer.Option("--id", help="Catalog plot id (see `analyze list`)")], + run_dir: Annotated[Path, typer.Option("--run-dir", help="Run directory from `analyze prep`")], ) -> None: """Merge one plot's chunk partials into its final reduced JSON. @@ -1450,14 +1491,10 @@ def analyze_list() -> None: @analyze_app.command("render") def analyze_render( - run_dir: Annotated[ - Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)") - ], + run_dir: Annotated[Path, typer.Argument(help="Run directory from `analyze prep` (holds reduced/)")], gallery: Annotated[ bool, - typer.Option( - "--gallery/--no-gallery", help="Run `gallery generate` after rendering" - ), + typer.Option("--gallery/--no-gallery", help="Run `gallery generate` after rendering"), ] = False, ) -> None: """Render reduced artifacts to styled PDFs + gallery metadata (local; needs LaTeX).""" @@ -1479,9 +1516,7 @@ def analyze_submit( help="Override the run directory (default: /analysis_runs/analysis_)", ), ] = None, - docker_image: Annotated[ - str, typer.Option("--docker-image") - ] = "cverstege/alma9-gridjob", + docker_image: Annotated[str, typer.Option("--docker-image")] = "cverstege/alma9-gridjob", request_memory: Annotated[int, typer.Option("--request-memory", help="MB")] = 8192, remote: Annotated[ bool, @@ -1497,9 +1532,7 @@ def analyze_submit( n_energy_bins: Annotated[int, typer.Option("--energy-bins")] = 4, n_marginal_bins: Annotated[int, typer.Option("--bins")] = 50, top_k_pdg: Annotated[int, typer.Option("--top-pdg")] = 6, - dry_run: Annotated[ - bool, typer.Option("--dry-run", help="Write files but don't condor_submit") - ] = False, + dry_run: Annotated[bool, typer.Option("--dry-run", help="Write files but don't condor_submit")] = False, ) -> None: """prep + write the HTCondor submit description (one job per plot x chunk), then submit.""" import subprocess diff --git a/giant/config.py b/giant/config.py index e4945e5..c7d10ec 100644 --- a/giant/config.py +++ b/giant/config.py @@ -1,8 +1,12 @@ +import copy +import difflib import hashlib import random +import re import subprocess import sys import tomllib +from dataclasses import dataclass, field from datetime import datetime, timezone from enum import Enum from pathlib import Path @@ -12,145 +16,712 @@ import torch class Conditioning(str, Enum): - """`model.conditioning` choices — shared by `giant.cli` and `scripts.dwarf`'s - Typer commands so the two CLIs can't silently drift apart on the option's - valid values (see DEFAULT_CONFIG["model"]["conditioning"] for what each - value means).""" + """`conditioning.particle.type` / `conditioning.material.type` choices — + shared by `giant.cli` and `scripts.dwarf`'s Typer commands so the two + CLIs can't silently drift apart on the option's valid values (see + DEFAULT_CONFIG["conditioning"] for what each value means).""" physical = "physical" embedding = "embedding" + onehot = "onehot" -DEFAULT_CONFIG: dict = { - "train": { - "mode": "flow", - "epochs": 100, - "batch_size": 4096, - "lr": 3e-4, - "weight_decay": 0.01, # AdamW default — exposed so it can be tuned - "ema_decay": 0.9999, # EMA of model weights for sampling; 0 disables - # per-epoch val loss (not the marginal/KL validate_every pass) is - # capped to this many batches; 0 = full val set every epoch - "max_val_batches": 200, - "val_fraction": 0.1, - "num_workers": 4, - "seed": 0, - "validate_every": 10, - "validate_steps": 10, - "warmup_epochs": 5, - "lambda_nsec": 0.1, - "lambda_s2": 1.0, - # WGAN-GP-only knobs (mode == "wgan"; ignored by flow/ddpm). n_critic: - # critic updates per generator update. gp_weight: gradient-penalty - # coefficient (Gulrajani et al. 2017). critic_lr: 0.0 means "use - # `lr`" — not None, since save_config's TOML writer has no null - # literal to round-trip. - "n_critic": 5, - "gp_weight": 10.0, - "critic_lr": 0.0, - # Weights & Biases per-epoch metric logging (opt-in; see giant.train). - # "" for wandb_run_name means "use the checkpoint out_dir name" — not - # None, since save_config's TOML writer has no null literal to - # round-trip. - "wandb": False, - "wandb_project": "giant", - "wandb_run_name": "", - # Batch-granularity metrics (loss/grad_norm/lr) are logged every N - # optimizer steps, not every batch — a single epoch can be tens of - # thousands of steps (see steps_per_epoch above), and logging every - # one of them would flood the run with points the UI has to downsample - # anyway. Per-epoch metrics (the metrics.csv row) always log in full. - "wandb_log_every": 50, - }, - "model": { - "hidden_dim": 256, - "n_blocks": 6, - "emb_dim": 16, - "dropout": 0.1, - # WGAN generator noise-vector width (mode == "wgan" only). - "noise_dim": 64, - # "physical" conditions on material/particle physical properties via - # a small MLP (giant.model.network.ConditionEncoder); "embedding" - # keeps the original learned pdg/material embedding tables — kept - # available as the generalization-comparison baseline. Checkpoints - # from before this option existed have no "conditioning" key and - # load as "embedding" (see giant.model.network.build_models). - "conditioning": "physical", - "router": { - "enabled": False, - "type": "energy", # selects the Router impl from ROUTER_REGISTRY - "n_experts": 4, - # 0 means "inherit model.hidden_dim/n_blocks" (see - # resolve_expert_dims below) — not a fixed 128/3, which silently - # ignored --hidden-dim/--n-blocks whenever routing was enabled. - # TOML has no null literal to round-trip (same pattern as - # critic_lr/wandb_run_name above), hence 0 rather than None. - "expert_hidden_dim": 0, - "expert_n_blocks": 0, - "temperature": 0.5, # energy/pdg-router kwarg - "learn_centers": True, # energy/pdg-router kwarg - # energy-router kwargs: mutually exclusive optional learnable - # gate-sharpness modes (see giant.model.network.EnergyRouter). - # learn_width generalizes the shared `temperature` to one - # learnable width per expert; learn_temperature instead makes - # the single shared `temperature` itself learnable. Both are - # bounded to [width_min_ratio, width_max_ratio] * temperature - # (sigmoid-parameterized, warm-started to reproduce `temperature` - # exactly at init) so gate sharpness can't run away to a - # collapse-inducing extreme during training. - "learn_width": False, - "learn_temperature": False, - "width_min_ratio": 0.1, - "width_max_ratio": 10.0, - "lambda_balance": 0.0, # optional load-balance aux loss weight - # optional entropy-regularization aux loss weight (generic - # Router.entropy_loss, penalizes uniform/collapsed gating) — a - # secondary guard against all experts' widths/temperature - # co-inflating together, which lambda_balance alone can't see - # since per-expert usage shares stay even throughout that - # failure mode. Off by default; bounding above is the primary - # defense. See giant.model.network.Router.entropy_loss. - "lambda_entropy": 0.0, - # Opt-in straight-through Gumbel-softmax train-time combine weights - # (see giant.model.network.Router.combine_weights): the training - # forward pass samples a hard one-hot combination — matching - # eval-time top-1 dispatch exactly — while the backward pass still - # flows a smooth gradient to every expert. Targets the train/eval - # mismatch identified as a likely contributor to experts - # overlapping instead of partitioning (see CLAUDE.md roadmap). - # gumbel_tau_start/_end are annealed linearly over training - # (giant.train._gumbel_tau); off by default, no effect unless - # gumbel = true. - "gumbel": False, - "gumbel_tau_start": 1.0, - "gumbel_tau_end": 0.1, - "emb_dim": 8, # process/pdg-router kwarg: own pdg(/mat) embedding width - "hidden_dim": 64, # process-router kwarg: its classifier's hidden width - "lambda_proc": 0.0, # process-router kwarg: supervised process-CE weight - # (0.0 still trains a working router — the gate gets gradient - # through the downstream flow loss like EnergyRouter's centers — - # but only lambda_proc > 0 grounds it in the true `process` label) - # type = "composed" routes on multiple axes at once (e.g. energy x - # pdg), each with its own expert count/hyperparameters. Axes are - # NOT in these defaults (there's no meaningful default axis list) - # — set them as flat axis{i}_{field} keys instead of "n_experts", - # e.g. axis0_type = "energy", axis0_n_experts = 4, axis1_type = - # "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See - # giant.model.network._parse_composed_axes / `--router-axis`. - }, - }, -} +# Tags a config dict (config.toml, or a checkpoint's model_config) as the new +# v0.3 nested format. Absence of `[meta].config_version == CONFIG_VERSION` is +# read as "this is a v0.2 dict" by migrate_config below. +CONFIG_VERSION = 3 + + +# --- Config dataclasses ----------------------------------------------------- +# +# These are the single source of truth for every default below. DEFAULT_CONFIG +# (a plain dict, for merge_cli_overrides/save_config/TOML round-tripping) is +# *generated* from GiantConfig().to_dict() rather than hand-maintained, so it +# cannot drift from the fallback defaults that build_models/build_critics +# (giant/model/network.py) and StageSpec.from_config (giant/training/trainers.py) +# read off these same dataclasses — see issues.md Issue 1. +# +# Each dataclass is frozen and carries an explicit from_dict/to_dict pair +# (mirroring StageSpec's established style in trainers.py) rather than a +# generic reflection-based helper, so every default is fully type-checkable. +# `lambda` is a Python keyword, so dict key "lambda" is always exposed as the +# field `lambda_weight`. +# +# Two sub-blocks — router and n_sec — carry genuinely dynamic keys that don't +# fit a fixed schema: composed-router `axis{i}_{field}` flags (see +# giant.model.network._parse_composed_axes) and pipeline.py's runtime-seeded +# `centers_init`, plus n_sec's `legacy_owner` (injected only by +# _migrate_legacy_model_config for v0.2 checkpoints). Both dataclasses carry +# an `extra: dict` catch-all so these keys round-trip losslessly without +# becoming named fields that would leak into every new run's config.toml. + + +@dataclass(frozen=True) +class ConditioningAxisConfig: + """One conditioning axis: `conditioning.particle` or `conditioning.material`.""" + + # "physical": a small MLP over log(mass)/charge, computable for any PDG + # code — generalizes beyond the training menu. + # "embedding": a learned nn.Embedding over a dense training-vocab index — + # memorizes the training menu; the generalization-comparison baseline, + # and the only mode compatible with stage2_model.particle_type.target = + # "embedding". + # "onehot": a fixed, unlearned vector — top (emb_dim - 1) PDG codes by + # training-set count, plus one "other" bin. NOT a reparameterization of + # "embedding": the vocabulary cap is the real difference. + type: str = "physical" + # Width of this axis's vector. Under "onehot" this also sets the class + # count. + emb_dim: int = 16 + # Depth of the sub-MLP under "physical". Ignored under "embedding"/"onehot". + n_layers: int = 1 + + @classmethod + def from_dict(cls, d: dict | None) -> "ConditioningAxisConfig": + d = d or {} + return cls( + type=d.get("type", "physical"), + emb_dim=d.get("emb_dim", 16), + n_layers=d.get("n_layers", 1), + ) + + def to_dict(self) -> dict: + return {"type": self.type, "emb_dim": self.emb_dim, "n_layers": self.n_layers} + + +@dataclass(frozen=True) +class ConditioningConfig: + # Width of the fused conditioning vector produced by the encoder's fusion + # MLP, consumed by every downstream trunk. + out_dim: int = 128 + # false: stage 1 and stage 2 each construct their own ConditionEncoder + # with identical config but independent weights. true: one instance, + # shared by reference (halves the conditioning parameter count, forces a + # common representation). + share_stages: bool = False + particle: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig) + material: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "ConditioningConfig": + d = d or {} + return cls( + out_dim=d.get("out_dim", 128), + share_stages=d.get("share_stages", False), + particle=ConditioningAxisConfig.from_dict(d.get("particle")), + material=ConditioningAxisConfig.from_dict(d.get("material")), + ) + + def to_dict(self) -> dict: + return { + "out_dim": self.out_dim, + "share_stages": self.share_stages, + "particle": self.particle.to_dict(), + "material": self.material.to_dict(), + } + + +@dataclass(frozen=True) +class FlowConfig: + # Width of the SinusoidalEmbedding for the flow time variable. + time_dim: int = 64 + + @classmethod + def from_dict(cls, d: dict | None) -> "FlowConfig": + d = d or {} + return cls(time_dim=d.get("time_dim", 64)) + + def to_dict(self) -> dict: + return {"time_dim": self.time_dim} + + +@dataclass(frozen=True) +class DdpmConfig: + time_dim: int = 64 + n_steps: int = 1000 + + @classmethod + def from_dict(cls, d: dict | None) -> "DdpmConfig": + d = d or {} + return cls(time_dim=d.get("time_dim", 64), n_steps=d.get("n_steps", 1000)) + + def to_dict(self) -> dict: + return {"time_dim": self.time_dim, "n_steps": self.n_steps} + + +@dataclass(frozen=True) +class Stage1WganConfig: + noise_dim: int = 64 + n_critic: int = 5 + gp_weight: float = 10.0 + # 0.0 means "inherit train.lr" — not None, since the TOML writer has no + # null literal to round-trip. + critic_lr: float = 0.0 + # 0 means "inherit stage1_model.hidden_dim/n_res_blocks" — same + # round-trip-friendly sentinel as critic_lr above. + critic_hidden_dim: int = 0 + critic_n_res_blocks: int = 0 + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage1WganConfig": + d = d or {} + return cls( + noise_dim=d.get("noise_dim", 64), + n_critic=d.get("n_critic", 5), + gp_weight=d.get("gp_weight", 10.0), + critic_lr=d.get("critic_lr", 0.0), + critic_hidden_dim=d.get("critic_hidden_dim", 0), + critic_n_res_blocks=d.get("critic_n_res_blocks", 0), + ) + + def to_dict(self) -> dict: + return { + "noise_dim": self.noise_dim, + "n_critic": self.n_critic, + "gp_weight": self.gp_weight, + "critic_lr": self.critic_lr, + "critic_hidden_dim": self.critic_hidden_dim, + "critic_n_res_blocks": self.critic_n_res_blocks, + } + + +@dataclass(frozen=True) +class Stage2WganConfig(Stage1WganConfig): + # Straight-through Gumbel temperature for the particle-type one-hot + # (distinct from router.gumbel_tau_start/_end, which anneal + # expert-combination weights). Read only under particle_type.target = + # "onehot". + gumbel_tau_start: float = 1.0 + gumbel_tau_end: float = 0.1 + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage2WganConfig": + d = d or {} + return cls( + noise_dim=d.get("noise_dim", 64), + n_critic=d.get("n_critic", 5), + gp_weight=d.get("gp_weight", 10.0), + critic_lr=d.get("critic_lr", 0.0), + critic_hidden_dim=d.get("critic_hidden_dim", 0), + critic_n_res_blocks=d.get("critic_n_res_blocks", 0), + gumbel_tau_start=d.get("gumbel_tau_start", 1.0), + gumbel_tau_end=d.get("gumbel_tau_end", 0.1), + ) + + def to_dict(self) -> dict: + return { + **super().to_dict(), + "gumbel_tau_start": self.gumbel_tau_start, + "gumbel_tau_end": self.gumbel_tau_end, + } + + +# Fixed router fields shared by stage1_model.router and stage2_model.router. +# Composed-router axis{i}_{field} keys and pipeline.py's runtime-seeded +# centers_init are NOT in this set — they land in RouterConfig.extra instead +# (see giant.model.network._parse_composed_axes). +_ROUTER_KNOWN_KEYS = frozenset( + { + "enabled", + "type", + "n_experts", + "temperature", + "learn_centers", + "learn_width", + "learn_temperature", + "width_min_ratio", + "width_max_ratio", + "lambda_balance", + "lambda_entropy", + "gumbel", + "gumbel_tau_start", + "gumbel_tau_end", + "emb_dim", + "hidden_dim", + "lambda_proc", + } +) + + +@dataclass(frozen=True) +class RouterConfig: + """`stage1_model.router`'s fixed fields. `extra` holds any key not named + below — composed-router `axis{i}_{field}` flags and pipeline.py's + runtime-seeded `centers_init` — so from_dict()/to_dict() round-trip + losslessly without this dataclass needing to know about them.""" + + enabled: bool = False + type: str = "energy" # selects the Router impl from ROUTER_REGISTRY + n_experts: int = 4 + temperature: float = 0.5 # energy/pdg-router kwarg + learn_centers: bool = True # energy/pdg-router kwarg + # energy-router kwargs: mutually exclusive optional learnable + # gate-sharpness modes. learn_width generalizes the shared `temperature` + # to one learnable width per expert; learn_temperature instead makes the + # single shared `temperature` itself learnable. Both are bounded to + # [width_min_ratio, width_max_ratio] * temperature so gate sharpness + # can't run away to a collapse-inducing extreme during training. + learn_width: bool = False + learn_temperature: bool = False + width_min_ratio: float = 0.1 + width_max_ratio: float = 10.0 + # Importance-CV^2 load-balancing aux loss weight (Shazeer et al. 2017). + lambda_balance: float = 0.0 + # Entropy-regularization weight penalizing uniform/collapsed gating — a + # secondary guard against all experts' widths co-inflating together, + # which lambda_balance alone can't see. + lambda_entropy: float = 0.0 + # Opt-in straight-through Gumbel-softmax train-time combine weights: the + # training forward pass samples a hard one-hot combination (matching + # eval-time top-1 dispatch exactly) while the backward pass still flows + # smooth gradient to every expert. + gumbel: bool = False + gumbel_tau_start: float = 1.0 + gumbel_tau_end: float = 0.1 + emb_dim: int = 8 # process/pdg-router kwarg: own pdg(/mat) embedding width + hidden_dim: int = 64 # process-router kwarg: its classifier's hidden width + lambda_proc: float = 0.0 # process-router kwarg: supervised process-CE weight + # type = "composed" routes on multiple axes at once (e.g. energy x pdg), + # each with its own expert count/hyperparameters. Axes are NOT in these + # defaults — set them as flat axis{i}_{field} keys instead of + # "n_experts", e.g. axis0_type = "energy", axis0_n_experts = 4, + # axis1_type = "pdg", axis1_n_experts = 3, axis1_emb_dim = 8. See + # giant.model.network._parse_composed_axes. + extra: dict = field(default_factory=dict) + + @classmethod + def from_dict(cls, d: dict | None) -> "RouterConfig": + d = d or {} + return cls( + enabled=d.get("enabled", False), + type=d.get("type", "energy"), + n_experts=d.get("n_experts", 4), + temperature=d.get("temperature", 0.5), + learn_centers=d.get("learn_centers", True), + learn_width=d.get("learn_width", False), + learn_temperature=d.get("learn_temperature", False), + width_min_ratio=d.get("width_min_ratio", 0.1), + width_max_ratio=d.get("width_max_ratio", 10.0), + lambda_balance=d.get("lambda_balance", 0.0), + lambda_entropy=d.get("lambda_entropy", 0.0), + gumbel=d.get("gumbel", False), + gumbel_tau_start=d.get("gumbel_tau_start", 1.0), + gumbel_tau_end=d.get("gumbel_tau_end", 0.1), + emb_dim=d.get("emb_dim", 8), + hidden_dim=d.get("hidden_dim", 64), + lambda_proc=d.get("lambda_proc", 0.0), + extra={k: v for k, v in d.items() if k not in _ROUTER_KNOWN_KEYS}, + ) + + def to_dict(self) -> dict: + return { + "enabled": self.enabled, + "type": self.type, + "n_experts": self.n_experts, + "temperature": self.temperature, + "learn_centers": self.learn_centers, + "learn_width": self.learn_width, + "learn_temperature": self.learn_temperature, + "width_min_ratio": self.width_min_ratio, + "width_max_ratio": self.width_max_ratio, + "lambda_balance": self.lambda_balance, + "lambda_entropy": self.lambda_entropy, + "gumbel": self.gumbel, + "gumbel_tau_start": self.gumbel_tau_start, + "gumbel_tau_end": self.gumbel_tau_end, + "emb_dim": self.emb_dim, + "hidden_dim": self.hidden_dim, + "lambda_proc": self.lambda_proc, + **self.extra, + } + + +@dataclass(frozen=True) +class Stage2RouterConfig(RouterConfig): + # true: stage 2 shares stage 1's Router module instance, so expert i in + # stage 1 and expert i in stage 2 gate on identical conditions by + # construction — every other key in this block is then ignored. false: + # an independent router. + tie_to_stage1: bool = False + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage2RouterConfig": + d = d or {} + known = _ROUTER_KNOWN_KEYS | {"tie_to_stage1"} + return cls( + tie_to_stage1=d.get("tie_to_stage1", False), + enabled=d.get("enabled", False), + type=d.get("type", "energy"), + n_experts=d.get("n_experts", 4), + temperature=d.get("temperature", 0.5), + learn_centers=d.get("learn_centers", True), + learn_width=d.get("learn_width", False), + learn_temperature=d.get("learn_temperature", False), + width_min_ratio=d.get("width_min_ratio", 0.1), + width_max_ratio=d.get("width_max_ratio", 10.0), + lambda_balance=d.get("lambda_balance", 0.0), + lambda_entropy=d.get("lambda_entropy", 0.0), + gumbel=d.get("gumbel", False), + gumbel_tau_start=d.get("gumbel_tau_start", 1.0), + gumbel_tau_end=d.get("gumbel_tau_end", 0.1), + emb_dim=d.get("emb_dim", 8), + hidden_dim=d.get("hidden_dim", 64), + lambda_proc=d.get("lambda_proc", 0.0), + extra={k: v for k, v in d.items() if k not in known}, + ) + + def to_dict(self) -> dict: + return {"tie_to_stage1": self.tie_to_stage1, **super().to_dict()} + + +@dataclass(frozen=True) +class NSecConfig: + # "head": a classifier over {0..k_max} on the condition encoding alone + # (no diffusion noise), callable independently at inference. + # "stop_token": an EOS-style implicit stop — accepted by the schema but + # not implemented in v0.3.0 (see validate_config). + # "truth": take n_sec from ground truth — standalone stage-2 evaluation + # only, never for rollout. + mode: str = "head" + lambda_weight: float = 0.1 # dict key "lambda" — cross-entropy weight for the head + # Holds "legacy_owner" when injected by _migrate_legacy_model_config + # (v0.2 checkpoints only) — not a user-facing config.toml key. + extra: dict = field(default_factory=dict) + + @property + def legacy_owner(self) -> str | None: + return self.extra.get("legacy_owner") + + @classmethod + def from_dict(cls, d: dict | None) -> "NSecConfig": + d = d or {} + known = {"mode", "lambda"} + return cls( + mode=d.get("mode", "head"), + lambda_weight=d.get("lambda", 0.1), + extra={k: v for k, v in d.items() if k not in known}, + ) + + def to_dict(self) -> dict: + return {"mode": self.mode, "lambda": self.lambda_weight, **self.extra} + + +@dataclass(frozen=True) +class ParticleTypeConfig: + # The three targets mirror the three conditioning.particle modes. + # "onehot": class logits over conditioning.particle.emb_dim classes. + # "physical": regressed (log mass, charge). "embedding": regressed + # against conditioning's own particle embedding table (requires + # conditioning.particle.type = "embedding"). + target: str = "onehot" + lambda_weight: float = 1.0 # dict key "lambda" + # How a predicted "other" class becomes a concrete PDG code at rollout. + # "sample": draw from the empirical within-bucket distribution recorded + # at map-build time. "modal": always the most common member. "drop": + # discard the secondary. Read only under target = "onehot". + other_policy: str = "sample" + + @classmethod + def from_dict(cls, d: dict | None) -> "ParticleTypeConfig": + d = d or {} + return cls( + target=d.get("target", "onehot"), + lambda_weight=d.get("lambda", 1.0), + other_policy=d.get("other_policy", "sample"), + ) + + def to_dict(self) -> dict: + return {"target": self.target, "lambda": self.lambda_weight, "other_policy": self.other_policy} + + +@dataclass(frozen=True) +class AutoregressiveConfig: + # Canonical generation order. Single-valued for now; the key exists so + # an alternative ordering is not a config break. + order: str = "energy_desc" + # How token i+1 sees tokens <= i. "markov": previous token plus running + # scalars (remaining energy budget, slot index) — a fixed-width summary. + # "attention": causal self-attention over all emitted tokens. + history: str = "markov" + # "always": condition on the ground-truth previous secondary throughout + # training. "scheduled": scheduled sampling — interpolate toward the + # model's own prediction. "never": free-running from the start. + teacher_forcing: str = "always" + tf_p_start: float = 1.0 + tf_p_end: float = 1.0 + attn_n_heads: int = 4 + attn_n_layers: int = 2 + + @classmethod + def from_dict(cls, d: dict | None) -> "AutoregressiveConfig": + d = d or {} + return cls( + order=d.get("order", "energy_desc"), + history=d.get("history", "markov"), + teacher_forcing=d.get("teacher_forcing", "always"), + tf_p_start=d.get("tf_p_start", 1.0), + tf_p_end=d.get("tf_p_end", 1.0), + attn_n_heads=d.get("attn_n_heads", 4), + attn_n_layers=d.get("attn_n_layers", 2), + ) + + def to_dict(self) -> dict: + return { + "order": self.order, + "history": self.history, + "teacher_forcing": self.teacher_forcing, + "tf_p_start": self.tf_p_start, + "tf_p_end": self.tf_p_end, + "attn_n_heads": self.attn_n_heads, + "attn_n_layers": self.attn_n_layers, + } + + +@dataclass(frozen=True) +class Stage1ModelConfig: + # false skips building/training stage 1 entirely. The resulting + # checkpoint holds only stage 2 and cannot be rolled out. + active: bool = True + # "flow": conditional flow matching (~10 ODE steps at inference). + # "ddpm": cosine-schedule diffusion baseline. + # "wgan": WGAN-GP, single forward pass at inference. + generator: str = "flow" + # Trunk width — also the width of every expert under a routed trunk. + hidden_dim: int = 256 + # Number of ResBlocks in the trunk, and in every expert under a routed + # trunk. + n_res_blocks: int = 6 + dropout: float = 0.0 + # Weight of this stage's loss in the total when both stages are active + # and non-adversarial. A WGAN stage's adversarial loss drives its own + # optimizer, so `lambda` scales only its non-adversarial auxiliary terms. + lambda_weight: float = 1.0 # dict key "lambda" + flow: FlowConfig = field(default_factory=FlowConfig) + ddpm: DdpmConfig = field(default_factory=DdpmConfig) + wgan: Stage1WganConfig = field(default_factory=Stage1WganConfig) + router: RouterConfig = field(default_factory=RouterConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage1ModelConfig": + d = d or {} + return cls( + active=d.get("active", True), + generator=d.get("generator", "flow"), + hidden_dim=d.get("hidden_dim", 256), + n_res_blocks=d.get("n_res_blocks", 6), + dropout=d.get("dropout", 0.0), + lambda_weight=d.get("lambda", 1.0), + flow=FlowConfig.from_dict(d.get("flow")), + ddpm=DdpmConfig.from_dict(d.get("ddpm")), + wgan=Stage1WganConfig.from_dict(d.get("wgan")), + router=RouterConfig.from_dict(d.get("router")), + ) + + def to_dict(self) -> dict: + return { + "active": self.active, + "generator": self.generator, + "hidden_dim": self.hidden_dim, + "n_res_blocks": self.n_res_blocks, + "dropout": self.dropout, + "lambda": self.lambda_weight, + "flow": self.flow.to_dict(), + "ddpm": self.ddpm.to_dict(), + "wgan": self.wgan.to_dict(), + "router": self.router.to_dict(), + } + + +@dataclass(frozen=True) +class Stage2ModelConfig: + # false trains stage 1 alone. giant rollout must then refuse the + # checkpoint; giant predict still works. + active: bool = True + # "one_shot": predict all k_max slots simultaneously with padded slots + # masked from the loss (v0.2 behaviour). + # "autoregressive": emit one secondary at a time in descending-energy + # order. + decoder: str = "autoregressive" + # As stage1_model.generator, but under "autoregressive" this is the + # objective for each token. + generator: str = "wgan" + hidden_dim: int = 256 + n_res_blocks: int = 6 + dropout: float = 0.0 + lambda_weight: float = 1.0 # dict key "lambda" + # Maximum secondary slots. Under "one_shot" this is the fixed output + # width; under "autoregressive" it is a safety cap on the generation loop. + k_max: int = 15 + # Width of the projected stage-1 outcome fed into stage 2's conditioning. + context_dim: int = 64 + # "truth": the ground-truth stage-1 target vector, detached — stage-level + # teacher forcing (v0.2 behaviour). "sampled": stage 1's own sampled + # output, closing the train/inference gap at the cost of a sampling pass + # per batch and a moving target early in training. + stage1_context: str = "truth" + n_sec: NSecConfig = field(default_factory=NSecConfig) + particle_type: ParticleTypeConfig = field(default_factory=ParticleTypeConfig) + autoregressive: AutoregressiveConfig = field(default_factory=AutoregressiveConfig) + flow: FlowConfig = field(default_factory=FlowConfig) + ddpm: DdpmConfig = field(default_factory=DdpmConfig) + wgan: Stage2WganConfig = field(default_factory=Stage2WganConfig) + router: Stage2RouterConfig = field(default_factory=Stage2RouterConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "Stage2ModelConfig": + d = d or {} + return cls( + active=d.get("active", True), + decoder=d.get("decoder", "autoregressive"), + generator=d.get("generator", "wgan"), + hidden_dim=d.get("hidden_dim", 256), + n_res_blocks=d.get("n_res_blocks", 6), + dropout=d.get("dropout", 0.0), + lambda_weight=d.get("lambda", 1.0), + k_max=d.get("k_max", 15), + context_dim=d.get("context_dim", 64), + stage1_context=d.get("stage1_context", "truth"), + n_sec=NSecConfig.from_dict(d.get("n_sec")), + particle_type=ParticleTypeConfig.from_dict(d.get("particle_type")), + autoregressive=AutoregressiveConfig.from_dict(d.get("autoregressive")), + flow=FlowConfig.from_dict(d.get("flow")), + ddpm=DdpmConfig.from_dict(d.get("ddpm")), + wgan=Stage2WganConfig.from_dict(d.get("wgan")), + router=Stage2RouterConfig.from_dict(d.get("router")), + ) + + def to_dict(self) -> dict: + return { + "active": self.active, + "decoder": self.decoder, + "generator": self.generator, + "hidden_dim": self.hidden_dim, + "n_res_blocks": self.n_res_blocks, + "dropout": self.dropout, + "lambda": self.lambda_weight, + "k_max": self.k_max, + "context_dim": self.context_dim, + "stage1_context": self.stage1_context, + "n_sec": self.n_sec.to_dict(), + "particle_type": self.particle_type.to_dict(), + "autoregressive": self.autoregressive.to_dict(), + "flow": self.flow.to_dict(), + "ddpm": self.ddpm.to_dict(), + "wgan": self.wgan.to_dict(), + "router": self.router.to_dict(), + } + + +@dataclass(frozen=True) +class TrainConfig: + epochs: int = 100 + batch_size: int = 4096 + lr: float = 3e-4 + weight_decay: float = 0.01 # AdamW default — exposed so it can be tuned + ema_decay: float = 0.9999 # EMA of model weights for sampling; 0 disables + warmup_epochs: int = 5 + val_fraction: float = 0.1 + # per-epoch val loss (not the marginal/KL validate_every pass) is capped + # to this many batches; 0 = full val set every epoch + max_val_batches: int = 200 + num_workers: int = 4 + seed: int = 0 + validate_every: int = 10 + validate_steps: int = 10 + # Weights & Biases per-epoch metric logging. Default true in v0.3.0 (was + # opt-in false): the v0.3.0 work is a sequence of architecture + # comparisons, and a run that wasn't logged isn't comparable. Set false + # for throwaway/debug runs. + wandb: bool = True + wandb_project: str = "giant" + # "" means "use the checkpoint out_dir name" — not None, since the TOML + # writer has no null literal to round-trip. + wandb_run_name: str = "" + # Batch-granularity metrics (loss/grad_norm/lr) are logged every N + # optimizer steps, not every batch — a single epoch can be tens of + # thousands of steps. Per-epoch metrics (the metrics.csv row) always log + # in full. + wandb_log_every: int = 50 + + @classmethod + def from_dict(cls, d: dict | None) -> "TrainConfig": + d = d or {} + return cls( + epochs=d.get("epochs", 100), + batch_size=d.get("batch_size", 4096), + lr=d.get("lr", 3e-4), + weight_decay=d.get("weight_decay", 0.01), + ema_decay=d.get("ema_decay", 0.9999), + warmup_epochs=d.get("warmup_epochs", 5), + val_fraction=d.get("val_fraction", 0.1), + max_val_batches=d.get("max_val_batches", 200), + num_workers=d.get("num_workers", 4), + seed=d.get("seed", 0), + validate_every=d.get("validate_every", 10), + validate_steps=d.get("validate_steps", 10), + wandb=d.get("wandb", True), + wandb_project=d.get("wandb_project", "giant"), + wandb_run_name=d.get("wandb_run_name", ""), + wandb_log_every=d.get("wandb_log_every", 50), + ) + + def to_dict(self) -> dict: + return { + "epochs": self.epochs, + "batch_size": self.batch_size, + "lr": self.lr, + "weight_decay": self.weight_decay, + "ema_decay": self.ema_decay, + "warmup_epochs": self.warmup_epochs, + "val_fraction": self.val_fraction, + "max_val_batches": self.max_val_batches, + "num_workers": self.num_workers, + "seed": self.seed, + "validate_every": self.validate_every, + "validate_steps": self.validate_steps, + "wandb": self.wandb, + "wandb_project": self.wandb_project, + "wandb_run_name": self.wandb_run_name, + "wandb_log_every": self.wandb_log_every, + } + + +@dataclass(frozen=True) +class GiantConfig: + """Root config dataclass — the single source of truth for every default + in DEFAULT_CONFIG below, which is generated from `GiantConfig().to_dict()` + rather than hand-maintained (see issues.md Issue 1).""" + + conditioning: ConditioningConfig = field(default_factory=ConditioningConfig) + stage1_model: Stage1ModelConfig = field(default_factory=Stage1ModelConfig) + stage2_model: Stage2ModelConfig = field(default_factory=Stage2ModelConfig) + train: TrainConfig = field(default_factory=TrainConfig) + + @classmethod + def from_dict(cls, d: dict | None) -> "GiantConfig": + d = d or {} + return cls( + conditioning=ConditioningConfig.from_dict(d.get("conditioning")), + stage1_model=Stage1ModelConfig.from_dict(d.get("stage1_model")), + stage2_model=Stage2ModelConfig.from_dict(d.get("stage2_model")), + train=TrainConfig.from_dict(d.get("train")), + ) + + def to_dict(self) -> dict: + return { + "conditioning": self.conditioning.to_dict(), + "stage1_model": self.stage1_model.to_dict(), + "stage2_model": self.stage2_model.to_dict(), + "train": self.train.to_dict(), + } + + +DEFAULT_CONFIG: dict = GiantConfig().to_dict() def git_hash() -> str: try: - return ( - subprocess.check_output( - ["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL - ) - .decode() - .strip() - ) + return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip() except Exception: return "unknown" @@ -169,6 +740,8 @@ def auto_device() -> torch.device: # Activation memory is assumed to scale linearly with # batch_size * hidden_dim * n_blocks (the ResBlock stack dominates), so this # is a rough estimate rather than a guaranteed bound. +# NOTE: not yet recalibrated for the v0.3.0 autoregressive stage-2 trunk — +# deliberately last in the implementation order. _REF_BYTES = 7683 * 1024**2 _REF_BATCH_SIZE = 29696 _REF_HIDDEN_DIM = 1024 @@ -200,12 +773,8 @@ def estimate_batch_size( calibration since there's no backward graph or optimizer state. """ if device.type != "cuda": - raise ValueError( - f"--batch-size auto is only supported on cuda devices, got {device.type!r}" - ) - device_index = ( - device.index if device.index is not None else torch.cuda.current_device() - ) + raise ValueError(f"--batch-size auto is only supported on cuda devices, got {device.type!r}") + device_index = device.index if device.index is not None else torch.cuda.current_device() free_bytes, _total_bytes = torch.cuda.mem_get_info(device_index) if training: ref_bytes, ref_batch_size, ref_hidden_dim, ref_n_blocks = ( @@ -254,13 +823,15 @@ def warn_if_git_hash_mismatch(file_cfg: dict, config_path: Path) -> None: def load_checkpoint_config(ckpt_path: str | Path) -> dict: - """Load the full ``[train]``/``[model]``/``[meta]`` config.toml written - alongside a checkpoint by ``save_config``. + """Load the full config.toml written alongside a checkpoint by + ``save_config``. Returns ``{}`` if no config.toml sits next to the checkpoint (older runs, or a checkpoint moved without its sidecar) — this is best-effort provenance for threading into a rollout's YAML sidecar, not a hard - requirement for using the checkpoint itself. + requirement for using the checkpoint itself. Returned as-loaded (v0.2 or + v0.3 shape); callers that need the v0.3 shape should run it through + `migrate_config` themselves. """ config_path = Path(ckpt_path).parent / "config.toml" if not config_path.exists(): @@ -283,130 +854,597 @@ def warn_if_checkpoint_config_mismatch(ckpt_path: str | Path) -> None: warn_if_git_hash_mismatch(load_toml(config_path), config_path) +def _get_path(d: dict, dotted: str): + """Read a dotted path (e.g. "stage1_model.router.enabled") out of a + nested dict. Returns None if any component along the path is missing.""" + cur = d + for part in dotted.split("."): + if not isinstance(cur, dict) or part not in cur: + return None + cur = cur[part] + return cur + + +def _set_path(d: dict, dotted: str, value) -> None: + """Write a dotted path into a nested dict, creating intermediate dicts as + needed.""" + parts = dotted.split(".") + cur = d + for part in parts[:-1]: + cur = cur.setdefault(part, {}) + cur[parts[-1]] = value + + +def _deep_merge(base: dict, override: dict) -> dict: + """Recursively merge `override` onto a copy of `base`. + + Dict-valued keys recurse instead of being replaced wholesale, so + overriding one leaf (e.g. stage1_model.router.enabled) never drops + untouched siblings — the rest of stage1_model.router, or of + stage1_model — the same property v0.2's router-only bespoke merge had, + generalized here to arbitrary depth. + """ + result = dict(base) + for k, v in override.items(): + if isinstance(v, dict) and isinstance(result.get(k), dict): + result[k] = _deep_merge(result[k], v) + else: + result[k] = v + return result + + +@dataclass(frozen=True) +class FlagSpec: + """One CLI flag's mapping into the config-overrides tree. + + `paths` lists every dotted config path this flag writes (>1 means fan-out + to multiple stages/axes, e.g. `--mode` -> both stages' `generator`). + `precedence` controls write order when two flags target the same path: + specs are applied in ascending precedence, so a higher-precedence (more + specific) flag overwrites a lower-precedence (shared/shorthand) one — + this is the "build a shared dict, then let a more specific dict win" + pattern `giant train`/`giant new-run` need (e.g. `--hidden-dim` vs + `--stage1-hidden-dim`, or `--n-critic` vs `--stage1-n-critic`), + generalized to one mechanism instead of three different ad hoc ones. + """ + + name: str + paths: tuple[str, ...] + precedence: int = 0 + + +# Flag -> config-path table shared by `giant train`/`giant new-run` +# (giant/cli.py) so both commands resolve CLI overrides identically. See +# issues.md Issue 3: this replaces ~140 lines of hand-written, imperative +# dict-building in cli.py with one declarative table plus +# `overrides_from_flags` below. +FLAG_SPECS: tuple[FlagSpec, ...] = ( + # train block -- flat pass-through, unique paths, precedence irrelevant. + FlagSpec("epochs", ("train.epochs",)), + FlagSpec("batch_size", ("train.batch_size",)), + FlagSpec("lr", ("train.lr",)), + FlagSpec("weight_decay", ("train.weight_decay",)), + FlagSpec("ema_decay", ("train.ema_decay",)), + FlagSpec("warmup_epochs", ("train.warmup_epochs",)), + FlagSpec("val_fraction", ("train.val_fraction",)), + FlagSpec("num_workers", ("train.num_workers",)), + FlagSpec("seed", ("train.seed",)), + FlagSpec("validate_every", ("train.validate_every",)), + FlagSpec("validate_steps", ("train.validate_steps",)), + FlagSpec("max_val_batches", ("train.max_val_batches",)), + FlagSpec("wandb", ("train.wandb",)), + FlagSpec("wandb_project", ("train.wandb_project",)), + FlagSpec("wandb_run_name", ("train.wandb_run_name",)), + FlagSpec("wandb_log_every", ("train.wandb_log_every",)), + # --hidden-dim/--n-blocks/--dropout are stage-1-only backward-compat + # shorthands (they predate stage2_model having its own flags); + # --stage1-* wins when both are given. + FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0), + FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1), + FlagSpec("n_blocks", ("stage1_model.n_res_blocks",), precedence=0), + FlagSpec("stage1_n_res_blocks", ("stage1_model.n_res_blocks",), precedence=1), + FlagSpec("dropout", ("stage1_model.dropout",), precedence=0), + FlagSpec("stage1_dropout", ("stage1_model.dropout",), precedence=1), + # stage2-only knobs. + FlagSpec("stage2_hidden_dim", ("stage2_model.hidden_dim",)), + FlagSpec("stage2_n_res_blocks", ("stage2_model.n_res_blocks",)), + FlagSpec("stage2_dropout", ("stage2_model.dropout",)), + FlagSpec("stage2_decoder", ("stage2_model.decoder",)), + FlagSpec("stage2_k_max", ("stage2_model.k_max",)), + FlagSpec("stage2_context_dim", ("stage2_model.context_dim",)), + FlagSpec("stage2_stage1_context", ("stage2_model.stage1_context",)), + # --mode applies to both stages by default (v0.2 had one shared + # mode/wgan config); --stage{1,2}-generator override a single stage. + FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator"), precedence=0), + FlagSpec("stage1_generator", ("stage1_model.generator",), precedence=1), + FlagSpec("stage2_generator", ("stage2_model.generator",), precedence=1), + # --emb-dim/--conditioning set both conditioning axes (v0.2 had one + # shared value for particle+material). + FlagSpec("conditioning", ("conditioning.particle.type", "conditioning.material.type")), + FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")), + # Pre-aggregated router override dict (built by `_router_cli_overrides` + # in cli.py from --router/--router-type/--n-experts/--router-axis). + # Router overrides only ever land on stage1_model -- this asymmetry is + # deliberate (see cli.py) and must not be "fixed" into a fan-out here. + FlagSpec("router_config", ("stage1_model.router",)), + # WGAN: shared knobs apply to both stages by default (v0.2 had one + # shared wgan config); --stage{1,2}-* override a single stage. + FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic"), precedence=0), + FlagSpec("stage1_n_critic", ("stage1_model.wgan.n_critic",), precedence=1), + FlagSpec("stage2_n_critic", ("stage2_model.wgan.n_critic",), precedence=1), + FlagSpec("gp_weight", ("stage1_model.wgan.gp_weight", "stage2_model.wgan.gp_weight"), precedence=0), + FlagSpec("stage1_gp_weight", ("stage1_model.wgan.gp_weight",), precedence=1), + FlagSpec("stage2_gp_weight", ("stage2_model.wgan.gp_weight",), precedence=1), + FlagSpec("noise_dim", ("stage1_model.wgan.noise_dim", "stage2_model.wgan.noise_dim"), precedence=0), + FlagSpec("stage1_noise_dim", ("stage1_model.wgan.noise_dim",), precedence=1), + FlagSpec("stage2_noise_dim", ("stage2_model.wgan.noise_dim",), precedence=1), + FlagSpec("critic_lr", ("stage1_model.wgan.critic_lr", "stage2_model.wgan.critic_lr"), precedence=0), + FlagSpec("stage1_critic_lr", ("stage1_model.wgan.critic_lr",), precedence=1), + FlagSpec("stage2_critic_lr", ("stage2_model.wgan.critic_lr",), precedence=1), +) + + +def overrides_from_flags(values: dict[str, object]) -> dict: + """Build the nested, section-keyed config-overrides dict + `merge_cli_overrides` expects, from `{flag_name: value}`. + + Flags absent from `values`, or mapped to `None` (= not given on the + CLI), are skipped. See `FlagSpec`/`FLAG_SPECS` above for the precedence + rule applied when two flags target the same path. + """ + overrides: dict = {} + for spec in sorted(FLAG_SPECS, key=lambda s: s.precedence): + if spec.name not in values or values[spec.name] is None: + continue + for path in spec.paths: + _set_path(overrides, path, values[spec.name]) + return overrides + + +# v0.2 [train] keys that pass through to v0.3 [train] unchanged (same name, +# same meaning) when present in the loaded file — everything model-shaped +# moved to the stage/conditioning blocks instead (see the rest of +# migrate_config below). +_V02_TRAIN_PASSTHROUGH = ( + "epochs", + "batch_size", + "lr", + "weight_decay", + "ema_decay", + "max_val_batches", + "val_fraction", + "num_workers", + "seed", + "validate_every", + "validate_steps", + "warmup_epochs", + "wandb", + "wandb_project", + "wandb_run_name", + "wandb_log_every", +) + +# v0.2 model.hidden_dim/n_blocks/dropout applied identically to both stages +# (there was only ever one trunk shape) -> copied to both stage{1,2}_model. +_V02_MODEL_TO_BOTH_STAGES = ( + ("hidden_dim", "hidden_dim"), + ("n_blocks", "n_res_blocks"), + ("dropout", "dropout"), +) + +# v0.2 train.{n_critic,gp_weight,critic_lr} applied identically to both +# stages' wgan sub-table (there was only ever one wgan objective, shared). +_V02_TRAIN_TO_BOTH_STAGES_WGAN = ( + ("n_critic", "n_critic"), + ("gp_weight", "gp_weight"), + ("critic_lr", "critic_lr"), +) + + +def migrate_config(cfg: dict) -> dict: + """Translate a v0.2 config dict (single [train] + [model]) into the v0.3 + nested format ([conditioning]/[stage1_model]/[stage2_model]/[train]). + + Called on every config.toml load (see merge_cli_overrides) so old + training configs on disk keep working under new code without hand- + editing. `[meta].config_version == CONFIG_VERSION` marks a dict as + already-v0.3; its absence is read as "this is v0.2", so an + already-migrated dict is returned unchanged (deep-copied). + + Only keys actually present in `cfg` are translated — `cfg` may be a + partial file (e.g. `[train]\\nepochs = 5\\n` with no [model] section at + all, relying on v0.2 defaults for everything else). Separately, a fixed + set of v0.2 architectural facts that were never exposed as config keys at + all (e.g. the conditioning MLP was always 2 layers deep, not the v0.3 + default of 1) are injected unconditionally whenever this function decides + it is migrating a v0.2 dict, regardless of which keys the file happened + to set. + + Operates on the config.toml shape. A checkpoint's `model_config` dict + (which additionally carries n_sec_head ownership and needs + `network.build_models`'s cooperation) is a separate migration surface, + deferred to the network.py refactor. + """ + if _get_path(cfg, "meta.config_version") == CONFIG_VERSION: + return copy.deepcopy(cfg) + + cfg = copy.deepcopy(cfg) + old_train = cfg.pop("train", {}) + old_model = cfg.pop("model", {}) + old_router = dict(old_model.pop("router", {})) + + new: dict = {} + + for key in _V02_TRAIN_PASSTHROUGH: + if key in old_train: + _set_path(new, f"train.{key}", old_train[key]) + + if "mode" in old_train: + _set_path(new, "stage1_model.generator", old_train["mode"]) + _set_path(new, "stage2_model.generator", old_train["mode"]) + if "lambda_nsec" in old_train: + _set_path(new, "stage2_model.n_sec.lambda", old_train["lambda_nsec"]) + if "lambda_s2" in old_train: + _set_path(new, "stage2_model.lambda", old_train["lambda_s2"]) + for old_key, new_key in _V02_TRAIN_TO_BOTH_STAGES_WGAN: + if old_key in old_train: + _set_path(new, f"stage1_model.wgan.{new_key}", old_train[old_key]) + _set_path(new, f"stage2_model.wgan.{new_key}", old_train[old_key]) + + for old_key, new_key in _V02_MODEL_TO_BOTH_STAGES: + if old_key in old_model: + _set_path(new, f"stage1_model.{new_key}", old_model[old_key]) + _set_path(new, f"stage2_model.{new_key}", old_model[old_key]) + if "emb_dim" in old_model: + _set_path(new, "conditioning.particle.emb_dim", old_model["emb_dim"]) + _set_path(new, "conditioning.material.emb_dim", old_model["emb_dim"]) + if "conditioning" in old_model: + _set_path(new, "conditioning.particle.type", old_model["conditioning"]) + _set_path(new, "conditioning.material.type", old_model["conditioning"]) + if "noise_dim" in old_model: + _set_path(new, "stage1_model.wgan.noise_dim", old_model["noise_dim"]) + _set_path(new, "stage2_model.wgan.noise_dim", old_model["noise_dim"]) + if "k_max" in old_model: + _set_path(new, "stage2_model.k_max", old_model["k_max"]) + + if old_router: + expert_hidden_dim = old_router.pop("expert_hidden_dim", 0) + expert_n_blocks = old_router.pop("expert_n_blocks", 0) + if expert_hidden_dim or expert_n_blocks: + raise ValueError( + "v0.2 config sets model.router.expert_hidden_dim/" + f"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 config's routed experts " + "have a different width/depth than the monolith and its " + "checkpoint can only be loaded by v0.2 code." + ) + _set_path(new, "stage1_model.router", dict(old_router)) + stage2_router = dict(old_router) + stage2_router["tie_to_stage1"] = False + _set_path(new, "stage2_model.router", stage2_router) + + # v0.2 architectural facts with no corresponding config key at all — + # always set once we've determined we're migrating a v0.2 dict, + # independent of what the file did/didn't specify. NOTE: n_layers here + # (2) differs from the v0.3 *default* (1) — this is not a typo, see the + # docstring above. + _set_path(new, "conditioning.out_dim", 128) + _set_path(new, "conditioning.particle.n_layers", 2) + _set_path(new, "conditioning.material.n_layers", 2) + _set_path(new, "stage1_model.active", True) + _set_path(new, "stage1_model.flow.time_dim", 64) + _set_path(new, "stage1_model.ddpm.time_dim", 64) + _set_path(new, "stage2_model.active", True) + _set_path(new, "stage2_model.flow.time_dim", 64) + _set_path(new, "stage2_model.ddpm.time_dim", 64) + _set_path(new, "stage2_model.context_dim", 64) + _set_path(new, "stage2_model.decoder", "one_shot") + _set_path(new, "stage2_model.particle_type.target", "physical") + + new_meta = dict(cfg.pop("meta", {})) + new_meta["config_version"] = CONFIG_VERSION + new["meta"] = new_meta + + # Anything else in the original dict (unrecognized top-level sections) + # carries through untouched rather than being silently dropped. + for k, v in cfg.items(): + new.setdefault(k, v) + + return new + + +# axis{i}_{field} composed-router keys (see network._parse_composed_axes) — field-name +# agnostic, matching network.py's own _AXIS_KEY_RE, since anything after axis{i}_ is +# passed straight through as a router kwarg there. +_AXIS_KEY_RE = re.compile(r"^axis\d+_.+$") + +# Paths (dotted, relative to the merged cfg root) that carry genuinely dynamic keys not +# in DEFAULT_CONFIG's fixed schema — composed-router axis{i}_{field} flags and +# pipeline.seed_router_centers's runtime-seeded centers_init (see RouterConfig.extra +# above). validate_config_keys allows any key under these paths through unconditionally +# apart from the axis-pattern/centers_init check below. +_DYNAMIC_ROUTER_PATHS = {"stage1_model.router", "stage2_model.router"} + + +def _unknown_key_error(full_path: str, key: str, valid_keys) -> ValueError: + hint = difflib.get_close_matches(key, list(valid_keys), n=1) + suggestion = f" — did you mean {hint[0]!r}?" if hint else "" + return ValueError(f"unknown config key {full_path!r}{suggestion}") + + +def _validate_keys(node: dict, default_node: dict, path: str) -> None: + for key, value in node.items(): + if path == "" and key == "meta": + continue + full_path = f"{path}.{key}" if path else key + if path in _DYNAMIC_ROUTER_PATHS and (key == "centers_init" or _AXIS_KEY_RE.match(key)): + continue + if key not in default_node: + raise _unknown_key_error(full_path, key, default_node.keys()) + if isinstance(value, dict) and isinstance(default_node[key], dict): + _validate_keys(value, default_node[key], full_path) + + +def validate_config_keys(cfg: dict) -> None: + """Reject any config key not part of the known v0.3 schema (DEFAULT_CONFIG's tree). + + Catches typos like `n_res_block` for `n_res_blocks` that would otherwise merge + cleanly, pass `validate_config`, and silently build the wrong model — see + issues.md Issue 2. + + Only exercised on the config.toml/CLI-overrides path (called from + `merge_cli_overrides` below). A checkpoint's `model_config` dict goes through + `network._migrate_legacy_model_config`/`build_models` instead and must keep + loading regardless of schema drift; old checkpoints predate this validator and + are never passed through here. + + Two areas are deliberately dynamic and excluded: `[meta]` (run provenance, no + DEFAULT_CONFIG counterpart), and `stage{1,2}_model.router`'s `axis{i}_{field}` + keys (composed-router axes, see `network._parse_composed_axes`) / + `centers_init` (runtime-seeded by `pipeline.seed_router_centers`). + """ + _validate_keys(cfg, DEFAULT_CONFIG, "") + + def merge_cli_overrides( defaults: dict, config_path: Path | None, - train_overrides: dict, - model_overrides: dict, + overrides: dict, ) -> dict: - """Resolve config as defaults -> TOML file -> explicit CLI flags. + """Resolve config as defaults -> TOML file -> explicit overrides. - `model.router` is deep-merged one level (rather than replaced wholesale) - at each stage, so a TOML file or CLI flag only overriding e.g. - `router.enabled` doesn't drop the rest of the router defaults. + `overrides` is keyed by top-level section name (e.g. "stage1_model", + "train"), each value an arbitrarily nested dict of overrides to + deep-merge (see `_deep_merge`) — the shape stage-prefixed CLI flags + naturally produce. A v0.2-shaped TOML file is transparently migrated + (`migrate_config`) before merging, so old configs on disk keep working + under the new schema. The result is checked against the known schema + (`validate_config_keys`) before being returned, so a typo'd key raises + here rather than silently building the wrong model. """ - cfg = {"train": dict(defaults["train"]), "model": dict(defaults["model"])} - cfg["model"]["router"] = dict(defaults["model"]["router"]) + cfg = copy.deepcopy(defaults) if config_path is not None: - file_cfg = load_toml(config_path) - cfg["train"].update(file_cfg.get("train", {})) - file_model = dict(file_cfg.get("model", {})) - file_router = file_model.pop("router", None) - cfg["model"].update(file_model) - if file_router: - cfg["model"]["router"].update(file_router) + file_cfg = migrate_config(load_toml(config_path)) + for section, values in file_cfg.items(): + if section == "meta": + continue + if isinstance(values, dict): + cfg[section] = _deep_merge(cfg.get(section, {}), values) + else: + cfg[section] = values warn_if_git_hash_mismatch(file_cfg, config_path) - model_overrides = dict(model_overrides) - router_overrides = model_overrides.pop("router", None) - cfg["train"].update(train_overrides) - cfg["model"].update(model_overrides) - if router_overrides: - cfg["model"]["router"].update(router_overrides) + for section, values in overrides.items(): + if isinstance(values, dict): + cfg[section] = _deep_merge(cfg.get(section, {}), values) + else: + cfg[section] = values + validate_config_keys(cfg) return cfg -def resolve_expert_dims( - router_cfg: dict, hidden_dim: int, n_blocks: int -) -> tuple[int, int]: - """Resolve a router's expert hidden_dim/n_blocks, inheriting from the - monolith's when left at the 0 ("unset") sentinel. +def validate_config(cfg: dict) -> None: + """Cross-block validation the per-block schema can't express on its own. - Used by both `giant.pipeline` (to build the checkpoint's `model_config`) - and `giant.cli`'s batch-size auto-estimate, so `--hidden-dim`/`--n-blocks` - size the experts the same way in both places unless - `router.expert_hidden_dim`/`expert_n_blocks` are explicitly overridden. + Raises ValueError with a clear message on the first violation found. Call + after `merge_cli_overrides` has produced a fully-merged v0.3 config — + these checks need to see across blocks, so they don't belong in + `migrate_config` (which only ever sees one dict's own keys) or in any + single block's defaults. """ - expert_hidden_dim = router_cfg.get("expert_hidden_dim") or hidden_dim - expert_n_blocks = router_cfg.get("expert_n_blocks") or n_blocks - return expert_hidden_dim, expert_n_blocks + particle_type = _get_path(cfg, "conditioning.particle.type") + + pt_target = _get_path(cfg, "stage2_model.particle_type.target") + if pt_target == "embedding" and particle_type != "embedding": + raise ValueError( + "stage2_model.particle_type.target = 'embedding' requires " + "conditioning.particle.type = 'embedding' (there is no embedding " + "table to regress against under conditioning.particle.type = " + f"{particle_type!r})" + ) + + for stage_name in ("stage1_model", "stage2_model"): + router = _get_path(cfg, f"{stage_name}.router") or {} + if router.get("enabled") and router.get("type") in ("pdg", "process") and particle_type == "physical": + raise ValueError( + f"{stage_name}.router.type = {router['type']!r} builds its " + "own training-vocab-scoped embedding, incompatible with " + "conditioning.particle.type = 'physical' (defeats " + "generalization beyond the training menu) — pick a " + "different router type or a different " + "conditioning.particle.type" + ) + + if _get_path(cfg, "stage2_model.router.tie_to_stage1") and not _get_path(cfg, "stage1_model.active"): + raise ValueError( + "stage2_model.router.tie_to_stage1 = true requires " + "stage1_model.active = true (there is no stage-1 router to tie to)" + ) + + if _get_path(cfg, "stage2_model.n_sec.mode") == "stop_token": + raise ValueError( + "stage2_model.n_sec.mode = 'stop_token' is accepted by the schema " + "but not implemented in v0.3.0 — use 'head' (default) or 'truth' " + "(standalone stage-2 evaluation only, never for rollout)" + ) + + if ( + _get_path(cfg, "stage2_model.n_sec.mode") == "truth" + and _get_path(cfg, "stage1_model.active") + and _get_path(cfg, "stage2_model.active") + ): + raise ValueError( + "stage2_model.n_sec.mode = 'truth' is invalid for a " + "rollout-capable checkpoint (both stage1_model.active and " + "stage2_model.active = true): " + "'truth' takes n_sec from ground truth, which giant rollout " + "doesn't have. 'truth' is for standalone stage-2 evaluation " + "only — set stage1_model.active = false for that, or use " + "'head' (default) for a rollout-capable checkpoint." + ) + + if _get_path(cfg, "stage2_model.decoder") == "autoregressive": + history = _get_path(cfg, "stage2_model.autoregressive.history") + if history not in ("markov", "attention"): + raise ValueError(f"stage2_model.autoregressive.history = {history!r} — must be 'markov' or 'attention'") + teacher_forcing = _get_path(cfg, "stage2_model.autoregressive.teacher_forcing") + if teacher_forcing not in ("always", "scheduled", "never"): + raise ValueError( + "stage2_model.autoregressive.teacher_forcing = " + f"{teacher_forcing!r} — must be 'always', 'scheduled' or " + "'never'" + ) -_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb"} +_CONDITIONING_CODE = {"physical": "phys", "embedding": "emb", "onehot": "oh"} -# Priority-ordered candidate fields for default_out_dir_name: (label, getter, -# formatter). `getter(train, model)` returns None when the field is at its -# default (and so should be omitted); otherwise formatter(value) renders the -# name token. The router is a single unit gated on `router.enabled` rather -# than one candidate per router key, since its type/n_experts are meaningless -# while disabled. -def _mode_candidate(train, model): - return None if train["mode"] == DEFAULT_CONFIG["train"]["mode"] else train["mode"] +def _path_candidate(dotted_path: str, prefix: str, formatter=str): + """Candidate factory: show `prefix + formatter(value)` when the value at + `dotted_path` differs from its DEFAULT_CONFIG value, else omit.""" - -def _router_candidate(train, model): - router = model["router"] - if router["enabled"] == DEFAULT_CONFIG["model"]["router"]["enabled"]: - return None - return f"r-{router['type']}{router['n_experts']}" - - -def _router_flag_candidate(field, token_map): - """Candidate factory for a boolean `model.router` sub-field. - - Gated on `router.enabled` like `_router_candidate` (a disabled router's - sub-fields are meaningless), then omitted unless `field` differs from - its DEFAULT_CONFIG value — same "only show non-default" rule as every - other candidate. `token_map` need only cover the non-default value(s), - since the default value always yields None. - """ - - def _candidate(train, model): - router = model["router"] - default_router = DEFAULT_CONFIG["model"]["router"] - if router["enabled"] == default_router["enabled"]: + def _candidate(cfg): + value = _get_path(cfg, dotted_path) + default = _get_path(DEFAULT_CONFIG, dotted_path) + if value == default: return None - value = router[field] - if value == default_router[field]: - return None - return token_map[value] + return f"{prefix}{formatter(value)}" return _candidate -def _conditioning_candidate(train, model): - if model["conditioning"] == DEFAULT_CONFIG["model"]["conditioning"]: - return None - code = _CONDITIONING_CODE.get(model["conditioning"], model["conditioning"]) - return f"c{code}" - - -def _default_field_candidate(section_key, field, prefix): - def _candidate(train, model): - section = train if section_key == "train" else model - value = section[field] - if value == DEFAULT_CONFIG[section_key][field]: +def _conditioning_candidate(axis: str, short: str): + def _candidate(cfg): + value = _get_path(cfg, f"conditioning.{axis}.type") + default = _get_path(DEFAULT_CONFIG, f"conditioning.{axis}.type") + if value == default: return None - return f"{prefix}{value}" + code = _CONDITIONING_CODE.get(value, value) + return f"{short}{code}" return _candidate +def _router_candidate(stage_key: str, short: str): + """Candidate for a stage's router as a single unit, gated on + `router.enabled` (a disabled router's type/n_experts are meaningless).""" + + def _candidate(cfg): + router = _get_path(cfg, f"{stage_key}.router") or {} + default_router = _get_path(DEFAULT_CONFIG, f"{stage_key}.router") or {} + if router.get("enabled") == default_router.get("enabled"): + return None + return f"{short}r-{router['type']}{router['n_experts']}" + + return _candidate + + +def _router_flag_candidate(stage_key: str, short: str, field: str, token_map: dict): + """Candidate factory for a boolean field inside a stage's router block. + + Gated on `router.enabled` like `_router_candidate`, then omitted unless + `field` differs from its DEFAULT_CONFIG value. `token_map` need only + cover the non-default value(s), since the default value always yields + None. + """ + + def _candidate(cfg): + router = _get_path(cfg, f"{stage_key}.router") or {} + default_router = _get_path(DEFAULT_CONFIG, f"{stage_key}.router") or {} + if router.get("enabled") == default_router.get("enabled"): + return None + value = router.get(field) + if value == default_router.get(field): + return None + return f"{short}{token_map[value]}" + + return _candidate + + +# Priority-ordered candidate fields for default_out_dir_name: (label, +# candidate(cfg) -> str | None). Beyond _OUT_DIR_NAME_MAX_FIELDS non-default +# fields, the remainder collapse into a hash suffix (see +# default_out_dir_name). Candidates read the whole nested cfg via dotted +# paths — there is no single "model" dict anymore now that architecture is +# split across conditioning/stage1_model/stage2_model. _OUT_DIR_NAME_CANDIDATES = [ - ("mode", _mode_candidate), - ("router", _router_candidate), - ("gumbel", _router_flag_candidate("gumbel", {True: "gum"})), - ("learn_centers", _router_flag_candidate("learn_centers", {False: "nolc"})), - ("learn_width", _router_flag_candidate("learn_width", {True: "lw"})), - ("learn_temperature", _router_flag_candidate("learn_temperature", {True: "lt"})), - ("conditioning", _conditioning_candidate), - ("hidden_dim", _default_field_candidate("model", "hidden_dim", "h")), - ("n_blocks", _default_field_candidate("model", "n_blocks", "b")), - ("emb_dim", _default_field_candidate("model", "emb_dim", "e")), - ("lr", _default_field_candidate("train", "lr", "lr")), - ("batch_size", _default_field_candidate("train", "batch_size", "bs")), - ("seed", _default_field_candidate("train", "seed", "seed")), - ("epochs", _default_field_candidate("train", "epochs", "ep")), + ("stage1_generator", _path_candidate("stage1_model.generator", "")), + ("stage2_generator", _path_candidate("stage2_model.generator", "s2-")), + ("stage2_decoder", _path_candidate("stage2_model.decoder", "dec-")), + ( + "stage2_history", + _path_candidate("stage2_model.autoregressive.history", "hist-"), + ), + ( + "particle_type_target", + _path_candidate("stage2_model.particle_type.target", "pt-"), + ), + ("stage1_router", _router_candidate("stage1_model", "s1")), + ("stage2_router", _router_candidate("stage2_model", "s2")), + ( + "stage1_gumbel", + _router_flag_candidate("stage1_model", "s1", "gumbel", {True: "gum"}), + ), + ( + "stage2_gumbel", + _router_flag_candidate("stage2_model", "s2", "gumbel", {True: "gum"}), + ), + ( + "stage1_learn_centers", + _router_flag_candidate("stage1_model", "s1", "learn_centers", {False: "nolc"}), + ), + ( + "stage2_learn_centers", + _router_flag_candidate("stage2_model", "s2", "learn_centers", {False: "nolc"}), + ), + ( + "stage1_learn_width", + _router_flag_candidate("stage1_model", "s1", "learn_width", {True: "lw"}), + ), + ( + "stage2_learn_width", + _router_flag_candidate("stage2_model", "s2", "learn_width", {True: "lw"}), + ), + ( + "stage1_learn_temperature", + _router_flag_candidate("stage1_model", "s1", "learn_temperature", {True: "lt"}), + ), + ( + "stage2_learn_temperature", + _router_flag_candidate("stage2_model", "s2", "learn_temperature", {True: "lt"}), + ), + ("particle_conditioning", _conditioning_candidate("particle", "c")), + ("material_conditioning", _conditioning_candidate("material", "m")), + ("stage1_hidden_dim", _path_candidate("stage1_model.hidden_dim", "h")), + ("stage2_hidden_dim", _path_candidate("stage2_model.hidden_dim", "s2h")), + ("stage1_n_res_blocks", _path_candidate("stage1_model.n_res_blocks", "b")), + ("stage2_n_res_blocks", _path_candidate("stage2_model.n_res_blocks", "s2b")), + ("particle_emb_dim", _path_candidate("conditioning.particle.emb_dim", "e")), + ("lr", _path_candidate("train.lr", "lr")), + ("batch_size", _path_candidate("train.batch_size", "bs")), + ("seed", _path_candidate("train.seed", "seed")), + ("epochs", _path_candidate("train.epochs", "ep")), ] _OUT_DIR_NAME_MAX_FIELDS = 6 @@ -421,14 +1459,13 @@ def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str: Beyond `_OUT_DIR_NAME_MAX_FIELDS` non-default fields, the remainder collapse into a short deterministic hash suffix rather than growing the name unboundedly. This name doubles as the run's W&B id (see - giant.train), which is the reason a timestamp is always included. + giant.training), which is the reason a timestamp is always included. """ now = now or datetime.now() - train, model = cfg["train"], cfg["model"] tokens = [] overflow = [] for label, candidate in _OUT_DIR_NAME_CANDIDATES: - token = candidate(train, model) + token = candidate(cfg) if token is None: continue if len(tokens) < _OUT_DIR_NAME_MAX_FIELDS: @@ -476,30 +1513,31 @@ def _toml_value(v) -> str: return str(v) -def save_config(cfg: dict, out_dir: Path, meta: dict) -> None: - lines = [] - # One-level-nested dict values (e.g. model.router) are rendered as their - # own [section.subsection] table after the parent section, since TOML - # doesn't accept a bare dict as a `key = value` scalar line. - nested_sections: list[tuple[str, dict]] = [] - for section, values in cfg.items(): - lines.append(f"[{section}]") - for k, v in values.items(): - if isinstance(v, dict): - nested_sections.append((f"{section}.{k}", v)) - continue - lines.append(f"{k:<14} = {_toml_value(v)}") - lines.append("") +def _write_section(lines: list[str], path: str, values: dict) -> None: + """Write one TOML table (`[path]`) and recurse depth-first into any + dict-valued keys as `[path.subkey]` — handles the v0.3 schema's 2-3 level + nesting (e.g. stage1_model.router, stage2_model.n_sec) with no depth + limit, unlike the one-level-only writer this replaces.""" + lines.append(f"[{path}]") + nested: list[tuple[str, dict]] = [] + for k, v in values.items(): + if isinstance(v, dict): + nested.append((f"{path}.{k}", v)) + else: + lines.append(f"{k:<18} = {_toml_value(v)}") + lines.append("") + for sub_path, sub_values in nested: + _write_section(lines, sub_path, sub_values) - for name, values in nested_sections: - lines.append(f"[{name}]") - for k, v in values.items(): - lines.append(f"{k:<14} = {_toml_value(v)}") - lines.append("") + +def save_config(cfg: dict, out_dir: Path, meta: dict) -> None: + lines: list[str] = [] + for section, values in cfg.items(): + _write_section(lines, section, values) lines.append("[meta]") for k, v in meta.items(): - lines.append(f"{k:<14} = {_toml_value(v)}") + lines.append(f"{k:<18} = {_toml_value(v)}") (out_dir / "config.toml").write_text("\n".join(lines)) @@ -514,6 +1552,7 @@ def build_run_meta( n_train_steps: int, ) -> dict: return { + "config_version": CONFIG_VERSION, "git_hash": git_hash(), "seed": seed, "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"), diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 3b56a79..6e2e282 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -6,6 +6,7 @@ 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 @@ -39,17 +40,27 @@ class StreamingStepsDataset(IterableDataset): numpy slicing instead of a per-row Python loop in the default collate. Each batch is a tuple: - (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx) where: cond_cont: (B, COND_DIM) float32 - cond_cat: (B, 2) int64 + 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, + 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) + 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 + + `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. """ def __init__( @@ -64,8 +75,13 @@ class StreamingStepsDataset(IterableDataset): shuffle_buffer: int = 65536, shuffle: bool = True, proc_map: dict[str, int] | None = None, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", + material_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)} @@ -79,8 +95,13 @@ class StreamingStepsDataset(IterableDataset): self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle self.proc_map = proc_map - self.conditioning = conditioning + self.particle_conditioning = particle_conditioning + self.material_conditioning = material_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() @@ -98,10 +119,11 @@ 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]): + for chunk in iter_file_chunks(path, offset=self._offsets[path], k_max=self.k_max): mask = sorted_membership(chunk["event_id"], self._events_arr) if not mask.any(): continue @@ -114,6 +136,7 @@ class StreamingStepsDataset(IterableDataset): n_sec, sec_cont, proc_idx, + sec_type_idx, _, _, ) = build_features( @@ -125,7 +148,12 @@ class StreamingStepsDataset(IterableDataset): sec_phys_normalizer=self.sec_phys_normalizer, proc_map=self.proc_map, require_secondaries=True, - conditioning=self.conditioning, + 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, ) buf_cont.append(cond_cont) buf_cat.append(cond_cat) @@ -133,6 +161,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec.append(n_sec) buf_sec.append(sec_cont) buf_proc.append(proc_idx) + buf_type.append(sec_type_idx) buf_n += len(cond_cont) if buf_n >= self.shuffle_buffer: @@ -143,6 +172,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec, buf_sec, buf_proc, + buf_type, buf_n, ) = yield from self._flush( buf_cont, @@ -151,6 +181,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec, buf_sec, buf_proc, + buf_type, final=False, ) @@ -162,6 +193,7 @@ class StreamingStepsDataset(IterableDataset): buf_nsec, buf_sec, buf_proc, + buf_type, final=True, ) @@ -173,6 +205,7 @@ 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) @@ -181,11 +214,12 @@ 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 = nsec[idx], sec[idx], proc[idx] + nsec, sec, proc, styp = nsec[idx], sec[idx], proc[idx], styp[idx] bs = self.batch_size n = len(cont) @@ -199,10 +233,11 @@ class StreamingStepsDataset(IterableDataset): torch.from_numpy(nsec[start:end]).long(), torch.from_numpy(sec[start:end]).float(), torch.from_numpy(proc[start:end]).long(), + torch.from_numpy(styp[start:end]).long(), ) if final: - return [], [], [], [], [], [], 0 + return [], [], [], [], [], [], [], 0 rem = n_full * bs return ( [cont[rem:]], @@ -211,5 +246,6 @@ class StreamingStepsDataset(IterableDataset): [nsec[rem:]], [sec[rem:]], [proc[rem:]], + [styp[rem:]], n - rem, ) diff --git a/giant/data/loader.py b/giant/data/loader.py index e76af49..94acfba 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from pathlib import Path from typing import Iterator @@ -5,6 +6,8 @@ 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 @@ -112,9 +115,7 @@ 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) -> dict[str, np.ndarray]: - from giant.constants import K_MAX - +def _df_to_dict(df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX) -> dict[str, np.ndarray]: has_sec_lists = "sec_E_list" in df.columns d: dict[str, np.ndarray] = { @@ -133,9 +134,7 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: # / 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), @@ -146,17 +145,15 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: } 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) -> dict[str, np.ndarray]: - return _df_to_dict(pd.read_parquet(path), offset=offset) +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_event_ids(path: str | Path, offset: int = 0) -> np.ndarray: @@ -165,13 +162,15 @@ 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 -) -> Iterator[dict[str, np.ndarray]]: - """Yield one parquet row-group at a time so a large file never fully loads.""" +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).""" 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) + yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max) _COND_COLS = [ @@ -205,15 +204,11 @@ 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( @@ -243,6 +238,44 @@ 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. @@ -253,16 +286,66 @@ 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: dict[str, int] = {} + 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 = {} for path in files: - 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 + 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) diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index b1c69d4..b1a613e 100644 --- a/giant/data/setup_cache.py +++ b/giant/data/setup_cache.py @@ -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 event_id_offset, load_event_ids +from giant.data.loader import TopNMap, 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,10 +96,50 @@ def fingerprint_files(files: list[Path]) -> list[list]: return out -def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str: +def normalizer_key( + val_fraction: float, + seed: int, + particle_conditioning: str, + material_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. - return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}" + # 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()}, + ) @dataclass @@ -119,9 +159,7 @@ 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 @@ -143,6 +181,8 @@ 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": @@ -156,6 +196,7 @@ 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 @@ -185,9 +226,8 @@ 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() - } + 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()} return cls( fingerprint=d["fingerprint"], git_hash=d.get("git_hash", "unknown"), @@ -195,6 +235,7 @@ class SetupCache: event_index=event_index, proc_maps=proc_maps, normalizers=normalizers, + topn_maps=topn_maps, ) def merge(self, other: "SetupCache") -> "SetupCache": @@ -209,17 +250,14 @@ 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 @@ -243,9 +281,7 @@ def load( 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: @@ -288,9 +324,7 @@ 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) @@ -298,9 +332,7 @@ 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: @@ -311,16 +343,12 @@ 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 diff --git a/giant/data/transforms.py b/giant/data/transforms.py index e41041e..4dfc672 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -2,6 +2,8 @@ import warnings 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 @@ -82,9 +84,7 @@ 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 @@ -116,9 +116,7 @@ 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 @@ -197,9 +195,7 @@ 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: @@ -339,22 +335,21 @@ 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 -) -> np.ndarray: +def _vectorized_map_lookup(values: np.ndarray, mapping: dict, strict: bool = True, default: int = 0) -> 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 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. + — 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. `build_cond_features` under `conditioning="physical"`, where - `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. + `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. """ keys = np.asarray(list(mapping.keys())) vals = np.asarray(list(mapping.values()), dtype=np.int64) @@ -366,7 +361,7 @@ def _vectorized_map_lookup( found = keys_sorted[pos] == values if not found.all(): if not strict: - out = np.zeros(values.shape, dtype=np.int64) + out = np.full(values.shape, default, dtype=np.int64) out[found] = vals_sorted[pos[found]] return out missing = np.unique(values[~found]) @@ -384,9 +379,7 @@ 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( @@ -405,9 +398,7 @@ 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. @@ -421,9 +412,7 @@ def inv_local_frame_rotation( 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 @@ -488,14 +477,10 @@ 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], @@ -522,9 +507,7 @@ 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 @@ -550,42 +533,69 @@ def encode_secondaries( return sec_cont.astype(np.float32) -def decode_secondaries( +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( 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. +) -> 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_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) + 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. 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). 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. + 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. """ - 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 + N, K = sec_cont.shape[0], sec_cont.shape[1] 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. @@ -625,9 +635,53 @@ def decode_secondaries( 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 @@ -639,52 +693,66 @@ def decode_secondaries( def _physical_cond_columns( - data: dict[str, np.ndarray], conditioning: str + data: dict[str, np.ndarray], + particle_conditioning: str, + material_conditioning: str, ) -> np.ndarray: """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns. - "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. + 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. """ from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM - 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}") + n = len(next(iter(data.values()))) - from giant.materials import material_properties_array - from giant.particles import particle_phys_array + if particle_conditioning == "physical": + 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) + 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) else: - mass, charge = particle_phys_array(data["pdg"]).T + raise ValueError(f"unknown conditioning.particle.type {particle_conditioning!r}") - z_eff, a_eff, density, x0, lambda_int = material_properties_array( - data["material"] - ).T + if material_conditioning == "physical": + from giant.materials import material_properties_array - 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) + 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) def build_cond_features( @@ -692,9 +760,24 @@ def build_cond_features( pdg_map: dict[int, int], mat_map: dict[str, int], cond_normalizer: "Normalizer | None" = None, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", + material_conditioning: str = "embedding", + pdg_topn_map: dict[int, int] | None = None, + mat_topn_map: dict[str, int] | None = None, ) -> tuple[np.ndarray, np.ndarray]: - """Build conditioning arrays only — no target, no post-step variables.""" + """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. + """ cond_cont = np.column_stack( [ data["pre_pos"], @@ -704,49 +787,69 @@ def build_cond_features( ] ).astype(np.float32) cond_cont = np.column_stack( - [cond_cont, _physical_cond_columns(data, conditioning)] + [ + cond_cont, + _physical_cond_columns(data, particle_conditioning, material_conditioning), + ] ).astype(np.float32) - # 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]) + # 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) if cond_normalizer is not None: - cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning) + cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, particle_conditioning, material_conditioning) return cond_cont, cond_cat def _cond_normalizer_transform( - cond_cont: np.ndarray, cond_normalizer: "Normalizer", conditioning: str + cond_cont: np.ndarray, + cond_normalizer: "Normalizer", + particle_conditioning: str, + material_conditioning: str, ) -> np.ndarray: """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed. Checkpoints trained before physical-property conditioning (``COND_DIM`` 8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond normalizer, fit before ``build_cond_features`` grew the extra physical - columns. In "embedding" mode those columns are never read by + columns. When NEITHER axis is "physical" those columns are never read by ``ConditionEncoder`` (``giant/model/network.py``), so padding the missing entries with mean=0/std=1 is a safe no-op that keeps such checkpoints - usable under the current, always-``COND_DIM``-wide contract. In - "physical" mode the physical columns are load-bearing, so a mismatch - there is a real incompatibility, not something to paper over. + 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. """ mean, std = cond_normalizer.mean, cond_normalizer.std assert mean is not None and std is not None, "Normalizer not fitted" width = cond_cont.shape[-1] if mean.shape[-1] < width: - if conditioning != "embedding": + physical_load_bearing = "physical" in ( + particle_conditioning, + material_conditioning, + ) + if physical_load_bearing: raise ValueError( f"cond normalizer has {mean.shape[-1]} columns, expected " - f"{width}, and conditioning={conditioning!r} reads the " + f"{width}, and particle_conditioning={particle_conditioning!r}/" + f"material_conditioning={material_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." @@ -767,8 +870,13 @@ def build_features( fit: bool = False, proc_map: dict[str, int] | None = None, require_secondaries: bool = False, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", + material_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, ) -> tuple[ np.ndarray, np.ndarray, @@ -776,10 +884,12 @@ def build_features( 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. + """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, + sec_type_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) @@ -787,9 +897,17 @@ def build_features( [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"`). 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 @@ -800,17 +918,27 @@ 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( [ @@ -831,30 +959,43 @@ def build_features( ] ).astype(np.float32) # (N, COND_DIM_BASE=8) cond_cont = np.column_stack( - [cond_cont, _physical_cond_columns(data, conditioning)] + [ + cond_cont, + _physical_cond_columns(data, particle_conditioning, material_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) - cond_cat = np.column_stack([pdg_idx, mat_idx]) # (N, 2) + 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) - 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,) + n_sec_raw = data["n_sec"].astype(np.int64) # (N,) unclamped, for the valid-slot mask # 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, @@ -863,7 +1004,12 @@ def build_features( data["pre_dir"], sec_pdg_list=sec_pdg_list, phys_only=sec_phys_only, - ) # (N, K_MAX, 6) + ) # (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) + ) 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 @@ -885,7 +1031,8 @@ 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_cont = np.zeros((N, k_max, 6), dtype=np.float32) + sec_type_idx = np.zeros((N, k_max), dtype=np.int64) if fit: cond_normalizer = Normalizer().fit(cond_cont) @@ -914,6 +1061,7 @@ def build_features( n_sec, sec_cont, proc_idx, + sec_type_idx, cond_normalizer, target_normalizer, ) diff --git a/giant/geometry.py b/giant/geometry.py index ffe27da..3d137e8 100644 --- a/giant/geometry.py +++ b/giant/geometry.py @@ -31,10 +31,7 @@ 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(): @@ -63,9 +60,7 @@ 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) @@ -75,11 +70,7 @@ 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 @@ -300,16 +291,12 @@ 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) @@ -344,12 +331,7 @@ 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] @@ -460,11 +442,7 @@ 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 diff --git a/giant/materials.py b/giant/materials.py index a15c8d2..1aed5de 100644 --- a/giant/materials.py +++ b/giant/materials.py @@ -64,18 +64,10 @@ 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, @@ -83,9 +75,7 @@ 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, @@ -108,20 +98,15 @@ 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 " - f"-- add it (known: {sorted(t)})" + f"material {name!r} is not in giant.materials.MATERIAL_PROPERTIES -- add it (known: {sorted(t)})" ) props = t[name] if any(v is None for v in props): @@ -134,9 +119,7 @@ def get_material_properties( 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)], diff --git a/giant/model/network.py b/giant/model/network.py index 85d145f..e946e20 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -7,28 +7,31 @@ import torch import torch.nn as nn import torch.nn.functional as F +from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig from giant.constants import ( COND_DIM, COND_DIM_BASE, + CONT_SLOT_DIM, EMB_DIM, K_MAX, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM, SEC_DIM, + SEC_SLOT_DIM, X_DIM, ) +# --------------------------------------------------------------------------- +# Building blocks +# --------------------------------------------------------------------------- + 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) - ) + 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: @@ -37,75 +40,157 @@ class SinusoidalEmbedding(nn.Module): return torch.cat([args.sin(), args.cos()], dim=-1) # (B, dim) +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 + + +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 ConditionEncoder(nn.Module): """Fuses continuous conditioning with particle/material identity. - Two mutually exclusive ways to turn (pdg, material) identity into the - two `emb_dim`-wide vectors concatenated with the base continuous - conditioning before the fusion MLP: - - "embedding": a learned `nn.Embedding` lookup table per axis, indexed - by `cond_cat`'s dense training-vocab index. Memorizes the training - menu; the original Phase-2 design. - - "physical": a small MLP per axis, mapping the axis's raw physical + 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) to an `emb_dim`-wide vector — - a drop-in replacement for the embedding lookup, computable for any - PDG code / material name rather than only ones seen in training. - Both modes produce the same `in_dim = COND_DIM_BASE + 2*emb_dim` for the - fusion MLP, so only how the two vectors are produced differs. + 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, - emb_dim: int = 16, out_dim: int = 128, - conditioning: str = "embedding", ) -> None: super().__init__() - if conditioning not in ("embedding", "physical"): - raise ValueError(f"unknown conditioning mode {conditioning!r}") - self.conditioning = conditioning - if conditioning == "embedding": - self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) - self.mat_emb = nn.Embedding(mat_vocab, emb_dim) - else: - self.particle_mlp = nn.Sequential( - nn.Linear(PARTICLE_PHYS_DIM, emb_dim), - nn.SiLU(), - nn.Linear(emb_dim, emb_dim), - ) - self.material_mlp = nn.Sequential( - nn.Linear(MATERIAL_PHYS_DIM, emb_dim), - nn.SiLU(), - nn.Linear(emb_dim, emb_dim), - ) - in_dim = COND_DIM_BASE + 2 * emb_dim + 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 forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - if self.conditioning == "embedding": - pdg_e = self.pdg_emb(cond_cat[:, 0]) - mat_e = self.mat_emb(cond_cat[:, 1]) - else: - particle_phys = cond_cont[ - :, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_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 :] - pdg_e = self.particle_mlp(particle_phys) - mat_e = self.material_mlp(material_phys) + 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) +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.1) -> None: + 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) @@ -123,406 +208,9 @@ class ResBlock(nn.Module): return x + h -class DenoisingMLP(nn.Module): - """Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits. - - The n_sec head runs on the condition encoding only (no diffusion noise), - so it can be called at inference time independently via `predict_n_sec`. - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - hidden_dim: int = 256, - n_blocks: int = 6, - emb_dim: int = 16, - time_dim: int = 64, - cond_out_dim: int = 128, - x_dim: int = X_DIM, - dropout: float = 0.1, - k_max: int = K_MAX, - conditioning: str = "embedding", - ) -> None: - super().__init__() - self.time_emb = SinusoidalEmbedding(time_dim) - self.cond_enc = ConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - out_dim=cond_out_dim, - conditioning=conditioning, - ) - merged_cond_dim = time_dim + cond_out_dim - self.input_proj = nn.Linear(x_dim, hidden_dim) - self.blocks = nn.ModuleList( - [ - ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) - for _ in range(n_blocks) - ] - ) - self.out_proj = nn.Linear(hidden_dim, x_dim) - # Predicts n_sec as classification over {0, 1, ..., k_max}. - # Applied to the condition encoding (not the diffused latent). - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, k_max + 1), - ) - - def forward( - self, - x_t: torch.Tensor, - t: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - ) -> torch.Tensor: - t_emb = self.time_emb(t) # (B, time_dim) - c_emb = self.cond_enc(cond_cont, cond_cat) # (B, cond_out_dim) - cond = torch.cat([t_emb, c_emb], dim=-1) - x = self.input_proj(x_t) - for block in self.blocks: - x = block(x, cond) - return self.out_proj(x) - - 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.""" - c_emb = self.cond_enc(cond_cont, cond_cat) - return self.n_sec_head(c_emb) - - -class SecondaryConditionEncoder(nn.Module): - """Encodes pre-step conditioning + Stage-1 output for the secondary decoder.""" - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - emb_dim: int = 16, - cond_out_dim: int = 128, - stage1_dim: int = X_DIM, - stage1_proj_dim: int = 64, - out_dim: int = 128, - conditioning: str = "embedding", - ) -> None: - super().__init__() - self.base = ConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - out_dim=cond_out_dim, - conditioning=conditioning, - ) - self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim) - fused_dim = cond_out_dim + stage1_proj_dim - self.fuse = nn.Sequential( - nn.Linear(fused_dim, out_dim), - nn.SiLU(), - ) - - def forward( - self, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - ) -> torch.Tensor: - base = self.base(cond_cont, cond_cat) # (B, cond_out_dim) - s1 = self.stage1_proj(stage1_out).tanh() # (B, stage1_proj_dim) - return self.fuse(torch.cat([base, s1], dim=-1)) # (B, out_dim) - - -class SecondaryDecoder(nn.Module): - """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously. - - Each slot encodes (stick_break_logit, local_dir_3D, log_mass, charge) for - one secondary ordered by descending energy — mass/charge are the - secondary's predicted physical identity, regressed directly against real - physics targets (see giant.data.transforms.encode_secondaries), used - as-is with no snapping to a discrete PDG code. Padded slots are masked - from loss. - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - hidden_dim: int = 256, - n_blocks: int = 6, - emb_dim: int = 16, - time_dim: int = 64, - cond_out_dim: int = 128, - stage1_proj_dim: int = 64, - sec_dim: int = SEC_DIM, - dropout: float = 0.1, - conditioning: str = "embedding", - ) -> None: - super().__init__() - self.time_emb = SinusoidalEmbedding(time_dim) - self.cond_enc = SecondaryConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - cond_out_dim=cond_out_dim, - stage1_proj_dim=stage1_proj_dim, - out_dim=cond_out_dim, - conditioning=conditioning, - ) - merged_cond_dim = time_dim + cond_out_dim - self.input_proj = nn.Linear(sec_dim, hidden_dim) - self.blocks = nn.ModuleList( - [ - ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) - for _ in range(n_blocks) - ] - ) - self.out_proj = nn.Linear(hidden_dim, sec_dim) - - def forward( - self, - x_t: torch.Tensor, - t: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - ) -> torch.Tensor: - t_emb = self.time_emb(t) - c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out) - cond = torch.cat([t_emb, c_emb], dim=-1) - x = self.input_proj(x_t) - for block in self.blocks: - x = block(x, cond) - return self.out_proj(x) - - -class WGANGenerator(nn.Module): - """Stage-1 WGAN-GP generator: single forward pass, no diffusion/flow time. - - Same `ConditionEncoder` + `ResBlock` trunk as `DenoisingMLP`, but the - input is a noise vector `z` (not a diffused/interpolated `x_t`) and the - ResBlocks condition on the condition encoding alone (no time embedding to - concatenate) — see `giant/model/wgan.py` for the adversarial losses, and - `giant.sample.sample_wgan` for single-pass sampling. - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - hidden_dim: int = 256, - n_blocks: int = 6, - emb_dim: int = 16, - cond_out_dim: int = 128, - x_dim: int = X_DIM, - noise_dim: int = 64, - dropout: float = 0.1, - k_max: int = K_MAX, - conditioning: str = "embedding", - ) -> None: - super().__init__() - self.noise_dim = noise_dim - self.cond_enc = ConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - out_dim=cond_out_dim, - conditioning=conditioning, - ) - self.input_proj = nn.Linear(noise_dim, hidden_dim) - self.blocks = nn.ModuleList( - [ - ResBlock(hidden_dim, cond_out_dim, dropout=dropout) - for _ in range(n_blocks) - ] - ) - self.out_proj = nn.Linear(hidden_dim, x_dim) - self.n_sec_head = nn.Sequential( - nn.Linear(cond_out_dim, hidden_dim // 2), - nn.SiLU(), - nn.Linear(hidden_dim // 2, k_max + 1), - ) - - def forward( - self, - z: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - ) -> torch.Tensor: - cond = self.cond_enc(cond_cont, cond_cat) - x = self.input_proj(z) - for block in self.blocks: - x = block(x, cond) - return self.out_proj(x) - - 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.""" - c_emb = self.cond_enc(cond_cont, cond_cat) - return self.n_sec_head(c_emb) - - -class Critic(nn.Module): - """Stage-1 WGAN-GP critic: scalar realism score, own `ConditionEncoder`. - - Kept structurally parallel to `WGANGenerator` (own condition encoder — - separate weights from the generator's, standard GAN practice) but has no - n_sec head: n_sec is never adversarial, it stays a plain classifier on - the generator side. - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - hidden_dim: int = 256, - n_blocks: int = 6, - emb_dim: int = 16, - cond_out_dim: int = 128, - x_dim: int = X_DIM, - dropout: float = 0.1, - conditioning: str = "embedding", - ) -> None: - super().__init__() - self.cond_enc = ConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - out_dim=cond_out_dim, - conditioning=conditioning, - ) - self.input_proj = nn.Linear(x_dim, hidden_dim) - self.blocks = nn.ModuleList( - [ - ResBlock(hidden_dim, cond_out_dim, dropout=dropout) - for _ in range(n_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, - ) -> torch.Tensor: - cond = self.cond_enc(cond_cont, cond_cat) - h = self.input_proj(x) - for block in self.blocks: - h = block(h, cond) - return self.out_proj(self.out_norm(h)).squeeze(-1) - - -class WGANSecondaryGenerator(nn.Module): - """Stage-2 WGAN-GP generator: single forward pass over all K_MAX slots. - - Mirrors `SecondaryDecoder` minus the time embedding, the same way - `WGANGenerator` mirrors `DenoisingMLP` — takes noise `z` instead of `x_t`, - conditions on `SecondaryConditionEncoder`'s output alone. - """ - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - hidden_dim: int = 256, - n_blocks: int = 6, - emb_dim: int = 16, - cond_out_dim: int = 128, - stage1_proj_dim: int = 64, - sec_dim: int = SEC_DIM, - noise_dim: int = 64, - dropout: float = 0.1, - conditioning: str = "embedding", - ) -> None: - super().__init__() - self.noise_dim = noise_dim - self.cond_enc = SecondaryConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - cond_out_dim=cond_out_dim, - stage1_proj_dim=stage1_proj_dim, - out_dim=cond_out_dim, - conditioning=conditioning, - ) - self.input_proj = nn.Linear(noise_dim, hidden_dim) - self.blocks = nn.ModuleList( - [ - ResBlock(hidden_dim, cond_out_dim, dropout=dropout) - for _ in range(n_blocks) - ] - ) - self.out_proj = nn.Linear(hidden_dim, sec_dim) - - def forward( - self, - z: torch.Tensor, - cond_cont: torch.Tensor, - cond_cat: torch.Tensor, - stage1_out: torch.Tensor, - ) -> torch.Tensor: - cond = self.cond_enc(cond_cont, cond_cat, stage1_out) - x = self.input_proj(z) - for block in self.blocks: - x = block(x, cond) - return self.out_proj(x) - - -class SecondaryCritic(nn.Module): - """Stage-2 WGAN-GP critic: scalar realism score over the flattened 90D slots.""" - - def __init__( - self, - pdg_vocab: int, - mat_vocab: int, - hidden_dim: int = 256, - n_blocks: int = 6, - emb_dim: int = 16, - cond_out_dim: int = 128, - stage1_proj_dim: int = 64, - sec_dim: int = SEC_DIM, - dropout: float = 0.1, - conditioning: str = "embedding", - ) -> None: - super().__init__() - self.cond_enc = SecondaryConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - cond_out_dim=cond_out_dim, - stage1_proj_dim=stage1_proj_dim, - out_dim=cond_out_dim, - conditioning=conditioning, - ) - self.input_proj = nn.Linear(sec_dim, hidden_dim) - self.blocks = nn.ModuleList( - [ - ResBlock(hidden_dim, cond_out_dim, dropout=dropout) - for _ in range(n_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, - ) -> torch.Tensor: - cond = self.cond_enc(cond_cont, cond_cat, stage1_out) - h = self.input_proj(x) - for block in self.blocks: - h = block(h, cond) - return self.out_proj(self.out_norm(h)).squeeze(-1) +# --------------------------------------------------------------------------- +# Routers — carried over unchanged from v0.2 +# --------------------------------------------------------------------------- class Router(nn.Module): @@ -537,13 +225,6 @@ class Router(nn.Module): def __init__(self, n_experts: int) -> None: super().__init__() self.n_experts = n_experts - # Opt-in straight-through Gumbel-softmax combine weights (see - # combine_weights below) — off by default, set from model.router.gumbel - # by _build_router_from_cfg. gumbel_tau is annealed per training step - # by giant.train (model.router.gumbel_tau_start/_end); neither is an - # nn.Parameter/buffer since neither is learned or needs checkpointing — - # the tau schedule is deterministic in global_step, so it recomputes - # correctly on resume. self.gumbel = False self.gumbel_tau = 1.0 @@ -551,28 +232,13 @@ class Router(nn.Module): """(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: + 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()` — the original dense - soft-mixture combination. Opt-in straight-through Gumbel-softmax - (`gumbel=True`, train mode only): samples a Gumbel-perturbed - categorical draw from the same distribution `gate()` defines - (`log(gate())` is a valid unnormalized-logit input to - `F.gumbel_softmax` since softmax is shift-invariant, so no subclass - needs to expose separate pre-softmax logits), then hardens it to a - one-hot vector on the forward pass while keeping the soft sample's - gradient on the backward pass. This makes the training-time forward - combination match eval-time top-1 dispatch exactly (one expert's - output, unweighted) instead of the smooth blend `gate()` gives — - intended to close the train/eval mismatch identified as a likely - cause of experts overlapping instead of partitioning (see the - router_gating write-up referenced in CLAUDE.md's roadmap). - `gate()` itself is untouched and still backs `balance_loss`/ - `entropy_loss`/`gate_stats`, so those diagnostics keep reading the - smooth distribution rather than a noisy sample. + 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): @@ -584,63 +250,27 @@ class Router(nn.Module): """(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: + 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: + 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), for routers like EnergyRouter that read a - quantity directly off cond_cont/cond_cat and need no label. Routers - gating on an unobservable pre-step quantity (e.g. ProcessRouter, - which predicts the physics process that will end the step) override - this to supervise their internal classifier against the true label. + 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. - - Reuses `gate_stats`'s `norm_entropy` (already in [0, 1], 1.0 = - uniform/collapsed) directly as the loss, so minimizing it pushes - every router's gate toward decisiveness. A generic base-class - default — works for any Router via gate_stats, no per-subclass - override needed. Off by default (see `lambda_entropy` in - giant.train): bounded width/temperature (EnergyRouter's - `learn_width`/`learn_temperature`) is the primary defense against - gate collapse; this is a secondary, use-with-caution lever, since - indiscriminately penalizing entropy can also suppress legitimate - soft ambiguity near a router's own decision boundary. - """ + 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 for catching a router that fails to specialize. - - Returns `(norm_entropy, importance)`: - - `norm_entropy`: scalar, the batch-mean of each row's gate entropy - divided by `log(n_experts)`, in [0, 1] and comparable across - routers with different `n_experts` (1.0 = uniform/collapsed - gating, 0.0 = fully hard routing). - - `importance`: (n_experts,) tensor, `gate(...).sum(dim=0)` — the - *unnormalized* per-expert weight mass for this batch. Callers - wanting a global utilization share across many batches must sum - this across batches first and normalize once at the end; - averaging per-batch shares instead would treat every batch as - equally important regardless of size and understate a - rarely-but-fully-used expert. - """ + 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) @@ -662,15 +292,13 @@ def register_router(name: str): 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 `model.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. + 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)}" - ) + 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} @@ -679,16 +307,13 @@ def build_router(name: str, n_experts: int, **kwargs) -> Router: def _bounded_interp(raw: torch.Tensor, lo: float, hi: float) -> torch.Tensor: """Sigmoid interpolation into `[lo, hi]` — smooth, always-positive-gradient - bound (unlike `clamp`, which zeroes gradient past the boundary) used for - EnergyRouter's `learn_width`/`learn_temperature` modes.""" + 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 `_bounded_interp(raw, lo, hi) == value` — lets `learn_width`/ - `learn_temperature` start out exactly reproducing the fixed-`temperature` - gate before any training moves them.""" + `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)) @@ -697,37 +322,9 @@ def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float: class EnergyRouter(Router): """Soft turn-on gate over normalized pre-step log-energy. - Reads `cond_cont[:, energy_idx]` (ignores cond_cat). Learnable (or - fixed) 1-D centers. By default initialized spread evenly across - [-2, 2] — an assumed-uniform z-normalized energy range that may not - match the true (often skewed) distribution and can leave experts - overlapping instead of partitioning the range; pass `centers_init` to - seed them from data (e.g. energy quantiles) instead. - `gate(e) = softmax_i(-(e - c_i)^2 / tau)`, differentiable in e; as - tau -> 0 this hardens to nearest-center (Voronoi) selection, which is - exactly what `top1` uses at eval. - - `temperature` is normally a single fixed scalar shared by every expert. - Two mutually exclusive optional modes generalize it: - - `learn_width`: each expert gets its own learnable width, so - `gate(e) = softmax_i(-(e - c_i)^2 / width_i)` — experts can learn - independently how much of the energy axis they cover. - - `learn_temperature`: the single shared `temperature` itself becomes - learnable (still one scalar for every expert). - Both parameterize their raw learnable value through a sigmoid bounded - into `[width_min_ratio, width_max_ratio] * temperature` (see - `_bounded_interp`), warm-started so the initial effective width/ - temperature exactly equals `temperature` — enabling either mode is a - no-op at init. The bound is deliberately not raw `softplus`/`exp` - (unbounded above): an unbounded width lets one expert's width run away - to infinity, making its logit `-d2/width -> 0` almost everywhere so it - wins nearly every row regardless of true distance to its center — the - same "experts overlap instead of partitioning" failure this whole - router design is trying to avoid, just via a new mechanism. See - `Router.entropy_loss`/`giant.train`'s `lambda_entropy` for a secondary, - optional guard against all experts' widths co-inflating together - (which bounding caps but doesn't forbid, and which the load-balance - loss alone can't see since usage shares stay even throughout). + 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__( @@ -752,8 +349,7 @@ class EnergyRouter(Router): 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 " - f"({width_max_ratio}) must bracket 1.0" + 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 @@ -766,10 +362,7 @@ class EnergyRouter(Router): centers = torch.linspace(-2.0, 2.0, n_experts) else: if len(centers_init) != n_experts: - raise ValueError( - f"centers_init has {len(centers_init)} values, " - f"expected n_experts={n_experts}" - ) + 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) @@ -777,9 +370,6 @@ class EnergyRouter(Router): self.register_buffer("centers", centers) def effective_width(self) -> torch.Tensor | float: - """Softmax denominator used by `gate()`: a fixed scalar `temperature` - (default), a per-expert `(n_experts,)` bounded width (`learn_width`), - or a single bounded learnable scalar (`learn_temperature`).""" if self.learn_width: return _bounded_interp(self.raw_width, self._width_lo, self._width_hi) if self.learn_temperature: @@ -794,20 +384,9 @@ class EnergyRouter(Router): @register_router("pdg") class PdgRouter(Router): - """Soft turn-on gate over a learned PDG embedding. - - Unlike ProcessRouter's process label, PDG code is already known at - pre-step time (it's a conditioning input, `cond_cat[:, 0]`), so no - supervision is needed — `classify_loss` falls back to the Router base - class's zero-loss default, same as EnergyRouter. Because PDG is - categorical rather than a scalar, this generalizes EnergyRouter's - soft-turn-on-then-Voronoi trick from a 1-D distance to a distance in a - small embedding space: its own embedding table (kept separate from the - trunk's ConditionEncoder, same reasoning as ProcessRouter's own - pdg/mat embeddings) maps each PDG code to a point, and `n_experts` - learnable (or fixed) centers partition that space. - `gate(pdg) = softmax_i(-||emb(pdg) - c_i||^2 / tau)`. - """ + """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, @@ -828,36 +407,18 @@ class PdgRouter(Router): 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) + 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. - - Unlike EnergyRouter (which reads a quantity that's already known at - pre-step time), the process — Compton, photoelectric, brems, ... — is a - *post-step outcome*: it can't be read off cond_cont/cond_cat directly. - Instead this router runs a small classifier over pre-step conditioning - (its own pdg/material embeddings, kept separate from the trunk's - ConditionEncoder) that predicts it, one class per expert slot - (`n_experts` doubles as the number of process classes — see - `build_process_map_from_files`, which caps the process vocabulary to - exactly this many classes, bucketing rare processes into a shared - "other" slot). - - The classifier is supervised by `classify_loss` against the true - `process` label (see `giant/train.py`) — a *training-time* signal only; - `gate`/`top1` never see it, so eval-time dispatch (rollout, predict) - needs no ground truth, same as every other Router. This sidesteps the - gradient/differentiability problem that sank the earlier - process-conditioned-flow proposal (see the archived decision doc): the - hard categorical choice only ever feeds a non-differentiable expert - *dispatch*, never the flow's own conditioning path. - """ + """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, @@ -877,7 +438,6 @@ class ProcessRouter(Router): ) def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - """(B, n_experts) raw process-classifier logits, one class per expert.""" 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) @@ -886,32 +446,13 @@ class ProcessRouter(Router): 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: + 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. - - Wraps N already-built sub-routers, each free to have its own - `n_experts` and hyperparameters (an `EnergyRouter(n_experts=4, ...)` - composed with a `PdgRouter(n_experts=3, ...)` needs no axis to match - the other's expert count). The joint gate is the outer product of the - per-axis softmax gates, flattened to `(B, prod(n_experts_i))` — still a - partition of unity, since each factor is one. Because the axes are - routed independently, the joint argmax factors into the per-axis - argmaxes, so `top1` (inherited from `Router`) costs no more than - routing each axis alone despite the multiplicative expert count; the - same is true of `balance_loss` (inherited, computed on the flattened - joint gate — now one importance term per *joint* expert cell). - - Not registered in `ROUTER_REGISTRY` / buildable via `build_router`, - since those assume one `n_experts` int shared by a single router type; - use `build_composed_router` instead, which resolves a list of per-axis - specs (each independently typed and sized) through `build_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: @@ -926,15 +467,10 @@ class ComposedRouter(Router): 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) + 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: - """Sum of each sub-router's own classify_loss (0 for unsupervised axes).""" + 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) @@ -942,15 +478,8 @@ class ComposedRouter(Router): def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter: - """Build a `ComposedRouter` from a list of per-axis router specs. - - Each spec is a `{"type": ..., "n_experts": ..., ...per-axis kwargs}` - dict resolved through `build_router` exactly like a single-axis router - config, so axes can differ in both expert count and hyperparameters - (e.g. an energy axis's `temperature` vs a pdg axis's `emb_dim`). - `shared_kwargs` (`pdg_vocab`, `mat_vocab`, ...) are merged under each - spec, with the spec's own keys taking precedence. - """ + """Build a `ComposedRouter` from a list of per-axis router specs — see + `_parse_composed_axes`.""" routers = [ build_router( spec["type"], @@ -965,30 +494,108 @@ def build_composed_router(specs: list[dict], **shared_kwargs) -> ComposedRouter: 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 + + +# --------------------------------------------------------------------------- +# Trunks +# --------------------------------------------------------------------------- + + class ExpertTrunk(nn.Module): """One small expert: `input_proj -> ResBlock stack -> out_proj`. - Same shape as the monolithic DenoisingMLP/SecondaryDecoder trunk, but - intended to be narrower/shallower (per-call cost is the whole point). + 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, - merged_cond_dim: int, - dropout: float = 0.1, + 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, merged_cond_dim, dropout=dropout) - for _ in range(n_blocks) - ] - ) - self.out_proj = nn.Linear(hidden_dim, in_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) @@ -1006,26 +613,23 @@ def _route_forward( cond_cat: torch.Tensor, training: bool, ) -> torch.Tensor: - """Shared dispatch for both Routed* trunks. + """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(...)`: the plain soft `gate()` by default, or (see - `Router.combine_weights`) a straight-through Gumbel-softmax one-hot sample - when `router.gumbel` is enabled — either way, no change to the compute - cost of this branch. Eval mode: grouped top-1 dispatch — each row runs - exactly one (small) expert, which is the actual source of the per-call - speedup this architecture is for. + 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_like(x) + 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 = torch.zeros_like(x) + 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(): @@ -1033,349 +637,1117 @@ def _route_forward( return out -class RoutedDenoisingMLP(nn.Module): - """Routed drop-in for `DenoisingMLP`. +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).""" - Shares the time embedding, `ConditionEncoder`, and `n_sec_head` (all - tiny) across experts and routes only the trunk (where the FLOPs are). - Same `forward`/`predict_n_sec` signatures as `DenoisingMLP`, so - sample.py/rollout.py/validate.py need no changes. + 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) + + +# --------------------------------------------------------------------------- +# History encoders — stage-2 autoregressive only +# --------------------------------------------------------------------------- + + +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 + + +# --------------------------------------------------------------------------- +# Stage models +# --------------------------------------------------------------------------- + + +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 + `conditioning.particle.emb_dim` wide).""" + 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 (legacy_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, - router: Router, - expert_hidden_dim: int = 128, - expert_n_blocks: int = 3, - emb_dim: int = EMB_DIM, - time_dim: int = 64, + 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.1, + dropout: float = 0.0, + generator: str = "wgan", + time_dim: int = 64, + noise_dim: int = 64, k_max: int = K_MAX, - conditioning: str = "embedding", + 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.router = router - self.time_emb = SinusoidalEmbedding(time_dim) - self.cond_enc = ConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - out_dim=cond_out_dim, - conditioning=conditioning, + 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, 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) ) - merged_cond_dim = time_dim + cond_out_dim - self.experts = nn.ModuleList( - [ - ExpertTrunk( - x_dim, - expert_hidden_dim, - expert_n_blocks, - merged_cond_dim, - dropout=dropout, - ) - for _ in range(router.n_experts) - ] - ) - self.n_sec_head = nn.Sequential( - nn.Linear(cond_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(), - nn.Linear(cond_out_dim, k_max + 1), ) + 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 = 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, - t: torch.Tensor, cond_cont: torch.Tensor, cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + t: torch.Tensor | None = None, ) -> torch.Tensor: - t_emb = self.time_emb(t) - c_emb = self.cond_enc(cond_cont, cond_cat) - cond = torch.cat([t_emb, c_emb], dim=-1) - return _route_forward( - self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training - ) + 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: - """Return n_sec logits (B, K_MAX+1) from conditioning alone.""" - c_emb = self.cond_enc(cond_cont, cond_cat) + if self.n_sec_head is None: + raise RuntimeError( + "this Stage2OneShot has no n_sec_head — it belongs to a " + "migrated v0.2 checkpoint (legacy_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 RoutedSecondaryDecoder(nn.Module): - """Routed drop-in for `SecondaryDecoder`. - Shares the time embedding and `SecondaryConditionEncoder` across - experts and routes only the trunk. Same `forward` signature as - `SecondaryDecoder`. +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, - router: Router, - expert_hidden_dim: int = 128, - expert_n_blocks: int = 3, - emb_dim: int = EMB_DIM, - time_dim: int = 64, + particle_cfg: dict, + material_cfg: dict, + hidden_dim: int = 256, + n_res_blocks: int = 6, cond_out_dim: int = 128, - stage1_proj_dim: int = 64, - sec_dim: int = SEC_DIM, - dropout: float = 0.1, - conditioning: str = "embedding", + 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__() - self.router = router - self.time_emb = SinusoidalEmbedding(time_dim) - self.cond_enc = SecondaryConditionEncoder( - pdg_vocab=pdg_vocab, - mat_vocab=mat_vocab, - emb_dim=emb_dim, - cond_out_dim=cond_out_dim, - stage1_proj_dim=stage1_proj_dim, - out_dim=cond_out_dim, - conditioning=conditioning, + 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 = 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) ) - merged_cond_dim = time_dim + cond_out_dim - self.experts = nn.ModuleList( - [ - ExpertTrunk( - sec_dim, - expert_hidden_dim, - expert_n_blocks, - merged_cond_dim, - dropout=dropout, - ) - for _ in range(router.n_experts) - ] + 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, - 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: - t_emb = self.time_emb(t) - c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out) - cond = torch.cat([t_emb, c_emb], dim=-1) - return _route_forward( - self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training + 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 (legacy_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) -_STAGE1_MODEL_KEYS = { - "pdg_vocab", - "mat_vocab", - "hidden_dim", - "n_blocks", - "emb_dim", - "dropout", - "k_max", - "conditioning", -} -_SEC_DECODER_MODEL_KEYS = { - "pdg_vocab", - "mat_vocab", - "hidden_dim", - "n_blocks", - "emb_dim", - "dropout", - "conditioning", -} -_WGAN_GENERATOR_MODEL_KEYS = _STAGE1_MODEL_KEYS | {"noise_dim"} -_WGAN_SEC_GENERATOR_MODEL_KEYS = _SEC_DECODER_MODEL_KEYS | {"noise_dim"} -# Critic has no n_sec head (n_sec is never adversarial), so it doesn't accept -# k_max the way DenoisingMLP/WGANGenerator do. -_CRITIC_MODEL_KEYS = _STAGE1_MODEL_KEYS - {"k_max"} +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) -_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$") +# --------------------------------------------------------------------------- +# v0.2 -> v0.3 checkpoint migration +# --------------------------------------------------------------------------- -def _parse_composed_axes(router_cfg: dict) -> list[dict]: - """Regroup `axis{i}_{field}` flat keys into a list of per-axis spec dicts. +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. - Flat keys (rather than a nested list-of-dicts) keep composed-router - config expressible in the same one-level-of-nesting TOML/CLI shape as - every other router option (`model.router` stays a flat table of - scalars) — 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; order follows the index, not - dict insertion order (TOML/CLI merging doesn't preserve it reliably). + Sets `stage2_model.n_sec.legacy_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. """ - 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 `conditioning` -# mode — see _check_router_conditioning_compat. -_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process") - - -def _check_router_conditioning_compat( - router_types: list[str], conditioning: str -) -> None: - """Reject a router axis that reintroduces a training-vocab PDG lookup - under `conditioning="physical"`. - - `PdgRouter`/`ProcessRouter` always build their own dataset-scoped - `nn.Embedding(pdg_vocab, ...)` (network.py's PdgRouter/ProcessRouter), - independent of `ConditionEncoder`'s `conditioning` mode. Pairing either - with `conditioning="physical"` would silently reintroduce a - training-menu-scoped lookup at the routing layer — defeating the entire - point of physical-property conditioning, which is to generalize to a - species/material outside that menu (see giant/rollout.py's - `build_cond_features(strict=...)` gate for the same concern on the - trunk side). Raised loudly at model-build time rather than left to - surface as a confusing rollout/generalization-benchmark result. - """ - bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES)) - if bad and conditioning == "physical": + 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 {}) + expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0) + expert_n_blocks = router_cfg.pop("expert_n_blocks", 0) + if expert_hidden_dim or expert_n_blocks: raise ValueError( - f"router type(s) {bad} always use a training-vocab PDG embedding, " - "which is incompatible with conditioning='physical' (whose whole " - "point is generalizing beyond that vocab) — pick a different " - "router type (e.g. 'energy') or use conditioning='embedding'." + "this checkpoint's model_config.router sets expert_hidden_dim/" + f"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 checkpoint's routed experts " + "have a different width/depth than the monolith — silently " + "dropping these keys would resize the experts instead of " + "refusing. This checkpoint can " + "only be loaded by v0.2 code." ) + router_cfg.setdefault("enabled", False) - -def _build_router_from_cfg( - router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding" -) -> Router: - """Resolve one `model.router` config into a `Router`, single-axis or composed. - - `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys - (see `_parse_composed_axes`) instead of a single `type`/`n_experts` pair. - - `gumbel` is set as a post-construction attribute here rather than a - per-subclass constructor kwarg, same reasoning as `lambda_balance`/ - `lambda_proc`/`lambda_entropy` living in `router_cfg` without being a - `Router` subclass constructor param: it's a training-time toggle shared by - every router type, not a per-type hyperparameter (`build_router`'s - kwarg-filtering would otherwise just silently drop it). - """ - 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], 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"]], conditioning) - router_kwargs = { - k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts") + return { + "pdg_vocab": m["pdg_vocab"], + "mat_vocab": m["mat_vocab"], + "conditioning": { + "out_dim": 128, + "share_stages": False, + "particle": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2}, + "material": {"type": conditioning_mode, "emb_dim": emb_dim, "n_layers": 2}, + }, + "stage1_model": { + "active": True, + "generator": generator, + "hidden_dim": hidden_dim, + "n_res_blocks": n_blocks, + "dropout": dropout, + "flow": {"time_dim": 64}, + "ddpm": {"time_dim": 64}, + "wgan": {"noise_dim": noise_dim}, + "router": dict(router_cfg), + }, + "stage2_model": { + "active": True, + "decoder": "one_shot", + "generator": generator, + "hidden_dim": hidden_dim, + "n_res_blocks": n_blocks, + "dropout": dropout, + "k_max": k_max, + "context_dim": 64, + "n_sec": {"mode": "head", "legacy_owner": "stage1"}, + "particle_type": {"target": "physical"}, + "flow": {"time_dim": 64}, + "ddpm": {"time_dim": 64}, + "wgan": {"noise_dim": noise_dim}, + "router": {**router_cfg, "tie_to_stage1": False}, + }, } - # Not every router needs these (EnergyRouter doesn't declare them, so - # build_router's kwarg filtering drops them silently) but ProcessRouter - # needs its own pdg/material embeddings sized to match the checkpoint's - # vocab, same as the trunk's ConditionEncoder. - 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 -def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: - """Construct (stage1, sec_decoder) from a persisted/CLI model_config dict. +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))`. - Dispatches to the routed pair when `model_config["router"]["enabled"]` - is truthy; a missing/absent "router" key (pre-routing checkpoints) - falls back to the monolithic pair unchanged, so this is a drop-in - replacement for the ad-hoc constructions it replaces. - - `model_config.get("conditioning", "embedding")` — old checkpoints have no - "conditioning" key and must keep loading with their original embedding - tables, so the default here is "embedding", not the training-time - default (which is "physical" — see giant.config.DEFAULT_CONFIG). Read - once and passed to both stage1/sec_decoder, so they structurally always - share one mode. + Only the monolithic (non-routed) trunk shape is handled. """ - if model_config.get("mode") == "wgan": - stage1 = WGANGenerator( - **{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS} - ) - sec_decoder = WGANSecondaryGenerator( - **{ - k: v - for k, v in model_config.items() - if k in _WGAN_SEC_GENERATOR_MODEL_KEYS - } - ) - return stage1, sec_decoder - router_cfg = model_config.get("router") - if router_cfg and router_cfg.get("enabled"): - pdg_vocab = model_config["pdg_vocab"] - mat_vocab = model_config["mat_vocab"] - shared = dict( + 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 (legacy_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 + + +# --------------------------------------------------------------------------- +# 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 + legacy_owner = s2_spec.n_sec.legacy_owner + n_sec_head_k_max = s2_spec.k_max if legacy_owner == "stage1" else None + result["stage1"] = Stage1Model( pdg_vocab=pdg_vocab, mat_vocab=mat_vocab, - expert_hidden_dim=model_config.get("expert_hidden_dim") - or model_config.get("hidden_dim", 128), - expert_n_blocks=model_config.get("expert_n_blocks") - or model_config.get("n_blocks", 3), - emb_dim=model_config.get("emb_dim", EMB_DIM), - dropout=model_config.get("dropout", 0.1), - conditioning=model_config.get("conditioning", "embedding"), + 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, ) - conditioning = shared["conditioning"] - stage1 = RoutedDenoisingMLP( - router=_build_router_from_cfg( - router_cfg, pdg_vocab, mat_vocab, conditioning - ), - k_max=model_config.get("k_max", K_MAX), - **shared, + + 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 + legacy_owner = s2_spec.n_sec.legacy_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=legacy_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, 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=legacy_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.hidden_dim, + n_res_blocks=s1_spec.n_res_blocks, + cond_out_dim=cond_out_dim, + dropout=s1_spec.dropout, + stage="stage1", ) - sec_decoder = RoutedSecondaryDecoder( - router=_build_router_from_cfg( - router_cfg, pdg_vocab, mat_vocab, conditioning - ), - **shared, + + 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, 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.hidden_dim, + n_res_blocks=s2_spec.n_res_blocks, + cond_out_dim=cond_out_dim, + dropout=s2_spec.dropout, + stage="stage2", + context_dim=s2_spec.context_dim, ) - return stage1, sec_decoder - stage1 = DenoisingMLP( - **{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS} - ) - sec_decoder = SecondaryDecoder( - **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS} - ) - return stage1, sec_decoder - - -def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]: - """Construct (critic, sec_critic) for `--mode wgan` training. - - Training-only — never persisted for inference the way `build_models`'s - pair is, since `predict`/`rollout` only ever run the generators. - """ - critic = Critic( - **{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS} - ) - sec_critic = SecondaryCritic( - **{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS} - ) - return critic, sec_critic + return result diff --git a/giant/model/schedule.py b/giant/model/schedule.py index 05d1e18..c63e018 100644 --- a/giant/model/schedule.py +++ b/giant/model/schedule.py @@ -11,9 +11,7 @@ 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)) @@ -48,7 +46,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, t_norm, cond_cont, cond_cat) + pred = model(x_t, cond_cont, cond_cat, t=t_norm) return F.mse_loss(pred, noise) @@ -67,7 +65,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, t, cond_cont, cond_cat) + v_t = model(x_t, cond_cont, cond_cat, t=t) return F.mse_loss(v_t, u_t) @@ -78,39 +76,125 @@ 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, SEC_DIM) — flattened secondary target (stick_logit, dir, log_mass, charge) + x1: (B, K_MAX * (CONT_SLOT_DIM + type_dim)) — flattened secondary target + (stick_logit, dir, then a `type_dim`-wide type slice) 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 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. + 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. """ - from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_SLOT_DIM + from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM + + if type_dim is None: + type_dim = PARTICLE_PHYS_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, t, cond_cont, cond_cat, stage1_out) + v_t = model(x_t, cond_cont, cond_cat, stage1_out, t=t) - 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) + 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 + 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 diff --git a/giant/particles.py b/giant/particles.py index 4e5ea17..d8fcdc1 100644 --- a/giant/particles.py +++ b/giant/particles.py @@ -7,16 +7,28 @@ 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 @@ -46,9 +58,7 @@ 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 " @@ -97,9 +107,7 @@ 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) @@ -111,3 +119,110 @@ 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: + """`conditioning.particle.type` / `stage2_model.particle_type.target = + "onehot"` inference decode: per-row top-N class index -> concrete 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: `conditioning.particle.emb_dim` — the class count; 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) diff --git a/giant/pipeline.py b/giant/pipeline.py index 21d6f6d..683053a 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -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, @@ -32,7 +32,7 @@ from giant.data.transforms import ( ) from giant.data.dataset import make_event_split, StreamingStepsDataset from giant.model.network import build_models, build_critics -from giant.train import train as run_training +from giant.training import train as run_training @dataclass @@ -48,6 +48,8 @@ class SetupStageResult: pdg_map: dict[int, int] mat_map: dict[str, int] proc_map: dict[str, int] | None + pdg_topn_map: TopNMap | None + mat_topn_map: TopNMap | None cond_norm: Normalizer tgt_norm: Normalizer sec_phys_norm: Normalizer @@ -56,26 +58,62 @@ 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, - conditioning: str, - router_cfg: dict, + 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 - (`router_cfg.type == "process"`), and the Stage-1/Stage-2 normalizers. + (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. 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). `router_cfg` may be mutated - in place: an active `EnergyRouter` (`router_cfg["type"] == "energy"`) + refreshes — any existing sidecar content). Each stage's `router` config + is mutated in place: an active `EnergyRouter` (`router.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)") @@ -97,23 +135,14 @@ 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 | " - f"{len(train_events)} train events | " - f"{len(val_events)} val events" - ) + echo(f" {int(counts.sum()):,} steps | {len(train_events)} train events | {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, " - f"{len(mat_map)} materials)" - ) + echo(f"vocabulary maps: cache hit ({len(pdg_map)} PDG codes, {len(mat_map)} materials)") else: echo("building vocabulary maps …") pdg_map, mat_map = build_index_maps_from_files(files) @@ -121,16 +150,23 @@ 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 - if router_cfg.get("enabled") and router_cfg.get("type") == "process": - n_experts = router_cfg["n_experts"] + 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"] 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, " - f"{n_experts} experts)" - ) + echo(f"process vocabulary: cache hit ({len(proc_map)} labels, {n_experts} experts)") else: echo("building process vocabulary …") proc_map = build_process_map_from_files(files, n_experts=n_experts) @@ -138,11 +174,49 @@ def run_setup_stage( if cache is not None: cache.proc_maps[n_experts] = proc_map - 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) + # Top-N-plus-other maps for onehot conditioning/type axes. + # The PDG axis is shared by + # conditioning.particle.type="onehot" and + # stage2_model.particle_type.target="onehot" (both key off + # conditioning.particle.emb_dim), so at most one PDG scan is needed even + # if both consumers are active. The material axis is independent. + particle_cfg = cfg["conditioning"]["particle"] + material_cfg = cfg["conditioning"]["material"] + particle_type_target = config.ParticleTypeConfig.from_dict(cfg["stage2_model"].get("particle_type")).target + + pdg_topn_map: TopNMap | None = None + if particle_cfg["type"] == "onehot" or particle_type_target == "onehot": + n_classes = particle_cfg["emb_dim"] + 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: + pdg_topn_map = cached + echo(f"pdg top-N map: cache hit ({len(pdg_topn_map.class_map)} codes, {n_classes} classes)") + else: + echo("building pdg top-N map …") + pdg_topn_map = build_pdg_topn_map_from_files(files, n_classes=n_classes) + echo(f" {len(pdg_topn_map.class_map)} pdg codes mapped to {n_classes} classes") + if cache is not None: + cache.topn_maps[cache_key] = pdg_topn_map + + 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) entry = cache.normalizers.get(norm_key) if cache is not None else None if entry is not None: @@ -163,33 +237,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 - # --router-type energy 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 an energy + # router 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)): + for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max): mask = sorted_membership(chunk["event_id"], events_arr) if not mask.any(): continue chunk_tr = {k: v[mask] for k, v in chunk.items()} - cond_cont, _, target_s1, n_sec, sec_cont, _proc, _, _ = 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, - conditioning=conditioning, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, sec_phys_only=True, + k_max=k_max, ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) if energy_sampler is not None: energy_sampler.update(cond_cont[:, energy_idx]) - sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None] + sec_valid = np.arange(sec_cont.shape[1])[None, :] < n_sec[:, None] sec_phys = sec_cont[:, :, 4:6][sec_valid] if len(sec_phys) > 0: sec_phys_acc.update(sec_phys) @@ -206,22 +280,8 @@ def run_setup_stage( cond_norm, tgt_norm, sec_phys_norm, n_train_steps, energy_quantiles ) - 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" - ) + _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 cache is not None: setup_cache.save(data, files, cache, echo=echo) @@ -231,6 +291,8 @@ def run_setup_stage( pdg_map=pdg_map, mat_map=mat_map, proc_map=proc_map, + pdg_topn_map=pdg_topn_map, + mat_topn_map=mat_topn_map, cond_norm=cond_norm, tgt_norm=tgt_norm, sec_phys_norm=sec_phys_norm, @@ -252,7 +314,7 @@ def run_train_job( rebuild_setup_cache: bool = False, echo=print, ) -> None: - t, m = cfg["train"], cfg["model"] + t = cfg["train"] config.seed_everything(t["seed"]) out_dir = Path(out_dir) @@ -271,20 +333,15 @@ def run_train_job( "section)" ) - 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"] + config.validate_config(cfg) + particle_conditioning = cfg["conditioning"]["particle"]["type"] + material_conditioning = cfg["conditioning"]["material"]["type"] + k_max = cfg["stage2_model"]["k_max"] setup = run_setup_stage( data, val_fraction=t["val_fraction"], seed=t["seed"], - conditioning=conditioning, - router_cfg=router_cfg, + cfg=cfg, cache_setup=cache_setup, rebuild_setup_cache=rebuild_setup_cache, echo=echo, @@ -302,6 +359,35 @@ 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.pdg_topn_map is not None + sec_type_class_map = setup.pdg_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") @@ -316,8 +402,13 @@ def run_train_job( shuffle_buffer=shuffle_buffer, shuffle=True, proc_map=proc_map, - conditioning=conditioning, + particle_conditioning=particle_conditioning, + material_conditioning=material_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, @@ -329,8 +420,13 @@ def run_train_job( batch_size=t["batch_size"], shuffle=False, proc_map=proc_map, - conditioning=conditioning, + particle_conditioning=particle_conditioning, + material_conditioning=material_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" @@ -347,60 +443,19 @@ 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), - "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), + "conditioning": cfg["conditioning"], + "stage1_model": cfg["stage1_model"], + "stage2_model": cfg["stage2_model"], } - 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" - ) + 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") out_dir.mkdir(parents=True, exist_ok=True) meta = config.build_run_meta( @@ -415,25 +470,13 @@ def run_train_job( config.save_config(cfg, out_dir, meta) run_training( - stage1_model=stage1_model, - sec_decoder=sec_decoder, + cfg=cfg, + models=models, + critics=critics, 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(), @@ -442,18 +485,12 @@ 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, + 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, - 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), + use_wandb=t.get("wandb", True), wandb_project=t.get("wandb_project", "giant"), wandb_run_name=t.get("wandb_run_name", ""), wandb_log_every=t.get("wandb_log_every", 50), diff --git a/giant/rollout.py b/giant/rollout.py index 8270a5e..087473d 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -17,7 +17,7 @@ treated as detector leakage and not deposited. from __future__ import annotations from collections import Counter -from typing import Callable, TypedDict +from typing import TYPE_CHECKING, Callable, TypedDict import numpy as np import torch @@ -33,18 +33,156 @@ 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 nearest_known_pdg, particle_phys_array -from giant.sample import ( - sample_flow, - sample_secondaries, - sample_wgan, - sample_secondaries_wgan, +from giant.particles import ( + decode_embedding_nearest, + decode_topn_class, + invert_dense_map, + nearest_known_pdg, + particle_phys_array, ) +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], + pdg_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 pdg_topn_map is None: + raise RuntimeError( + "particle_type.target='onehot' rollout needs pdg_topn_map " + "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']" + ) + class_idx = sec_type_np.argmax(axis=-1) + sec_pdg = decode_topn_class( + class_idx, + pdg_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 = [ @@ -153,13 +291,9 @@ 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() @@ -167,10 +301,7 @@ 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 != ""] @@ -187,8 +318,7 @@ 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(): @@ -205,7 +335,7 @@ def make_seed_frontier( pre_pos: np.ndarray, pre_E: np.ndarray, pre_dir: np.ndarray, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", ) -> tuple[dict[str, np.ndarray], dict[int, int]]: """Build the initial frontier from primary entry states. @@ -225,17 +355,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 conditioning == "physical": + if particle_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" mode never reads mass/charge (see + # "embedding"/"onehot" never read mass/charge (see # _physical_cond_columns), so resolving them here would only risk - # crashing an embedding-mode rollout on a PDG code giant.particles - # can't resolve, for a value that's never used. + # crashing a 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) @@ -310,8 +440,15 @@ def rollout( max_tracks_per_event: int | None = None, escape_threshold: float | None = None, on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None, - conditioning: str = "embedding", - mode: str = "flow", + particle_conditioning: str = "embedding", + material_conditioning: str = "embedding", + pdg_topn_map: "TopNMap | None" = None, + mat_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, ) -> dict[str, np.ndarray] | RolloutSummary: """Run showers to completion. @@ -324,12 +461,43 @@ 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` serve two independent purposes that happen + to share `pdg_topn_map` (one PDG map, not two): they're required + whenever `particle_conditioning`/`material_conditioning` is `"onehot"` + (feeds `build_cond_features`'s extra `cond_cat` top-N columns), and + `pdg_topn_map`/`other_policy` are additionally read under + `stage2_model.particle_type.target = "onehot"` (secondary-species + decode). `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 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"], @@ -337,7 +505,7 @@ def rollout( seeds["pre_pos"], seeds["pre_E"], seeds["pre_dir"], - conditioning=conditioning, + particle_conditioning=particle_conditioning, ) rec = _Recorder(sink=on_chunk) @@ -364,8 +532,15 @@ def rollout( steps, device, max_tracks_per_event, - conditioning, - mode, + particle_conditioning, + material_conditioning, + pdg_topn_map, + mat_topn_map, + other_policy, + rng, + stage1_ddpm_steps, + stage2_ddpm_steps, + l1_dist_collector, ) ) frontier = _concat_frontiers(next_parts) @@ -395,8 +570,15 @@ def _step_chunk( steps, device, max_tracks_per_event, - conditioning, - mode="flow", + particle_conditioning, + material_conditioning, + pdg_topn_map, + mat_topn_map, + other_policy, + rng, + stage1_ddpm_steps, + stage2_ddpm_steps, + l1_dist_collector, ) -> dict[str, np.ndarray]: """Advance one chunk of tracks by a single step; return the next frontier.""" n = len(tr["event_id"]) @@ -407,7 +589,7 @@ def _step_chunk( tr["_material"] = material tr["_layer_id"] = layer_id - if conditioning == "physical": + if particle_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 @@ -421,33 +603,19 @@ 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 @@ -475,60 +643,60 @@ def _step_chunk( "charge": tr["charge"], } cond_cont, cond_cat = build_cond_features( - cond_dict, pdg_map, mat_map, cond_norm, conditioning=conditioning + 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, ) cc = torch.from_numpy(cond_cont).float().to(device) ck = torch.from_numpy(cond_cat).long().to(device) - 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) + stage1_norm, n_sec_pred_stage1 = sample_stage1(stage1_model, cc, ck, steps, stage1_ddpm_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: 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, + # 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, n_sec_np, e_sec, tr["pre_dir"], - sec_phys_normalizer=sec_phys_norm, + sec_phys_norm, + pdg_map, + pdg_topn_map, + other_policy, + rng, ) - sec_pdg_code = nearest_known_pdg( - sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys() - ).reshape(sec_mass.shape) + 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) edep = edep.astype(np.float64) post_E = post_E.astype(np.float64) diff --git a/giant/sample.py b/giant/sample.py index 271897b..d4486b2 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -1,26 +1,22 @@ import torch +import torch.nn.functional as F -from giant.constants import K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM +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 -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 +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) @torch.no_grad() @@ -29,12 +25,14 @@ def sample_flow( cond_cont: torch.Tensor, cond_cat: torch.Tensor, steps: int = 10, -) -> tuple[torch.Tensor, torch.Tensor]: +) -> tuple[torch.Tensor, torch.Tensor | None]: """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 + 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`). """ model.eval() B = cond_cont.size(0) @@ -43,11 +41,138 @@ def sample_flow( dt = 1.0 / steps for i in range(steps): t = torch.full((B,), i * dt, device=device) - v = model(x, t, cond_cont, cond_cat) + v = model(x, cond_cont, cond_cat, t=t) x = x + v * dt - 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 + 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 @torch.no_grad() @@ -59,76 +184,25 @@ def sample_secondaries( n_sec_pred: torch.Tensor, steps: int = 10, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Euler integration of the Stage-2 vector field; return raw slot outputs. + """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. 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, SEC_DIM, device=device) + x = torch.randn(B, flat_width, device=device) dt = 1.0 / steps for i in range(steps): t = torch.full((B,), i * dt, device=device) - v = sec_decoder(x, t, cond_cont, cond_cat, stage1_out) + v = sec_decoder(x, cond_cont, cond_cat, stage1_out, t=t) x = x + v * dt - 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 + return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred) @torch.no_grad() @@ -139,41 +213,228 @@ def sample_secondaries_wgan( stage1_out: torch.Tensor, n_sec_pred: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: - """Single-pass Stage-2 WGAN generator sample; see `sample_secondaries`'s - docstring for the returned (sec_cont, sec_phys, sec_valid) shapes.""" + """Single-pass `Stage2OneShot` WGAN generator sample; see + `_decode_stage2_flat`'s docstring for the returned (sec_cont, sec_type, + sec_valid) shapes/meaning.""" 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 _slots_from_flat(x, n_sec_pred) + return _decode_stage2_flat(sec_decoder, x, cond_cont, cond_cat, stage1_out, n_sec_pred) @torch.no_grad() -def sample_ddim( - model: torch.nn.Module, +def sample_secondaries_ar( + sec_decoder: torch.nn.Module, cond_cont: torch.Tensor, cond_cat: torch.Tensor, - schedule, - steps: int = 50, -) -> tuple[torch.Tensor, torch.Tensor]: - """DDIM deterministic sampling (Song et al. 2020). Returns (sample, n_sec_pred).""" - model.eval() + 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() 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, 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())] + 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, + ) 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 - 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 + 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) diff --git a/giant/train.py b/giant/train.py deleted file mode 100644 index a6324bc..0000000 --- a/giant/train.py +++ /dev/null @@ -1,1155 +0,0 @@ -import copy -import csv -import math -import os -import signal -import time -from pathlib import Path -from types import FrameType -from typing import Callable - -import numpy as np -import torch -import torch.nn.functional as F -import torch.optim as optim -from torch.utils.data import DataLoader -from tqdm import tqdm - -from giant.constants import K_MAX, SEC_SLOT_DIM -from giant.model.schedule import ( - CosineSchedule, - flow_matching_loss, - flow_matching_loss_secondary, -) -from giant.model.wgan import gradient_penalty, generator_loss -from giant.validate import validate_marginals - -_METRICS_FIELDS = [ - "epoch", - "train_loss", - "train_loss_s1", - "train_loss_nsec", - "train_loss_s2", - "train_loss_balance", - "train_loss_proc", - "train_loss_entropy", - "train_nsec_acc", - "d_loss", - "g_loss", - "wasserstein_estimate", - "gp_loss", - "val_loss", - "val_loss_s1", - "val_loss_nsec", - "val_loss_s2", - "val_loss_balance", - "val_loss_proc", - "val_loss_entropy", - "val_nsec_acc", - "val_marginal_kl", - "router_s1_entropy", - "router_s1_util_min", - "router_s1_util_max", - "router_s1_util_std", - "router_s2_entropy", - "router_s2_util_min", - "router_s2_util_max", - "router_s2_util_std", - "lr", - "critic_lr", - "grad_norm", - "grad_norm_d", - "grad_norm_g", - "gpu_mem_mb", - "samples_per_sec", - "is_best", - "epoch_time_s", -] - -_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)" - ) - - -@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 _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 _wandb_run_config( - *, - mode: str, - epochs: int, - lr: float, - warmup_epochs: int, - weight_decay: float, - ema_decay: float, - lambda_nsec: float, - lambda_s2: float, - lambda_balance: float, - lambda_proc: float, - lambda_entropy: float, - gumbel_tau_start: float, - gumbel_tau_end: float, - n_critic: int, - gp_weight: float, - model_config: dict | None, - stage1_params: int, - sec_decoder_params: int, - critic_params: int, - sec_critic_params: int, - total_params: int, -) -> dict: - """Build the dict logged as a wandb run's `config`. - - Router-only knobs (`lambda_balance`/`lambda_proc`/`lambda_entropy`/ - `gumbel_tau_start`/`gumbel_tau_end`) and WGAN-only knobs (`n_critic`/ - `gp_weight`) are omitted unless actually active, so a run's wandb config - doesn't imply hyperparameters from an inactive code path (a disabled - router's fine-tuning knobs, or GAN critic settings for a flow/DDPM run). - The full `model_config` (including its `router` sub-dict, whatever the - router type/state) is always included, so no information is lost — this - only trims the flattened top-level convenience duplicates. - """ - router_enabled = bool((model_config or {}).get("router", {}).get("enabled", False)) - cfg = { - "mode": mode, - "epochs": epochs, - "lr": lr, - "warmup_epochs": warmup_epochs, - "weight_decay": weight_decay, - "ema_decay": ema_decay, - "lambda_nsec": lambda_nsec, - "lambda_s2": lambda_s2, - "model": model_config or {}, - "stage1_params": stage1_params, - "sec_decoder_params": sec_decoder_params, - "critic_params": critic_params, - "sec_critic_params": sec_critic_params, - "total_params": total_params, - } - if router_enabled: - cfg.update( - { - "lambda_balance": lambda_balance, - "lambda_proc": lambda_proc, - "lambda_entropy": lambda_entropy, - "gumbel_tau_start": gumbel_tau_start, - "gumbel_tau_end": gumbel_tau_end, - } - ) - if mode == "wgan": - cfg.update({"n_critic": n_critic, "gp_weight": gp_weight}) - return cfg - - -def _compute_losses( - stage1_model: torch.nn.Module, - sec_decoder: torch.nn.Module, - batch: tuple, - mode: str, - ddpm_schedule, - device: torch.device, - lambda_nsec: float, - lambda_s2: float, - lambda_balance: float = 0.0, - lambda_proc: float = 0.0, - lambda_entropy: float = 0.0, -) -> tuple[ - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, - torch.Tensor, -]: - """Compute (total_loss, L_s1, L_nsec, L_s2, L_balance, L_proc, L_entropy, nsec_acc) for one batch.""" - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx = batch - cond_cont = cond_cont.to(device) - cond_cat = cond_cat.to(device) - x1_s1 = x1_s1.to(device) - n_sec = n_sec.to(device) - sec_cont = sec_cont.to(device) - proc_idx = proc_idx.to(device) - - # Stage-1 flow loss - if mode == "flow": - l_s1 = flow_matching_loss(stage1_model, x1_s1, cond_cont, cond_cat) - else: - assert ddpm_schedule is not None - l_s1 = ddpm_schedule.loss(stage1_model, x1_s1, cond_cont, cond_cat) - - # n_sec classification loss - n_sec_logits = stage1_model.predict_n_sec(cond_cont, cond_cat) - l_nsec = F.cross_entropy(n_sec_logits, n_sec) - nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean() - - # Stage-2 secondary flow loss - # Use a noiseless Stage-1 target as context (detach to avoid back-prop - # coupling between the two flow paths). sec_cont's log_mass/charge - # columns are already a fixed physics-derived regression target (see - # giant.data.transforms.encode_secondaries) rather than a learned/moving - # one, so — unlike the embedding-table target this replaced — no - # detaching is needed to keep the target from chasing the decoder. - x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM) - - sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) - l_s2 = flow_matching_loss_secondary( - sec_decoder, - x1_s2, - cond_cont, - cond_cat, - x1_s1.detach(), - sec_mask, - ) - - # Optional MoE load-balance auxiliary loss: only present when both stages - # are routed (RoutedDenoisingMLP/RoutedSecondaryDecoder carry `.router`, - # the monolith models don't), computed on cond_cont alone (cheap — no - # trunk compute) so it's reported even when lambda_balance == 0. - if hasattr(stage1_model, "router") and hasattr(sec_decoder, "router"): - l_balance = stage1_model.router.balance_loss( - cond_cont, cond_cat - ) + sec_decoder.router.balance_loss(cond_cont, cond_cat) - # Supervised router auxiliary loss (e.g. ProcessRouter's process - # classifier); a scalar 0 for routers with no such loss (EnergyRouter). - l_proc = stage1_model.router.classify_loss( - cond_cont, cond_cat, proc_idx - ) + sec_decoder.router.classify_loss(cond_cont, cond_cat, proc_idx) - # Optional entropy-regularization aux loss (see Router.entropy_loss): - # penalizes uniform/collapsed gating, a secondary guard against - # gate-sharpness collapse that lambda_balance alone can't see. - l_entropy = stage1_model.router.entropy_loss( - cond_cont, cond_cat - ) + sec_decoder.router.entropy_loss(cond_cont, cond_cat) - else: - l_balance = torch.zeros((), device=device) - l_proc = torch.zeros((), device=device) - l_entropy = torch.zeros((), device=device) - - total = l_s1 + lambda_nsec * l_nsec + lambda_s2 * l_s2 - if lambda_balance > 0: - total = total + lambda_balance * l_balance - if lambda_proc > 0: - total = total + lambda_proc * l_proc - if lambda_entropy > 0: - total = total + lambda_entropy * l_entropy - return total, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc - - -def _wgan_train_step( - generator: torch.nn.Module, - sec_generator: torch.nn.Module, - critic: torch.nn.Module, - sec_critic: torch.nn.Module, - batch: tuple, - device: torch.device, - optimizer_g: optim.Optimizer, - optimizer_d: optim.Optimizer, - g_params: list, - d_params: list, - step_count: int, - n_critic: int, - gp_weight: float, - lambda_nsec: float, - lambda_s2: float, -) -> dict: - """One WGAN-GP training step, both stages (see giant/model/wgan.py for the losses). - - Both critics update every batch. Every `n_critic`-th batch additionally - updates both generators. The (non-adversarial) n_sec classifier updates - every batch regardless — folded into whichever `optimizer_g` step happens - this batch (full adversarial g_loss on generator batches, n_sec-only in - between) rather than throttled to the generator's cadence, since n_sec - accuracy is a headline flow-vs-wgan comparison metric and shares the - generator's ConditionEncoder. - - Stage 2's real/fake target is a flattened (B, SEC_DIM) vector with - `K_MAX - n_sec` padded slots per row; both critic's input and its - gradient-penalty gradient are masked to the valid slots (see - `giant.model.wgan.gradient_penalty`) so the critic can't key on padding - instead of genuine content. Stage 2 is conditioned on the *real* - ground-truth Stage-1 target (`x1_s1`, detached) rather than the - generator's own fake Stage-1 output — same precedent as the flow-matching - path's `flow_matching_loss_secondary` call, avoiding compounding errors - during training. - """ - cond_cont, cond_cat, x1_s1, n_sec, sec_cont, _proc_idx = batch - cond_cont = cond_cont.to(device) - cond_cat = cond_cat.to(device) - x1_s1 = x1_s1.to(device) - n_sec = n_sec.to(device) - sec_cont = sec_cont.to(device) - - B = x1_s1.size(0) - x1_s2 = sec_cont.flatten(1) # (B, SEC_DIM) - sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) - mask_flat = ( - sec_mask.unsqueeze(-1).expand(-1, -1, SEC_SLOT_DIM).reshape(B, -1).float() - ) - stage1_ctx = x1_s1.detach() - - def critic_fn1(x: torch.Tensor) -> torch.Tensor: - return critic(x, cond_cont, cond_cat) - - def critic_fn2(x: torch.Tensor) -> torch.Tensor: - return sec_critic(x, cond_cont, cond_cat, stage1_ctx) - - z1 = torch.randn(B, generator.noise_dim, device=device) - fake1 = generator(z1, cond_cont, cond_cat) - z2 = torch.randn(B, sec_generator.noise_dim, device=device) - fake2 = sec_generator(z2, cond_cont, cond_cat, stage1_ctx) - fake2_masked = fake2 * mask_flat - real2_masked = x1_s2 * mask_flat - - # --- Critic step (every batch) --- - fake1_detached = fake1.detach() - real1_score = critic_fn1(x1_s1) - fake1_score = critic_fn1(fake1_detached) - gp1 = gradient_penalty(critic_fn1, x1_s1, fake1_detached) - d1 = fake1_score.mean() - real1_score.mean() + gp_weight * gp1 - wasserstein_estimate = (real1_score.mean() - fake1_score.mean()).detach() - - fake2_detached_masked = fake2_masked.detach() - real2_score = critic_fn2(real2_masked) - fake2_score = critic_fn2(fake2_detached_masked) - gp2 = gradient_penalty( - critic_fn2, real2_masked, fake2_detached_masked, mask=mask_flat - ) - d2 = fake2_score.mean() - real2_score.mean() + gp_weight * gp2 - - d_loss = d1 + lambda_s2 * d2 - optimizer_d.zero_grad() - d_loss.backward() - grad_norm_d = torch.nn.utils.clip_grad_norm_(d_params, 1.0) - optimizer_d.step() - - # --- Generator (+ n_sec) step --- - did_g_step = step_count % n_critic == 0 - n_sec_logits = generator.predict_n_sec(cond_cont, cond_cat) - l_nsec = F.cross_entropy(n_sec_logits, n_sec) - nsec_acc = (n_sec_logits.argmax(dim=-1) == n_sec).float().mean() - optimizer_g.zero_grad() - if did_g_step: - g1 = generator_loss(critic_fn1, fake1) - g2 = generator_loss(critic_fn2, fake2_masked) - g_loss = g1 + lambda_nsec * l_nsec + lambda_s2 * g2 - else: - g1 = torch.zeros((), device=device) - g2 = torch.zeros((), device=device) - g_loss = lambda_nsec * l_nsec - g_loss.backward() - grad_norm_g = torch.nn.utils.clip_grad_norm_(g_params, 1.0) - optimizer_g.step() - - return { - "d_loss": d_loss.detach(), - "g_loss": (g1 + lambda_s2 * g2).detach(), - "wasserstein_estimate": wasserstein_estimate, - "gp_loss": (gp1 + lambda_s2 * gp2).detach(), - "l_nsec": l_nsec.detach(), - "nsec_acc": nsec_acc.detach(), - "did_g_step": did_g_step, - "grad_norm": grad_norm_d.item() + grad_norm_g.item(), - "grad_norm_d": grad_norm_d.item(), - "grad_norm_g": grad_norm_g.item(), - } - - -def train( - stage1_model: torch.nn.Module, - sec_decoder: torch.nn.Module, - train_loader: DataLoader, - val_loader: DataLoader, - mode: str, - epochs: int, - lr: float, - warmup_epochs: int, - device: torch.device, - out_dir: str | Path, - weight_decay: float = 0.01, - ema_decay: float = 0.9999, - lambda_nsec: float = 0.1, - lambda_s2: float = 1.0, - 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, - normalizer_dict: dict | None = None, - pdg_map: dict | None = None, - mat_map: dict | None = None, - proc_map: dict | None = None, - model_config: dict | None = None, - resume_path: str | Path | None = None, - validate_every: int = 0, - validate_steps: int = 10, - max_val_batches: int = 0, - total_train_batches: int = 0, - critic: torch.nn.Module | None = None, - sec_critic: torch.nn.Module | None = None, - n_critic: int = 5, - gp_weight: float = 10.0, - critic_lr: float | None = None, - use_wandb: bool = False, - wandb_project: str = "giant", - wandb_run_name: str = "", - wandb_log_every: int = 50, -) -> None: - out_dir = Path(out_dir) - out_dir.mkdir(parents=True, exist_ok=True) - - stage1_params = sum(p.numel() for p in stage1_model.parameters()) - sec_decoder_params = sum(p.numel() for p in sec_decoder.parameters()) - critic_params = ( - sum(p.numel() for p in critic.parameters()) if critic is not None else 0 - ) - sec_critic_params = ( - sum(p.numel() for p in sec_critic.parameters()) if sec_critic is not None else 0 - ) - total_params = ( - stage1_params + sec_decoder_params + critic_params + sec_critic_params - ) - - 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 - # `id` is derived from out_dir so resuming a run (--resume) reattaches - # to the same wandb run instead of starting a new one. - 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( - mode=mode, - epochs=epochs, - lr=lr, - warmup_epochs=warmup_epochs, - weight_decay=weight_decay, - ema_decay=ema_decay, - lambda_nsec=lambda_nsec, - lambda_s2=lambda_s2, - lambda_balance=lambda_balance, - lambda_proc=lambda_proc, - lambda_entropy=lambda_entropy, - gumbel_tau_start=gumbel_tau_start, - gumbel_tau_end=gumbel_tau_end, - n_critic=n_critic, - gp_weight=gp_weight, - model_config=model_config, - stage1_params=stage1_params, - sec_decoder_params=sec_decoder_params, - critic_params=critic_params, - sec_critic_params=sec_critic_params, - total_params=total_params, - ), - ) - - stage1_model = stage1_model.to(device) - sec_decoder = sec_decoder.to(device) - if mode == "wgan": - assert critic is not None and sec_critic is not None, ( - "mode='wgan' requires critic/sec_critic (see giant.model.network.build_critics)" - ) - critic = critic.to(device) - sec_critic = sec_critic.to(device) - - # MoE routing trunk (RoutedDenoisingMLP/RoutedSecondaryDecoder) is - # optional and orthogonal to `mode` — both stages carry a `.router` - # when enabled. Each router is an independent instance (their - # `n_experts` need not match), used both for the batch-level gate - # entropy snapshot below and the val-level gate stats further down. - has_router = hasattr(stage1_model, "router") and hasattr(sec_decoder, "router") - - # Flow-matching/diffusion models sample noticeably better from an EMA of - # the weights than from the raw SGD-noisy ones — buffers (e.g. the fixed - # sinusoidal-embedding freqs, or non-learned router centers) never change - # after this initial copy, so only parameters need the running average. - ema_stage1_model: torch.nn.Module | None = None - ema_sec_decoder: torch.nn.Module | None = None - if ema_decay > 0: - ema_stage1_model = copy.deepcopy(stage1_model).eval() - ema_sec_decoder = copy.deepcopy(sec_decoder).eval() - for p in ema_stage1_model.parameters(): - p.requires_grad_(False) - for p in ema_sec_decoder.parameters(): - p.requires_grad_(False) - - all_params = list(stage1_model.parameters()) + list(sec_decoder.parameters()) - all_params_d: list = [] - optimizer_d: optim.Optimizer | None = None - if mode == "wgan": - assert critic is not None and sec_critic is not None - # Standard WGAN-GP recipe (Gulrajani et al. 2017): Adam with - # beta1=0 (momentum destabilizes critic training) and no weight - # decay, rather than the AdamW(weight_decay=...) used for flow/ddpm. - optimizer = optim.Adam(all_params, lr=lr, betas=(0.0, 0.9)) - all_params_d = list(critic.parameters()) + list(sec_critic.parameters()) - optimizer_d = optim.Adam( - all_params_d, - lr=critic_lr if critic_lr is not None else lr, - betas=(0.0, 0.9), - ) - else: - optimizer = optim.AdamW(all_params, lr=lr, weight_decay=weight_decay) - - # Warmup/decay in units of optimizer steps rather than epochs: at large - # dataset sizes a single epoch can be tens of thousands of steps, and an - # epoch-granularity schedule would leave warmup/cosine decay unable to - # move within it. Requires an accurate `total_train_batches` (steps per - # epoch); the only caller, run_train_job, always supplies one. - # - # In wgan mode, `lr_sched.step()`/EMA only fire on generator steps (see - # the per-batch loop below) — 1 in every `n_critic` batches — so the - # schedule's own step-counting must be in those same units, or warmup - # would never finish and cosine decay would barely move. - steps_per_epoch = max(total_train_batches, 1) - if mode == "wgan": - # Generator steps fire every n_critic-th batch (did_g_step = - # step_count % n_critic == 0 in _wgan_train_step), not n_critic + 1. - steps_per_epoch = max(total_train_batches // n_critic, 1) - warmup_steps = warmup_epochs * steps_per_epoch - total_steps = max(epochs * steps_per_epoch, 1) - - 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)) - - lr_sched = optim.lr_scheduler.LambdaLR(optimizer, _lr_lambda) - - ddpm_schedule = CosineSchedule().to(device) if mode == "ddpm" else None - - def _build_checkpoint(epoch: int, global_step: int, best_val_loss: float) -> dict: - ckpt: dict = { - "model": stage1_model.state_dict(), - "sec_decoder": sec_decoder.state_dict(), - "optimizer": optimizer.state_dict(), - "lr_sched": lr_sched.state_dict(), - "epoch": epoch, - "best_val_loss": best_val_loss, - "global_step": global_step, - } - if mode == "wgan": - assert ( - critic is not None - and sec_critic is not None - and optimizer_d is not None - ) - ckpt["critic"] = critic.state_dict() - ckpt["sec_critic"] = sec_critic.state_dict() - ckpt["optimizer_d"] = optimizer_d.state_dict() - if ema_decay > 0: - assert ema_stage1_model is not None and ema_sec_decoder is not None - ckpt["model_ema"] = ema_stage1_model.state_dict() - ckpt["sec_decoder_ema"] = ema_sec_decoder.state_dict() - if normalizer_dict is not None: - ckpt["normalizer"] = normalizer_dict - if pdg_map is not None: - ckpt["pdg_map"] = pdg_map - if mat_map is not None: - ckpt["mat_map"] = mat_map - if proc_map is not None: - ckpt["proc_map"] = proc_map - if model_config is not None: - ckpt["model_config"] = model_config - return ckpt - - start_epoch = 1 - best_val_loss = float("inf") - resumed_global_step = 0 - if resume_path is not None: - ckpt = torch.load(resume_path, map_location=device, weights_only=False) - stage1_model.load_state_dict(ckpt["model"]) - sec_decoder.load_state_dict(ckpt["sec_decoder"]) - if ema_decay > 0: - assert ema_stage1_model is not None and ema_sec_decoder is not None - ema_stage1_model.load_state_dict(ckpt.get("model_ema", ckpt["model"])) - ema_sec_decoder.load_state_dict( - ckpt.get("sec_decoder_ema", ckpt["sec_decoder"]) - ) - if mode == "wgan": - assert ( - critic is not None - and sec_critic is not None - and optimizer_d is not None - ) - critic.load_state_dict(ckpt["critic"]) - sec_critic.load_state_dict(ckpt["sec_critic"]) - optimizer_d.load_state_dict(ckpt["optimizer_d"]) - # Mirrors the `lr` fixup below for the generator optimizer: - # optimizer_d.load_state_dict() above restores the checkpoint's - # own critic LR, which would otherwise silently override an - # explicit `--critic-lr` passed on this resume. optimizer_d has - # no LR scheduler (unlike `optimizer`/`lr_sched`), so this is a - # flat set rather than a schedule-relative one. - resumed_critic_lr = critic_lr if critic_lr is not None else lr - for group in optimizer_d.param_groups: - group["lr"] = resumed_critic_lr - optimizer.load_state_dict(ckpt["optimizer"]) - lr_sched.load_state_dict(ckpt["lr_sched"]) - start_epoch = ckpt.get("epoch", 0) + 1 - best_val_loss = ckpt.get("best_val_loss", float("inf")) - resumed_global_step = ckpt.get("global_step", 0) - - # optimizer/lr_sched.load_state_dict() above restore the checkpoint's - # own base LR, which would otherwise silently override an explicit - # `lr` argument. Make `lr` authoritative again, applied at whatever - # point the cosine/warmup schedule has already reached. - lr_sched.base_lrs = [lr for _ in lr_sched.base_lrs] - resumed_lr = lr * _lr_lambda(lr_sched.last_epoch) - for group in optimizer.param_groups: - group["lr"] = resumed_lr - - if start_epoch > epochs: - print( - f"checkpoint already completed epoch {start_epoch - 1} " - f"(>= --epochs {epochs}) — nothing to train" - ) - return - - metrics_path = out_dir / "metrics.csv" - resuming_existing_metrics = resume_path is not None and metrics_path.exists() - write_header = not resuming_existing_metrics - metrics_file = open( - metrics_path, "a" if resuming_existing_metrics else "w", newline="" - ) - metrics_writer = csv.DictWriter(metrics_file, fieldnames=_METRICS_FIELDS) - if write_header: - metrics_writer.writeheader() - - epoch_w = len(str(epochs)) - last_completed_epoch = start_epoch - 1 - # Restored from the checkpoint on --resume so wandb_run.log(..., step=...) - # keeps advancing monotonically instead of restarting at 0 mid-run (a - # reattached wandb run — see wandb.init(id=..., resume="allow") below — - # would otherwise silently drop every post-resume point). - global_step = resumed_global_step - 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) - stage1_model.train() - sec_decoder.train() - if mode == "wgan": - assert critic is not None and sec_critic is not None - critic.train() - sec_critic.train() - train_loss_sum = 0.0 - train_s1_sum = 0.0 - train_nsec_sum = 0.0 - train_s2_sum = 0.0 - train_balance_sum = 0.0 - train_proc_sum = 0.0 - train_entropy_sum = 0.0 - train_d_sum = 0.0 - train_g_sum = 0.0 - train_wasserstein_sum = 0.0 - train_gp_sum = 0.0 - train_nsec_acc_sum = 0.0 - train_grad_norm_d_sum = 0.0 - train_grad_norm_g_sum = 0.0 - train_n = 0 - train_batches = 0 - grad_norm_sum = 0.0 - ema_loss = 0.0 - ema_grad_norm = 0.0 - 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: - if has_router: - gumbel_tau = _gumbel_tau( - global_step, total_steps, gumbel_tau_start, gumbel_tau_end - ) - stage1_model.router.gumbel_tau = gumbel_tau - sec_decoder.router.gumbel_tau = gumbel_tau - - if mode == "wgan": - assert ( - critic is not None - and sec_critic is not None - and optimizer_d is not None - ) - stats = _wgan_train_step( - stage1_model, - sec_decoder, - critic, - sec_critic, - batch, - device, - optimizer, - optimizer_d, - all_params, - all_params_d, - global_step, - n_critic, - gp_weight, - lambda_nsec, - lambda_s2, - ) - if stats["did_g_step"]: - lr_sched.step() - if ema_decay > 0: - assert ( - ema_stage1_model is not None - and ema_sec_decoder is not None - ) - _update_ema(ema_stage1_model, stage1_model, ema_decay) - _update_ema(ema_sec_decoder, sec_decoder, ema_decay) - - B = batch[0].size(0) - batch_loss = stats["d_loss"].item() + stats["g_loss"].item() - batch_grad_norm = stats["grad_norm"] - train_loss_sum += batch_loss * B - train_nsec_sum += stats["l_nsec"].item() * B - train_d_sum += stats["d_loss"].item() * B - train_g_sum += stats["g_loss"].item() * B - train_wasserstein_sum += stats["wasserstein_estimate"].item() * B - train_gp_sum += stats["gp_loss"].item() * B - train_nsec_acc_sum += stats["nsec_acc"].item() * B - train_grad_norm_d_sum += stats["grad_norm_d"] * B - train_grad_norm_g_sum += stats["grad_norm_g"] * B - else: - loss, l_s1, l_nsec, l_s2, l_balance, l_proc, l_entropy, nsec_acc = ( - _compute_losses( - stage1_model, - sec_decoder, - batch, - mode, - ddpm_schedule, - device, - lambda_nsec, - lambda_s2, - lambda_balance, - lambda_proc, - lambda_entropy, - ) - ) - optimizer.zero_grad() - loss.backward() - grad_norm = torch.nn.utils.clip_grad_norm_(all_params, 1.0) - optimizer.step() - lr_sched.step() - if ema_decay > 0: - assert ( - ema_stage1_model is not None and ema_sec_decoder is not None - ) - _update_ema(ema_stage1_model, stage1_model, ema_decay) - _update_ema(ema_sec_decoder, sec_decoder, ema_decay) - - B = batch[0].size(0) - batch_loss = loss.item() - batch_grad_norm = grad_norm.item() - train_loss_sum += batch_loss * B - train_s1_sum += l_s1.item() * B - train_nsec_sum += l_nsec.item() * B - train_s2_sum += l_s2.item() * B - train_balance_sum += l_balance.item() * B - train_proc_sum += l_proc.item() * B - train_entropy_sum += l_entropy.item() * B - train_nsec_acc_sum += nsec_acc.item() * B - - train_n += B - train_batches += 1 - grad_norm_sum += batch_grad_norm - ema_loss = ( - batch_loss if train_n == B else 0.95 * ema_loss + 0.05 * batch_loss - ) - ema_grad_norm = ( - batch_grad_norm - if train_batches == 1 - else 0.95 * ema_grad_norm + 0.05 * batch_grad_norm - ) - bar.set_postfix_str( - f"loss={ema_loss:.4f} gnorm={ema_grad_norm:.3f}", refresh=False - ) - - global_step += 1 - if ( - wandb_run is not None - and wandb_log_every > 0 - and global_step % wandb_log_every == 0 - ): - log_payload = { - "batch/epoch": epoch, - "batch/loss": batch_loss, - "batch/loss_ema": ema_loss, - "batch/grad_norm": batch_grad_norm, - "batch/lr": optimizer.param_groups[0]["lr"], - "batch/critic_lr": ( - optimizer_d.param_groups[0]["lr"] - if optimizer_d is not None - else 0.0 - ), - } - if has_router: - # Cheap re-use of the batch already in hand — no - # extra data loading, just a small forward through - # each router's own gate function. Only entropy is - # logged at this granularity (not per-expert - # utilization): a single batch's importance sum is - # too noisy as a "global share" estimate, whereas - # the val-loop aggregate (below) sums over the - # whole val set for that. Batch-level entropy alone - # is still enough to see a router collapsing in - # real time, mid-epoch, rather than only at the - # next validation pass. - with torch.no_grad(): - cond_cont_b = batch[0].to(device) - cond_cat_b = batch[1].to(device) - s1_entropy, _ = stage1_model.router.gate_stats( - cond_cont_b, cond_cat_b - ) - s2_entropy, _ = sec_decoder.router.gate_stats( - cond_cont_b, cond_cat_b - ) - log_payload["batch/router_s1_entropy"] = s1_entropy.item() - log_payload["batch/router_s2_entropy"] = s2_entropy.item() - log_payload["batch/gumbel_tau"] = gumbel_tau - wandb_run.log(log_payload, step=global_step) - - if shutdown.requested: - break - bar.close() - - if shutdown.requested: - # Epoch was interrupted mid-loop, so there's no val_loss to - # weigh a "best" checkpoint against — save the in-progress - # weights as last.pt only, under the last *fully completed* - # epoch number so --resume restarts this epoch from scratch - # rather than skipping it (weights/optimizer state are still - # kept, so those partial-epoch batches aren't wasted work). - ckpt = _build_checkpoint(epoch - 1, global_step, best_val_loss) - 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 - - train_loss = train_loss_sum / max(train_n, 1) - train_grad_norm = grad_norm_sum / max(train_batches, 1) - train_nsec_acc = train_nsec_acc_sum / max(train_n, 1) - train_grad_norm_d = train_grad_norm_d_sum / max(train_n, 1) - train_grad_norm_g = train_grad_norm_g_sum / max(train_n, 1) - current_lr = optimizer.param_groups[0]["lr"] - critic_lr_value = ( - optimizer_d.param_groups[0]["lr"] if optimizer_d is not None else 0.0 - ) - - stage1_model.eval() - sec_decoder.eval() - if mode == "wgan": - assert critic is not None and sec_critic is not None - critic.eval() - sec_critic.eval() - - val_marginal_kl = float("nan") - if mode == "wgan": - # WGANGenerator.forward(z, cond_cont, cond_cat) has no - # diffusion/flow `t` argument, so the usual _compute_losses - # val loop below (which calls flow_matching_loss -> - # stage1_model(x_t, t, ...)) doesn't apply — and a WGAN - # critic loss isn't a monotone quality signal fit for - # best-checkpoint selection anyway. Select on marginal KL - # against the EMA generators instead (matches what - # predict/rollout sample from by default, --weights ema). - eval_stage1 = ( - ema_stage1_model if ema_stage1_model is not None else stage1_model - ) - eval_sec_decoder = ( - ema_sec_decoder if ema_sec_decoder is not None else sec_decoder - ) - marginal_result = validate_marginals( - eval_stage1, - val_loader, - mode=mode, - device=device, - sec_decoder=eval_sec_decoder, - ) - val_marginal_kl = float(np.mean(marginal_result["kl_divergence"])) - val_loss = val_marginal_kl - val_s1_sum = val_nsec_sum = val_s2_sum = val_balance_sum = ( - val_proc_sum - ) = val_entropy_sum = val_nsec_acc_sum = 0.0 - val_n = 1 - val_nsec_acc = 0.0 - router_s1_entropy = router_s2_entropy = 0.0 - router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0 - router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0 - else: - val_loss_sum = 0.0 - val_s1_sum = 0.0 - val_nsec_sum = 0.0 - val_s2_sum = 0.0 - val_balance_sum = 0.0 - val_proc_sum = 0.0 - val_entropy_sum = 0.0 - val_nsec_acc_sum = 0.0 - val_n = 0 - if has_router: - n_experts_s1 = stage1_model.router.n_experts - n_experts_s2 = sec_decoder.router.n_experts - val_router_s1_entropy_sum = 0.0 - val_router_s2_entropy_sum = 0.0 - val_router_s1_importance_sum = torch.zeros( - n_experts_s1, device=device - ) - val_router_s2_importance_sum = torch.zeros( - n_experts_s2, device=device - ) - 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 - ( - loss, - l_s1, - l_nsec, - l_s2, - l_balance, - l_proc, - l_entropy, - nsec_acc, - ) = _compute_losses( - stage1_model, - sec_decoder, - batch, - mode, - ddpm_schedule, - device, - lambda_nsec, - lambda_s2, - lambda_balance, - lambda_proc, - lambda_entropy, - ) - B = batch[0].size(0) - val_loss_sum += loss.item() * B - val_s1_sum += l_s1.item() * B - val_nsec_sum += l_nsec.item() * B - val_s2_sum += l_s2.item() * B - val_balance_sum += l_balance.item() * B - val_proc_sum += l_proc.item() * B - val_entropy_sum += l_entropy.item() * B - val_nsec_acc_sum += nsec_acc.item() * B - if has_router: - cond_cont = batch[0].to(device) - cond_cat = batch[1].to(device) - s1_entropy, s1_importance = stage1_model.router.gate_stats( - cond_cont, cond_cat - ) - s2_entropy, s2_importance = sec_decoder.router.gate_stats( - cond_cont, cond_cat - ) - val_router_s1_entropy_sum += s1_entropy.item() * B - val_router_s2_entropy_sum += s2_entropy.item() * B - val_router_s1_importance_sum += s1_importance - val_router_s2_importance_sum += s2_importance - val_n += B - val_loss = val_loss_sum / max(val_n, 1) - val_nsec_acc = val_nsec_acc_sum / max(val_n, 1) - - if has_router: - router_s1_entropy = val_router_s1_entropy_sum / max(val_n, 1) - router_s2_entropy = val_router_s2_entropy_sum / max(val_n, 1) - s1_util = val_router_s1_importance_sum / ( - val_router_s1_importance_sum.sum().clamp_min(1e-8) - ) - s2_util = val_router_s2_importance_sum / ( - val_router_s2_importance_sum.sum().clamp_min(1e-8) - ) - router_s1_util_min = s1_util.min().item() - router_s1_util_max = s1_util.max().item() - router_s1_util_std = ( - s1_util.std().item() if n_experts_s1 > 1 else 0.0 - ) - router_s2_util_min = s2_util.min().item() - router_s2_util_max = s2_util.max().item() - router_s2_util_std = ( - s2_util.std().item() if n_experts_s2 > 1 else 0.0 - ) - else: - router_s1_entropy = router_s2_entropy = 0.0 - router_s1_util_min = router_s1_util_max = router_s1_util_std = 0.0 - router_s2_util_min = router_s2_util_max = router_s2_util_std = 0.0 - - if validate_every > 0 and epoch % validate_every == 0: - print(f"[epoch {epoch}] marginal validation:") - marginal_result = validate_marginals( - stage1_model, - val_loader, - mode=mode, - schedule=ddpm_schedule, - device=device, - steps=validate_steps, - sec_decoder=sec_decoder, - ) - val_marginal_kl = float(np.mean(marginal_result["kl_divergence"])) - - epoch_time = time.monotonic() - epoch_start - gpu_mem_mb = ( - torch.cuda.max_memory_allocated(device) / (1024 * 1024) - if device.type == "cuda" - else 0.0 - ) - - is_best = val_loss < best_val_loss - marker = " [best]" if is_best else "" - print( - f"epoch {epoch:{epoch_w}d}/{epochs}" - f" train {train_loss:.4f}" - f" (s1={train_s1_sum / max(train_n, 1):.3f}" - f" nsec={train_nsec_sum / max(train_n, 1):.3f}" - f" s2={train_s2_sum / max(train_n, 1):.3f}" - f" bal={train_balance_sum / max(train_n, 1):.3f}" - f" proc={train_proc_sum / max(train_n, 1):.3f}" - f" entropy={train_entropy_sum / max(train_n, 1):.3f}" - f" d={train_d_sum / max(train_n, 1):.3f}" - f" g={train_g_sum / max(train_n, 1):.3f})" - f" val {val_loss:.4f}" - f" lr {current_lr:.2e} gnorm {train_grad_norm:.3f}" - f" {epoch_time:.1f}s{marker}" - ) - metrics_row = { - "epoch": epoch, - "train_loss": train_loss, - "train_loss_s1": train_s1_sum / max(train_n, 1), - "train_loss_nsec": train_nsec_sum / max(train_n, 1), - "train_loss_s2": train_s2_sum / max(train_n, 1), - "train_loss_balance": train_balance_sum / max(train_n, 1), - "train_loss_proc": train_proc_sum / max(train_n, 1), - "train_loss_entropy": train_entropy_sum / max(train_n, 1), - "train_nsec_acc": train_nsec_acc, - "d_loss": train_d_sum / max(train_n, 1), - "g_loss": train_g_sum / max(train_n, 1), - "wasserstein_estimate": train_wasserstein_sum / max(train_n, 1), - "gp_loss": train_gp_sum / max(train_n, 1), - "val_loss": val_loss, - "val_loss_s1": val_s1_sum / max(val_n, 1), - "val_loss_nsec": val_nsec_sum / max(val_n, 1), - "val_loss_s2": val_s2_sum / max(val_n, 1), - "val_loss_balance": val_balance_sum / max(val_n, 1), - "val_loss_proc": val_proc_sum / max(val_n, 1), - "val_loss_entropy": val_entropy_sum / max(val_n, 1), - "val_nsec_acc": val_nsec_acc, - "val_marginal_kl": val_marginal_kl, - "router_s1_entropy": router_s1_entropy, - "router_s1_util_min": router_s1_util_min, - "router_s1_util_max": router_s1_util_max, - "router_s1_util_std": router_s1_util_std, - "router_s2_entropy": router_s2_entropy, - "router_s2_util_min": router_s2_util_min, - "router_s2_util_max": router_s2_util_max, - "router_s2_util_std": router_s2_util_std, - "lr": current_lr, - "critic_lr": critic_lr_value, - "grad_norm": train_grad_norm, - "grad_norm_d": train_grad_norm_d, - "grad_norm_g": train_grad_norm_g, - "gpu_mem_mb": gpu_mem_mb, - "samples_per_sec": train_n / max(epoch_time, 1e-8), - "is_best": int(is_best), - "epoch_time_s": epoch_time, - } - metrics_writer.writerow(metrics_row) - metrics_file.flush() - if wandb_run is not None: - # Shares the same monotonic step axis as the per-batch - # `batch/*` logs above (global_step) rather than `epoch`, - # since a wandb run's `step` argument across `log()` calls - # must never decrease. - wandb_run.log(metrics_row, step=global_step) - - ckpt = _build_checkpoint(epoch, global_step, best_val_loss) - - if val_loss < best_val_loss: - 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 - - metrics_file.close() - if wandb_run is not None: - wandb_run.finish() - - if shutdown.requested: - print( - f"stopped after epoch {last_completed_epoch} due to shutdown signal — " - f"resume with --resume {out_dir / 'last.pt'}" - ) diff --git a/giant/training/__init__.py b/giant/training/__init__.py new file mode 100644 index 0000000..ed50b5f --- /dev/null +++ b/giant/training/__init__.py @@ -0,0 +1,30 @@ +"""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", +] diff --git a/giant/training/checkpoint.py b/giant/training/checkpoint.py new file mode 100644 index 0000000..c946314 --- /dev/null +++ b/giant/training/checkpoint.py @@ -0,0 +1,68 @@ +"""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_` / `optimizer_d_` / `lr_sched_` +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) diff --git a/giant/training/loop.py b/giant/training/loop.py new file mode 100644 index 0000000..ca43651 --- /dev/null +++ b/giant/training/loop.py @@ -0,0 +1,295 @@ +"""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, + 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, + "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'}" + ) diff --git a/giant/training/metrics.py b/giant/training/metrics.py new file mode 100644 index 0000000..d1fb0dd --- /dev/null +++ b/giant/training/metrics.py @@ -0,0 +1,384 @@ +"""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: `/train/`, `/val/`, +`/` for point-in-time values (`lr`, `critic_lr`), +`/router/` 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 `/`.""" + 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) diff --git a/giant/training/stage2_inputs.py b/giant/training/stage2_inputs.py new file mode 100644 index 0000000..e2f8367 --- /dev/null +++ b/giant/training/stage2_inputs.py @@ -0,0 +1,318 @@ +"""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=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) diff --git a/giant/training/trainers.py b/giant/training/trainers.py new file mode 100644 index 0000000..d6ffd6d --- /dev/null +++ b/giant/training/trainers.py @@ -0,0 +1,952 @@ +"""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 + +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.model.network import Router, 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: tuple, device: torch.device) -> tuple: + return tuple(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_emb_dim: 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_emb_dim=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 batch tuple `(cond_cont, cond_cat, x1_s1, n_sec, + sec_cont, proc_idx, sec_type_idx)` — 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_emb_dim = spec.particle_type_emb_dim + 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: tuple, device: torch.device, global_step: int) -> dict: + raise NotImplementedError + + def val_loss(self, batch: tuple, 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_emb_dim, + 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_emb_dim, + ) + 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: tuple, 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: + l_balance = self.router.balance_loss(cond_cont, cond_cat) + l_proc = self.router.classify_loss(cond_cont, cond_cat, proc_idx) + 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: tuple, 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: tuple, 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 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, 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_emb_dim) + 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: tuple, 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( + (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_emb_dim), + 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 diff --git a/giant/validate.py b/giant/validate.py index 62a88df..1d3b5a0 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -2,26 +2,13 @@ import numpy as np import torch from torch.utils.data import DataLoader -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, -) +from giant.constants import LOCAL_TARGET_NAMES +from giant.sample import resolve_n_sec, sample_stage1, sample_stage2 _SEC_PHYS_NAMES = ["log_mass", "charge"] -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: +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()) @@ -44,15 +31,40 @@ 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( - model: torch.nn.Module, + stage1_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 | None = None, + steps: int = 10, + ddpm_steps: int = 1000, sec_decoder: torch.nn.Module | None = None, ) -> dict[str, np.ndarray | float]: """Compare per-dimension marginals of generated vs. real steps. @@ -61,51 +73,55 @@ def validate_marginals( normalised space. `kl_divergence[j]` is KL(real || generated) for dimension j, estimated from a shared histogram over both samples. - `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. + 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. - 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. + 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). """ if device is None: - device = next(model.parameters()).device - model.eval() + device = next(stage1_model.parameters()).device + stage1_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_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)] - all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)] + 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)] for i, batch in enumerate(val_loader): if n_batches is not None and i >= n_batches: break - # 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 + # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, + # sec_type_idx). + cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) - 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) - ) + gen, n_sec_pred = sample_stage1(stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps) all_real.append(x1.numpy()) all_gen.append(gen.cpu().numpy()) @@ -113,54 +129,46 @@ 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] - 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)) + sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2( + sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps ) - gen_phys = sec_phys_pred.cpu().numpy() + gen_frac = 1.0 / (1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64))) gen_valid = sec_valid_pred.cpu().numpy() - all_phys_real.append(real_phys[real_valid]) - all_phys_gen.append(gen_phys[gen_valid]) - for j in range(K_MAX): + 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_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} " - f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}" - ) + header = f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} {'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): @@ -181,24 +189,9 @@ 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) - 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}" - ) + 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}") 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)) @@ -208,46 +201,70 @@ 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} " - f"{r.std():>10.4f} {g.std():>10.4f} {kl:>14.4f}" - ) + print(f"{j:<24} {r.mean():>10.4f} {g.mean():>10.4f} {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 diff --git a/issues.md b/issues.md new file mode 100644 index 0000000..7d5a085 --- /dev/null +++ b/issues.md @@ -0,0 +1,1281 @@ +# GIANT — architecture issues + +Software-engineering review of the `giant` codebase, conducted 2026-08-12 on branch +`v0.3.0-stage2-autoregressive` at commit `55332db`. + +This document is written to be read **without the context of the review conversation**. +Each issue states what the code does today, why that is a problem, how to verify the +claim independently, what the fix looks like, and — importantly — what the fix should +*not* touch. Line numbers are accurate as of `55332db`; if they have drifted, the +accompanying code excerpts and `grep` commands will still locate the code. + +## Baseline: what is already good + +Read this first, so the issues below are calibrated correctly. **This is a healthy +codebase.** The problems listed are about structural resilience as the v0.3.0 +architecture matrix grows, not about rot or breakage. + +- `pytest` — **725 tests pass in ~18 s**. Fast enough that there is no excuse for not + running it on every change. +- `uv run ruff check .`, `uv run ruff format --check .`, `uv run ty check .` — all clean, + and all three are enforced in CI (`.gitea/workflows/ci.yml`). +- Zero `TODO` / `FIXME` / `XXX` / `HACK` markers in `giant/` or `scripts/`. +- The internal import graph is **acyclic** with clean layering: + `constants → data → model → training → pipeline → cli`, and `analysis` almost fully + independent of the rest. +- Line coverage is **96–99.6 % on every core module** (`network.py` 98.6 %, + `trainers.py` 99.5 %, `rollout.py` 99.6 %, `transforms.py` 96.2 %, `pipeline.py` + 97.1 %, `analysis/catalog.py` 99.6 %). The single exception is `cli.py` at 35.8 % — + see Issue 4. +- Comment and docstring quality is unusually high. Docstrings routinely explain *why* + a decision was made, not just what the code does, and several call out their own + known limitations honestly. **Preserve this when refactoring.** A refactor that + deletes the reasoning in `giant/config.py`'s `DEFAULT_CONFIG` comments or + `giant/analysis/router_gating.py`'s module docstring is a net loss even if the code + gets shorter. +- `giant/analysis/` is the best-designed subsystem in the repo and should be treated as + the template the rest of the codebase moves toward. See Issue 11. + +## Issue index + +| # | Issue | Severity | Effort | Status | +|---|---|---|---|---| +| 1 | Config defaults are declared twice; `DEFAULT_CONFIG` and consumers already disagree | **High** | Medium | **Fixed** (`9bf5874`) | +| 2 | No unknown-key validation — a typo in `config.toml` silently trains the wrong model | **High** | Small | **Fixed** | +| 3 | `cli.py:train()` is a 58-parameter, 510-line fat controller | **High** | Medium | **Fixed** (`2bfb1ab`) | +| 4 | `cli.py` is at 35.8 % coverage and holds untested override-precedence logic | **High** | Medium | **Fixed** (`2bfb1ab`, partial — see status note) | +| 5 | Inference bootstrap is duplicated verbatim between `predict` and `rollout` | **High** | Small | **Fixed** | +| 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | Open | +| 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | Open | +| 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | Open | +| 9 | `scripts` is published as a top-level distribution package | Medium | Small | Open | +| 10 | `torch.load(weights_only=False)` — checkpoints are arbitrary pickles | Low | Medium | Open | +| 11 | Minor: `echo=print` threading, `particles → data.loader` layering | Low | Small | Open | + +--- + +## Issue 1 — Config defaults are declared twice, and the two declarations already disagree + +> **Status: Fixed, commit `9bf5874` on `v0.3.0-stage2-autoregressive`.** `giant/config.py` now +> declares a hierarchy of frozen dataclasses (`GiantConfig` and its nested blocks, following +> `StageSpec`'s existing style) as the single source of truth; `DEFAULT_CONFIG` is generated from +> `GiantConfig().to_dict()` rather than hand-maintained, and `build_models`/`build_critics` +> (`network.py`) and `StageSpec.from_config` (`trainers.py`) convert their dict input into these +> dataclasses at the top of each function and read attributes instead of duplicating +> `.get(key, literal)` defaults. Both drifted keys (`decoder`, `particle_type.target`) now resolve +> to `DEFAULT_CONFIG`'s documented v0.3.0 values (`"autoregressive"`/`"onehot"`) unconditionally. +> Router/`n_sec` sub-blocks keep a dict-shaped `extra` catch-all for their genuinely dynamic keys +> (composed-router `axis{i}_*`, runtime-seeded `centers_init`, legacy `legacy_owner`) rather than +> being fully typed — see the "Recommended fix" section below, which this diverges from slightly +> on that one point. Fixing the fallback surfaced two existing callers that had been silently +> depending on the old (wrong) default — a `tests/test_train.py` fixture and +> `scripts/warm_setup_cache.py`'s hand-built minimal config (now merged against `DEFAULT_CONFIG` +> instead of hand-rolled) — both fixed in the same commit. Regression tests added in +> `tests/test_config.py`, `tests/test_network.py`, `tests/test_train.py`, including one pinning +> `GiantConfig().to_dict() == DEFAULT_CONFIG` so this class of drift can't recur silently. +> Everything below this point describes the pre-fix state and is kept for historical context. + +**Severity: High. Effort: Medium. Risk if unfixed: silent wrong-model training.** + +**Location:** `giant/config.py:32` (`DEFAULT_CONFIG`) versus `giant/model/network.py:1570-1693` +(`build_models`), `giant/training/trainers.py:130-175` (`StageSpec.from_config`), +`giant/pipeline.py:381`. + +### What the code does today + +`DEFAULT_CONFIG` in `giant/config.py:32` is a fully-specified nested dict. Every key has +a value and most have an explanatory comment. It is, on paper, the single source of +truth for configuration defaults. + +The consumers do not treat it that way. They re-declare the same defaults inline via +`dict.get(key, default)`: + +```python +# giant/model/network.py, inside build_models +cond_out_dim = conditioning.get("out_dim", 128) +... +hidden_dim=s1cfg.get("hidden_dim", 256), +n_res_blocks=s1cfg.get("n_res_blocks", 6), +dropout=s1cfg.get("dropout", 0.0), +time_dim=gen_sub.get("time_dim", 64), +``` + +```python +# giant/training/trainers.py, inside StageSpec.from_config +lambda_weight=stage_cfg.get("lambda", 1.0), +n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1), +ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000), +n_critic=wgan_cfg.get("n_critic", 5), +gp_weight=wgan_cfg.get("gp_weight", 10.0), +``` + +There are **132 `.get("key", default)` call sites across `giant/`**, concentrated in +`network.py` (55) and `trainers.py` (25). Each one is an independent copy of a value +that `DEFAULT_CONFIG` also declares. + +### Why it's a problem + +Because `cfg` is typed `dict` (i.e. `dict[str, Any]`), neither `ty` nor the tests can +see that the two declarations are supposed to agree. When they drift, nothing fails — +the model is simply built differently than the config file says. + +**This has already happened.** Two defaults currently disagree: + +| Key | `DEFAULT_CONFIG` says | `build_models` fallback says | +|---|---|---| +| `stage2_model.decoder` | `"autoregressive"` | `"one_shot"` (`network.py:1632`) | +| `stage2_model.particle_type.target` | `"onehot"` | `"physical"` (`network.py:1644`, `trainers.py:143`, `pipeline.py:381`) | + +Both are *architecturally load-bearing*: `decoder` selects between `Stage2Autoregressive` +and `Stage2OneShot` — two different networks — and `particle_type.target` selects between +categorical class logits and continuous `(log-mass, charge)` regression, which is the +exact axis the v0.3.0 redesign was created to change (see `CLAUDE.md`, "v0.3.0 — +Stage-2 autoregressive redesign"). + +### Is it currently a live bug? + +**No — it is currently latent.** Two things save it today: + +1. The training path always passes a fully-merged config. `giant/pipeline.py:446` builds + `model_config` from `cfg["stage1_model"]` / `cfg["stage2_model"]`, and `cfg` always + originates from `merge_cli_overrides(DEFAULT_CONFIG, ...)`, so every key is present + and no fallback fires. +2. The legacy path also populates the keys explicitly. + `_migrate_legacy_model_config` (`network.py:1446`) hard-codes + `"decoder": "one_shot"` and `"particle_type": {"target": "physical"}` into the dict it + returns, so a v0.2 checkpoint does not rely on the fallbacks either. +3. The tests build their configs by deep-copying `DEFAULT_CONFIG` (e.g. + `tests/test_network.py:611` `_minimal_model_config`), so they never exercise the + fallback branch. + +That is precisely what makes this dangerous rather than harmless. The fallbacks are +**unreachable by every current caller, untested, and wrong.** The first caller that +hand-builds a partial `model_config` — a new test, a notebook, a debugging script, a +future `giant train --stage2-only` path — silently gets the v0.2 architecture while the +config documentation promises the v0.3 one. + +### How to verify + +```bash +uv run python - <<'EOF' +import copy +from giant import config as gc +from giant.model.network import build_models + +full = copy.deepcopy(gc.DEFAULT_CONFIG) +print("DEFAULT_CONFIG stage2 decoder :", full["stage2_model"]["decoder"]) +print("DEFAULT_CONFIG particle_type.target:", full["stage2_model"]["particle_type"]["target"]) + +partial = { + "pdg_vocab": 3, "mat_vocab": 2, + "conditioning": full["conditioning"], + "stage1_model": {"hidden_dim": 8, "n_res_blocks": 1}, + "stage2_model": {"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3}, +} +print("build_models fallback stage2 class :", type(build_models(partial)["stage2"]).__name__) +EOF +``` + +Observed output on `55332db`: + +``` +DEFAULT_CONFIG stage2 decoder : autoregressive +DEFAULT_CONFIG particle_type.target: onehot +build_models fallback stage2 class : Stage2OneShot +``` + +### Recommended fix + +Make the config **typed at the boundary**, so the defaults exist exactly once and the +type checker can see them. + +1. Define a frozen dataclass per config block in `giant/config.py`: + `ConditioningAxisConfig`, `ConditioningConfig`, `RouterConfig`, `WganConfig`, + `Stage1ModelConfig`, `Stage2ModelConfig`, `TrainConfig`, `GiantConfig`. The dataclass + field defaults become the *only* declaration of each default. Move the existing + `DEFAULT_CONFIG` comments onto the fields — they are the most valuable thing in that + file and must not be lost. +2. Give each a `from_dict(d)` classmethod that constructs from the merged TOML dict and + **rejects unknown keys** (that is Issue 2, and the same constructor solves both). +3. Change `build_models`, `build_critics` and `StageSpec.from_config` to take the + dataclasses instead of dicts. Every `.get(key, default)` becomes `cfg.key`. The 132 + duplicate defaults disappear by construction. +4. Keep `DEFAULT_CONFIG` as a derived artifact if anything still needs the dict shape + (`save_config`, `default_out_dir_name`, the setup-cache sidecar) — generate it from + the dataclasses via `dataclasses.asdict`, do not maintain it by hand. + +Precedent exists in this codebase and should be followed: +`@dataclass(frozen=True) StageSpec` (`trainers.py:79`) and `PlotSpec` +(`analysis/catalog.py:65`) already do exactly this, well. + +### Interim fix if the full change is too large + +If the dataclass migration is deferred, at minimum add a test that asserts the inline +fallbacks agree with `DEFAULT_CONFIG`, so the drift is caught: + +```python +def test_build_models_fallbacks_match_default_config(): + """Every `.get(key, default)` in build_models must use the same default + DEFAULT_CONFIG declares — see issues.md Issue 1.""" +``` + +and fix the two divergent values. Decide deliberately which way: if the fallbacks are +meant to encode *v0.2 legacy* behaviour, they are redundant (the legacy migration at +`network.py:1446` already sets both keys explicitly) and should be **deleted**, letting +a missing key raise `KeyError` instead of silently choosing an architecture. + +### Scope guard + +Do **not** change any actual default *value* as part of this refactor. The goal is to +make the two declarations agree and then have only one; changing what the model does is +a separate, physics-relevant decision that belongs in its own commit with its own +retraining. When resolving the two divergences above, the correct target is whatever +`DEFAULT_CONFIG` says (`autoregressive` / `onehot`), because that is the documented +v0.3.0 intent per `CLAUDE.md`. + +--- + +## Issue 2 — No unknown-key validation: a typo in `config.toml` silently trains the wrong model + +> **Status: Fixed.** `giant/config.py` now has `validate_config_keys(cfg)`, which walks a +> merged config dict recursively against `DEFAULT_CONFIG`'s tree (itself generated from +> the `GiantConfig` dataclasses added for Issue 1) and raises `ValueError` on any key not +> present there, with a `difflib`-based "did you mean" suggestion. `[meta]` is skipped +> unconditionally, and `stage{1,2}_model.router`'s `axis{i}_{field}` composed-router keys +> and runtime-seeded `centers_init` are allowed through explicitly. `merge_cli_overrides` +> calls it right before returning, so all three real entry points (`giant train`, +> `giant new-run`, `scripts/warm_setup_cache.py`) are covered automatically; checkpoint +> `model_config` loading (a separate code path, `network._migrate_legacy_model_config`) +> is untouched, so old checkpoints keep loading regardless of schema drift. This is the +> "standalone fallback" option from the recommended fix below rather than the +> dataclass-`from_dict`-rejects-unknown-keys option, specifically to keep the check scoped +> to the config.toml/CLI-overrides path without touching `validate_config`'s many direct +> unit-test callers or `run_train_job`'s redundant internal `validate_config` call. +> Regression tests added in `tests/test_config.py` cover top-level and nested typos, the +> axis/`centers_init` allowances, `[meta]` skipping, and both the file and CLI-overrides +> paths, plus the two real `configs/*.toml` fixtures that don't already fail +> `migrate_config` for unrelated reasons. Everything below this point describes the +> pre-fix state and is kept for historical context. + +**Severity: High. Effort: Small. Risk if unfixed: wasted GPU-days on a run that did not +use the setting you thought it did.** + +**Location:** `giant/config.py:456` (`_deep_merge`), `giant/config.py:628` +(`merge_cli_overrides`), `giant/config.py:661` (`validate_config`). + +### What the code does today + +`merge_cli_overrides` resolves configuration as `DEFAULT_CONFIG → TOML file → CLI +overrides`, deep-merging at each step. `_deep_merge` accepts any key: + +```python +def _deep_merge(base: dict, override: dict) -> dict: + result = dict(base) + for k, v in override.items(): + if isinstance(v, dict) and isinstance(result.get(k), dict): + result[k] = _deep_merge(result[k], v) + else: + result[k] = v # <-- unknown k is accepted and stored + return result +``` + +`migrate_config` explicitly preserves unrecognised sections +(`giant/config.py:612-615`: *"Anything else in the original dict (unrecognized top-level +sections) carries through untouched rather than being silently dropped"*). + +`validate_config` (`giant/config.py:661`) exists and is *good* — it catches cross-block +contradictions with genuinely excellent, actionable error messages (e.g. "`router.type = +'pdg'` builds its own training-vocab-scoped embedding, incompatible with +`conditioning.particle.type = 'physical'`"). But by design it only inspects keys it knows +about. It cannot detect a key that should not exist. + +### Why it's a problem + +A config file containing + +```toml +[stage1_model] +n_res_block = 12 # typo: should be n_res_blocks +``` + +merges cleanly, validates cleanly, trains to completion, and reports success — having +built a 6-block model. The stray `n_res_block` key is faithfully written back into the +run's saved `config.toml` by `save_config`, so the artifact of record also claims the +typo was a real setting. Nothing anywhere in the pipeline will ever say otherwise. + +The blast radius scales with the config surface. v0.3.0 introduced a large nested schema +(`conditioning` × `stage1_model` × `stage2_model` × per-generator sub-tables × router +sub-tables), which is a great deal of surface for typos, and the v0.3.0 workflow is +explicitly *"a sequence of architecture comparisons"* (`config.py`, the `wandb` default +comment). A comparison in which one arm silently ignored its distinguishing setting is +worse than no comparison — it produces a confident, wrong conclusion. + +There is a partial mitigation already: `warn_if_git_hash_mismatch` and +`warn_if_checkpoint_config_mismatch` warn about *provenance* drift. Neither looks at key +names. + +### How to verify + +```bash +uv run python - <<'EOF' +import tempfile, pathlib +from giant import config as gc + +p = pathlib.Path(tempfile.mkdtemp()) / "config.toml" +p.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n") +cfg = gc.merge_cli_overrides(gc.DEFAULT_CONFIG, p, {}) +gc.validate_config(cfg) # passes +print("typo key survived :", cfg["stage1_model"]["n_res_block"]) +print("real key unchanged:", cfg["stage1_model"]["n_res_blocks"]) +EOF +``` + +### Recommended fix + +Add strict key validation against the known schema. Two options, in order of preference: + +1. **Preferred — fold into Issue 1.** If each config block becomes a dataclass with a + `from_dict` that raises on unknown keys, this issue is solved for free and cannot + regress. Raise a `ValueError` in the established house style, naming the offending + key, its section, and the closest valid key by edit distance + (`difflib.get_close_matches`) — that last part matters, because "unknown key + `n_res_block`; did you mean `n_res_blocks`?" is the message that actually saves the + run. + +2. **Standalone fallback.** Add `validate_config_keys(cfg)` to `giant/config.py` that + walks `cfg` alongside `DEFAULT_CONFIG` and raises on any key not present in the + defaults tree. Call it from `merge_cli_overrides` right before returning, so every + entry point gets it. + +Either way, the following need explicit allowances, since they legitimately carry keys +not in `DEFAULT_CONFIG`: + +- **`meta`** — holds `config_version`, `git_hash`, and run metadata from + `build_run_meta`, none of which are in the defaults tree. +- **Composed-router axis keys** — `stage{1,2}_model.router.axis{i}_{field}` are + deliberately dynamic and deliberately absent from `DEFAULT_CONFIG` (see the comment at + the end of the `router` block in `config.py`, and `_parse_composed_axes` in + `network.py:499`). Validate these against the `axis(\d+)_(\w+)` pattern plus the known + per-axis field names rather than against a fixed key list. +- **`stage2_model.n_sec.legacy_owner`** — injected by `_migrate_legacy_model_config`, not + a user-facing key. It appears in checkpoint `model_config` dicts, not in `config.toml`, + so it should not reach this validator; confirm that before assuming. + +### Scope guard + +Strict validation must apply to the **`config.toml` path only**, not to loading old +checkpoints. A v0.2 checkpoint's `model_config` is migrated by a different function +(`network._migrate_legacy_model_config`, see Issue 6) and must keep loading. Adding +strictness that rejects historical checkpoints would break `giant predict` / `giant +rollout` against every model trained so far. + +--- + +## Issue 3 — `cli.py:train()` is a 58-parameter, 510-line fat controller + +> **Status: Fixed, commit `2bfb1ab` on `v0.3.0-stage2-autoregressive`.** `giant/config.py` +> now declares `FlagSpec` (a frozen dataclass: `name`, `paths` — one or more dotted config +> paths, `precedence`) and a `FLAG_SPECS` table covering every flag `train`/`new-run` map +> into the config tree, plus `overrides_from_flags(values: dict[str, object]) -> dict`, +> which applies specs in ascending precedence order (so a more-specific flag overwrites a +> shared/shorthand one written earlier at the same path) — one mechanism replacing the +> three different ad hoc "more specific wins" patterns identified below +> (`.update()`-call-order, `setdefault(...)[...] =` overwrite-order, and +> `{**shared, **specific}` merge). `train()`'s ~140-line override-building block +> (`cli.py:656-790` pre-fix) is now a single flat `flag_values` dict (mostly enum `.value` +> unwrapping) plus one call to `overrides_from_flags`; `new_run()`'s near-verbatim copy +> (`cli.py:915-983` pre-fix) collapsed the same way, reusing the identical table. Router +> overrides (`_router_cli_overrides`, unchanged, still shared by both commands) feed into +> the table as a single pre-aggregated `router_config` entry mapped only to +> `stage1_model.router` — the stage1-only asymmetry noted below is preserved exactly, with +> a regression test. No flag was added, removed, or renamed, and no precedence semantics +> changed: `giant train --help`/`giant new-run --help` are byte-identical before and after, +> verified by diffing both. Everything below this point describes the pre-fix state and is +> kept for historical context. + +**Severity: High. Effort: Medium.** + +**Location:** `giant/cli.py:329-839`. + +### What the code does today + +Measured on `55332db`: + +| Command | Parameters | Lines | +|---|---|---| +| `train` | **58** | **510** | +| `new_run` | 31 | 188 | +| `predict` | 9 | 387 | +| `rollout` | 14 | 237 | + +Of `train`'s 510 lines, roughly 430 are the Typer signature (one `Annotated[...]` block +per flag) and roughly 80 are hand-written translation from flags into the nested +overrides dict. That translation follows the same shape five times over — for `train`, +`stage1_model`, `stage2_model`, `conditioning`, and the WGAN sub-tables: + +```python +cli_stage1_model: dict[str, object] = { + k: v for k, v in { + "hidden_dim": hidden_dim, + "n_res_blocks": n_blocks, + "dropout": dropout, + }.items() if v is not None +} +cli_stage1_model.update( + {k: v for k, v in { + "hidden_dim": stage1_hidden_dim, + "n_res_blocks": stage1_n_res_blocks, + "dropout": stage1_dropout, + }.items() if v is not None} +) +``` + +On top of that sits genuinely intricate precedence logic, correctly implemented but +expressed imperatively: + +- `--mode` sets `generator` on **both** stages; `--stage1-generator` / `--stage2-generator` + then override a single stage (`cli.py:713-720`). +- `--hidden-dim` / `--n-blocks` / `--dropout` are stage-1-only backward-compatible + shorthands; `--stage1-*` wins when both are given (`cli.py:676-700`). +- `--n-critic` / `--gp-weight` / `--noise-dim` / `--critic-lr` fan out to + `stage{1,2}_model.wgan.*`, with `--stage{1,2}-*` variants overriding per stage — built + by merging a shared dict under a stage-specific one (`cli.py:722-753`). +- `--emb-dim` and `--conditioning` each set **two** places + (`conditioning.particle.*` and `conditioning.material.*`) (`cli.py:702-710`). + +### Why it's a problem + +1. **The mapping is data, written as code.** Flag → dotted config path → fan-out rule is + a table. Written as 80 lines of copy-adapted dict comprehensions, correctness has to + be re-verified by reading, every time a flag is added. `giant/config.py:446` already + provides `_set_path(d, "stage1_model.wgan.n_critic", v)` — the primitive a table-driven + version needs. +2. **Adding a flag touches three places at once** (signature, the right comprehension, + the right precedence rule) with no mechanism forcing them to stay consistent. v0.3.0's + design explicitly anticipates more per-stage flags. +3. **It is the least-tested code in the repo** (Issue 4), because unit-testing a + 58-parameter Typer command means going through `CliRunner` with argv strings. +4. It obscures the ~15 lines that actually do something (`cli.py:825-839`): resolve the + out-dir, echo two lines, call `run_train_job`. + +### Recommended fix + +Extract the mapping into a pure, table-driven function in `giant/config.py` (not +`cli.py`), so it is importable and directly testable without Typer: + +```python +@dataclass(frozen=True) +class FlagSpec: + """One CLI flag's mapping into the config tree.""" + name: str # "stage1_hidden_dim" + paths: tuple[str, ...] # ("stage1_model.hidden_dim",) — >1 means fan-out + precedence: int = 0 # higher wins; per-stage flags outrank shorthands + +FLAG_SPECS: tuple[FlagSpec, ...] = ( + FlagSpec("hidden_dim", ("stage1_model.hidden_dim",), precedence=0), + FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1), + FlagSpec("mode", ("stage1_model.generator", "stage2_model.generator")), + FlagSpec("n_critic", ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic")), + FlagSpec("emb_dim", ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")), + ... +) + +def overrides_from_flags(values: dict[str, object]) -> dict: + """Build the nested overrides dict from {flag_name: value}, dropping + None (= flag not given) and applying precedence.""" +``` + +`train()` then becomes: parse `--batch-size auto`, collect `locals()`-style flag values, +call `overrides_from_flags`, `merge_cli_overrides`, `validate_config`, resolve out-dir, +call `run_train_job`. The 430-line Typer signature stays — that is irreducible, it *is* +the user interface — but it becomes the only bulk in the function. + +The precedence table also becomes self-documenting, which today requires reading three +separate inline comment blocks to reconstruct. + +### Scope guard + +- **Do not remove or rename any flag.** The backward-compatible shorthands + (`--hidden-dim`, `--n-blocks`, `--dropout`, `--mode`) exist because they predate the + per-stage flags and are in people's shell history and job scripts. Their surprising + stage-1-only scoping is documented behaviour, not a bug to fix. +- **Do not change precedence semantics.** Reproduce the current rules exactly, then + test them (Issue 4). Any intentional change is a separate commit. +- Preserve the explanatory comments at `cli.py:676-680` and `cli.py:713-719` — they + record *why* the precedence is the way it is. + +--- + +## Issue 4 — `cli.py` sits at 35.8 % coverage and holds untested override-precedence logic + +> **Status: Fixed for the override-precedence logic (Issue 3's scope); the +> `predict`/`rollout` inference-bootstrap portion described below is still open — that is +> Issue 5, deliberately not attempted here.** Commit `2bfb1ab` on +> `v0.3.0-stage2-autoregressive` adds direct, `CliRunner`-free unit tests for +> `overrides_from_flags` in `tests/test_config.py` — one per precedence rule, including +> several with previously **zero** coverage: both legs of `--stage{1,2}-generator` +> overriding `--mode` (only the stage1 leg had a test before), all three stage1 +> shorthand-vs-`--stage1-*` pairs (previously only `hidden_dim`), all four WGAN knobs' +> shared-vs-per-stage precedence (previously only `n_critic`/`gp_weight`), the +> `--emb-dim`/`--conditioning` dual-axis fan-out (previously untested via `train` at all), +> and an explicit regression test that `router_config` only ever writes +> `stage1_model.router`. `tests/test_cli_train_overrides.py` keeps its original 4 +> `CliRunner` smoke tests unmodified, plus 4 new ones covering the `--batch-size auto` +> parse-error and success paths and all three `out_dir` resolution branches +> (`--out`/`--resume`/default) — previously entirely uncovered. The percentage barely moves +> (36 % → 35 %), because the extraction *deleted* more statements from `cli.py` (539→478) +> than the new tests cover elsewhere in the file, but covered statements rose in absolute +> terms (157→169) and every override-precedence line the original issue called out by +> number is now covered. Coverage of `cli.py`'s other listed gap — the `predict`/ +> `rollout` inference bootstrap (`cli.py:1104-1423`/`1528-1702` in the pre-fix numbering) — +> is unchanged, since fixing that requires the `load_for_inference` extraction described in +> Issue 5, which is explicitly out of scope for this fix (done separately, if at all). +> Everything below this point describes the pre-fix state and is kept for historical +> context. + +**Severity: High. Effort: Medium.** + +**Location:** `giant/cli.py` (1865 lines). + +### What the code does today + +Per-module line coverage from `coverage.xml`: + +| Module | Coverage | +|---|---| +| `rollout.py` | 99.6 % | +| `analysis/catalog.py` | 99.6 % | +| `training/trainers.py` | 99.5 % | +| `model/network.py` | 98.6 % | +| `pipeline.py` | 97.1 % | +| `data/transforms.py` | 96.2 % | +| `config.py` | 91.8 % | +| **`cli.py`** | **35.8 %** | + +Project total is 81.8 %. `cli.py` is not merely the lowest — it is a **60-point outlier** +against an otherwise uniformly high standard, and it is simultaneously the largest module +in the repo. + +Existing CLI tests (`tests/test_cli_train_overrides.py`, `tests/test_cli_new_run.py`, +`tests/test_cli_predict.py`) are the right idea and well-written; they simply cover a +fraction of the surface. + +### Why it's a problem + +The uncovered code is not boilerplate. It includes: + +- The entire flag → config override translation and its precedence rules (Issue 3). Every + "`--stage2-hidden-dim` beats `--hidden-dim`", "`--stage2-generator` beats `--mode`", + "`--n-critic` fans out to both stages unless `--stage1-n-critic` is given" rule is + currently **unverified by any test**. +- The complete inference bootstrap in `predict` and `rollout` — checkpoint loading, + normalizer reconstruction, conditioning-axis resolution, top-N map wiring (Issue 5). + This is the code where a divergence between the two commands produces silent + train/inference skew rather than a crash. +- The `--batch-size auto` parsing and estimation paths, and the six distinct + "retrain with the current code" checkpoint-compatibility guards. + +A failure here is expensive in a way a unit-test failure is not: a mis-scoped flag means +a multi-hour GPU run trains a model that differs from the one the experiment log claims. +Given `CLAUDE.md` describes v0.3.0 as a sequence of architecture comparisons, this is the +worst possible place for silent divergence. + +### How to verify + +```bash +uv run pytest -q --cov=giant --cov-report=term-missing:skip-covered 2>&1 | grep "cli.py" +``` + +### Recommended fix + +Coverage here is a *consequence* of Issue 3, not an independent goal. Sequence the work: + +1. **First, extract the testable logic** (Issue 3): `overrides_from_flags` in + `config.py`, and `load_for_inference` in a new `giant/checkpoint_io.py` (Issue 5). +2. **Then test the extracted functions directly** — no Typer, no `CliRunner`, no + filesystem. One test per precedence rule: + + ```python + def test_stage1_hidden_dim_beats_hidden_dim_shorthand(): ... + def test_mode_sets_both_stage_generators(): ... + def test_stage2_generator_overrides_mode_for_stage2_only(): ... + def test_n_critic_fans_out_to_both_wgan_subtables(): ... + def test_stage1_n_critic_overrides_shared_n_critic_for_stage1_only(): ... + def test_emb_dim_sets_both_conditioning_axes(): ... + ``` + +3. **Keep a thin layer of `CliRunner` smoke tests** over `train --help`, `predict`, + `rollout` to catch signature/wiring breakage, but do not try to reach high coverage + through the CLI surface — that is slow and brittle. + +Target: get `cli.py` to roughly the repo norm by *moving code out of it*, not by writing +elaborate CLI-invocation tests. A `cli.py` that is genuinely just argument declaration +plus delegation can sit at modest coverage without concern, because there will be +nothing in it left to get wrong. + +### Scope guard + +Do not add `# pragma: no cover` to close the gap. The gap is a real signal and the fix is +extraction. + +--- + +## Issue 5 — The inference bootstrap is duplicated verbatim between `predict` and `rollout` + +> **Status: Fixed.** `giant/checkpoint_io.py` now holds a single +> `load_for_inference(checkpoint, device, command_name, weights="raw", require_stage2=True)` +> plus an `InferenceContext` dataclass (stage1/stage2 models, all three normalizers, both +> vocab maps, both top-N maps, both conditioning axes, `k_max`, both stages' ddpm step +> counts, `other_policy`, the raw `model_config`, and `epoch`/`best_val_loss` for +> `rollout`'s YAML sidecar) — exactly the design this issue proposed, verified against the +> current (not `55332db`-era) code before writing it. `predict`/`rollout` in `cli.py` +> each shrink to one `try/except CheckpointCompatibilityError` call plus a block of +> `ctx.` unpacks; `cli.py` lost `_conditioning_axes`, `_stage_cfg`, `_ddpm_steps`, +> `_particle_type_other_policy`, `_load_pdg_topn_map`/`_load_mat_topn_map`, and +> `_load_model_weights` entirely (net ~180 lines off `cli.py`). The independent third copy +> in `giant/analysis/router_gating.py` was deleted in favour of a lazy +> `from giant.checkpoint_io import conditioning_axes` inside `load_router`'s existing +> lazy-import block, preserving that module's "polars/numpy only at module scope" +> contract (`checkpoint_io.py` imports torch eagerly, so it must never be imported at +> `router_gating.py` module scope). Two guard orderings from the two commands' drifted +> copies were consolidated into one (`warn_if_checkpoint_config_mismatch` now always +> runs right after the existence guards, and `predict`'s `--batch-size auto` estimate now +> reads `ctx.model_config` after the checkpoint loads rather than before) — both are +> console-output-order changes only, no error text or model behavior changed, confirmed by +> diffing `giant predict --help`/`giant rollout --help` byte-for-byte before and after (no +> flag touched) and re-reading both rewritten command bodies field-by-field against the +> original. `require_stage2` exists as a real parameter, exercised by a new test, even +> though both current callers pass the default `True`. New `tests/test_checkpoint_io.py` +> (17 tests: happy path, every guard individually with exact message-text assertions, the +> `require_stage2=False`/inactive-stage2 path, `conditioning_axes`/`stage_cfg` directly) +> plus thin `CliRunner` smoke tests in `tests/test_cli_predict.py` and the new +> `tests/test_cli_rollout.py` confirming `CheckpointCompatibilityError` actually surfaces +> as `typer.Exit(1)` through the CLI — previously this entire code path had zero test +> coverage. `uv run pytest -q` (803 passed, up from 784), `ruff check`, `ruff format +> --check`, and `ty check` all clean. Everything below this point describes the pre-fix +> state and is kept for historical context. + +**Severity: High. Effort: Small. Risk if unfixed: silent train/inference skew.** + +**Location:** `giant/cli.py:1121-1201` (`predict`), `giant/cli.py:1534-1593` (`rollout`), +plus a third partial copy at `giant/analysis/router_gating.py:86-120`. + +### What the code does today + +`predict` and `rollout` each contain ~65 lines that are near-identical line-for-line: + +1. `ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)` +2. Guard: no `model_config` → "retrain with the current code" → `Exit(1)` +3. Guard: no `sec_decoder` → same +4. Guard: no `normalizer.sec_phys` → same +5. `_conditioning_axes(model_cfg)` +6. `_load_pdg_topn_map(ckpt)` / `_load_mat_topn_map(ckpt)` +7. Guard: `particle_conditioning == "onehot"` but no `pdg_topn_map` → `Exit(1)` +8. Guard: `material_conditioning == "onehot"` but no `mat_topn_map` → `Exit(1)` +9. `pdg_map` / `mat_map` key-type coercion (`{int(k): v ...}` / `{str(k): v ...}`) +10. Three `Normalizer.from_dict(ckpt["normalizer"][...])` calls +11. `build_models(model_cfg)`, unpack `stage1` / `stage2` +12. Guard: either stage `None` → "needs both" → `Exit(1)` +13. `_load_model_weights(...)`, `.to(device).eval()` on both +14. `typer.echo(f"loaded checkpoint: ...")` +15. `gconfig.warn_if_checkpoint_config_mismatch(checkpoint)` + +Additionally, `_conditioning_axes` is **duplicated verbatim into a second module**. +`giant/analysis/router_gating.py:69` carries a byte-for-byte copy of +`giant/cli.py:74`, with the docstring openly acknowledging it: *"Mirrors +`giant.cli._conditioning_axes`."* + +```bash +grep -rn "def _conditioning_axes" giant/ +# giant/cli.py:74 +# giant/analysis/router_gating.py:69 +``` + +### Why it's a problem + +This is not cosmetic duplication. This code decides **how input features are assembled at +inference time** — which conditioning mode is used, which top-N vocabulary maps are +applied, which normalizer statistics are restored. If `predict` and `rollout` ever +diverge on any of those, the result is not an exception. It is two commands producing +subtly different physics from the same checkpoint, with no error and no warning. That +class of bug is found by noticing that a plot looks wrong, weeks later. + +The duplication is also **entirely untested** — it lives in the 35.8 %-covered region of +`cli.py` (Issue 4). And it is *already* growing: the third copy in `router_gating.py` +shows the pattern spreading into a package that is otherwise carefully isolated. + +The self-aware "Mirrors `giant.cli._conditioning_axes`" comment is the tell. When a +developer documents a copy rather than removing it, the missing abstraction has been +identified but not yet built. + +### Recommended fix + +Create `giant/checkpoint_io.py` (name it whatever fits; the point is that it is **not** +`cli.py`) exposing one function and one result object: + +```python +@dataclass(frozen=True) +class InferenceContext: + """Everything needed to run a trained checkpoint forward, resolved once.""" + stage1: nn.Module + stage2: nn.Module + 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 + particle_conditioning: str + material_conditioning: str + k_max: int + stage1_ddpm_steps: int + stage2_ddpm_steps: int + other_policy: str + model_config: dict + +def load_for_inference( + checkpoint: Path, + device: torch.device, + weights: str = "raw", + require_stage2: bool = True, +) -> InferenceContext: ... +``` + +Raise a dedicated `CheckpointCompatibilityError` (carrying the current, genuinely helpful +message text) instead of calling `typer.echo` + `typer.Exit`; the CLI catches it and does +the echo/exit. That keeps the module free of Typer and makes it unit-testable. + +`predict` and `rollout` each shrink by ~60 lines to a single call. `router_gating.py` +drops its `_conditioning_axes` copy and imports from the new module — note that +`router_gating` deliberately imports torch **lazily inside the function** +(`router_gating.py:88`) to keep the analysis workers' "polars/numpy only" contract; the +new module must be imported the same way there, and `checkpoint_io.py` must not be +imported at `analysis/` module scope. + +### Scope guard + +- Preserve every error message verbatim. They are specific and actionable + ("checkpoint's `conditioning.particle.type='onehot'` but has no `pdg_topn_map` — retrain + with the current code"), and users have seen them before. +- Preserve the `require_stage2` distinction: `giant predict` and `giant rollout` both need + both stages today, but `stage1_model.active = false` / `stage2_model.active = false` are + real config options (`DEFAULT_CONFIG`), so the parameter should exist rather than + hard-coding the requirement. +- Do not fold `giant/analysis/router_gating.py`'s router-specific loading + (`load_router`, which returns `None` for non-MoE checkpoints) into the shared function. + It has different semantics — absence is a normal outcome there, not an error. + +--- + +## Issue 6 — Two independent v0.2→v0.3 migration surfaces encode the same knowledge + +**Severity: Medium. Effort: Medium.** + +**Location:** `giant/config.py:514` (`migrate_config`) and +`giant/model/network.py:1446` (`_migrate_legacy_model_config`). + +### What the code does today + +v0.3.0 broke the config format: the v0.2 single `[train]` + `[model]` layout became +`[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]`. That break has to be +absorbed in two different places, and today it is absorbed by two unrelated functions: + +- **`config.migrate_config`** translates a v0.2 **`config.toml`** on load. +- **`network._migrate_legacy_model_config`** translates a v0.2 **checkpoint's + `model_config` dict** on `build_models`. + +The split is acknowledged as deferred work in `migrate_config`'s own docstring: + +> *"Operates on the config.toml shape. A checkpoint's `model_config` dict (which +> additionally carries n_sec_head ownership and needs `network.build_models`'s +> cooperation) is a separate migration surface, deferred to the network.py refactor."* + +Both functions independently encode the same translation facts: + +| v0.2 concept | v0.3 destination | in `migrate_config` | in `_migrate_legacy_model_config` | +|---|---|---|---| +| `train.mode` / `model.mode` | `stage{1,2}_model.generator` | ✓ | ✓ | +| `model.emb_dim` | `conditioning.{particle,material}.emb_dim` | ✓ | ✓ | +| `model.conditioning` | `conditioning.{particle,material}.type` | ✓ | ✓ | +| `model.n_blocks` | `stage{1,2}_model.n_res_blocks` | ✓ | ✓ | +| `model.noise_dim` | `stage{1,2}_model.wgan.noise_dim` | ✓ | ✓ | +| conditioning MLP depth was always 2 | `conditioning.*.n_layers = 2` | ✓ | ✓ | +| `router.expert_hidden_dim` set → hard error | — | ✓ | ✓ (near-identical message) | +| n_sec head lived on stage 1 | `stage2_model.n_sec.legacy_owner` | ✗ | ✓ (only here) | + +The `expert_hidden_dim` rejection is the clearest symptom: the same policy, the same +reasoning, two hand-maintained copies of a ~10-line error message +(`config.py:581-591` and `network.py:1475-1489`). + +### Why it's a problem + +1. **Drift.** A future correction to the v0.2 interpretation must be applied to both. + Applying it to one produces a checkpoint that loads with different architecture than + its own config file describes. +2. **`legacy_owner` leaks into the builder.** Because the checkpoint migration is + downstream of the config migration rather than sharing it, `build_models` has to carry + legacy-specific branching at four sites (`network.py:1612`, `1642`, `1663`, `1688`): + + ```python + legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner") + n_sec_head_k_max = s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None + ... + build_n_sec_head=legacy_owner != "stage1", + ``` + + A builder for the current architecture should not need to know where v0.2 put its + `n_sec` head. +3. **The retention policy is undeclared.** `tests/legacy/network_v02_snapshot.py` is + **1009 lines** of frozen v0.2 network code kept purely so migration can be tested + against it. That is a substantial maintenance surface with no stated expiry. + +### Recommended fix + +1. **Extract the shared translation into one table.** A single + `_V02_TO_V03_TRANSLATION` mapping (old dotted key → tuple of new dotted keys) plus one + set of "v0.2 architectural facts with no config key" constants, consumed by both + functions. `config.py:475-511` already has the beginnings of this + (`_V02_TRAIN_PASSTHROUGH`, `_V02_MODEL_TO_BOTH_STAGES`, + `_V02_TRAIN_TO_BOTH_STAGES_WGAN`) — extend that pattern and share it. Put the shared + table in a neutral module (e.g. `giant/_migration.py`) so neither `config.py` nor + `network.py` has to import the other. +2. **Move `legacy_owner` handling out of `build_models`.** The migration function should + emit a `model_config` that `build_models` can consume without legacy branching — for + example by emitting an explicit `stage2_model.n_sec.owner` key that the *current* + schema also carries (with value `"stage2"` for new runs), so the builder reads one key + with two valid values rather than a nullable legacy sentinel. +3. **Write down the retention policy.** Add to `CLAUDE.md`: which v0.2 checkpoints must + remain loadable, until when, and what triggers dropping the shim and the 1009-line + snapshot. Without that, nobody will ever feel authorised to delete it. + +### Scope guard + +- **v0.2 checkpoints must keep loading** until the policy above says otherwise. There are + trained models on `/ceph` that predate v0.3.0 and analysis runs referencing them. +- Keep `tests/legacy/network_v02_snapshot.py` as a **frozen snapshot** — do not "clean it + up", reformat it, or make it share code with current `network.py`. Its entire value is + that it is an independent, unchanging record of what v0.2 did. If ruff/ty complain about + it, exclude it rather than edit it (it is already excluded from coverage via + `pyproject.toml`'s `omit = ["*/legacy/*"]`). + +--- + +## Issue 7 — Positional tuple contracts between the data, model and training layers + +**Severity: Medium. Effort: Small.** + +**Location:** `giant/data/transforms.py:863-890` (`build_features`), +`giant/data/dataset.py:229-237` (`StreamingStepsDataset` yield), +`giant/training/trainers.py:574-582` (unpack), `giant/cli.py:1233` (unpack). + +### What the code does today + +`build_features` returns a **bare 9-tuple**: + +```python +) -> tuple[ + np.ndarray, 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, + sec_type_idx) arrays. ...""" +``` + +The element *meanings* live only in the docstring. Call sites re-derive them positionally: + +```python +# giant/cli.py:1233 — predict +cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(...) +``` + +`StreamingStepsDataset` yields a **7-tuple** of tensors (`dataset.py:229`), unpacked +positionally in the trainer: + +```python +# giant/training/trainers.py:574 +( + cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx, sec_type_idx, +) = _batch_to_device(batch, device) +``` + +and again, with a *different arity*, in the WGAN path: + +```python +# giant/training/trainers.py:737 +cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors +``` + +`_batch_to_device` is typed `(batch: tuple, device) -> tuple` — no element information at +all. + +### Why it's a problem + +1. **Invisible to the type checker.** `ty check` passes on all of it. Reordering two + `np.ndarray` elements — or inserting a new one in the middle — type-checks cleanly and + produces a model trained on scrambled features. Several of these arrays are + shape-compatible with each other (`n_sec`, `proc_idx` and `sec_type_idx` are all + integer arrays), so a swap may not even produce a shape error at runtime. +2. **This is the hottest contract in the codebase.** It crosses three layer boundaries + (`data` → `training`, `data` → `cli`) and is the mechanism by which the model receives + physics. A silent corruption here is the most expensive possible failure: it does not + crash, it just trains something wrong. +3. **The nine-underscore unpack** at `cli.py:1233` is unreadable and does not survive any + change to the tuple. +4. **Two different batch arities** for the same conceptual batch (7 in `_compute`, 5 in + `_stage2_real_and_fake`) means the reader must track which slice is in play. + +### Recommended fix + +Convert both to `NamedTuple`. Zero runtime cost, full `ty` visibility, tuple-unpacking +still works so migration is incremental: + +```python +class StepFeatures(NamedTuple): + """Output of build_features. Field order is load-bearing for existing + positional unpacking — append only, never insert or reorder.""" + 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 + +class StepBatch(NamedTuple): + 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 +``` + +Then `cli.py:1233` becomes `feats = build_features(...)` / `feats.cond_cont`, and +`_batch_to_device` gets the real signature `(batch: StepBatch, device) -> StepBatch`. + +Precedent exists in this codebase: `RolloutSummary` (`rollout.py:276`, a `TypedDict`) and +`SetupStageResult` (`pipeline.py:39`, a dataclass) already do this correctly. + +Note that `StreamingStepsDataset` passes batches through a PyTorch `DataLoader` with +`batch_size=None` and `num_workers > 0`, so the batch type must survive worker-process +pickling and `pin_memory`. `NamedTuple` does — it is a plain tuple subclass, and +`pin_memory` recurses into tuples — but **verify this with a real multi-worker run**, not +just the test suite, since the tests may run single-process. + +### Scope guard + +Do not reorder any existing fields while converting. The whole point of the change is to +make future reordering safe; performing one during the conversion, when nothing yet +protects against it, is the single riskiest version of this change. + +--- + +## Issue 8 — `network.py` is 1745 lines holding three distinct modules + +**Severity: Medium. Effort: Small (mechanical).** + +**Location:** `giant/model/network.py`. + +### What the code does today + +One file contains six unrelated concerns: + +| Lines (approx.) | Concern | +|---|---| +| 28-213 | Primitives: `SinusoidalEmbedding`, `cat_col_layout`, `_make_axis_mlp`, `ConditionEncoder`, `ContextAdapter`, `ResBlock` | +| 215-575 | **The entire router subsystem**: `Router` base, `register_router`/`build_router` registry, `EnergyRouter`, `PdgRouter`, `ProcessRouter`, `ComposedRouter`, `build_composed_router`, `_parse_composed_axes`, `_check_router_conditioning_compat`, `_build_router_from_cfg` | +| 577-727 | Trunks: `ExpertTrunk`, `_route_forward`, `Trunk`, `MonolithicTrunk`, `RoutedTrunk`, `build_trunk` | +| 729-881 | **History encoders**: `HistoryEncoder`, `MarkovHistory`, `_CausalAttnBlock`, `AttentionHistory` | +| 883-1443 | Top-level models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel` | +| 1446-1568 | **Legacy migration**: `_migrate_legacy_model_config`, `migrate_legacy_state_dict` | +| 1570-1745 | Builders: `build_models`, `build_critics` | + +The v0.3.0 refactor described in `CLAUDE.md` — *"`network.py` is refactored from ten +permutation classes into composable parts (encoder × trunk × objective)"* — clearly +landed and is a genuine improvement. The composition is visible and clean. The file +simply was not split to match. + +### Why it's a problem + +- The seams are already there in the class layout; the file just does not honour them. A + reader looking for `AttentionHistory` has no reason to expect it 800 lines into a file + whose name suggests "the network". +- The router subsystem alone is ~370 lines with its own registry, its own config parsing, + and its own compatibility validation. It is a subsystem, not a section. +- Legacy migration (Issue 6) sitting in the same file as the current builders is exactly + what lets `legacy_owner` bleed into `build_models`. +- It is a merge-conflict magnet on a repo with parallel feature branches + (`condor-gpu-train-rollout`, `v0.3.0-stage2-autoregressive`). + +### Recommended fix + +Split along the existing seams into `giant/model/`: + +``` +giant/model/ + layers.py # SinusoidalEmbedding, ResBlock, ContextAdapter, _make_axis_mlp + encoders.py # ConditionEncoder, cat_col_layout + routers.py # Router base + registry + all 4 router types + config parsing + trunks.py # Trunk, MonolithicTrunk, RoutedTrunk, ExpertTrunk, build_trunk + history.py # HistoryEncoder, MarkovHistory, AttentionHistory, _CausalAttnBlock + models.py # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel + builders.py # build_models, build_critics + _legacy.py # _migrate_legacy_model_config, migrate_legacy_state_dict + network.py # re-export shim: `from giant.model.layers import *` etc. +``` + +Keep `network.py` as a **re-export shim** so no import site outside `giant/model/` has to +change in the same commit. `network.py` is imported by `cli.py`, `sample.py`, +`trainers.py`, `pipeline.py`, `analysis/router_gating.py`, and heavily by the test suite — +`grep -rn "from giant.model.network import" giant/ tests/ | wc -l` before starting, to +size the blast radius if you do decide to update call sites (21 sites on `55332db`). + +This is a pure file-move refactor: no logic changes, and the 725-test suite plus `ty +check` is a strong safety net for it. + +### Scope guard + +Do it as its own commit, containing **only** moves and imports. Do not combine with Issue 1 +or Issue 6, both of which change behaviour in this file — a diff mixing moves with logic +changes is effectively unreviewable. + +--- + +## Issue 9 — `scripts` is published as a top-level distribution package + +**Severity: Medium. Effort: Small.** + +**Location:** `pyproject.toml`: + +```toml +[project.scripts] +giant = "giant.cli:app" +dwarf = "scripts.dwarf:app" + +[tool.hatch.build.targets.wheel] +packages = ["giant", "scripts"] +``` + +### What the code does today + +The repo's tooling CLI (`dwarf`) lives in a directory called `scripts/`, which is declared +as a wheel package and referenced by the console-script entry point as `scripts.dwarf:app`. +Installing `giant` therefore creates a top-level importable module named **`scripts`** in +`site-packages`. + +`scripts/` holds ten real modules: `dwarf.py`, `bump_dataset_version.py`, +`create_root_files.py`, `geometry_oracle.py`, `hparam_scan.py`, `migrate_geant_steps.py`, +`profile_analysis_costs.py`, `steps_to_parquet.py`, `steps_to_parquet_parallel.py`, +`warm_setup_cache.py`. + +### Why it's a problem + +`scripts` is one of the most generic names possible in the Python ecosystem. Consequences: + +1. **Collision.** Any other installed distribution that also ships a top-level `scripts` + package silently shadows or is shadowed by this one, depending on `sys.path` order. + The failure mode is an `ImportError` or — worse — importing someone else's `dwarf`-less + `scripts` and getting `AttributeError` at CLI startup. +2. **Environment-order fragility.** On the portal machines (`/work/lbogner`, shared with + other users, per `CLAUDE.md`), a stray `scripts/` directory in the CWD shadows the + installed package, because CWD precedes `site-packages` on `sys.path`. Running `dwarf` + from a directory that happens to contain a `scripts/` folder can break in a confusing + way. +3. **`scripts/` reads as "not part of the product"**, yet it is installed, has an entry + point, is covered by `[tool.coverage.run] source`, and has a full test suite + (`tests/test_dwarf.py` at 81.7 % coverage, plus `test_steps_to_parquet*.py`, + `test_create_root_files.py`, `test_bump_dataset_version.py`). Its name misrepresents + its status. + +### Recommended fix + +Move it under the `giant` namespace, where it cannot collide: + +``` +giant/tools/ # was scripts/ + dwarf.py + ... +``` + +```toml +[project.scripts] +giant = "giant.cli:app" +dwarf = "giant.tools.dwarf:app" + +[tool.hatch.build.targets.wheel] +packages = ["giant"] +``` + +Then update: + +- `[tool.coverage.run] source = ["giant"]` (drop the now-redundant `"scripts"`). +- All `from scripts.X import Y` in `tests/` (`grep -rn "scripts\." tests/`). +- Any `python scripts/foo.py` invocations in docs, `README.md`, `CLAUDE.md`, HTCondor + submit files, or shell history/job scripts on the portal machines. + +The `dwarf` **command name does not change**, so anything invoking `dwarf ...` keeps +working — only Python-level imports and direct `python scripts/...` paths move. + +### Scope guard + +Check `giant/analysis/condor.py` and any generated HTCondor submit descriptions for +hard-coded `scripts/` paths before moving. Jobs submitted to the cluster may reference the +path as it exists on `/work` or `/ceph`, and a rename that lands mid-flight breaks queued +jobs. Search for the literal string `scripts/` across the repo, not just Python imports. + +--- + +## Issue 10 — `torch.load(weights_only=False)`: checkpoints are arbitrary pickles + +**Severity: Low (given the threat model). Effort: Medium.** + +**Location:** `giant/cli.py:1122`, `giant/cli.py:1535`, +`giant/analysis/router_gating.py:93`, `giant/training/loop.py:153`. + +### What the code does today + +All four checkpoint loads pass `weights_only=False`: + +```python +ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) +``` + +This is *necessary* today, not careless: the checkpoint dict carries non-tensor objects +beyond raw state dicts — `model_config`, `pdg_map` / `mat_map`, `pdg_topn_map` / +`mat_topn_map`, and the `normalizer` sub-dict, assembled in +`giant/training/checkpoint.py:19` (`build_checkpoint`) from the `extras` argument. +`weights_only=True` would reject them. + +### Why it's worth recording + +`weights_only=False` means loading a checkpoint executes arbitrary pickle opcodes. The +practical threat model here is mild — checkpoints are produced by this codebase and live +on `/ceph/lbogner` — but the portal machines and `/ceph` are explicitly **shared with +other users** (`CLAUDE.md`, "Compute environment"), and `giant predict --checkpoint ` +will happily load a path someone else wrote. It is also the kind of thing that becomes a +blocker later, if a model is ever shared outside the group or published alongside a paper. + +Secondary practical cost: a pickled `model_config` cannot be inspected without importing +torch and unpickling. `dwarf status`-style tooling, or a human answering "what conditioning +mode was this trained with?", has to load the whole checkpoint. + +### Recommended fix + +Split the checkpoint into a tensor part and a JSON sidecar: + +- `ckpt.pt` — state dicts only, loadable with `weights_only=True`. +- `ckpt.meta.json` — `model_config`, vocab maps, top-N maps, normalizer statistics, + `epoch` / `global_step` / `best_val_loss`. + +This has real secondary benefits: the sidecar is greppable and diffable, so comparing two +runs' architectures becomes `diff` rather than a Python session, and it composes well +with the config-provenance machinery already in `config.py` +(`warn_if_checkpoint_config_mismatch`, `git_hash`). + +Note that `giant/data/setup_cache.py` already establishes the "JSON sidecar next to the +data" pattern for exactly this kind of non-tensor state — follow its conventions. + +### Scope guard + +This is a **format change**, so it needs a compatibility path: `load_for_inference` +(Issue 5) reads the sidecar if present and falls back to the pickled keys if not. Do not +attempt this before Issue 5 lands — with three separate copies of the loading code, a +format change means three separate compatibility paths. After Issue 5, it is one. + +Given the mild threat model, this is correctly the **lowest-priority** item here. It is +recorded so the decision is deliberate rather than accidental. + +--- + +## Issue 11 — Minor items + +### 11a — `echo=print` threaded through the pipeline instead of `logging` + +**Location:** `giant/pipeline.py:305` (`run_train_job(..., echo=print)`), +`giant/pipeline.py:86` (`run_setup_stage(..., echo=...)`); 52 `typer.echo` calls in +`cli.py`. + +Passing the output function as a parameter is *better* than hard-coding `typer.echo` deep +in the pipeline — it keeps `pipeline.py` free of Typer and makes output capturable in +tests. But at this scale it has costs: no severity levels (the `--num-workers` quota +warning at `pipeline.py:329` and routine progress output are indistinguishable), no +timestamps on long training runs, no way to route to a file without wrapping, and the +parameter has to be threaded through every function that might print. + +Suggested: standard `logging` with a Typer/rich handler configured once in `cli.py`. +Module-level `log = logging.getLogger(__name__)` replaces the threaded `echo`. Keep +`typer.echo` for genuine CLI output (the `device:` / `out_dir:` banners, error messages +before `Exit(1)`) — those are interface, not logs. + +Low priority; the current approach is defensible. Worth doing opportunistically if +`pipeline.py` is being touched anyway. + +### 11b — `giant.particles` imports from `giant.data.loader` + +**Location:** `giant/particles.py:30`. + +```python +if TYPE_CHECKING: + from giant.data.loader import TopNMap +``` + +`giant/particles.py` is a physics-domain module (PDG code decoding, mass/charge lookup). +`giant/data/loader.py` is data I/O. A domain module depending on an I/O module is a +layering inversion. + +Mitigating factors, which is why this is minor: the import is `TYPE_CHECKING`-guarded, so +there is **no runtime cycle**, and `ty check` passes. It is a design smell, not a defect. + +Suggested: move `TopNMap` to a neutral location — `giant/constants.py` or a new +`giant/types.py` — so both `particles.py` and `data/loader.py` depend on it rather than on +each other. Do this only if `TopNMap` is being touched for other reasons; it is not worth +a standalone commit. + +--- + +## Recommended sequence + +The issues are interdependent. This order minimises rework: + +1. **Issue 5** — extract `load_for_inference` into `giant/checkpoint_io.py`. Small, + self-contained, removes the highest-risk duplication, and unblocks Issue 10. +2. **Issue 3** — extract the flag→config mapping table into `config.py`. +3. **Issue 4** — test the two extracted units directly. Issues 3 and 5 make this cheap; + attempting it first means testing through `CliRunner`, which is slow and brittle. +4. **Issue 1 + Issue 2 together** — typed config dataclasses. One migration solves both, + and doing them separately means touching the same 132 call sites twice. +5. **Issue 6** — unify the migration surfaces (easier once config is typed). +6. **Issue 8** — split `network.py`. Pure file moves; do it as an isolated commit, after + the logic changes in Issues 1 and 6 have settled, to keep both diffs readable. +7. **Issue 7** — `NamedTuple` the feature/batch contracts. +8. **Issue 9** — move `scripts/` → `giant/tools/`. Independent of everything else; can be + done at any point. +9. **Issue 10** — checkpoint format split, if the threat model or a publication makes it + worthwhile. +10. **Issue 11** — opportunistically. + +## General guidance for whoever picks this up + +- **Run the full suite on every change.** `uv run pytest -q` takes 18 seconds. There is no + reason to batch up unverified changes. +- **Run all three checks before committing:** `uv run ruff check .`, + `uv run ruff format .`, `uv run ty check .`. CI enforces all three. +- **Do not change model behaviour while refactoring.** Several issues touch code that + determines what network gets built (Issues 1, 6) or what features it receives (Issues 5, + 7). A refactor that also changes a default silently invalidates every prior benchmark in + `/home/lars/knowledge-base/experiments/`. If a behaviour change is warranted, it is a + separate commit with a separate message saying so. +- **Preserve the comments.** This codebase's docstrings explain reasoning, trade-offs, and + known limitations to an unusually high standard. When code moves, the comments move with + it. When code is deleted, check whether the comment records a decision that still needs + recording elsewhere. +- **`giant/analysis/` is the template.** It is declarative where the rest of the codebase + is imperative (`PlotSpec` registry), explicit about its own contracts (the + `compute_partial` / `finalize` split, `chunkable=False`, the "polars/numpy only on + workers" boundary), and honest about its one exception (`router_gating.py`'s module + docstring explains precisely why it is allowed to import torch and why that is still + safe). When deciding how a refactored `config.py` or `cli.py` should look, look there. diff --git a/pyproject.toml b/pyproject.toml index e7fad25..1d7538b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "giant" -version = "0.2.0" +version = "0.3.0" description = "Geant4 step-function surrogate via conditional flow matching" readme = "README.md" requires-python = ">=3.12" @@ -23,6 +23,7 @@ cuda = [ ] dev = [ "pytest>=8,<10", + "pytest-cov>=5,<8", "ruff>=0.15,<1", "ty>=0.0.50,<0.1", "giant[convert,analysis,geometry,wandb]", @@ -51,6 +52,19 @@ analysis = [ giant = "giant.cli:app" dwarf = "scripts.dwarf:app" +[tool.ruff] +line-length = 120 + +[tool.coverage.run] +source = ["giant", "scripts"] +omit = ["*/legacy/*"] + +[tool.coverage.report] +exclude_also = [ + "if TYPE_CHECKING:", + "raise NotImplementedError", +] + [build-system] requires = ["hatchling"] build-backend = "hatchling.build" diff --git a/scripts/bump_dataset_version.py b/scripts/bump_dataset_version.py index 216748a..27ecabe 100644 --- a/scripts/bump_dataset_version.py +++ b/scripts/bump_dataset_version.py @@ -82,9 +82,7 @@ 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 @@ -142,9 +140,7 @@ 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}") @@ -154,9 +150,7 @@ 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 @@ -212,9 +206,7 @@ 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 @@ -243,9 +235,7 @@ def _referenced_root_count( 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 @@ -342,11 +332,7 @@ 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 @@ -359,25 +345,18 @@ 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 @@ -425,9 +404,7 @@ 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 @@ -526,13 +503,8 @@ 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") @@ -552,9 +524,7 @@ 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] = [] @@ -569,9 +539,7 @@ def plan_create_manifest( 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 @@ -641,9 +609,7 @@ 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:") diff --git a/scripts/check_migration_v02_v03.py b/scripts/check_migration_v02_v03.py new file mode 100644 index 0000000..3386a91 --- /dev/null +++ b/scripts/check_migration_v02_v03.py @@ -0,0 +1,161 @@ +"""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 scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt + uv run python scripts/check_migration_v02_v03.py /ceph/lbogner/.../best.pt --ema + uv run python scripts/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()) diff --git a/scripts/create_root_files.py b/scripts/create_root_files.py index d960c78..66ba11a 100644 --- a/scripts/create_root_files.py +++ b/scripts/create_root_files.py @@ -64,9 +64,7 @@ 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 @@ -94,9 +92,7 @@ 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: @@ -124,9 +120,7 @@ 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: @@ -147,10 +141,7 @@ 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) @@ -174,20 +165,12 @@ def run_job( job, False, None, - f"expected exactly one .root output in {workdir}, found {len(produced)}: " - f"{[p.name for p in produced]}", + f"expected exactly one .root output in {workdir}, found {len(produced)}: {[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, @@ -279,14 +262,7 @@ 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}") diff --git a/scripts/dwarf.py b/scripts/dwarf.py index 200b1ec..674e252 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -80,8 +80,7 @@ def convert( typer.Option( "--output", "-o", - help="Output Parquet file (default: .parquet). Only valid " - "with a single input file and --jobs 1.", + help="Output Parquet file (default: .parquet). Only valid with a single input file and --jobs 1.", ), ] = None, batch_size: Annotated[ @@ -91,9 +90,7 @@ 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, @@ -129,15 +126,11 @@ 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: @@ -150,10 +143,7 @@ def convert( ) total_orphaned += n_orphaned if total_orphaned: - typer.echo( - f"\n{total_orphaned} orphaned child track(s) dropped across " - f"{len(root_files)} file(s)." - ) + typer.echo(f"\n{total_orphaned} orphaned child track(s) dropped across {len(root_files)} file(s).") return if output is not None: @@ -176,9 +166,7 @@ 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( @@ -190,8 +178,7 @@ 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: @@ -202,12 +189,8 @@ 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)"), @@ -220,12 +203,8 @@ 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( @@ -243,12 +222,8 @@ 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)"), @@ -261,12 +236,8 @@ 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( @@ -283,9 +254,7 @@ 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)) @@ -293,9 +262,7 @@ 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( @@ -306,9 +273,7 @@ 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, @@ -316,9 +281,7 @@ 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") @@ -333,22 +296,15 @@ def create_manifest( typer.Option( "--pool", metavar="DETECTOR", - help="Detector name; combined with --type and --root to form " - "/pools//.manifest", + help="Detector name; combined with --type and --root to form /pools//.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"), @@ -368,9 +324,7 @@ 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( @@ -382,15 +336,9 @@ 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//, 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//, e.g. gen1")], energy_gev: Annotated[ float | None, typer.Option( @@ -401,20 +349,12 @@ 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.""" @@ -441,9 +381,7 @@ 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, @@ -456,9 +394,7 @@ 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"), @@ -504,9 +440,7 @@ 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, @@ -518,49 +452,52 @@ def warm_cache( ] = 0.1, seed: Annotated[ int, - typer.Option( - "--seed", "-s", help="Must match the `giant train` run(s) to warm for" - ), + typer.Option("--seed", "-s", help="Must match the `giant train` run(s) to warm for"), ] = 0, - conditioning: Annotated[ + particle_conditioning: Annotated[ Conditioning, typer.Option( - "--conditioning", help="Must match the `giant train` run(s) to warm for" + "--particle-conditioning", + help="Must match the `giant train` run(s)' conditioning.particle.type to warm for", + ), + ] = Conditioning.physical, + material_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.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/--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/--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. """ run_warm_setup_cache( data=str(data), val_fraction=val_fraction, seed=seed, - conditioning=conditioning.value, + particle_conditioning=particle_conditioning.value, + material_conditioning=material_conditioning.value, router_enabled=router, router_type=router_type, n_experts=n_experts, diff --git a/scripts/geometry_oracle.py b/scripts/geometry_oracle.py index e5a36fb..762d288 100644 --- a/scripts/geometry_oracle.py +++ b/scripts/geometry_oracle.py @@ -38,9 +38,7 @@ 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}") diff --git a/scripts/hparam_scan.py b/scripts/hparam_scan.py index dfb5ca2..0c64f6c 100644 --- a/scripts/hparam_scan.py +++ b/scripts/hparam_scan.py @@ -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,9 +154,7 @@ 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, { @@ -171,11 +169,6 @@ def run_hparam_scan( "wall_time_s": round(wall_time_s, 1), }, ) - print( - f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} " - f"({wall_time_s:.1f}s)" - ) + print(f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} ({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") diff --git a/scripts/migrate_geant_steps.py b/scripts/migrate_geant_steps.py index 27e96c9..447e0dc 100644 --- a/scripts/migrate_geant_steps.py +++ b/scripts/migrate_geant_steps.py @@ -50,12 +50,8 @@ PREDICTED_RE = re.compile( r"^(?P[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P\d+)" r"_predicted(?P_local)?\.parquet$" ) -SHARD_RE = re.compile( - r"^(?P[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P\d+)\.(?Proot|parquet)$" -) -LEGACY_PREDICTED_RE = re.compile( - r"^pbwo4_10000events_hits_predicted(?P_local)?\.parquet$" -) +SHARD_RE = re.compile(r"^(?P[a-z0-9]+(?:_[a-z0-9]+)*)_10k_(?P\d+)\.(?Proot|parquet)$") +LEGACY_PREDICTED_RE = re.compile(r"^pbwo4_10000events_hits_predicted(?P_local)?\.parquet$") LEGACY_RE = re.compile(r"^pbwo4_10000events_hits\.(?Proot|parquet)$") @@ -110,24 +106,9 @@ 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 @@ -135,19 +116,9 @@ 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 @@ -165,20 +136,9 @@ 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: diff --git a/scripts/profile_analysis_costs.py b/scripts/profile_analysis_costs.py index de97a2b..60dda89 100644 --- a/scripts/profile_analysis_costs.py +++ b/scripts/profile_analysis_costs.py @@ -94,9 +94,7 @@ 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), @@ -163,9 +161,7 @@ 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, diff --git a/scripts/steps_to_parquet_parallel.py b/scripts/steps_to_parquet_parallel.py index b2ee41e..1817c36 100644 --- a/scripts/steps_to_parquet_parallel.py +++ b/scripts/steps_to_parquet_parallel.py @@ -49,9 +49,7 @@ 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////.root (relative to *dataset_root*) to processed/////.parquet. @@ -66,12 +64,7 @@ def resolve_destination( 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////.root " f"under {dataset_root} (got relative path: {rel})" @@ -216,13 +209,7 @@ 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.") diff --git a/scripts/warm_setup_cache.py b/scripts/warm_setup_cache.py index 3eeebb3..1feed1c 100644 --- a/scripts/warm_setup_cache.py +++ b/scripts/warm_setup_cache.py @@ -9,6 +9,8 @@ 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 @@ -16,7 +18,8 @@ def run_warm_setup_cache( data: str, val_fraction: float = 0.1, seed: int = 0, - conditioning: str = "physical", + particle_conditioning: str = "physical", + material_conditioning: str = "physical", router_enabled: bool = False, router_type: str = "energy", n_experts: int = 4, @@ -25,9 +28,11 @@ def run_warm_setup_cache( ) -> None: """Populate (or refresh) the setup cache sidecar for `data`. - `val_fraction`/`seed`/`conditioning` select the normalizer cache entry + `val_fraction`/`seed`/`particle_conditioning`/`material_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. + `giant train` invocation will use so it hits this warmed entry. The two + conditioning axes are independent and may differ. `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 @@ -39,12 +44,30 @@ 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, - conditioning=conditioning, - router_cfg=router_cfg, + cfg=cfg, cache_setup=True, rebuild_setup_cache=rebuild, echo=echo, diff --git a/tests/legacy/__init__.py b/tests/legacy/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/legacy/network_v02_snapshot.py b/tests/legacy/network_v02_snapshot.py new file mode 100644 index 0000000..a58992d --- /dev/null +++ b/tests/legacy/network_v02_snapshot.py @@ -0,0 +1,1009 @@ +"""Frozen snapshot of `giant/model/network.py` as it stood at the v0.3.0 +"step 1" commit (eb6dd27), i.e. the last commit before the step-2 +composable-parts decomposition. + +This is a deliberate verbatim copy, not an import of the live module — the +whole point is that this file's classes keep behaving exactly as v0.2 did +even after `giant/model/network.py` itself is rewritten, so +`tests/test_migration_v02_v03.py` has a stable "old" side to diff the new +`build_models`/`Stage1Model`/`Stage2OneShot` against (the bit-identical +acceptance test). Do not edit this file to track future +`network.py` changes — it exists specifically to stop tracking them. +""" + +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, + COND_DIM_BASE, + EMB_DIM, + K_MAX, + MATERIAL_PHYS_DIM, + PARTICLE_PHYS_DIM, + SEC_DIM, + X_DIM, +) + + +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) + + +class ConditionEncoder(nn.Module): + """Fuses continuous conditioning with particle/material identity. + + Two mutually exclusive ways to turn (pdg, material) identity into the + two `emb_dim`-wide vectors concatenated with the base continuous + conditioning before the fusion MLP: + - "embedding": a learned `nn.Embedding` lookup table per axis, indexed + by `cond_cat`'s dense training-vocab index. Memorizes the training + menu; the original Phase-2 design. + - "physical": a small MLP per axis, mapping the axis's raw physical + properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see + giant.data.transforms.build_features) to an `emb_dim`-wide vector — + a drop-in replacement for the embedding lookup, computable for any + PDG code / material name rather than only ones seen in training. + Both modes produce the same `in_dim = COND_DIM_BASE + 2*emb_dim` for the + fusion MLP, so only how the two vectors are produced differs. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + cont_dim: int = COND_DIM, + emb_dim: int = 16, + out_dim: int = 128, + conditioning: str = "embedding", + ) -> None: + super().__init__() + if conditioning not in ("embedding", "physical"): + raise ValueError(f"unknown conditioning mode {conditioning!r}") + self.conditioning = conditioning + if conditioning == "embedding": + self.pdg_emb = nn.Embedding(pdg_vocab, emb_dim) + self.mat_emb = nn.Embedding(mat_vocab, emb_dim) + else: + self.particle_mlp = nn.Sequential( + nn.Linear(PARTICLE_PHYS_DIM, emb_dim), + nn.SiLU(), + nn.Linear(emb_dim, emb_dim), + ) + self.material_mlp = nn.Sequential( + nn.Linear(MATERIAL_PHYS_DIM, emb_dim), + nn.SiLU(), + nn.Linear(emb_dim, emb_dim), + ) + in_dim = COND_DIM_BASE + 2 * emb_dim + self.mlp = nn.Sequential( + nn.Linear(in_dim, out_dim), + nn.SiLU(), + nn.Linear(out_dim, out_dim), + ) + + def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + if self.conditioning == "embedding": + pdg_e = self.pdg_emb(cond_cat[:, 0]) + mat_e = self.mat_emb(cond_cat[:, 1]) + else: + particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM] + material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :] + pdg_e = self.particle_mlp(particle_phys) + mat_e = self.material_mlp(material_phys) + x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1) + return self.mlp(x) + + +class ResBlock(nn.Module): + def __init__(self, dim: int, cond_dim: int, dropout: float = 0.1) -> 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 + + +class DenoisingMLP(nn.Module): + """Stage-1 model: predicts the 9D primary post-step vector field + n_sec logits. + + The n_sec head runs on the condition encoding only (no diffusion noise), + so it can be called at inference time independently via `predict_n_sec`. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + time_dim: int = 64, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + dropout: float = 0.1, + k_max: int = K_MAX, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + merged_cond_dim = time_dim + cond_out_dim + self.input_proj = nn.Linear(x_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)]) + self.out_proj = nn.Linear(hidden_dim, x_dim) + # Predicts n_sec as classification over {0, 1, ..., k_max}. + # Applied to the condition encoding (not the diffused latent). + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, k_max + 1), + ) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) # (B, time_dim) + c_emb = self.cond_enc(cond_cont, cond_cat) # (B, cond_out_dim) + cond = torch.cat([t_emb, c_emb], dim=-1) + x = self.input_proj(x_t) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + 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.""" + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + +class SecondaryConditionEncoder(nn.Module): + """Encodes pre-step conditioning + Stage-1 output for the secondary decoder.""" + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + emb_dim: int = 16, + cond_out_dim: int = 128, + stage1_dim: int = X_DIM, + stage1_proj_dim: int = 64, + out_dim: int = 128, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.base = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.stage1_proj = nn.Linear(stage1_dim, stage1_proj_dim) + fused_dim = cond_out_dim + stage1_proj_dim + self.fuse = nn.Sequential( + nn.Linear(fused_dim, out_dim), + nn.SiLU(), + ) + + def forward( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + base = self.base(cond_cont, cond_cat) # (B, cond_out_dim) + s1 = self.stage1_proj(stage1_out).tanh() # (B, stage1_proj_dim) + return self.fuse(torch.cat([base, s1], dim=-1)) # (B, out_dim) + + +class SecondaryDecoder(nn.Module): + """Stage-2 model: predicts vector field over K_MAX secondary slots simultaneously. + + Each slot encodes (stick_break_logit, local_dir_3D, log_mass, charge) for + one secondary ordered by descending energy — mass/charge are the + secondary's predicted physical identity, regressed directly against real + physics targets (see giant.data.transforms.encode_secondaries), used + as-is with no snapping to a discrete PDG code. Padded slots are masked + from loss. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + time_dim: int = 64, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + merged_cond_dim = time_dim + cond_out_dim + self.input_proj = nn.Linear(sec_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)]) + self.out_proj = nn.Linear(hidden_dim, sec_dim) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out) + cond = torch.cat([t_emb, c_emb], dim=-1) + x = self.input_proj(x_t) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + +class WGANGenerator(nn.Module): + """Stage-1 WGAN-GP generator: single forward pass, no diffusion/flow time. + + Same `ConditionEncoder` + `ResBlock` trunk as `DenoisingMLP`, but the + input is a noise vector `z` (not a diffused/interpolated `x_t`) and the + ResBlocks condition on the condition encoding alone (no time embedding to + concatenate) — see `giant/model/wgan.py` for the adversarial losses, and + `giant.sample.sample_wgan` for single-pass sampling. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + noise_dim: int = 64, + dropout: float = 0.1, + k_max: int = K_MAX, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.noise_dim = noise_dim + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(noise_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)]) + self.out_proj = nn.Linear(hidden_dim, x_dim) + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, hidden_dim // 2), + nn.SiLU(), + nn.Linear(hidden_dim // 2, k_max + 1), + ) + + def forward( + self, + z: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat) + x = self.input_proj(z) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + 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.""" + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + +class Critic(nn.Module): + """Stage-1 WGAN-GP critic: scalar realism score, own `ConditionEncoder`. + + Kept structurally parallel to `WGANGenerator` (own condition encoder — + separate weights from the generator's, standard GAN practice) but has no + n_sec head: n_sec is never adversarial, it stays a plain classifier on + the generator side. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(x_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_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, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat) + h = self.input_proj(x) + for block in self.blocks: + h = block(h, cond) + return self.out_proj(self.out_norm(h)).squeeze(-1) + + +class WGANSecondaryGenerator(nn.Module): + """Stage-2 WGAN-GP generator: single forward pass over all K_MAX slots. + + Mirrors `SecondaryDecoder` minus the time embedding, the same way + `WGANGenerator` mirrors `DenoisingMLP` — takes noise `z` instead of `x_t`, + conditions on `SecondaryConditionEncoder`'s output alone. + """ + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + noise_dim: int = 64, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.noise_dim = noise_dim + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(noise_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_blocks)]) + self.out_proj = nn.Linear(hidden_dim, sec_dim) + + def forward( + self, + z: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat, stage1_out) + x = self.input_proj(z) + for block in self.blocks: + x = block(x, cond) + return self.out_proj(x) + + +class SecondaryCritic(nn.Module): + """Stage-2 WGAN-GP critic: scalar realism score over the flattened 90D slots.""" + + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + hidden_dim: int = 256, + n_blocks: int = 6, + emb_dim: int = 16, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + self.input_proj = nn.Linear(sec_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, cond_out_dim, dropout=dropout) for _ in range(n_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, + ) -> torch.Tensor: + cond = self.cond_enc(cond_cont, cond_cat, stage1_out) + h = self.input_proj(x) + for block in self.blocks: + h = block(h, cond) + return self.out_proj(self.out_norm(h)).squeeze(-1) + + +class Router(nn.Module): + """Contract for a pluggable mixture-of-experts routing axis.""" + + 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: + raise NotImplementedError + + def combine_weights(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + 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: + 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 = 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: + return torch.zeros((), device=cond_cont.device) + + def entropy_loss(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: + 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]: + 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: + 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: + return lo + (hi - lo) * torch.sigmoid(raw) + + +def _inverse_bounded_interp(value: float, lo: float, hi: float) -> float: + p = min(max((value - lo) / (hi - lo), 1e-6), 1 - 1e-6) + return math.log(p / (1 - p)) + + +@register_router("energy") +class EnergyRouter(Router): + 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): + 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): + 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): + 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: + 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) + + +class ExpertTrunk(nn.Module): + def __init__( + self, + in_dim: int, + hidden_dim: int, + n_blocks: int, + merged_cond_dim: int, + dropout: float = 0.1, + ) -> None: + super().__init__() + self.input_proj = nn.Linear(in_dim, hidden_dim) + self.blocks = nn.ModuleList([ResBlock(hidden_dim, merged_cond_dim, dropout=dropout) for _ in range(n_blocks)]) + self.out_proj = nn.Linear(hidden_dim, in_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: + if training: + weights = router.combine_weights(cond_cont, cond_cat) # (B, n_experts) + out = torch.zeros_like(x) + 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 = torch.zeros_like(x) + for i, expert in enumerate(experts): + mask = idx == i + if mask.any(): + out[mask] = expert(x[mask], cond[mask]) + return out + + +class RoutedDenoisingMLP(nn.Module): + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + router: Router, + expert_hidden_dim: int = 128, + expert_n_blocks: int = 3, + emb_dim: int = EMB_DIM, + time_dim: int = 64, + cond_out_dim: int = 128, + x_dim: int = X_DIM, + dropout: float = 0.1, + k_max: int = K_MAX, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.router = router + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = ConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + merged_cond_dim = time_dim + cond_out_dim + self.experts = nn.ModuleList( + [ + ExpertTrunk( + x_dim, + expert_hidden_dim, + expert_n_blocks, + merged_cond_dim, + dropout=dropout, + ) + for _ in range(router.n_experts) + ] + ) + self.n_sec_head = nn.Sequential( + nn.Linear(cond_out_dim, cond_out_dim), + nn.SiLU(), + nn.Linear(cond_out_dim, k_max + 1), + ) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat) + cond = torch.cat([t_emb, c_emb], dim=-1) + return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training) + + def predict_n_sec( + self, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + ) -> torch.Tensor: + c_emb = self.cond_enc(cond_cont, cond_cat) + return self.n_sec_head(c_emb) + + +class RoutedSecondaryDecoder(nn.Module): + def __init__( + self, + pdg_vocab: int, + mat_vocab: int, + router: Router, + expert_hidden_dim: int = 128, + expert_n_blocks: int = 3, + emb_dim: int = EMB_DIM, + time_dim: int = 64, + cond_out_dim: int = 128, + stage1_proj_dim: int = 64, + sec_dim: int = SEC_DIM, + dropout: float = 0.1, + conditioning: str = "embedding", + ) -> None: + super().__init__() + self.router = router + self.time_emb = SinusoidalEmbedding(time_dim) + self.cond_enc = SecondaryConditionEncoder( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + emb_dim=emb_dim, + cond_out_dim=cond_out_dim, + stage1_proj_dim=stage1_proj_dim, + out_dim=cond_out_dim, + conditioning=conditioning, + ) + merged_cond_dim = time_dim + cond_out_dim + self.experts = nn.ModuleList( + [ + ExpertTrunk( + sec_dim, + expert_hidden_dim, + expert_n_blocks, + merged_cond_dim, + dropout=dropout, + ) + for _ in range(router.n_experts) + ] + ) + + def forward( + self, + x_t: torch.Tensor, + t: torch.Tensor, + cond_cont: torch.Tensor, + cond_cat: torch.Tensor, + stage1_out: torch.Tensor, + ) -> torch.Tensor: + t_emb = self.time_emb(t) + c_emb = self.cond_enc(cond_cont, cond_cat, stage1_out) + cond = torch.cat([t_emb, c_emb], dim=-1) + return _route_forward(self.experts, self.router, x_t, cond, cond_cont, cond_cat, self.training) + + +_STAGE1_MODEL_KEYS = { + "pdg_vocab", + "mat_vocab", + "hidden_dim", + "n_blocks", + "emb_dim", + "dropout", + "k_max", + "conditioning", +} +_SEC_DECODER_MODEL_KEYS = { + "pdg_vocab", + "mat_vocab", + "hidden_dim", + "n_blocks", + "emb_dim", + "dropout", + "conditioning", +} +_WGAN_GENERATOR_MODEL_KEYS = _STAGE1_MODEL_KEYS | {"noise_dim"} +_WGAN_SEC_GENERATOR_MODEL_KEYS = _SEC_DECODER_MODEL_KEYS | {"noise_dim"} +_CRITIC_MODEL_KEYS = _STAGE1_MODEL_KEYS - {"k_max"} + + +_AXIS_KEY_RE = re.compile(r"^axis(\d+)_(.+)$") + + +def _parse_composed_axes(router_cfg: dict) -> list[dict]: + 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))] + + +_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process") + + +def _check_router_conditioning_compat(router_types: list[str], conditioning: str) -> None: + bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES)) + if bad and conditioning == "physical": + raise ValueError( + f"router type(s) {bad} always use a training-vocab PDG embedding, " + "which is incompatible with conditioning='physical' (whose whole " + "point is generalizing beyond that vocab) — pick a different " + "router type (e.g. 'energy') or use conditioning='embedding'." + ) + + +def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding") -> Router: + 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], 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"]], 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 + + +def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: + if model_config.get("mode") == "wgan": + stage1 = WGANGenerator(**{k: v for k, v in model_config.items() if k in _WGAN_GENERATOR_MODEL_KEYS}) + sec_decoder = WGANSecondaryGenerator( + **{k: v for k, v in model_config.items() if k in _WGAN_SEC_GENERATOR_MODEL_KEYS} + ) + return stage1, sec_decoder + + router_cfg = model_config.get("router") + if router_cfg and router_cfg.get("enabled"): + pdg_vocab = model_config["pdg_vocab"] + mat_vocab = model_config["mat_vocab"] + shared = dict( + pdg_vocab=pdg_vocab, + mat_vocab=mat_vocab, + expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128), + expert_n_blocks=model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3), + emb_dim=model_config.get("emb_dim", EMB_DIM), + dropout=model_config.get("dropout", 0.1), + conditioning=model_config.get("conditioning", "embedding"), + ) + conditioning = shared["conditioning"] + stage1 = RoutedDenoisingMLP( + router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning), + k_max=model_config.get("k_max", K_MAX), + **shared, + ) + sec_decoder = RoutedSecondaryDecoder( + router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning), + **shared, + ) + return stage1, sec_decoder + + stage1 = DenoisingMLP(**{k: v for k, v in model_config.items() if k in _STAGE1_MODEL_KEYS}) + sec_decoder = SecondaryDecoder(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}) + return stage1, sec_decoder + + +def build_critics(model_config: dict) -> tuple[nn.Module, nn.Module]: + critic = Critic(**{k: v for k, v in model_config.items() if k in _CRITIC_MODEL_KEYS}) + sec_critic = SecondaryCritic(**{k: v for k, v in model_config.items() if k in _SEC_DECODER_MODEL_KEYS}) + return critic, sec_critic diff --git a/tests/test_bump_dataset_version.py b/tests/test_bump_dataset_version.py index 3f169a5..83f7dd9 100644 --- a/tests/test_bump_dataset_version.py +++ b/tests/test_bump_dataset_version.py @@ -22,9 +22,7 @@ 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", @@ -56,9 +54,7 @@ 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 @@ -66,18 +62,14 @@ 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"] @@ -91,9 +83,7 @@ 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 @@ -125,9 +115,7 @@ 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 @@ -164,15 +152,7 @@ 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" @@ -194,15 +174,7 @@ 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" @@ -231,15 +203,7 @@ 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" @@ -254,15 +218,7 @@ 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" @@ -278,15 +234,7 @@ 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" @@ -303,15 +251,7 @@ 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" @@ -328,15 +268,7 @@ 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" @@ -358,24 +290,8 @@ 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) @@ -419,9 +335,7 @@ 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 @@ -435,9 +349,7 @@ 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" diff --git a/tests/test_catalog.py b/tests/test_catalog.py index c8a47f0..92ddb24 100644 --- a/tests/test_catalog.py +++ b/tests/test_catalog.py @@ -13,9 +13,7 @@ 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") @@ -142,10 +140,7 @@ 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 diff --git a/tests/test_checkpoint_io.py b/tests/test_checkpoint_io.py new file mode 100644 index 0000000..8107945 --- /dev/null +++ b/tests/test_checkpoint_io.py @@ -0,0 +1,236 @@ +"""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 {} + 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_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") == {} diff --git a/tests/test_cli_new_run.py b/tests/test_cli_new_run.py index 54788c4..a4f9450 100644 --- a/tests/test_cli_new_run.py +++ b/tests/test_cli_new_run.py @@ -37,13 +37,14 @@ def test_writes_config_with_overrides_applied(tmp_path: Path): with open(config_path, "rb") as f: cfg = tomllib.load(f) - assert cfg["train"]["mode"] == "ddpm" + assert cfg["stage1_model"]["generator"] == "ddpm" + assert cfg["stage2_model"]["generator"] == "ddpm" assert cfg["train"]["lr"] == 0.0005 - assert cfg["model"]["hidden_dim"] == 128 - assert cfg["model"]["n_blocks"] == 4 + assert cfg["stage1_model"]["hidden_dim"] == 128 + assert cfg["stage1_model"]["n_res_blocks"] == 4 # untouched defaults still present assert cfg["train"]["epochs"] == 100 - assert "router" in cfg["model"] + assert "router" in cfg["stage1_model"] assert str(out_dir) in result.output assert "" in result.output @@ -100,9 +101,7 @@ 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() diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py index a6f74fb..7c8e237 100644 --- a/tests/test_cli_predict.py +++ b/tests/test_cli_predict.py @@ -1,13 +1,18 @@ 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 @@ -116,9 +121,7 @@ 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" @@ -133,9 +136,7 @@ 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). @@ -150,9 +151,24 @@ 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 diff --git a/tests/test_cli_rollout.py b/tests/test_cli_rollout.py new file mode 100644 index 0000000..b5bd2a6 --- /dev/null +++ b/tests/test_cli_rollout.py @@ -0,0 +1,33 @@ +"""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 diff --git a/tests/test_cli_train_overrides.py b/tests/test_cli_train_overrides.py new file mode 100644 index 0000000..7df258f --- /dev/null +++ b/tests/test_cli_train_overrides.py @@ -0,0 +1,176 @@ +"""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): + cfg = _invoke_and_capture_cfg( + monkeypatch, + tmp_path, + [ + "--stage2-decoder", + "one_shot", + "--stage2-k-max", + "8", + "--stage2-hidden-dim", + "32", + "--stage2-context-dim", + "16", + "--stage2-stage1-context", + "sampled", + ], + ) + 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 + assert cfg["stage2_model"]["stage1_context"] == "sampled" + # 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 diff --git a/tests/test_condor.py b/tests/test_condor.py index 4027a50..02912ed 100644 --- a/tests/test_condor.py +++ b/tests/test_condor.py @@ -62,9 +62,7 @@ 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, @@ -224,16 +222,12 @@ def test_write_submit_description(tmp_path: Path): assert "--chunk" in body and "--run-dir" in body -def test_write_submit_requires_synced_venv( - tmp_path: Path, monkeypatch: pytest.MonkeyPatch -): +def test_write_submit_requires_synced_venv(tmp_path: Path, monkeypatch: pytest.MonkeyPatch): run_dir = _prep(_write_inputs(tmp_path)) cfg = SubmitConfig(run_dir=run_dir, accounting_group="cms", repo_dir=tmp_path) # No `giant` next to the (fake) active interpreter, so this falls through # to repo_dir/.venv/bin/giant, which _write_inputs/_prep also didn't create. - monkeypatch.setattr( - sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python") - ) + monkeypatch.setattr(sys, "executable", str(tmp_path / "not-a-venv" / "bin" / "python")) with pytest.raises(FileNotFoundError, match="uv sync"): write_submit(cfg) @@ -241,9 +235,7 @@ def test_write_submit_requires_synced_venv( 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 @@ -253,9 +245,7 @@ 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] = {} @@ -272,9 +262,7 @@ 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) @@ -297,16 +285,9 @@ 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 diff --git a/tests/test_config.py b/tests/test_config.py index 30f6e41..58d4d1c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,92 +1,1023 @@ from datetime import datetime +from pathlib import Path + +import pytest from giant import config as gconfig +_CONFIGS_DIR = Path(__file__).resolve().parents[1] / "configs" -def _write_config(path, git_hash): - path.write_text( - f""" -[train] -epochs = 5 -[model] -hidden_dim = 64 +def _write_toml(path, git_hash=None, extra=""): + meta = f'\n[meta]\ngit_hash = "{git_hash}"\n' if git_hash is not None else "" + path.write_text(extra + meta) -[meta] -git_hash = "{git_hash}" -""" + +# --------------------------------------------------------------------------- +# Conditioning enum +# --------------------------------------------------------------------------- + + +def test_conditioning_enum_has_onehot(): + assert gconfig.Conditioning.onehot == "onehot" + assert {c.value for c in gconfig.Conditioning} == { + "physical", + "embedding", + "onehot", + } + + +# --------------------------------------------------------------------------- +# Config dataclasses (issues.md Issue 1) +# --------------------------------------------------------------------------- + + +def test_giant_config_to_dict_matches_default_config(): + """DEFAULT_CONFIG is generated from GiantConfig().to_dict() (not + hand-maintained), so the two cannot structurally drift apart — but this + pins the *equality* too, catching e.g. a stray in-place mutation of + DEFAULT_CONFIG added elsewhere after import.""" + assert gconfig.GiantConfig().to_dict() == gconfig.DEFAULT_CONFIG + + +@pytest.mark.parametrize( + "cls", + [ + gconfig.ConditioningAxisConfig, + gconfig.ConditioningConfig, + gconfig.FlowConfig, + gconfig.DdpmConfig, + gconfig.Stage1WganConfig, + gconfig.Stage2WganConfig, + gconfig.RouterConfig, + gconfig.Stage2RouterConfig, + gconfig.NSecConfig, + gconfig.ParticleTypeConfig, + gconfig.AutoregressiveConfig, + gconfig.Stage1ModelConfig, + gconfig.Stage2ModelConfig, + gconfig.TrainConfig, + gconfig.GiantConfig, + ], +) +def test_config_dataclass_from_dict_round_trips_through_to_dict(cls): + assert cls.from_dict(cls().to_dict()) == cls() + assert cls.from_dict(None) == cls() + + +def test_stage2_model_config_defaults_match_documented_v030_intent(): + """The two keys issues.md Issue 1 found drifted between DEFAULT_CONFIG + and build_models/StageSpec.from_config's own .get(key, default) + fallbacks — pinned directly against the dataclass that is now their + shared single source of truth.""" + spec = gconfig.Stage2ModelConfig() + assert spec.decoder == "autoregressive" + assert spec.particle_type.target == "onehot" + + +def test_router_config_extra_round_trips_composed_axis_keys(): + d = {"enabled": True, "type": "composed", "axis0_type": "energy", "axis0_n_experts": 4} + router = gconfig.RouterConfig.from_dict(d) + assert router.enabled is True + assert router.extra == {"axis0_type": "energy", "axis0_n_experts": 4} + assert router.to_dict()["axis0_type"] == "energy" + + +def test_stage2_router_config_tie_to_stage1_not_leaked_into_extra(): + router = gconfig.Stage2RouterConfig.from_dict({"tie_to_stage1": True}) + assert router.tie_to_stage1 is True + assert "tie_to_stage1" not in router.extra + + +def test_stage1_router_config_has_no_tie_to_stage1_key(): + """Stage 1's router schema must not gain stage 2's tie_to_stage1 key — + that would change every future run's saved config.toml shape.""" + assert "tie_to_stage1" not in gconfig.RouterConfig().to_dict() + + +def test_n_sec_config_extra_round_trips_legacy_owner(): + n_sec = gconfig.NSecConfig.from_dict({"mode": "head", "legacy_owner": "stage1"}) + assert n_sec.legacy_owner == "stage1" + assert n_sec.to_dict() == {"mode": "head", "lambda": 0.1, "legacy_owner": "stage1"} + + +# --------------------------------------------------------------------------- +# _deep_merge +# --------------------------------------------------------------------------- + + +def test_deep_merge_leaf_override_keeps_untouched_siblings(): + base = {"a": 1, "b": {"c": 2, "d": 3}} + result = gconfig._deep_merge(base, {"b": {"c": 99}}) + assert result == {"a": 1, "b": {"c": 99, "d": 3}} + + +def test_deep_merge_recurses_at_multiple_levels(): + base = { + "stage1_model": { + "hidden_dim": 256, + "router": {"enabled": False, "type": "energy", "n_experts": 4}, + } + } + result = gconfig._deep_merge(base, {"stage1_model": {"router": {"enabled": True}}}) + assert result["stage1_model"]["hidden_dim"] == 256 + assert result["stage1_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 4, + } + + +def test_deep_merge_does_not_mutate_base(): + base = {"a": {"b": 1}} + gconfig._deep_merge(base, {"a": {"b": 2}}) + assert base == {"a": {"b": 1}} + + +def test_deep_merge_non_dict_override_replaces_wholesale(): + base = {"a": {"b": 1}} + result = gconfig._deep_merge(base, {"a": 5}) + assert result == {"a": 5} + + +# --------------------------------------------------------------------------- +# migrate_config +# --------------------------------------------------------------------------- + + +def test_migrate_config_already_v3_returned_unchanged(): + cfg = {"meta": {"config_version": 3}, "stage1_model": {"generator": "flow"}} + result = gconfig.migrate_config(cfg) + assert result == cfg + result["stage1_model"]["generator"] = "wgan" + assert cfg["stage1_model"]["generator"] == "flow" # deep-copied, not aliased + + +def test_migrate_config_empty_dict_still_injects_hardcoded_v02_facts(): + # No [train]/[model] at all still counts as "v0.2" (config_version + # absent) — the hardcoded architectural facts fire unconditionally. + new = gconfig.migrate_config({}) + assert "train" not in new + assert new["conditioning"]["out_dim"] == 128 + assert new["conditioning"]["particle"]["n_layers"] == 2 + assert new["conditioning"]["material"]["n_layers"] == 2 + assert new["stage1_model"]["active"] is True + assert new["stage1_model"]["flow"]["time_dim"] == 64 + assert new["stage1_model"]["ddpm"]["time_dim"] == 64 + assert new["stage2_model"]["active"] is True + assert new["stage2_model"]["flow"]["time_dim"] == 64 + assert new["stage2_model"]["ddpm"]["time_dim"] == 64 + assert new["stage2_model"]["context_dim"] == 64 + assert new["stage2_model"]["decoder"] == "one_shot" + assert new["stage2_model"]["particle_type"]["target"] == "physical" + assert new["meta"] == {"config_version": 3} + + +def test_migrate_config_mode_maps_to_both_stage_generators(): + new = gconfig.migrate_config({"train": {"mode": "wgan"}}) + assert new["stage1_model"]["generator"] == "wgan" + assert new["stage2_model"]["generator"] == "wgan" + + +def test_migrate_config_lambda_nsec_and_lambda_s2(): + new = gconfig.migrate_config({"train": {"lambda_nsec": 0.2, "lambda_s2": 2.0}}) + assert new["stage2_model"]["n_sec"]["lambda"] == 0.2 + assert new["stage2_model"]["lambda"] == 2.0 + + +def test_migrate_config_wgan_knobs_map_to_both_stages(): + new = gconfig.migrate_config({"train": {"n_critic": 3, "gp_weight": 5.0, "critic_lr": 1e-4}}) + for stage in ("stage1_model", "stage2_model"): + assert new[stage]["wgan"]["n_critic"] == 3 + assert new[stage]["wgan"]["gp_weight"] == 5.0 + assert new[stage]["wgan"]["critic_lr"] == 1e-4 + + +def test_migrate_config_model_hidden_dim_n_blocks_dropout_map_to_both_stages(): + new = gconfig.migrate_config({"model": {"hidden_dim": 128, "n_blocks": 4, "dropout": 0.2}}) + for stage in ("stage1_model", "stage2_model"): + assert new[stage]["hidden_dim"] == 128 + assert new[stage]["n_res_blocks"] == 4 + assert new[stage]["dropout"] == 0.2 + + +def test_migrate_config_emb_dim_and_conditioning_map_to_both_axes(): + new = gconfig.migrate_config({"model": {"emb_dim": 32, "conditioning": "embedding"}}) + for axis in ("particle", "material"): + assert new["conditioning"][axis]["emb_dim"] == 32 + assert new["conditioning"][axis]["type"] == "embedding" + + +def test_migrate_config_noise_dim_maps_to_both_stages_wgan(): + new = gconfig.migrate_config({"model": {"noise_dim": 128}}) + assert new["stage1_model"]["wgan"]["noise_dim"] == 128 + assert new["stage2_model"]["wgan"]["noise_dim"] == 128 + + +def test_migrate_config_k_max_maps_to_stage2_only(): + new = gconfig.migrate_config({"model": {"k_max": 20}}) + assert new["stage2_model"]["k_max"] == 20 + assert "k_max" not in new.get("stage1_model", {}) + + +def test_migrate_config_router_copied_to_both_stages_with_tie_to_stage1_false(): + new = gconfig.migrate_config( + { + "model": { + "router": { + "enabled": True, + "type": "energy", + "n_experts": 10, + "temperature": 0.05, + } + } + } ) + assert new["stage1_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 10, + "temperature": 0.05, + } + assert new["stage2_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 10, + "temperature": 0.05, + "tie_to_stage1": False, + } -def test_merge_cli_overrides_applies_file_then_cli(tmp_path, monkeypatch): +def test_migrate_config_router_nonzero_expert_dims_raises(): + cfg = {"model": {"router": {"enabled": True, "expert_hidden_dim": 128, "expert_n_blocks": 0}}} + try: + gconfig.migrate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "expert_hidden_dim" in str(e) + + +def test_migrate_config_router_zero_expert_dims_dropped_silently(): + new = gconfig.migrate_config( + { + "model": { + "router": { + "enabled": True, + "type": "energy", + "n_experts": 4, + "expert_hidden_dim": 0, + "expert_n_blocks": 0, + } + } + } + ) + assert "expert_hidden_dim" not in new["stage1_model"]["router"] + assert "expert_n_blocks" not in new["stage1_model"]["router"] + + +def test_migrate_config_train_passthrough_is_exact(): + new = gconfig.migrate_config({"train": {"epochs": 7, "batch_size": 999, "seed": 3}}) + assert new["train"] == {"epochs": 7, "batch_size": 999, "seed": 3} + + +def test_migrate_config_preserves_meta_git_hash(): + new = gconfig.migrate_config({"meta": {"git_hash": "abc123"}}) + assert new["meta"] == {"git_hash": "abc123", "config_version": 3} + + +def test_migrate_config_real_router_fixture_raises_on_nonzero_expert_dims(): + cfg = gconfig.load_toml(_CONFIGS_DIR / "router_energy_n10_embedding.toml") + try: + gconfig.migrate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "expert_hidden_dim" in str(e) + + +def test_migrate_config_real_wgan_fixture(): + cfg = gconfig.load_toml(_CONFIGS_DIR / "wgan_h128_b4_physical.toml") + new = gconfig.migrate_config(cfg) + assert new["stage1_model"]["generator"] == "wgan" + assert new["stage2_model"]["generator"] == "wgan" + for stage in ("stage1_model", "stage2_model"): + assert new[stage]["hidden_dim"] == 128 + assert new[stage]["n_res_blocks"] == 4 + assert new[stage]["dropout"] == 0.0 + assert new["conditioning"]["particle"]["type"] == "physical" + assert new["conditioning"]["material"]["type"] == "physical" + assert new["train"]["epochs"] == 30 + assert new["train"]["warmup_epochs"] == 3 + + +# --------------------------------------------------------------------------- +# merge_cli_overrides +# --------------------------------------------------------------------------- + + +def test_merge_cli_overrides_defaults_only_matches_default_config(): + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, None, {}) + assert cfg == gconfig.DEFAULT_CONFIG + assert cfg is not gconfig.DEFAULT_CONFIG + + +def test_merge_cli_overrides_nested_override_keeps_siblings(): + cfg = gconfig.merge_cli_overrides( + gconfig.DEFAULT_CONFIG, + None, + {"stage1_model": {"router": {"enabled": True}}}, + ) + assert cfg["stage1_model"]["router"]["enabled"] is True + assert cfg["stage1_model"]["router"]["type"] == "energy" # default preserved + assert cfg["stage1_model"]["hidden_dim"] == 256 # untouched sibling section + + +def test_merge_cli_overrides_file_then_explicit_override_precedence(tmp_path, monkeypatch): monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123") path = tmp_path / "config.toml" - _write_config(path, "abc123") + _write_toml( + path, + git_hash="abc123", + extra="[train]\nepochs = 5\n\n[model]\nhidden_dim = 64\n", + ) cfg = gconfig.merge_cli_overrides( gconfig.DEFAULT_CONFIG, path, - train_overrides={}, - model_overrides={"hidden_dim": 128}, + {"stage1_model": {"hidden_dim": 128}}, ) assert cfg["train"]["epochs"] == 5 # from file - assert cfg["model"]["hidden_dim"] == 128 # CLI override wins over file + assert cfg["stage1_model"]["hidden_dim"] == 128 # explicit override wins over file + assert cfg["stage2_model"]["hidden_dim"] == 64 # migrated from file, not overridden + + +def test_merge_cli_overrides_migrates_v2_file_transparently(tmp_path, monkeypatch): + monkeypatch.setattr(gconfig, "git_hash", lambda: "abc123") + path = tmp_path / "config.toml" + _write_toml( + path, + git_hash="abc123", + extra='[train]\nmode = "wgan"\n\n[model]\nconditioning = "embedding"\n', + ) + + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) + assert cfg["stage1_model"]["generator"] == "wgan" + assert cfg["stage2_model"]["generator"] == "wgan" + assert cfg["conditioning"]["particle"]["type"] == "embedding" + # hardcoded v0.2 fact still applied even though it's not a CLI-settable key + assert cfg["conditioning"]["particle"]["n_layers"] == 2 def test_merge_cli_overrides_warns_on_git_hash_mismatch(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") path = tmp_path / "config.toml" - _write_config(path, "old111") + _write_toml(path, git_hash="old111", extra="[train]\nepochs = 5\n") - cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) - - assert cfg["train"]["epochs"] == 5 # does not fail, config still applied + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) captured = capsys.readouterr() assert "warning" in captured.err assert "old111" in captured.err assert "current999" in captured.err -def test_merge_cli_overrides_no_warning_on_matching_git_hash( - tmp_path, monkeypatch, capsys -): +def test_merge_cli_overrides_no_warning_on_matching_git_hash(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "same123") path = tmp_path / "config.toml" - _write_config(path, "same123") + _write_toml(path, git_hash="same123", extra="[train]\nepochs = 5\n") - gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" -def test_merge_cli_overrides_no_warning_when_git_hash_unknown( - tmp_path, monkeypatch, capsys -): +def test_merge_cli_overrides_no_warning_when_git_hash_unknown(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "unknown") path = tmp_path / "config.toml" - _write_config(path, "abc123") + _write_toml(path, git_hash="abc123", extra="[train]\nepochs = 5\n") - gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" -def test_merge_cli_overrides_no_warning_when_meta_section_absent( - tmp_path, monkeypatch, capsys -): +def test_merge_cli_overrides_no_warning_when_meta_section_absent(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") path = tmp_path / "config.toml" path.write_text("[train]\nepochs = 5\n") - gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}, {}) + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) assert capsys.readouterr().err == "" -def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml( - tmp_path, monkeypatch, capsys -): +def test_merge_cli_overrides_real_default_toml_fixture(monkeypatch): + monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243") + cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / "default.toml", {}) + assert cfg["stage1_model"]["generator"] == "flow" + assert cfg["stage1_model"]["hidden_dim"] == 256 + assert cfg["stage2_model"]["hidden_dim"] == 256 + assert cfg["conditioning"]["particle"]["emb_dim"] == 16 + assert cfg["conditioning"]["particle"]["n_layers"] == 2 # migrated hardcoded fact + assert cfg["train"]["epochs"] == 100 + assert cfg["stage2_model"]["decoder"] == "one_shot" + + +# --------------------------------------------------------------------------- +# save_config +# --------------------------------------------------------------------------- + + +def test_save_config_round_trips_multi_level_nesting(tmp_path): + cfg = { + "stage1_model": { + "hidden_dim": 256, + "router": { + "enabled": True, + "type": "energy", + "n_experts": 4, + }, + }, + "train": {"epochs": 100}, + } + meta = {"config_version": 3, "git_hash": "abc123"} + + gconfig.save_config(cfg, tmp_path, meta) + loaded = gconfig.load_toml(tmp_path / "config.toml") + + assert loaded["stage1_model"]["hidden_dim"] == 256 + assert loaded["stage1_model"]["router"] == { + "enabled": True, + "type": "energy", + "n_experts": 4, + } + assert loaded["train"] == {"epochs": 100} + assert loaded["meta"] == meta + + +def test_save_config_round_trips_three_level_nesting(tmp_path): + cfg = { + "stage2_model": { + "decoder": "autoregressive", + "n_sec": {"mode": "head", "lambda": 0.1}, + "router": {"tie_to_stage1": True}, + } + } + gconfig.save_config(cfg, tmp_path, {"config_version": 3}) + loaded = gconfig.load_toml(tmp_path / "config.toml") + + assert loaded["stage2_model"]["decoder"] == "autoregressive" + assert loaded["stage2_model"]["n_sec"] == {"mode": "head", "lambda": 0.1} + assert loaded["stage2_model"]["router"] == {"tie_to_stage1": True} + + +# --------------------------------------------------------------------------- +# default_out_dir_name +# --------------------------------------------------------------------------- + +_NOW = datetime(2026, 7, 29, 14, 30) + + +def _cfg_with(**dotted_overrides): + """Build a full DEFAULT_CONFIG-shaped dict with dotted-path overrides + applied via _deep_merge, e.g. _cfg_with(**{"stage1_model.hidden_dim": 512}).""" + overrides: dict = {} + for dotted, value in dotted_overrides.items(): + gconfig._set_path(overrides, dotted, value) + return gconfig._deep_merge(gconfig.DEFAULT_CONFIG, overrides) + + +def test_default_out_dir_name_all_defaults_is_just_the_timestamp(): + assert gconfig.default_out_dir_name(gconfig.DEFAULT_CONFIG, now=_NOW) == "20260729_1430" + + +def test_default_out_dir_name_stage1_generator_shown_bare_no_prefix(): + cfg = _cfg_with(**{"stage1_model.generator": "wgan"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan" + + +def test_default_out_dir_name_stage2_decoder_shown(): + cfg = _cfg_with(**{"stage2_model.decoder": "one_shot"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_dec-one_shot" + + +def test_default_out_dir_name_particle_type_target_shown(): + cfg = _cfg_with(**{"stage2_model.particle_type.target": "physical"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_pt-physical" + + +def test_default_out_dir_name_particle_conditioning_embedding_abbreviated(): + cfg = _cfg_with(**{"conditioning.particle.type": "embedding"}) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb" + + +def test_default_out_dir_name_stage1_router_shown_as_unit(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + } + ) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8" + + +def test_default_out_dir_name_stage2_router_shown_as_unit_distinct_from_stage1(): + cfg = _cfg_with( + **{ + "stage2_model.router.enabled": True, + "stage2_model.router.type": "pdg", + "stage2_model.router.n_experts": 3, + } + ) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s2r-pdg3" + + +def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": False, + "stage1_model.router.type": "pdg", + "stage1_model.router.n_experts": 8, + } + ) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" + + +def test_default_out_dir_name_router_gumbel_shown_when_enabled(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + "stage1_model.router.gumbel": True, + } + ) + assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_s1r-energy8_s1gum" + + +def test_default_out_dir_name_overflow_caps_and_hashes_remainder(): + cfg = _cfg_with( + **{ + "stage1_model.generator": "wgan", + "stage2_model.generator": "flow", + "stage2_model.decoder": "one_shot", + "stage2_model.autoregressive.history": "attention", + "stage2_model.particle_type.target": "physical", + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + "stage2_model.router.enabled": True, + "stage2_model.router.type": "pdg", + "stage2_model.router.n_experts": 3, + } + ) + name = gconfig.default_out_dir_name(cfg, now=_NOW) + # First 6 by priority: stage1_generator, stage2_generator, stage2_decoder, + # stage2_history, particle_type_target, stage1_router — stage2_router + # overflows into the hash suffix. + assert name.startswith("20260729_1430_wgan_s2-flow_dec-one_shot_hist-attention_pt-physical_s1r-energy8_+") + + +def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive(): + overrides = { + "stage1_model.generator": "wgan", + "stage2_model.generator": "flow", + "stage2_model.decoder": "one_shot", + "stage2_model.autoregressive.history": "attention", + "stage2_model.particle_type.target": "physical", + "stage1_model.router.enabled": True, + "stage1_model.router.type": "energy", + "stage1_model.router.n_experts": 8, + "train.seed": 3, + } + name_a = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW) + name_b = gconfig.default_out_dir_name(_cfg_with(**overrides), now=_NOW) + assert name_a == name_b + + changed = dict(overrides, **{"train.seed": 99}) + name_c = gconfig.default_out_dir_name(_cfg_with(**changed), now=_NOW) + assert name_c != name_a + + +# --------------------------------------------------------------------------- +# validate_config +# --------------------------------------------------------------------------- + + +def test_validate_config_default_config_passes(): + gconfig.validate_config(gconfig.DEFAULT_CONFIG) # must not raise + + +def test_validate_config_embedding_target_requires_embedding_conditioning(): + cfg = _cfg_with( + **{ + "stage2_model.particle_type.target": "embedding", + "conditioning.particle.type": "physical", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "embedding" in str(e) + + +def test_validate_config_embedding_target_passes_with_embedding_conditioning(): + cfg = _cfg_with( + **{ + "stage2_model.particle_type.target": "embedding", + "conditioning.particle.type": "embedding", + "conditioning.material.type": "embedding", + } + ) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_mixed_particle_material_conditioning_is_valid(): + """The particle and material conditioning axes are configured + independently and may mix freely — e.g. material + "physical" with particle "embedding" — and the data pipeline + (giant/data/transforms.py) now implements that end-to-end, so + validate_config must not reject it.""" + cfg = _cfg_with( + **{ + "conditioning.particle.type": "physical", + "conditioning.material.type": "embedding", + } + ) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_pdg_router_incompatible_with_physical_conditioning(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": True, + "stage1_model.router.type": "pdg", + "conditioning.particle.type": "physical", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "pdg" in str(e) + + +def test_validate_config_tie_to_stage1_requires_stage1_active(): + cfg = _cfg_with( + **{ + "stage2_model.router.tie_to_stage1": True, + "stage1_model.active": False, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "tie_to_stage1" in str(e) + + +def test_validate_config_stop_token_not_implemented(): + cfg = _cfg_with(**{"stage2_model.n_sec.mode": "stop_token"}) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "stop_token" in str(e) + + +def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint(): + """'n_sec.mode = "truth" is invalid for a rollout-capable checkpoint' — + both stages active means giant rollout + could load this checkpoint, but 'truth' has no ground truth to draw + n_sec from at rollout time.""" + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "truth", + "stage1_model.active": True, + "stage2_model.active": True, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "n_sec.mode" in str(e) and "truth" in str(e) + + +def test_validate_config_n_sec_truth_allowed_for_stage2_only_checkpoint(): + """'truth' is exactly the standalone stage-2 evaluation mode the design + doc carves out — stage1_model.active = false must still pass.""" + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "truth", + "stage1_model.active": False, + "stage2_model.active": True, + } + ) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_n_sec_truth_allowed_when_stage2_inactive(): + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "truth", + "stage1_model.active": True, + "stage2_model.active": False, + } + ) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_ar_default_markov_always_passes(): + """DEFAULT_CONFIG already has decoder='autoregressive', + history='markov', teacher_forcing='always' — must not raise (v0.3.0 + step 5; see also test_validate_config_default_config_passes).""" + cfg = _cfg_with(**{"stage2_model.decoder": "autoregressive"}) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_ar_history_attention_passes(): + """v0.3.0 step 7 implements history='attention' — must not raise.""" + cfg = _cfg_with( + **{ + "stage2_model.decoder": "autoregressive", + "stage2_model.autoregressive.history": "attention", + } + ) + gconfig.validate_config(cfg) # must not raise + + +@pytest.mark.parametrize("teacher_forcing", ["scheduled", "never"]) +def test_validate_config_ar_teacher_forcing_scheduled_or_never_passes(teacher_forcing): + """v0.3.0 step 7 implements teacher_forcing in {'scheduled', 'never'} — + must not raise.""" + cfg = _cfg_with( + **{ + "stage2_model.decoder": "autoregressive", + "stage2_model.autoregressive.teacher_forcing": teacher_forcing, + } + ) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_ar_history_invalid_value_rejected(): + cfg = _cfg_with( + **{ + "stage2_model.decoder": "autoregressive", + "stage2_model.autoregressive.history": "bogus", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "history" in str(e) + + +def test_validate_config_ar_teacher_forcing_invalid_value_rejected(): + cfg = _cfg_with( + **{ + "stage2_model.decoder": "autoregressive", + "stage2_model.autoregressive.teacher_forcing": "bogus", + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "teacher_forcing" in str(e) + + +def test_validate_config_ar_checks_skipped_under_one_shot(): + """history/teacher_forcing values that would fail under AR are irrelevant + (and unchecked) when decoder='one_shot'.""" + cfg = _cfg_with( + **{ + "stage2_model.decoder": "one_shot", + "stage2_model.autoregressive.history": "attention", + "stage2_model.autoregressive.teacher_forcing": "scheduled", + } + ) + gconfig.validate_config(cfg) # must not raise + + +# --------------------------------------------------------------------------- +# validate_config_keys / merge_cli_overrides unknown-key rejection +# --------------------------------------------------------------------------- + + +def test_validate_config_keys_default_config_passes(): + gconfig.validate_config_keys(gconfig.DEFAULT_CONFIG) # must not raise + + +def test_validate_config_keys_rejects_unknown_top_level_key(): + cfg = _cfg_with(**{"bogus_section.foo": 1}) + try: + gconfig.validate_config_keys(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "bogus_section" in str(e) + + +def test_validate_config_keys_rejects_unknown_nested_key_with_close_match_hint(): + cfg = _cfg_with(**{"stage1_model.n_res_block": 12}) + try: + gconfig.validate_config_keys(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "stage1_model.n_res_block" in str(e) + assert "n_res_blocks" in str(e) + + +def test_validate_config_keys_allows_composed_router_axis_keys(): + cfg = _cfg_with( + **{ + "stage1_model.router.enabled": True, + "stage1_model.router.type": "composed", + "stage1_model.router.axis0_type": "energy", + "stage1_model.router.axis0_n_experts": 4, + "stage1_model.router.axis1_type": "pdg", + "stage1_model.router.axis1_emb_dim": 8, + } + ) + gconfig.validate_config_keys(cfg) # must not raise + + +def test_validate_config_keys_allows_centers_init(): + cfg = _cfg_with(**{"stage1_model.router.centers_init": [-1.0, 0.0, 1.0]}) + gconfig.validate_config_keys(cfg) # must not raise + + +def test_validate_config_keys_rejects_unrelated_unknown_router_key(): + cfg = _cfg_with(**{"stage1_model.router.n_expert": 4}) # typo for n_experts + try: + gconfig.validate_config_keys(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "stage1_model.router.n_expert" in str(e) + assert "n_experts" in str(e) + + +def test_validate_config_keys_skips_meta_section(): + cfg = _cfg_with() + cfg["meta"] = {"config_version": 3, "git_hash": "abc123"} + gconfig.validate_config_keys(cfg) # must not raise + + +def test_merge_cli_overrides_rejects_typo_in_toml_file(tmp_path): + path = tmp_path / "config.toml" + path.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n") + try: + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, path, {}) + assert False, "expected ValueError" + except ValueError as e: + assert "n_res_block" in str(e) + + +def test_merge_cli_overrides_rejects_typo_in_cli_overrides(): + try: + gconfig.merge_cli_overrides( + gconfig.DEFAULT_CONFIG, + None, + {"stage1_model": {"n_res_block": 12}}, + ) + assert False, "expected ValueError" + except ValueError as e: + assert "n_res_block" in str(e) + + +@pytest.mark.parametrize("fixture_name", ["default.toml", "wgan_h128_b4_physical.toml"]) +def test_merge_cli_overrides_real_config_fixtures_pass_key_validation(fixture_name, monkeypatch): + monkeypatch.setattr(gconfig, "git_hash", lambda: "c3bf3abebfe29a10fe42b9cbafbb3460ab78d243") + gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, _CONFIGS_DIR / fixture_name, {}) # must not raise + + +# --------------------------------------------------------------------------- +# overrides_from_flags (issues.md Issue 3): the flag -> config-path table +# shared by `giant train`/`giant new-run`. Each test below pins one +# precedence rule directly, without CliRunner — see also +# tests/test_cli_train_overrides.py for the thin end-to-end smoke coverage. +# --------------------------------------------------------------------------- + + +def test_overrides_from_flags_empty_values_yield_empty_overrides(): + assert gconfig.overrides_from_flags({}) == {} + assert gconfig.overrides_from_flags({"epochs": None, "hidden_dim": None}) == {} + + +def test_overrides_from_flags_train_block_passthrough(): + overrides = gconfig.overrides_from_flags({"epochs": 5, "lr": 1e-3, "hidden_dim": None}) + assert overrides == {"train": {"epochs": 5, "lr": 1e-3}} + + +@pytest.mark.parametrize( + ("shorthand", "explicit", "path_key"), + [ + ("hidden_dim", "stage1_hidden_dim", "hidden_dim"), + ("n_blocks", "stage1_n_res_blocks", "n_res_blocks"), + ("dropout", "stage1_dropout", "dropout"), + ], +) +def test_overrides_from_flags_stage1_explicit_overrides_shorthand(shorthand, explicit, path_key): + overrides = gconfig.overrides_from_flags({shorthand: 1, explicit: 2}) + assert overrides["stage1_model"][path_key] == 2 + + +@pytest.mark.parametrize( + ("shorthand", "path_key"), + [("hidden_dim", "hidden_dim"), ("n_blocks", "n_res_blocks"), ("dropout", "dropout")], +) +def test_overrides_from_flags_stage1_shorthand_alone(shorthand, path_key): + overrides = gconfig.overrides_from_flags({shorthand: 7}) + assert overrides["stage1_model"][path_key] == 7 + + +def test_overrides_from_flags_stage2_only_knobs(): + overrides = gconfig.overrides_from_flags( + { + "stage2_hidden_dim": 32, + "stage2_n_res_blocks": 4, + "stage2_dropout": 0.1, + "stage2_decoder": "one_shot", + "stage2_k_max": 8, + "stage2_context_dim": 16, + "stage2_stage1_context": "sampled", + } + ) + assert overrides["stage2_model"] == { + "hidden_dim": 32, + "n_res_blocks": 4, + "dropout": 0.1, + "decoder": "one_shot", + "k_max": 8, + "context_dim": 16, + "stage1_context": "sampled", + } + assert "stage1_model" not in overrides + + +def test_overrides_from_flags_mode_fans_to_both_stages(): + overrides = gconfig.overrides_from_flags({"mode": "wgan"}) + assert overrides["stage1_model"]["generator"] == "wgan" + assert overrides["stage2_model"]["generator"] == "wgan" + + +def test_overrides_from_flags_stage1_generator_overrides_mode_for_stage1_only(): + overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage1_generator": "flow"}) + assert overrides["stage1_model"]["generator"] == "flow" + assert overrides["stage2_model"]["generator"] == "wgan" + + +def test_overrides_from_flags_stage2_generator_overrides_mode_for_stage2_only(): + overrides = gconfig.overrides_from_flags({"mode": "wgan", "stage2_generator": "flow"}) + assert overrides["stage1_model"]["generator"] == "wgan" + assert overrides["stage2_model"]["generator"] == "flow" + + +def test_overrides_from_flags_emb_dim_sets_both_conditioning_axes(): + overrides = gconfig.overrides_from_flags({"emb_dim": 24}) + assert overrides["conditioning"]["particle"]["emb_dim"] == 24 + assert overrides["conditioning"]["material"]["emb_dim"] == 24 + + +def test_overrides_from_flags_conditioning_sets_both_axes_type(): + overrides = gconfig.overrides_from_flags({"conditioning": "onehot"}) + assert overrides["conditioning"]["particle"]["type"] == "onehot" + assert overrides["conditioning"]["material"]["type"] == "onehot" + + +def test_overrides_from_flags_router_config_only_touches_stage1(): + overrides = gconfig.overrides_from_flags({"router_config": {"enabled": True, "type": "energy"}}) + assert overrides["stage1_model"]["router"] == {"enabled": True, "type": "energy"} + assert "stage2_model" not in overrides + + +@pytest.mark.parametrize( + ("shared", "stage1_specific", "stage2_specific", "path_key"), + [ + ("n_critic", "stage1_n_critic", "stage2_n_critic", "n_critic"), + ("gp_weight", "stage1_gp_weight", "stage2_gp_weight", "gp_weight"), + ("noise_dim", "stage1_noise_dim", "stage2_noise_dim", "noise_dim"), + ("critic_lr", "stage1_critic_lr", "stage2_critic_lr", "critic_lr"), + ], +) +def test_overrides_from_flags_wgan_knobs_split_per_stage(shared, stage1_specific, stage2_specific, path_key): + overrides = gconfig.overrides_from_flags({shared: 5.0, stage1_specific: 3.0}) + assert overrides["stage1_model"]["wgan"][path_key] == 3.0 + assert overrides["stage2_model"]["wgan"][path_key] == 5.0 + + overrides = gconfig.overrides_from_flags({shared: 5.0, stage2_specific: 2.5}) + assert overrides["stage1_model"]["wgan"][path_key] == 5.0 + assert overrides["stage2_model"]["wgan"][path_key] == 2.5 + + +# --------------------------------------------------------------------------- +# checkpoint config-mismatch warnings (unchanged surface, still exercised) +# --------------------------------------------------------------------------- + + +def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") ckpt_path = tmp_path / "best.pt" - ckpt_path.write_bytes(b"") # contents irrelevant, only its directory is used - _write_config(tmp_path / "config.toml", "old111") + ckpt_path.write_bytes(b"") + _write_toml(tmp_path / "config.toml", git_hash="old111", extra="[train]\nepochs = 5\n") gconfig.warn_if_checkpoint_config_mismatch(ckpt_path) @@ -96,9 +1027,7 @@ def test_warn_if_checkpoint_config_mismatch_finds_sibling_toml( assert "current999" in captured.err -def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent( - tmp_path, monkeypatch, capsys -): +def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "current999") ckpt_path = tmp_path / "best.pt" ckpt_path.write_bytes(b"") @@ -107,208 +1036,11 @@ def test_warn_if_checkpoint_config_mismatch_no_warning_when_toml_absent( assert capsys.readouterr().err == "" -def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match( - tmp_path, monkeypatch, capsys -): +def test_warn_if_checkpoint_config_mismatch_no_warning_when_hashes_match(tmp_path, monkeypatch, capsys): monkeypatch.setattr(gconfig, "git_hash", lambda: "same123") ckpt_path = tmp_path / "best.pt" ckpt_path.write_bytes(b"") - _write_config(tmp_path / "config.toml", "same123") + _write_toml(tmp_path / "config.toml", git_hash="same123", extra="[train]\nepochs = 5\n") gconfig.warn_if_checkpoint_config_mismatch(ckpt_path) assert capsys.readouterr().err == "" - - -def test_resolve_expert_dims_default_config_inherits_hidden_dim_and_n_blocks(): - # The unset sentinel (expert_hidden_dim/n_blocks == 0 in DEFAULT_CONFIG) - # is exactly the bug fixed by resolve_expert_dims: it must not silently - # fall back to some other hardcoded default, only to the monolith's own - # hidden_dim/n_blocks, so --hidden-dim/--n-blocks reach the experts too. - router_cfg = dict(gconfig.DEFAULT_CONFIG["model"]["router"]) - assert router_cfg["expert_hidden_dim"] == 0 - assert router_cfg["expert_n_blocks"] == 0 - - hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6) - assert (hidden_dim, n_blocks) == (512, 6) - - -def test_resolve_expert_dims_missing_keys_also_inherit(): - hidden_dim, n_blocks = gconfig.resolve_expert_dims({}, 512, 6) - assert (hidden_dim, n_blocks) == (512, 6) - - -def test_default_config_gumbel_router_defaults_off(): - # Straight-through Gumbel-softmax combine weights (giant.model.network. - # Router.combine_weights) must be opt-in — existing routed configs and - # checkpoints should be unaffected unless gumbel is explicitly enabled. - router_cfg = gconfig.DEFAULT_CONFIG["model"]["router"] - assert router_cfg["gumbel"] is False - assert router_cfg["gumbel_tau_start"] == 1.0 - assert router_cfg["gumbel_tau_end"] == 0.1 - - -def test_resolve_expert_dims_explicit_override_wins(): - router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 3} - hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6) - assert (hidden_dim, n_blocks) == (128, 3) - - -def test_resolve_expert_dims_partial_override_mixes_explicit_and_inherited(): - router_cfg = {"expert_hidden_dim": 128, "expert_n_blocks": 0} - hidden_dim, n_blocks = gconfig.resolve_expert_dims(router_cfg, 512, 6) - assert (hidden_dim, n_blocks) == (128, 6) - - -def _default_cfg(**overrides): - train_overrides = { - k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["train"] - } - model_overrides = { - k: v for k, v in overrides.items() if k in gconfig.DEFAULT_CONFIG["model"] - } - router_overrides = overrides.get("router") - if router_overrides: - model_overrides["router"] = router_overrides - return gconfig.merge_cli_overrides( - gconfig.DEFAULT_CONFIG, None, train_overrides, model_overrides - ) - - -_NOW = datetime(2026, 7, 29, 14, 30) - - -def test_default_out_dir_name_all_defaults_is_just_the_timestamp(): - cfg = _default_cfg() - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_single_non_default_field(): - cfg = _default_cfg(hidden_dim=512) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_h512" - - -def test_default_out_dir_name_conditioning_embedding_shown_abbreviated(): - cfg = _default_cfg(conditioning="embedding") - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_cemb" - - -def test_default_out_dir_name_conditioning_default_omitted(): - cfg = _default_cfg(conditioning="physical") - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_router_enabled_shown_as_unit(): - cfg = _default_cfg(router={"enabled": True, "type": "energy", "n_experts": 8}) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8" - - -def test_default_out_dir_name_router_disabled_omitted_even_if_subfields_nondefault(): - cfg = _default_cfg(router={"enabled": False, "type": "pdg", "n_experts": 8}) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_router_gumbel_shown_when_enabled(): - cfg = _default_cfg( - router={"enabled": True, "type": "energy", "n_experts": 8, "gumbel": True} - ) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_gum" - - -def test_default_out_dir_name_router_gumbel_omitted_when_router_disabled(): - cfg = _default_cfg( - router={"enabled": False, "type": "energy", "n_experts": 8, "gumbel": True} - ) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430" - - -def test_default_out_dir_name_router_learn_centers_shown_only_when_disabled(): - cfg_default = _default_cfg( - router={"enabled": True, "type": "energy", "n_experts": 8} - ) - assert ( - gconfig.default_out_dir_name(cfg_default, now=_NOW) == "20260729_1430_r-energy8" - ) - - cfg_off = _default_cfg( - router={ - "enabled": True, - "type": "energy", - "n_experts": 8, - "learn_centers": False, - } - ) - assert ( - gconfig.default_out_dir_name(cfg_off, now=_NOW) - == "20260729_1430_r-energy8_nolc" - ) - - -def test_default_out_dir_name_router_learn_width_and_temperature_shown(): - cfg = _default_cfg( - router={ - "enabled": True, - "type": "energy", - "n_experts": 8, - "learn_width": True, - } - ) - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_r-energy8_lw" - - cfg2 = _default_cfg( - router={ - "enabled": True, - "type": "energy", - "n_experts": 8, - "learn_temperature": True, - } - ) - assert gconfig.default_out_dir_name(cfg2, now=_NOW) == "20260729_1430_r-energy8_lt" - - -def test_default_out_dir_name_mode_shown_bare_no_prefix(): - cfg = _default_cfg(mode="wgan") - assert gconfig.default_out_dir_name(cfg, now=_NOW) == "20260729_1430_wgan" - - -def test_default_out_dir_name_overflow_caps_and_hashes_remainder(): - cfg = _default_cfg( - mode="wgan", - router={"enabled": True, "type": "energy", "n_experts": 8}, - conditioning="embedding", - hidden_dim=512, - n_blocks=8, - emb_dim=32, - lr=1e-3, - batch_size=2048, - seed=3, - epochs=200, - ) - name = gconfig.default_out_dir_name(cfg, now=_NOW) - # First 6 by priority: mode, router, conditioning, hidden_dim, n_blocks, emb_dim. - assert name.startswith("20260729_1430_wgan_r-energy8_cemb_h512_b8_e32_+4more-") - digest = name.split("-")[-1] - assert len(digest) == 6 - - -def test_default_out_dir_name_overflow_hash_is_deterministic_and_value_sensitive(): - base = dict( - mode="wgan", - router={"enabled": True, "type": "energy", "n_experts": 8}, - conditioning="embedding", - hidden_dim=512, - n_blocks=8, - emb_dim=32, - lr=1e-3, - batch_size=2048, - seed=3, - epochs=200, - ) - name_a = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW) - name_b = gconfig.default_out_dir_name(_default_cfg(**base), now=_NOW) - assert name_a == name_b # stable across calls with the same overflow set - - changed = dict(base, epochs=999) - name_c = gconfig.default_out_dir_name(_default_cfg(**changed), now=_NOW) - assert ( - name_c != name_a - ) # differs when an overflowed value changes # n_blocks inherited, hidden_dim not diff --git a/tests/test_create_root_files.py b/tests/test_create_root_files.py index f07453c..686e198 100644 --- a/tests/test_create_root_files.py +++ b/tests/test_create_root_files.py @@ -16,9 +16,7 @@ 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 @@ -89,17 +87,13 @@ 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): @@ -108,9 +102,7 @@ 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) @@ -320,9 +312,7 @@ 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: diff --git a/tests/test_dataset.py b/tests/test_dataset.py index e311dae..5d2e71e 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -126,7 +126,8 @@ def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path): target_normalizer=tgt_norm, batch_size=4, shuffle=False, - conditioning="embedding", + particle_conditioning="embedding", + material_conditioning="embedding", ) return sum(len(batch[0]) for batch in ds) diff --git a/tests/test_dwarf.py b/tests/test_dwarf.py index 14b2673..2d9aeae 100644 --- a/tests/test_dwarf.py +++ b/tests/test_dwarf.py @@ -51,9 +51,7 @@ 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 @@ -103,7 +101,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_cond=physical" in loaded.normalizers + assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers def test_warm_cache_second_run_hits_cache(tmp_path): @@ -167,5 +165,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_cond=physical" in loaded.normalizers - assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers + 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 diff --git a/tests/test_flow.py b/tests/test_flow.py index b9e7e8c..aae4a93 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -1,12 +1,23 @@ import torch from giant.constants import COND_DIM -from giant.model.network import DenoisingMLP +from giant.model.network import Stage1Model 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 DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2) + 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, + ) def _batch(B=8): @@ -41,7 +52,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.shape == (B,) + assert n_sec is not None and n_sec.shape == (B,) def test_ddpm_loss_nonneg(): @@ -58,4 +69,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.shape == (B,) + assert n_sec is not None and n_sec.shape == (B,) diff --git a/tests/test_geometry.py b/tests/test_geometry.py index d967315..aa67110 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -4,6 +4,7 @@ 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 @@ -12,6 +13,70 @@ 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) @@ -137,9 +202,7 @@ 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] @@ -168,9 +231,7 @@ 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() diff --git a/tests/test_loader.py b/tests/test_loader.py index f0900c7..5d57656 100644 --- a/tests/test_loader.py +++ b/tests/test_loader.py @@ -6,7 +6,9 @@ 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, @@ -178,6 +180,63 @@ 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) ──────────────────────────────────────────── @@ -256,9 +315,7 @@ 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] @@ -339,9 +396,7 @@ 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): diff --git a/tests/test_materials.py b/tests/test_materials.py index c0d0c8e..233bf93 100644 --- a/tests/test_materials.py +++ b/tests/test_materials.py @@ -23,11 +23,7 @@ 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 diff --git a/tests/test_migration_v02_v03.py b/tests/test_migration_v02_v03.py new file mode 100644 index 0000000..35f4db5 --- /dev/null +++ b/tests/test_migration_v02_v03.py @@ -0,0 +1,283 @@ +"""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) + # legacy_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 legacy_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"]["legacy_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, no legacy_owner: n_sec lives on stage 2. + assert models["stage1"].n_sec_head is None + assert models["stage2"].n_sec_head is not None diff --git a/tests/test_network.py b/tests/test_network.py index 397aa79..7857d39 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -1,6 +1,28 @@ +import copy + +import pytest import torch -from giant.constants import COND_DIM -from giant.model.network import DenoisingMLP, SinusoidalEmbedding +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} def test_sinusoidal_embedding_shape(): @@ -15,9 +37,15 @@ def test_sinusoidal_embedding_batch_1(): assert emb(t).shape == (1, 32) -def test_denoising_mlp_output_shape(): +def test_stage1_model_output_shape(): B = 8 - model = DenoisingMLP(pdg_vocab=5, mat_vocab=3) + model = Stage1Model( + pdg_vocab=5, + mat_vocab=3, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + n_sec_head_k_max=15, + ) x_t = torch.randn(B, 9) t = torch.rand(B) cond_cont = torch.randn(B, COND_DIM) @@ -28,20 +56,655 @@ def test_denoising_mlp_output_shape(): ], dim=1, ) - out = model(x_t, t, cond_cont, cond_cat) + out = model(x_t, cond_cont, cond_cat, t=t) assert out.shape == (B, 9) -def test_denoising_mlp_gradients_flow(): +def test_stage1_model_gradients_flow(): B = 4 - model = DenoisingMLP(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2) + 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, + ) 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, t, cond_cont, cond_cat).sum() + flow_loss = model(x_t, cond_cont, cond_cat, t=t).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) + + +# --- 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") + + +@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()} + + +# ── 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 diff --git a/tests/test_particles.py b/tests/test_particles.py index 224e2c8..c80e724 100644 --- a/tests/test_particles.py +++ b/tests/test_particles.py @@ -1,7 +1,11 @@ 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, @@ -45,9 +49,7 @@ 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(): @@ -126,3 +128,99 @@ 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) diff --git a/tests/test_phase2.py b/tests/test_phase2.py index c8e7fd1..7a8efec 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -4,44 +4,64 @@ import numpy as np import pytest import torch -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.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.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"): - return DenoisingMLP( + particle_cfg, material_cfg = _particle_material_cfg(conditioning) + return Stage1Model( pdg_vocab=pdg, mat_vocab=mat, + particle_cfg=particle_cfg, + material_cfg=material_cfg, hidden_dim=32, - n_blocks=2, - conditioning=conditioning, + n_res_blocks=2, + n_sec_head_k_max=K_MAX, ) def _sec_decoder(pdg=3, mat=2, conditioning="embedding"): - return SecondaryDecoder( + particle_cfg, material_cfg = _particle_material_cfg(conditioning) + return Stage2OneShot( pdg_vocab=pdg, mat_vocab=mat, + particle_cfg=particle_cfg, + material_cfg=material_cfg, hidden_dim=32, - n_blocks=2, - conditioning=conditioning, + n_res_blocks=2, + generator="flow", + time_dim=16, ) 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 -# ── DenoisingMLP Phase-2 additions ─────────────────────────────────────────── +# ── Stage1Model Phase-2 additions ──────────────────────────────────────────── def test_predict_n_sec_shape(): @@ -82,7 +102,7 @@ def test_condition_encoder_embedding_mode_has_embedding_tables(): assert not hasattr(model.cond_enc, "particle_mlp") -# ── SecondaryDecoder ────────────────────────────────────────────────────────── +# ── Stage2OneShot ───────────────────────────────────────────────────────────── @pytest.mark.parametrize("conditioning", ["embedding", "physical"]) @@ -93,7 +113,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, t, cond_cont, cond_cat, stage1_out) + out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t) assert out.shape == (B, SEC_DIM) @@ -104,7 +124,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, t, cond_cont, cond_cat, stage1_out) + out = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t) assert torch.isfinite(out).all() @@ -115,7 +135,9 @@ def test_sec_decoder_gradients(): t = torch.rand(B) cond_cont, cond_cat = _cond(B) stage1_out = torch.randn(B, X_DIM) - decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward() + 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() for name, p in decoder.named_parameters(): assert p.grad is not None, f"no grad for {name}" @@ -130,9 +152,7 @@ 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 @@ -145,9 +165,7 @@ 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) @@ -158,8 +176,102 @@ 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 + 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, ).backward() assert any(p.grad is not None for p in decoder.parameters()) @@ -173,9 +285,7 @@ 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) @@ -188,9 +298,7 @@ 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() @@ -224,9 +332,7 @@ 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() @@ -273,9 +379,7 @@ 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(): @@ -387,9 +491,7 @@ 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) @@ -423,9 +525,7 @@ 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 @@ -448,9 +548,7 @@ 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) @@ -470,9 +568,7 @@ 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: @@ -496,12 +592,8 @@ 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] @@ -523,20 +615,14 @@ 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) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 4c84c45..9c25e67 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -6,8 +6,10 @@ import pytest import torch from giant import config as gconfig +from giant.constants import COND_DIM from giant.data import setup_cache -from giant.pipeline import run_train_job +from giant.data.transforms import Normalizer +from giant.pipeline import _seed_energy_router, run_train_job def _unit(v): @@ -46,9 +48,7 @@ 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,10 +99,19 @@ def _tiny_cfg(**train_overrides): "warmup_epochs": 0, "validate_every": 0, "max_val_batches": 1, + "wandb": False, } ) cfg["train"].update(train_overrides) - cfg["model"].update({"hidden_dim": 8, "n_blocks": 1, "emb_dim": 4, "dropout": 0.0}) + 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 return cfg @@ -143,17 +152,74 @@ def test_run_train_job_second_run_hits_cache(tmp_path, data, monkeypatch): assert "normalizer: cache hit" in joined -def test_run_train_job_warns_when_num_workers_exceeds_shared_quota( - tmp_path, data, monkeypatch -): +def test_run_train_job_builds_caches_and_persists_pdg_topn_map(tmp_path, data): + """DEFAULT_CONFIG's stage2_model.particle_type.target defaults to + "onehot" — a plain _tiny_cfg() run must + build the shared pdg top-N map, cache it in the setup-cache sidecar, and + persist it into the checkpoint, with no extra config needed.""" + 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 + key = setup_cache.topn_key("pdg", 4) # conditioning.particle.emb_dim = 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 "pdg_topn_map" in ckpt + assert set(ckpt["pdg_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_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. 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) -def test_run_train_job_no_warning_when_num_workers_within_shared_quota( - tmp_path, data, monkeypatch -): +@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): 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) @@ -193,6 +259,70 @@ 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) @@ -202,11 +332,63 @@ 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] diff --git a/tests/test_render.py b/tests/test_render.py index 49bb84c..064cb32 100644 --- a/tests/test_render.py +++ b/tests/test_render.py @@ -4,10 +4,12 @@ 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 @@ -19,6 +21,152 @@ 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( @@ -90,3 +238,133 @@ 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 diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 5227109..1cdb799 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -8,10 +8,16 @@ import numpy as np import pytest import torch -from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG +from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX +from giant.data.loader import TopNMap from giant.data.transforms import Normalizer -from giant.model.network import DenoisingMLP, SecondaryDecoder -from giant.rollout import make_seed_frontier, rollout +from giant.model.network import ( + Stage1Model, + Stage2Autoregressive, + Stage2OneShot, + stage2_trunk_sec_dim, +) +from giant.rollout import L1DistCollector, make_seed_frontier, rollout pytest.importorskip("sklearn") from giant import geometry as g # noqa: E402 @@ -21,11 +27,26 @@ MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1} def _models(conditioning="embedding"): - s1 = DenoisingMLP( - pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, conditioning=conditioning + 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, ) - s2 = SecondaryDecoder( - 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, ) return s1.eval(), s2.eval() @@ -90,7 +111,8 @@ def _run( batch_size=128, max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, - conditioning=conditioning, + particle_conditioning=conditioning, + material_conditioning=conditioning, ) @@ -99,12 +121,8 @@ 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 @@ -151,7 +169,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, conditioning="embedding") + fr, _counts = make_seed_frontier(**seeds, particle_conditioning="embedding") np.testing.assert_array_equal(fr["mass"], 0.0) np.testing.assert_array_equal(fr["charge"], 0.0) @@ -325,3 +343,341 @@ 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, ], 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", + pdg_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, + pdg_topn_map=pdg_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, pdg_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="pdg_topn_map"): + _run_v3(s1, s2, pdg_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) + + +@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 diff --git a/tests/test_router.py b/tests/test_router.py index a8e9174..d733218 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -6,47 +6,55 @@ 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, - RoutedDenoisingMLP, - RoutedSecondaryDecoder, - SecondaryDecoder, + RoutedTrunk, + Stage1Model, + Stage2OneShot, 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 RoutedDenoisingMLP( + return Stage1Model( pdg_vocab=pdg, mat_vocab=mat, + particle_cfg=PARTICLE_CFG, + material_cfg=MATERIAL_CFG, + hidden_dim=16, + n_res_blocks=2, router=router, - expert_hidden_dim=16, - expert_n_blocks=2, + n_sec_head_k_max=K_MAX, ) def _routed_sec_decoder(n_experts=4, pdg=3, mat=2, **router_kwargs): router = build_router("energy", n_experts, **router_kwargs) - return RoutedSecondaryDecoder( + return Stage2OneShot( 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, ) @@ -68,9 +76,7 @@ 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(): @@ -92,7 +98,7 @@ def test_energy_router_balance_loss_is_nonnegative_scalar(): def test_build_router_ignores_unrecognized_kwargs(): - # lambda_balance is a model_config.router key but not an EnergyRouter kwarg + # lambda_balance is a router config 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 @@ -118,12 +124,8 @@ 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) @@ -150,9 +152,7 @@ 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,21 +185,15 @@ 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(): @@ -236,9 +230,7 @@ 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) @@ -254,9 +246,7 @@ 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 @@ -270,9 +260,7 @@ 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) @@ -375,18 +363,18 @@ def test_build_router_from_cfg_sets_gumbel_for_composed_router(): assert router.gumbel is True -def test_routed_denoising_mlp_forward_runs_with_gumbel_enabled(): +def test_routed_stage1_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.router.gumbel = True - model.router.gumbel_tau = 0.5 + model.trunk.router.gumbel = True + model.trunk.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, t, cond_cont, cond_cat) + out = model(x_t, cond_cont, cond_cat, t=t) assert out.shape == (B, X_DIM) assert torch.isfinite(out).all() @@ -409,9 +397,7 @@ 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(): @@ -463,59 +449,89 @@ 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(): - model_config = dict( + cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, - 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, - }, + particle_type="embedding", + material_type="embedding", + stage1_router={"enabled": True, "type": "pdg", "n_experts": 3}, ) - 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 + 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 def test_build_models_rejects_pdg_router_with_physical_conditioning(): - """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( + """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( pdg_vocab=4, mat_vocab=2, - 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}, + stage1_router={"enabled": True, "type": "pdg", "n_experts": 3}, ) with pytest.raises(ValueError, match="physical"): - build_models(model_config) + build_models(cfg) def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditioning(): - model_config = dict( + cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, - emb_dim=16, - dropout=0.1, - k_max=K_MAX, - expert_hidden_dim=16, - expert_n_blocks=2, - conditioning="physical", - router={ + stage1_router={ "enabled": True, "type": "composed", "axis0_type": "energy", @@ -525,7 +541,7 @@ def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditi }, ) with pytest.raises(ValueError, match="physical"): - build_models(model_config) + build_models(cfg) # ── ProcessRouter ──────────────────────────────────────────────────────────── @@ -546,9 +562,7 @@ 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(): @@ -597,43 +611,38 @@ def test_build_router_process_type_uses_pdg_mat_vocab(): def test_build_models_routed_with_process_router(): - model_config = dict( + cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, - emb_dim=16, - dropout=0.1, - k_max=K_MAX, - expert_hidden_dim=16, - expert_n_blocks=2, - router={ + particle_type="embedding", + material_type="embedding", + stage1_router={ "enabled": True, "type": "process", "n_experts": 3, "lambda_proc": 1.0, }, ) - 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 + 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 # ── 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) @@ -670,9 +679,7 @@ 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) @@ -680,9 +687,7 @@ 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) @@ -777,15 +782,12 @@ def test_build_composed_router_resolves_per_axis_specs(): def test_build_models_routed_with_composed_router(): - model_config = dict( + cfg = _nested_cfg( pdg_vocab=5, mat_vocab=2, - emb_dim=16, - dropout=0.1, - k_max=K_MAX, - expert_hidden_dim=16, - expert_n_blocks=2, - router={ + particle_type="embedding", + material_type="embedding", + stage1_router={ "enabled": True, "type": "composed", "axis0_type": "energy", @@ -793,29 +795,58 @@ 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, + }, ) - 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 + 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 def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): from giant.sample import sample_flow, sample_secondaries - model_config = dict( + cfg = _nested_cfg( pdg_vocab=3, mat_vocab=2, - emb_dim=16, - dropout=0.1, - k_max=K_MAX, - expert_hidden_dim=8, - expert_n_blocks=1, - router={ + # 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={ "enabled": True, "type": "composed", "axis0_type": "energy", @@ -824,24 +855,31 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): "axis1_n_experts": 2, }, ) - stage1, sec_decoder = build_models(model_config) + models = build_models(cfg) + stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None 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( - sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2 + stage2, 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) -# ── RoutedDenoisingMLP ─────────────────────────────────────────────────────── +# ── Stage1Model with a routed trunk ───────────────────────────────────────── -def test_routed_denoising_mlp_output_shape_train_and_eval(): +def test_routed_stage1_output_shape_train_and_eval(): B = 8 model = _routed_stage1() x_t = torch.randn(B, X_DIM) @@ -849,16 +887,16 @@ def test_routed_denoising_mlp_output_shape_train_and_eval(): cond_cont, cond_cat = _cond(B) model.train() - out_train = model(x_t, t, cond_cont, cond_cat) + out_train = model(x_t, cond_cont, cond_cat, t=t) assert out_train.shape == (B, X_DIM) model.eval() with torch.no_grad(): - out_eval = model(x_t, t, cond_cont, cond_cat) + out_eval = model(x_t, cond_cont, cond_cat, t=t) assert out_eval.shape == (B, X_DIM) -def test_routed_denoising_mlp_gradients_flow_in_train_mode(): +def test_routed_stage1_gradients_flow_in_train_mode(): """Soft mixture in train mode should touch every expert's parameters.""" B = 8 model = _routed_stage1(n_experts=3) @@ -866,14 +904,14 @@ def test_routed_denoising_mlp_gradients_flow_in_train_mode(): t = torch.rand(B) cond_cont, cond_cat = _cond(B) model.train() - flow_loss = model(x_t, t, cond_cont, cond_cat).sum() + flow_loss = model(x_t, cond_cont, cond_cat, t=t).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_denoising_mlp_eval_dispatch_matches_manual_grouping(): +def test_routed_stage1_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 @@ -884,20 +922,20 @@ def test_routed_denoising_mlp_eval_dispatch_matches_manual_grouping(): cond_cont, cond_cat = _cond(B) with torch.no_grad(): - batched = model(x_t, t, cond_cont, cond_cat) + batched = model(x_t, cond_cont, cond_cat, t=t) 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.router.top1(cond_cont, cond_cat) + idx = model.trunk.router.top1(cond_cont, cond_cat) manual = torch.zeros_like(x_t) for i in range(B): - manual[i] = model.experts[int(idx[i])](x_t[i : i + 1], cond[i : i + 1])[0] + manual[i] = model.trunk.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_denoising_mlp_predict_n_sec_shape(): +def test_routed_stage1_predict_n_sec_shape(): B = 6 model = _routed_stage1() cond_cont, cond_cat = _cond(B) @@ -905,15 +943,15 @@ def test_routed_denoising_mlp_predict_n_sec_shape(): assert logits.shape == (B, K_MAX + 1) -def test_routed_denoising_mlp_has_no_pdg_embedding_weight_method(): +def test_routed_stage1_has_no_pdg_embedding_weight_method(): model = _routed_stage1(pdg=5, mat=2) assert not hasattr(model, "pdg_embedding_weight") -# ── RoutedSecondaryDecoder ─────────────────────────────────────────────────── +# ── Stage2OneShot with a routed trunk ──────────────────────────────────────── -def test_routed_secondary_decoder_output_shape_train_and_eval(): +def test_routed_stage2_output_shape_train_and_eval(): B = 8 decoder = _routed_sec_decoder() x_t = torch.randn(B, SEC_DIM) @@ -922,16 +960,16 @@ def test_routed_secondary_decoder_output_shape_train_and_eval(): stage1_out = torch.randn(B, X_DIM) decoder.train() - out_train = decoder(x_t, t, cond_cont, cond_cat, stage1_out) + out_train = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t) assert out_train.shape == (B, SEC_DIM) decoder.eval() with torch.no_grad(): - out_eval = decoder(x_t, t, cond_cont, cond_cat, stage1_out) + out_eval = decoder(x_t, cond_cont, cond_cat, stage1_out, t=t) assert out_eval.shape == (B, SEC_DIM) -def test_routed_secondary_decoder_gradients_flow(): +def test_routed_stage2_gradients_flow(): B = 4 decoder = _routed_sec_decoder(n_experts=3) x_t = torch.randn(B, SEC_DIM) @@ -939,7 +977,9 @@ def test_routed_secondary_decoder_gradients_flow(): cond_cont, cond_cat = _cond(B) stage1_out = torch.randn(B, X_DIM) decoder.train() - decoder(x_t, t, cond_cont, cond_cat, stage1_out).sum().backward() + 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() for name, p in decoder.named_parameters(): assert p.grad is not None, f"no grad for {name}" @@ -948,46 +988,33 @@ def test_routed_secondary_decoder_gradients_flow(): def test_build_models_monolith_when_router_absent(): - 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) + 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) def test_build_models_monolith_when_router_disabled(): - model_config = dict( + cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, - hidden_dim=32, - n_blocks=2, - emb_dim=16, - dropout=0.1, - k_max=K_MAX, - router={"enabled": False, "type": "energy", "n_experts": 4}, + stage1_router={"enabled": False, "type": "energy", "n_experts": 4}, ) - stage1, sec_decoder = build_models(model_config) - assert isinstance(stage1, DenoisingMLP) - assert isinstance(sec_decoder, SecondaryDecoder) + 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) def test_build_models_routed_when_enabled(): - model_config = dict( + cfg = _nested_cfg( pdg_vocab=4, mat_vocab=2, - emb_dim=16, - dropout=0.1, - k_max=K_MAX, - expert_hidden_dim=16, - expert_n_blocks=2, - router={ + stage1_router={ "enabled": True, "type": "energy", "n_experts": 4, @@ -995,37 +1022,49 @@ 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, + }, ) - 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 + 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 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 - model_config = dict( + cfg = _nested_cfg( pdg_vocab=3, mat_vocab=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}, + stage1_router={"enabled": True, "type": "energy", "n_experts": 2}, ) - stage1, sec_decoder = build_models(model_config) + models = build_models(cfg) + stage1, stage2 = models["stage1"], models["stage2"] + assert stage1 is not None and stage2 is not None 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( - sec_decoder, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2 + stage2, 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) diff --git a/tests/test_router_gating.py b/tests/test_router_gating.py index 51123d3..18be8c8 100644 --- a/tests/test_router_gating.py +++ b/tests/test_router_gating.py @@ -36,7 +36,8 @@ def _model_cfg() -> dict: def _write_checkpoint(tmp_path) -> str: cfg = _model_cfg() - stage1, _ = build_models(cfg) + stage1 = build_models(cfg)["stage1"] + assert stage1 is not None norm = Normalizer() norm.mean = np.zeros(15, dtype=np.float32) norm.std = np.ones(15, dtype=np.float32) diff --git a/tests/test_sample.py b/tests/test_sample.py new file mode 100644 index 0000000..7906069 --- /dev/null +++ b/tests/test_sample.py @@ -0,0 +1,224 @@ +"""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]] diff --git a/tests/test_setup_cache.py b/tests/test_setup_cache.py index a243a53..584b5e8 100644 --- a/tests/test_setup_cache.py +++ b/tests/test_setup_cache.py @@ -7,15 +7,14 @@ 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 @@ -28,9 +27,7 @@ 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 ───────────────────────────────────────────────────────── @@ -38,9 +35,7 @@ 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): @@ -50,10 +45,7 @@ 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 ──────────────────────────────────────────────────── @@ -101,6 +93,37 @@ 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 @@ -149,6 +172,25 @@ 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] @@ -325,3 +367,11 @@ 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 diff --git a/tests/test_steps_to_parquet.py b/tests/test_steps_to_parquet.py index 8e839c1..169bda8 100644 --- a/tests/test_steps_to_parquet.py +++ b/tests/test_steps_to_parquet.py @@ -24,18 +24,14 @@ 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 @@ -69,9 +65,7 @@ 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]] diff --git a/tests/test_steps_to_parquet_parallel.py b/tests/test_steps_to_parquet_parallel.py index 55e0cce..101d36f 100644 --- a/tests/test_steps_to_parquet_parallel.py +++ b/tests/test_steps_to_parquet_parallel.py @@ -124,31 +124,13 @@ 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): diff --git a/tests/test_train.py b/tests/test_train.py index 627fe38..fc5781d 100644 --- a/tests/test_train.py +++ b/tests/test_train.py @@ -1,6 +1,47 @@ -"""Tests for giant/train.py helpers.""" +"""Tests for giant/training/.""" -from giant.train import _gumbel_tau, _wandb_run_config +import copy +import csv +import math +import tempfile +from pathlib import Path + +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.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 def test_gumbel_tau_at_step_zero_is_start(): @@ -26,72 +67,627 @@ def test_gumbel_tau_handles_zero_total_steps(): assert abs(_gumbel_tau(1, 0, 1.0, 0.1) - 0.1) < 1e-9 -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_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 test_wandb_run_config_handles_missing_model_config(): - cfg = _wandb_run_config(**_base_wandb_kwargs(model_config=None)) - assert cfg["model"] == {} - assert "lambda_balance" not in cfg + 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((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" + + +# --- 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 diff --git a/tests/test_transforms.py b/tests/test_transforms.py index dcda475..7f19326 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -162,9 +162,7 @@ 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) @@ -200,14 +198,10 @@ 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) @@ -221,12 +215,8 @@ 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) @@ -319,7 +309,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]) @@ -356,24 +346,20 @@ 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]) @@ -396,9 +382,7 @@ 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() @@ -413,11 +397,7 @@ 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 @@ -426,7 +406,13 @@ 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, conditioning="embedding") + cond_cont, *_ = build_features( + data, + pdg_map, + mat_map, + particle_conditioning="embedding", + material_conditioning="embedding", + ) assert cond_cont.shape[1] == COND_DIM np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0) @@ -438,14 +424,18 @@ 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, conditioning="physical") + cond_cont, *_ = build_features( + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_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 @@ -461,7 +451,13 @@ 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, conditioning="physical") + build_features( + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_conditioning="physical", + ) def test_build_cond_features_mass_charge_override(fake_material_props): @@ -474,11 +470,15 @@ 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, conditioning="physical") - - np.testing.assert_allclose( - cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])) + cond_cont, _ = build_cond_features( + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_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 + 1], [2.0, -2.0]) @@ -494,7 +494,12 @@ 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, conditioning="embedding" + data, + pdg_map, + mat_map, + cond_normalizer=legacy_norm, + particle_conditioning="embedding", + material_conditioning="embedding", ) assert cond_cont.shape[-1] == COND_DIM @@ -521,7 +526,8 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode( pdg_map, mat_map, cond_normalizer=legacy_norm, - conditioning="physical", + particle_conditioning="physical", + material_conditioning="physical", ) @@ -610,13 +616,23 @@ 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, conditioning="physical" + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_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, conditioning="embedding") + build_cond_features( + data, + pdg_map, + mat_map, + particle_conditioning="embedding", + material_conditioning="embedding", + ) # ── _WelfordAccumulator ────────────────────────────────────────────────────── @@ -674,9 +690,7 @@ 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: diff --git a/tests/test_type_embedding_distance.py b/tests/test_type_embedding_distance.py new file mode 100644 index 0000000..6b3b8dd --- /dev/null +++ b/tests/test_type_embedding_distance.py @@ -0,0 +1,43 @@ +"""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 diff --git a/tests/test_validate.py b/tests/test_validate.py index ae0025f..ca220ae 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -1,30 +1,63 @@ import numpy as np import torch -from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM -from giant.model.network import DenoisingMLP, SecondaryDecoder +from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM +from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim 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(): - 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) + +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 return s1.eval(), s2.eval() -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.""" +def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8): + """A val_loader matching StreamingStepsDataset's 7-tuple batch shape: + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx).""" 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.zeros(B, dtype=torch.long) - sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM) + n_sec = torch.full((B,), n_sec_value, dtype=torch.long) + sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM) proc_idx = torch.zeros(B, dtype=torch.long) - batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx)) + sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long) + batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx)) return batches @@ -32,22 +65,54 @@ 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 -- a regression the old species/bincount code this - replaced explicitly guarded against.""" + _histogram_kl.""" s1, s2 = _tiny_models() - loader = _zero_secondaries_loader() + loader = _loader(n_sec_value=0) - # 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) + 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) - monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow) + monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec) - result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2) + result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2, steps=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"} diff --git a/tests/test_wgan.py b/tests/test_wgan.py index 4b85b0c..5b17aa8 100644 --- a/tests/test_wgan.py +++ b/tests/test_wgan.py @@ -1,15 +1,13 @@ import torch from giant.constants import COND_DIM, K_MAX, SEC_DIM, SEC_SLOT_DIM, X_DIM -from giant.model.network import ( - Critic, - SecondaryCritic, - WGANGenerator, - WGANSecondaryGenerator, -) +from giant.model.network import CriticModel, Stage1Model, Stage2OneShot 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) @@ -18,23 +16,56 @@ def _cond(B=8): def _small_generator(): - return WGANGenerator( - pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8 + 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, ) def _small_critic(): - return Critic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2) + 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", + ) def _small_sec_generator(): - return WGANSecondaryGenerator( - pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2, noise_dim=8 + 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, ) def _small_sec_critic(): - return SecondaryCritic(pdg_vocab=3, mat_vocab=2, hidden_dim=32, n_blocks=2) + 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", + ) def _mask(B, n_sec): @@ -89,7 +120,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.shape == (B,) + assert n_sec is not None and n_sec.shape == (B,) # --- Stage-2 generator/critic --- @@ -121,9 +152,7 @@ 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) @@ -152,9 +181,7 @@ 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 @@ -164,9 +191,7 @@ 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()) diff --git a/uv.lock b/uv.lock index bd8cb35..97f81e7 100644 --- a/uv.lock +++ b/uv.lock @@ -357,6 +357,105 @@ 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" @@ -534,7 +633,7 @@ wheels = [ [[package]] name = "giant" -version = "0.2.0" +version = "0.3.0" source = { editable = "." } dependencies = [ { name = "numpy" }, @@ -572,6 +671,7 @@ dev = [ { name = "plotstyle" }, { name = "polars" }, { name = "pytest" }, + { name = "pytest-cov" }, { name = "ruff" }, { name = "scikit-learn" }, { name = "ty" }, @@ -599,6 +699,7 @@ 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" }, @@ -1713,6 +1814,20 @@ 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"