From 93b19911f8e6370999f7eade04b1e1532bc392d3 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 7 Aug 2026 10:37:57 +0200 Subject: [PATCH] v0.3.0 step 6: sample.py/rollout.py AR generation + class->PDG decode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - giant/sample.py: fix every sampler's call convention against Stage1Model/Stage2OneShot's actual forward signatures (was still calling model(x, t, cond_cont, cond_cat) positionally); add sample_secondaries_ar (free-running AR loop, unsnapped history feature) and sample_stage1/sample_stage2/resolve_n_sec dispatch helpers that read each stage's generator_kind/decoder off the model instance itself. - giant/particles.py: decode_topn_class (argmax + other_policy) and decode_embedding_nearest (L1-snap + distance) turn a secondary's "onehot"/"embedding" type prediction into a concrete PDG. - giant/rollout.py: decode_secondary_identity routes all three particle_type.target values to real mass/charge; per-stage generator dispatch (drops the single shared `mode` string, adds ddpm support); L1DistCollector accumulates the §11.3 embedding-distance diagnostic. - giant/cli.py: drop the onehot/embedding-target rejection gate (narrowed to the still-unimplemented conditioning.particle/material.type=onehot axis); fix the dead model_cfg.get("mode") bug in predict/rollout. - giant/analysis/: new type_embedding_l1_distance PlotSpec, wired through the rollout YAML sidecar (no live-model call needed, unlike router_gating -- the histogram is already pre-aggregated at rollout time). - Un-xfail every test that was blocked on this step (test_rollout.py, test_flow.py, test_wgan.py, test_phase2.py, test_router.py, test_validate.py); add test_sample.py, test_type_embedding_distance.py. Known follow-up: giant/validate.py still unpacks the training val-batch as a stale 6-tuple and doesn't use the new per-stage dispatch, so marginal validation during training degrades gracefully with a warning rather than working -- not in this step's scope. Co-Authored-By: Claude Sonnet 5 --- giant/analysis/catalog.py | 18 + giant/analysis/condor.py | 15 +- giant/analysis/render.py | 2 + giant/analysis/type_embedding_distance.py | 70 ++++ giant/cli.py | 173 +++++--- giant/data/transforms.py | 91 ++-- giant/model/network.py | 2 + giant/particles.py | 122 ++++++ giant/rollout.py | 259 ++++++++++-- giant/sample.py | 489 +++++++++++++++++----- giant/train.py | 28 +- tests/test_cli_predict.py | 37 +- tests/test_flow.py | 10 - tests/test_particles.py | 102 +++++ tests/test_phase2.py | 9 - tests/test_rollout.py | 265 +++++++++++- tests/test_router.py | 33 +- tests/test_sample.py | 237 +++++++++++ tests/test_type_embedding_distance.py | 43 ++ tests/test_validate.py | 10 - 20 files changed, 1696 insertions(+), 319 deletions(-) create mode 100644 giant/analysis/type_embedding_distance.py create mode 100644 tests/test_sample.py create mode 100644 tests/test_type_embedding_distance.py diff --git a/giant/analysis/catalog.py b/giant/analysis/catalog.py index fc665e3..8f5c751 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 + # §11.3 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, ) @@ -658,6 +666,9 @@ _router_share_pdg_partial, _router_share_pdg_finalize = _unchunkable( _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)) +) # --------------------------------------------------------------------------- @@ -831,6 +842,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..ec2d3e0 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", + # §11.3 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", ) @@ -247,6 +253,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. @@ -265,7 +272,12 @@ def compute_reduced( f"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 +303,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"), ) diff --git a/giant/analysis/render.py b/giant/analysis/render.py index b746942..ec166ca 100644 --- a/giant/analysis/render.py +++ b/giant/analysis/render.py @@ -144,6 +144,8 @@ def _render_single(r: Reduced, params: dict): ) 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") diff --git a/giant/analysis/type_embedding_distance.py b/giant/analysis/type_embedding_distance.py new file mode 100644 index 0000000..adb39c6 --- /dev/null +++ b/giant/analysis/type_embedding_distance.py @@ -0,0 +1,70 @@ +"""Secondary-type embedding-distance diagnostic (docs/v0.3.0-design.md §11.3). + +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/cli.py b/giant/cli.py index 99dfdba..d476390 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -34,24 +34,22 @@ 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.data.setup_cache import topnmap_from_json 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) @@ -86,41 +84,66 @@ def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str: return raw -def _check_v030_onehot_support(model_cfg: dict, command: str) -> None: - """`giant predict`/`giant rollout` don't yet support conditioning - `"onehot"` mode or `stage2_model.particle_type.target` in `("onehot", - "embedding")` — full decode (class index -> concrete PDG, `other_policy` - sampling) is v0.3.0 step 6 (docs/v0.3.0-design.md), which also lands the - autoregressive decoder these targets are meant to pair with. Without - this guard, predict/rollout would crash later on a `cond_cat`/`sec_dim` - shape mismatch (onehot) or silently read a meaningless raw (log_mass, - charge) pair (embedding) instead of failing clearly. A v0.2 (flat) - `model_config` never has these, so this is a no-op there. +def _check_conditioning_onehot_support(model_cfg: dict, command: str) -> None: + """`giant predict`/`giant rollout` don't yet support + `conditioning.particle/material.type = "onehot"` — that needs + `pdg_topn_map.class_map`/`mat_topn_map.class_map` threaded into every + `build_cond_features` call in this module and `giant/rollout.py`'s + `_step_chunk`, which hasn't been wired (a separate axis from + `stage2_model.particle_type.target`, handled below — see + docs/v0.3.0-design.md §3.1 vs §3.3). Without this guard, predict/rollout + would silently build a `cond_cat` missing the topN columns + `ConditionEncoder`'s `"onehot"` mode expects. A v0.2 (flat) `model_config` + never has this, so this is a no-op there. """ conditioning = model_cfg.get("conditioning") if not isinstance(conditioning, dict): return particle_type = conditioning.get("particle", {}).get("type") material_type = conditioning.get("material", {}).get("type") - particle_type_target = ( - model_cfg.get("stage2_model", {}).get("particle_type", {}).get("target") - ) - if ( - particle_type == "onehot" - or material_type == "onehot" - or particle_type_target in ("onehot", "embedding") - ): + if particle_type == "onehot" or material_type == "onehot": typer.echo( - f"error: giant {command} does not yet support conditioning " - "onehot mode or stage2_model.particle_type.target in " - "('onehot', 'embedding') — full decode (class index -> concrete " - "PDG) lands in v0.3.0 step 6 alongside the autoregressive " - "decoder these targets are meant to pair with.", + f"error: giant {command} does not yet support " + "conditioning.particle/material.type = 'onehot' — only " + "stage2_model.particle_type.target in ('onehot', 'embedding') " + "is implemented (docs/v0.3.0-design.md step 6).", err=True, ) raise typer.Exit(1) +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): + """`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 _batch_size_estimate_dims( model_cfg: dict, training: bool, stage: str = "stage1" ) -> tuple[int, int]: @@ -965,7 +988,10 @@ def predict( raise typer.Exit(1) model_cfg = ckpt["model_config"] - _check_v030_onehot_support(model_cfg, "predict") + _check_conditioning_onehot_support(model_cfg, "predict") + pdg_topn_map = _load_pdg_topn_map(ckpt) + other_policy = _particle_type_other_policy(model_cfg) + stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1") if batch_size_auto: est_hidden_dim, est_n_blocks = _batch_size_estimate_dims( @@ -1046,23 +1072,21 @@ def predict( 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 (decision 1) — + # 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 @@ -1119,19 +1143,26 @@ def predict( 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, - n_sec_pred_np, - e_sec_pred, - piece["pre_dir"], - sec_phys_normalizer=sec_phys_norm, + # 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_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) ] @@ -1378,8 +1409,12 @@ def rollout( training_cfg = gconfig.load_checkpoint_config(checkpoint) model_cfg = ckpt["model_config"] - _check_v030_onehot_support(model_cfg, "rollout") + _check_conditioning_onehot_support(model_cfg, "rollout") conditioning = _conditioning_str(model_cfg) + pdg_topn_map = _load_pdg_topn_map(ckpt) + other_policy = _particle_type_other_policy(model_cfg) + stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1") + stage2_ddpm_steps = _ddpm_steps(model_cfg, "stage2") 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"]) @@ -1432,6 +1467,10 @@ def rollout( writer = pq.ParquetWriter(out, table.schema) writer.write_table(table) + # Only meaningful under particle_type.target="embedding" (§11.3) — a + # no-op collector otherwise, cheaper than branching the call itself. + l1_dist_collector = L1DistCollector() + summary = run_rollout( model, sec_decoder, @@ -1451,11 +1490,18 @@ def rollout( escape_threshold=escape_threshold, on_chunk=_write_chunk, conditioning=conditioning, - mode=model_cfg.get("mode", "flow"), + pdg_topn_map=pdg_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( @@ -1475,6 +1521,13 @@ def rollout( "rollout_seed": seed, "n_rows": summary["n_rows"], "termination_reason_counts": summary["termination_reason_counts"], + # §11.3 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, ...) diff --git a/giant/data/transforms.py b/giant/data/transforms.py index a605cab..562abf0 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -587,42 +587,37 @@ def encode_secondary_type_idx( return np.where(sec_valid, idx, 0).astype(np.int64) -def decode_secondaries( +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 (docs/v0.3.0-design.md §6.1) 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. @@ -666,6 +661,54 @@ def decode_secondaries( pre_dir[valid], dir_local[valid, i] ) + return sec_E, sec_dir_world, sec_valid + + +def decode_secondaries( + sec_cont: np.ndarray, + n_sec: np.ndarray, + e_sec: np.ndarray, + pre_dir: np.ndarray, + sec_phys_normalizer: "Normalizer | None" = None, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Inverse of encode_secondaries: continuous targets → physical secondary attrs. + + `particle_type.target = "physical"` only (the type slice is a raw + (log_mass, charge) regression target folded straight into `sec_cont`) — + `"onehot"`/`"embedding"` decode through `decode_secondary_cont` + + `giant.particles.decode_topn_class`/`decode_embedding_nearest` instead, + since their type slice isn't (log_mass, charge) at all. See + `decode_secondary_cont`'s docstring for why the two share the energy/ + direction logic below. + + sec_cont: (N, K_MAX, 6) — [stick_logit, local_dir_x, local_dir_y, + local_dir_z, log_mass, charge] (log_mass/charge normalised iff + `sec_phys_normalizer` was applied when this was produced — e.g. a + raw model prediction; pass the same normalizer here to invert it) + n_sec: (N,) integer secondary counts + e_sec: (N,) total secondary energy budget [MeV] + pre_dir: (N, 3) pre-step world-frame direction + + Returns (sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid) each + shape (N, K_MAX). mass/charge are the model's raw predicted physical + identity for each secondary, used as-is (no snapping to a discrete PDG + code) — see giant/particles.py for the separate, reporting-only + nearest-PDG lookup callers may apply on top of this for display/ + bookkeeping purposes. + """ + if sec_phys_normalizer is not None: + N_, K_, _ = sec_cont.shape + phys = sec_phys_normalizer.inverse_transform(sec_cont[:, :, 4:6].reshape(-1, 2)) + sec_cont = sec_cont.copy() + sec_cont[:, :, 4:6] = phys.reshape(N_, K_, 2) + + sec_E, sec_dir_world, sec_valid = decode_secondary_cont( + sec_cont, n_sec, e_sec, pre_dir + ) + + log_mass = sec_cont[:, :, 4] # (N, K) + charge = sec_cont[:, :, 5] # (N, K) + # mass is non-negative by construction (inv_log_transform of a real # number is always > 0); clip to 0 for padded/invalid slots rather than # leaving a spurious small positive floor from the log inverse. diff --git a/giant/model/network.py b/giant/model/network.py index 581e1dd..0624e13 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -982,7 +982,9 @@ class Stage2OneShot(nn.Module): super().__init__() self.generator_kind = generator self.noise_dim = noise_dim + self.k_max = k_max self.particle_type_cfg = dict(particle_type_cfg or {"target": "physical"}) + self.type_dim = stage2_type_dim(self.particle_type_cfg, particle_cfg["emb_dim"]) self.cond_enc = ConditionEncoder( pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim ) diff --git a/giant/particles.py b/giant/particles.py index 4e5ea17..7e2d5ac 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 (§3.3/§8/§11.3 +of docs/v0.3.0-design.md): `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 @@ -111,3 +123,113 @@ 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 (docs/v0.3.0-design.md §3.3): 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 + (docs/v0.3.0-design.md §3.3): 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 the §11.3 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/rollout.py b/giant/rollout.py index 8270a5e..3bca136 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,165 @@ 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 §11.3 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 (docs/v0.3.0-design.md §11.3). + + 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` + (docs/v0.3.0-design.md §3.3): + + - `"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 (§11.3 + diagnostic — see `giant/rollout.py`'s L1-distance accumulator). + + 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 = [ @@ -311,7 +458,12 @@ def rollout( escape_threshold: float | None = None, on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None, conditioning: str = "embedding", - mode: str = "flow", + pdg_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 +476,30 @@ 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` + (docs/v0.3.0-design.md decision 2: 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`/`other_policy` are only read under + `stage2_model.particle_type.target = "onehot"` (§3.3); `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 §11.3 embedding-distance + diagnostic across the whole run — see `L1DistCollector`. Only populated + under `particle_type.target = "embedding"`; a no-op otherwise. """ 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"], @@ -365,7 +535,12 @@ def rollout( device, max_tracks_per_event, conditioning, - mode, + pdg_topn_map, + other_policy, + rng, + stage1_ddpm_steps, + stage2_ddpm_steps, + l1_dist_collector, ) ) frontier = _concat_frontiers(next_parts) @@ -396,7 +571,12 @@ def _step_chunk( device, max_tracks_per_event, conditioning, - mode="flow", + pdg_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"]) @@ -480,10 +660,9 @@ def _step_chunk( 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]) @@ -503,32 +682,40 @@ def _step_chunk( 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, - n_sec_np, - e_sec, - tr["pre_dir"], - sec_phys_normalizer=sec_phys_norm, + # 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_pdg_code = nearest_known_pdg( - sec_mass.reshape(-1), sec_charge.reshape(-1), pdg_map.keys() - ).reshape(sec_mass.shape) + 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_norm, + pdg_map, + pdg_topn_map, + other_policy, + rng, + ) + ) + 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..f9ba019 100644 --- a/giant/sample.py +++ b/giant/sample.py @@ -1,26 +1,23 @@ 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 + (docs/v0.3.0-design.md decision 1 moves it to stage 2 for fresh runs — + 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 +26,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 +42,142 @@ 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, docs/v0.3.0-design.md decision 2).""" + 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 (docs/v0.3.0-design.md + decision 2/3): 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 +189,27 @@ 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 +220,231 @@ 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 (docs/v0.3.0-design.md §6.4): + one token at a time, in descending-energy slot order, `k_max` sequential + calls. Unlike training (teacher forcing, §6.2 point 3 — a single + parallel pass over ground-truth tokens, see + `giant.train._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 + §6.2 point 4 explicitly flags as the cost of markov history's + expressiveness. + + A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass — + §6.4's "K sequential forwards" cost note applies per-token here, not + once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per + physics step. + + 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` was 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) + + 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 + ) + + 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, + ) 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, + ) + 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, + ).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 (docs/v0.3.0-design.md +# decision 2: 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 under + decision 1 (docs/v0.3.0-design.md §2). 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 (§3.3, §9).""" + 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 " + "(docs/v0.3.0-design.md §3.3/§9)" + ) + 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 index a117d2f..a504056 100644 --- a/giant/train.py +++ b/giant/train.py @@ -991,27 +991,31 @@ _WARNED_MARGINAL_VALIDATION_BROKEN = False def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs): - """`validate_marginals` delegates to `giant.sample`'s samplers, which - still assume Stage 1 always owns `n_sec_head` (v0.2 behaviour) — under - decision 1 (docs/v0.3.0-design.md §2) a fresh Stage2OneShot owns it by - default instead, so this currently raises for any non-legacy checkpoint. - `giant/sample.py` needs a per-stage generator/n_sec update (design doc - §10, deferred to step 6); until then this degrades gracefully with a - one-time warning instead of crashing the whole training run. + """`giant.sample` itself has a per-stage generator/n_sec-ownership + update (v0.3.0 step 6 — `sample_stage1`/`sample_stage2`/`resolve_n_sec`, + docs/v0.3.0-design.md §10); `giant/validate.py` hasn't been updated to + call it yet, and also still unpacks the val batch as a 6-tuple, which no + longer matches `StreamingStepsDataset`'s 7-tuple shape (the `sec_type_idx` + column step 4 added). Both are `giant/validate.py`-side gaps, not + `giant/sample.py`'s — tracked, not yet scheduled. Until fixed, this + degrades gracefully with a one-time warning instead of crashing the + whole training run. """ global _WARNED_MARGINAL_VALIDATION_BROKEN model = trainer.sampling_model() try: return validate_marginals(model, val_loader, device=device, **kwargs) except Exception as exc: # noqa: BLE001 — see docstring: any failure here - # is expected until giant/sample.py's per-stage n_sec update lands. + # is expected until giant/validate.py is updated for the v0.3.0 + # per-stage dispatch + 7-tuple batch shape. if not _WARNED_MARGINAL_VALIDATION_BROKEN: warnings.warn( "marginal validation unavailable this run " - f"({type(exc).__name__}: {exc}) — giant/sample.py doesn't yet " - "support a Stage2-owned n_sec head (docs/v0.3.0-design.md " - "step 6); val_marginal_kl stays NaN, and wgan best-checkpoint " - "selection falls back to the Wasserstein-distance magnitude.", + f"({type(exc).__name__}: {exc}) — giant/validate.py hasn't " + "been updated for v0.3.0's per-stage sample dispatch / " + "7-tuple val batch shape yet; val_marginal_kl stays NaN, and " + "wgan best-checkpoint selection falls back to the " + "Wasserstein-distance magnitude.", stacklevel=2, ) _WARNED_MARGINAL_VALIDATION_BROKEN = True diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py index 8bae93f..1ee8af6 100644 --- a/tests/test_cli_predict.py +++ b/tests/test_cli_predict.py @@ -6,14 +6,14 @@ import yaml from giant.cli import ( _CEPH_PREDICTIONS, - _check_v030_onehot_support, + _check_conditioning_onehot_support, _resolve_prediction_output, _write_prediction_ref, ) # --------------------------------------------------------------------------- -# _check_v030_onehot_support +# _check_conditioning_onehot_support # --------------------------------------------------------------------------- @@ -29,42 +29,43 @@ def _nested_model_cfg( } -def test_check_v030_onehot_support_allows_physical(): - _check_v030_onehot_support(_nested_model_cfg(), "predict") # no raise +def test_check_conditioning_onehot_support_allows_physical(): + _check_conditioning_onehot_support(_nested_model_cfg(), "predict") # no raise -def test_check_v030_onehot_support_rejects_onehot_particle_conditioning(): +def test_check_conditioning_onehot_support_rejects_onehot_particle_conditioning(): cfg = _nested_model_cfg(particle_type="onehot") with pytest.raises(typer.Exit): - _check_v030_onehot_support(cfg, "predict") + _check_conditioning_onehot_support(cfg, "predict") -def test_check_v030_onehot_support_rejects_onehot_material_conditioning(): +def test_check_conditioning_onehot_support_rejects_onehot_material_conditioning(): cfg = _nested_model_cfg(material_type="onehot") with pytest.raises(typer.Exit): - _check_v030_onehot_support(cfg, "rollout") + _check_conditioning_onehot_support(cfg, "rollout") -def test_check_v030_onehot_support_rejects_onehot_particle_type_target(): +def test_check_conditioning_onehot_support_allows_onehot_particle_type_target(): + """stage2_model.particle_type.target="onehot" is implemented (v0.3.0 + step 6, giant.rollout.decode_secondary_identity) — it's a separate axis + from conditioning.particle.type, which this guard doesn't gate at all.""" cfg = _nested_model_cfg(target="onehot") - with pytest.raises(typer.Exit): - _check_v030_onehot_support(cfg, "predict") + _check_conditioning_onehot_support(cfg, "predict") # no raise -def test_check_v030_onehot_support_rejects_embedding_particle_type_target(): +def test_check_conditioning_onehot_support_allows_embedding_particle_type_target(): cfg = _nested_model_cfg( particle_type="embedding", material_type="embedding", target="embedding" ) - with pytest.raises(typer.Exit): - _check_v030_onehot_support(cfg, "predict") + _check_conditioning_onehot_support(cfg, "predict") # no raise -def test_check_v030_onehot_support_is_noop_for_v02_flat_model_config(): +def test_check_conditioning_onehot_support_is_noop_for_v02_flat_model_config(): """A v0.2 checkpoint's flat model_config has conditioning as a plain - string, not a dict — never onehot/embedding-target, so this must be a - silent no-op rather than crash on `.get("particle")` against a string.""" + string, not a dict — never onehot, so this must be a silent no-op rather + than crash on `.get("particle")` against a string.""" cfg = {"conditioning": "embedding", "mode": "flow"} - _check_v030_onehot_support(cfg, "predict") # no raise + _check_conditioning_onehot_support(cfg, "predict") # no raise # --------------------------------------------------------------------------- diff --git a/tests/test_flow.py b/tests/test_flow.py index 6a6730e..f1276f0 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -1,4 +1,3 @@ -import pytest import torch from giant.constants import COND_DIM from giant.model.network import Stage1Model @@ -8,13 +7,6 @@ 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} -_SAMPLE_XFAIL_REASON = ( - "giant/sample.py isn't updated yet — its sample_flow/sample_ddim call " - "models positionally as model(x, t, cond_cont, cond_cat), which doesn't " - "match Stage1Model's new forward signature. Deferred to " - "docs/v0.3.0-design.md step 6." -) - def _small_model(): return Stage1Model( @@ -54,7 +46,6 @@ def test_flow_matching_loss_has_grad(): assert any(p.grad is not None for p in model.parameters()) -@pytest.mark.xfail(reason=_SAMPLE_XFAIL_REASON, strict=False) def test_sample_flow_shape(): B = 6 cond_cont = torch.randn(B, COND_DIM) @@ -71,7 +62,6 @@ def test_ddpm_loss_nonneg(): assert loss.item() >= 0.0 -@pytest.mark.xfail(reason=_SAMPLE_XFAIL_REASON, strict=False) def test_sample_ddim_shape(): B = 4 schedule = CosineSchedule(T=50) diff --git a/tests/test_particles.py b/tests/test_particles.py index 224e2c8..bcfbfdf 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, @@ -126,3 +130,101 @@ 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 f548f7a..afb12bd 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -19,13 +19,6 @@ from giant.model.schedule import ( ) from giant.sample import sample_secondaries -_SAMPLE_SECONDARIES_XFAIL_REASON = ( - "giant/sample.py isn't updated yet — sample_secondaries calls the " - "decoder positionally as decoder(x, t, cond_cont, cond_cat, stage1_out), " - "which doesn't match Stage2OneShot's new forward signature. Deferred to " - "docs/v0.3.0-design.md step 6." -) - # ── helpers ────────────────────────────────────────────────────────────────── @@ -300,7 +293,6 @@ def test_flow_matching_loss_secondary_ar_has_grad(): # ── sampling ────────────────────────────────────────────────────────────────── -@pytest.mark.xfail(reason=_SAMPLE_SECONDARIES_XFAIL_REASON, strict=False) def test_sample_secondaries_shapes(): B, pdg, mat = 6, 3, 2 decoder = _sec_decoder(pdg, mat) @@ -316,7 +308,6 @@ def test_sample_secondaries_shapes(): assert sec_valid.dtype == torch.bool -@pytest.mark.xfail(reason=_SAMPLE_SECONDARIES_XFAIL_REASON, strict=False) def test_sample_secondaries_valid_mask_matches_n_sec(): B, pdg, mat = 4, 3, 2 decoder = _sec_decoder(pdg, mat) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index 0e71b90..f03eab8 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -9,23 +9,19 @@ import pytest import torch from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX +from giant.data.loader import TopNMap from giant.data.transforms import Normalizer -from giant.model.network import Stage1Model, Stage2OneShot -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 -# giant/rollout.py isn't updated yet — it drives Stage1Model/Stage2OneShot -# through giant.sample's sample_flow/sample_secondaries, which still call -# models with the pre-refactor positional convention -# (model(x, t, cond_cont, cond_cat)) that no longer matches these classes' -# forward signatures. Deferred to docs/v0.3.0-design.md step 6/§10. -pytestmark = pytest.mark.xfail( - reason="giant/rollout.py not updated for v0.3.0 network.py yet (step 6)", - strict=False, -) - PDG_MAP = {22: 0, 11: 1, -11: 2} MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1} @@ -350,3 +346,248 @@ 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 + # (decision 1 moves n_sec ownership 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} + common = dict( + 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, + ) + if decoder == "one_shot": + sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator2, k_max, emb_dim) + s2 = Stage2OneShot(sec_dim=sec_dim, **common) + else: + s2 = Stage2Autoregressive(**common) + 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, + 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 (decision 1) — 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 — the matrix docs/v0.3.0-design.md §7 calls out for comparison.""" + 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) + + +@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(): + """§11.3: 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 96b8c5b..8d791d7 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -869,21 +869,17 @@ def test_build_models_routed_stage2_ties_to_stage1_router(): assert stage1.trunk.router is stage2.trunk.router -@pytest.mark.xfail( - reason=( - "giant/sample.py isn't updated yet — its sample_flow/sample_secondaries " - "call models positionally as model(x, t, cond_cont, cond_cat), which " - "doesn't match Stage1Model/Stage2OneShot's new forward signature. " - "Deferred to docs/v0.3.0-design.md step 6." - ), - strict=False, -) def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): from giant.sample import sample_flow, sample_secondaries cfg = _nested_cfg( pdg_vocab=3, mat_vocab=2, + # PdgRouter (axis1) always builds its own training-vocab embedding, + # incompatible with conditioning.particle.type="physical" (the + # _nested_cfg default) — see _check_router_conditioning_compat. + particle_type="embedding", + material_type="embedding", stage1_router={ "enabled": True, "type": "composed", @@ -900,6 +896,11 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow(): 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 (decision 1 moves it 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( @@ -1075,15 +1076,6 @@ def test_build_models_routed_when_enabled(): assert len(stage2.trunk.experts) == 4 -@pytest.mark.xfail( - reason=( - "giant/sample.py isn't updated yet — its sample_flow/sample_secondaries " - "call models positionally as model(x, t, cond_cont, cond_cat), which " - "doesn't match Stage1Model/Stage2OneShot's new forward signature. " - "Deferred to docs/v0.3.0-design.md step 6." - ), - strict=False, -) 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 @@ -1100,6 +1092,11 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow(): 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 (decision 1 moves it 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( diff --git a/tests/test_sample.py b/tests/test_sample.py new file mode 100644 index 0000000..670942a --- /dev/null +++ b/tests/test_sample.py @@ -0,0 +1,237 @@ +"""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 (docs/v0.3.0-design.md step 6).""" + +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 (docs/v0.3.0-design.md §3.3) — 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, +) -> 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}, + ).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 (decision 1) ──────────────────────────────────── + + +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.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("generator", ["flow", "wgan"]) +@pytest.mark.parametrize("target", ["physical", "onehot", "embedding"]) +def test_sample_secondaries_ar_shapes(target, generator): + B, k_max, emb_dim = 4, 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.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_type_embedding_distance.py b/tests/test_type_embedding_distance.py new file mode 100644 index 0000000..c17f965 --- /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) — the §11.3 diagnostic.""" + +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 5c316fa..61fa72f 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -1,5 +1,4 @@ import numpy as np -import pytest import torch from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM @@ -49,15 +48,6 @@ def _zero_secondaries_loader(B=4, n_batches=2): return batches -@pytest.mark.xfail( - reason=( - "giant/validate.py isn't updated yet — it calls the stage models " - "(sample_secondaries et al.) with the old positional convention, " - "which doesn't match Stage1Model/Stage2OneShot's new forward " - "signature. Deferred to docs/v0.3.0-design.md step 6/§10." - ), - strict=False, -) 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