From da7cde3ef9670d9fd9f25a3ec6463c0e8b0e1d45 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Fri, 7 Aug 2026 16:12:58 +0200 Subject: [PATCH] v0.3.0 post-implementation audit: resolve all 9 tracked discrepancies MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Works through docs/v0.3.0-followups.md item by item, closing the gap between the design doc and the shipped v0.3.0-stage2-autoregressive code: 1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2 dispatch, stage-2 particle-type-class marginal. 2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run. 3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/ train instead of the hardcoded K_MAX constant. 4. Mixed conditioning.particle.type / conditioning.material.type support end-to-end (data pipeline + dwarf warm-cache). 5. conditioning.share_stages = true: one shared ConditionEncoder instance across both stages. 6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2 (was silently unimplemented). 7. giant predict/rollout: implement conditioning.*.type = "onehot" via the checkpoint's saved pdg_topn_map/mat_topn_map. 8. network.py's checkpoint-path model_config migration now fails loudly on non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's TOML-load path (§4.2). 9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a rollout-capable checkpoint (§9). Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics), mostly a test-helper dict-unpack pattern that made every unrelated constructor keyword look like a type error. Co-Authored-By: Claude Sonnet 5 --- docs/v0.3.0-design.md | 7 + docs/v0.3.0-followups.md | 37 +++ giant/analysis/router_gating.py | 32 ++- giant/cli.py | 425 +++++++++++++++++++++++++----- giant/config.py | 27 +- giant/data/dataset.py | 25 +- giant/data/loader.py | 30 ++- giant/data/setup_cache.py | 18 +- giant/data/transforms.py | 177 ++++++++----- giant/model/network.py | 74 +++++- giant/model/schedule.py | 7 +- giant/pipeline.py | 57 ++-- giant/rollout.py | 58 +++- giant/train.py | 67 ++--- giant/validate.py | 251 +++++++++++------- scripts/dwarf.py | 25 +- scripts/warm_setup_cache.py | 25 +- tests/test_cli_predict.py | 59 ----- tests/test_cli_train_overrides.py | 97 +++++++ tests/test_config.py | 59 ++++- tests/test_dataset.py | 3 +- tests/test_dwarf.py | 6 +- tests/test_flow.py | 4 +- tests/test_migration_v02_v03.py | 50 ++++ tests/test_network.py | 53 ++++ tests/test_pipeline.py | 69 +++++ tests/test_rollout.py | 140 ++++++++-- tests/test_sample.py | 2 +- tests/test_transforms.py | 56 +++- tests/test_validate.py | 93 +++++-- tests/test_wgan.py | 2 +- 31 files changed, 1536 insertions(+), 499 deletions(-) create mode 100644 docs/v0.3.0-followups.md create mode 100644 tests/test_cli_train_overrides.py diff --git a/docs/v0.3.0-design.md b/docs/v0.3.0-design.md index b5ebe64..7a47aee 100644 --- a/docs/v0.3.0-design.md +++ b/docs/v0.3.0-design.md @@ -892,6 +892,13 @@ All in `giant/config.py`: undesigned:** the mechanism should be worked out on its own terms when it is taken up, not pre-shaped by choices made for this refactor. Adding the block later is a config addition, not a break. +- **`stage2_model.generator = "ddpm"`.** The value is **accepted by the + schema** (§3.3 lists `"flow" | "ddpm" | "wgan"` with no caveat) but + `FlowDDPMStageTrainer.__init__` (`giant/train.py`) raises + `NotImplementedError` for stage 2 — only `"flow"` and `"wgan"` have a + stage-2 secondary-decoder loss implemented. `stage1_model.generator = + "ddpm"` is unaffected; this restriction is stage-2-only. Landing stage-2 + ddpm later is a trainer addition, not a config break. ### 11.3 Tracked as implementation work diff --git a/docs/v0.3.0-followups.md b/docs/v0.3.0-followups.md new file mode 100644 index 0000000..48ea25b --- /dev/null +++ b/docs/v0.3.0-followups.md @@ -0,0 +1,37 @@ +# v0.3.0 — post-implementation audit: open discrepancies + +**Status:** steps 1–7 of `docs/v0.3.0-design.md` §12 are implemented (branch +`v0.3.0-stage2-autoregressive`, commits `eb6dd27`..`200c6d2`). This document +tracked discrepancies found between that implementation and the design contract +during a 2026-08-07 audit, as concrete work items. **All items (1-9) are now +resolved** — either implemented (1-5, 7-9) or explicitly deferred into +`docs/v0.3.0-design.md` §11.2 (6: `stage2_model.generator = "ddpm"`). Step 8 +(`estimate_batch_size` recalibration) was intentionally still outstanding per +§12 and was never tracked here. + +--- + +## Confirmed correct during the audit (no action needed) + +For reference — these were explicitly checked against the design doc and +match it, including two spots the doc itself flagged as likely stale that +turned out fine: + +- Config schema, `DEFAULT_CONFIG`, `migrate_config` table (§4), `save_config`/ + `merge_cli_overrides` recursion, `default_out_dir_name`, `Conditioning` enum, + `n_sec.mode = "stop_token"` error (§9, §11.2). +- `network.py`'s full class decomposition (§5.3), dict-returning + `build_models`/`build_critics` (§5.4), `ExpertTrunk` separate in/out dims, + ST-Gumbel wiring (§2.1), AR token layout (§6.1), Markov/Attention history + encoders (§6.2). +- Shared PDG top-N type map (one map, not two, per §8), `other_policy` + sample/modal/drop (§11.1), embedding L1-nearest decode, and the L1-distance + diagnostic surfaced in `giant analyze` (§11.3). +- `analysis/render.py`/`analysis/router_gating.py` correctly branch + old-flat vs new-nested `model_config["router"]` location — doc flagged this + as a likely stale spot (§10) but it's actually fine. +- `train.py`'s per-stage trainers, mixed flow+wgan runs, WGAN critic cadence, + per-stage router auxiliary losses, stage-prefixed metrics, stage-2-only + training via ground-truth `x1_s1` (§7). +- `pipeline.py`'s deleted wgan+router rejection, per-stage `centers_init` + seeding, removed stale expert-size warning (§9, §10). diff --git a/giant/analysis/router_gating.py b/giant/analysis/router_gating.py index 9dd6a96..edfb826 100644 --- a/giant/analysis/router_gating.py +++ b/giant/analysis/router_gating.py @@ -61,21 +61,26 @@ class _RouterHandle: pdg_map: dict[int, int] mat_map: dict[str, int] cond_normalizer: "Normalizer" - conditioning: str + particle_conditioning: str + material_conditioning: str router_type: str -def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str: - """The single conditioning-mode string `giant.data.transforms. - build_cond_features` still expects — from either a v0.2 checkpoint's - flat `model_config["conditioning"]` (already a string) or a new-format - one (`model_config["conditioning"]["particle"]["type"]`; `validate_config` - guarantees particle/material agree, see its mixed-conditioning check). - Mirrors `giant.cli._conditioning_str`.""" +def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]: + """(particle_conditioning, material_conditioning) for + `giant.data.transforms.build_cond_features` — from either a v0.2 + checkpoint's flat `model_config["conditioning"]` (one shared string, same + for both axes) or a new-format one (independent + `model_config["conditioning"]["particle"/"material"]["type"]` — + docs/v0.3.0-design.md §3.1: the two axes may differ). Mirrors + `giant.cli._conditioning_axes`.""" raw = model_cfg.get("conditioning", default) if isinstance(raw, dict): - return raw.get("particle", {}).get("type", default) - return raw + return ( + raw.get("particle", {}).get("type", default), + raw.get("material", {}).get("type", default), + ) + return raw, raw def load_router(checkpoint: str | Path) -> _RouterHandle | None: @@ -107,12 +112,14 @@ def load_router(checkpoint: str | Path) -> _RouterHandle | None: if router is None: return None + particle_conditioning, material_conditioning = _conditioning_axes(model_cfg) return _RouterHandle( router=router, pdg_map={int(k): v for k, v in ckpt["pdg_map"].items()}, mat_map={str(k): v for k, v in ckpt["mat_map"].items()}, cond_normalizer=Normalizer.from_dict(ckpt["normalizer"]["cond"]), - conditioning=_conditioning_str(model_cfg), + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, router_type=router_cfg["type"], ) @@ -169,7 +176,8 @@ def _gate_for_df( handle.pdg_map, handle.mat_map, cond_normalizer=handle.cond_normalizer, - conditioning=handle.conditioning, + particle_conditioning=handle.particle_conditioning, + material_conditioning=handle.material_conditioning, ) with torch.no_grad(): gate = handle.router.gate( diff --git a/giant/cli.py b/giant/cli.py index d476390..81ac407 100644 --- a/giant/cli.py +++ b/giant/cli.py @@ -19,6 +19,7 @@ from tqdm import tqdm from giant import config as gconfig from giant.constants import ( + K_MAX, LOCAL_TARGET_NAMES, PREDICT_COORD_METADATA_KEY, PREDICT_SCHEMA_VERSION, @@ -72,44 +73,21 @@ def _router_total_experts(router_cfg: dict) -> int: return int(router_cfg.get("n_experts", 1)) -def _conditioning_str(model_cfg: dict, default: str = "embedding") -> str: - """The single conditioning-mode string `giant.data.transforms. - build_cond_features` still expects — from either a v0.2 checkpoint's - flat `model_config["conditioning"]` (already a string) or a new-format - one (`model_config["conditioning"]["particle"]["type"]`; `validate_config` - guarantees particle/material agree, see its mixed-conditioning check).""" +def _conditioning_axes(model_cfg: dict, default: str = "embedding") -> tuple[str, str]: + """(particle_conditioning, material_conditioning) for + `giant.data.transforms.build_cond_features`/`build_features` — from + either a v0.2 checkpoint's flat `model_config["conditioning"]` (one + shared string, same for both axes) or a new-format one (independent + `model_config["conditioning"]["particle"/"material"]["type"]` — + docs/v0.3.0-design.md §3.1: the two axes are configured independently and + may differ).""" raw = model_cfg.get("conditioning", default) if isinstance(raw, dict): - return raw.get("particle", {}).get("type", default) - return raw - - -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") - if particle_type == "onehot" or material_type == "onehot": - typer.echo( - 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, + return ( + raw.get("particle", {}).get("type", default), + raw.get("material", {}).get("type", default), ) - raise typer.Exit(1) + return raw, raw def _stage_cfg(model_cfg: dict, stage: str) -> dict: @@ -144,6 +122,14 @@ def _load_pdg_topn_map(ckpt: dict): return topnmap_from_json(raw, axis="pdg") if raw is not None else None +def _load_mat_topn_map(ckpt: dict): + """`ckpt["mat_topn_map"]` as a `giant.data.loader.TopNMap`, or `None` if + this checkpoint's `conditioning.material.type` was never `"onehot"` (see + `giant.pipeline.run_setup_stage`).""" + raw = ckpt.get("mat_topn_map") + return topnmap_from_json(raw, axis="material") if raw is not None else None + + def _batch_size_estimate_dims( model_cfg: dict, training: bool, stage: str = "stage1" ) -> tuple[int, int]: @@ -295,6 +281,16 @@ class Mode(str, Enum): wgan = "wgan" +class Decoder(str, Enum): + one_shot = "one_shot" + autoregressive = "autoregressive" + + +class Stage1Context(str, Enum): + truth = "truth" + sampled = "sampled" + + # Conditioning itself lives in giant.config (imported below as gconfig) — # shared with scripts/dwarf.py's Typer commands so the two CLIs can't # silently drift apart on the option's valid values. @@ -389,6 +385,89 @@ def train( "--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)" ), ] = None, + stage1_generator: Annotated[ + Optional[Mode], + typer.Option( + "--stage1-generator", + help="Stage 1's generative objective — overrides --mode for stage 1 " + "only (see docs/v0.3.0-design.md decision 7)", + ), + ] = None, + stage1_hidden_dim: Annotated[ + Optional[int], + typer.Option( + "--stage1-hidden-dim", + help="Overrides --hidden-dim for stage 1 only (same effect today; " + "--hidden-dim is kept as a shorthand since stage 1 was the only " + "target before stage2_model got its own flags)", + ), + ] = None, + stage1_n_res_blocks: Annotated[ + Optional[int], + typer.Option( + "--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only" + ), + ] = None, + stage1_dropout: Annotated[ + Optional[float], + typer.Option("--stage1-dropout", help="Overrides --dropout for stage 1 only"), + ] = None, + stage2_generator: Annotated[ + Optional[Mode], + typer.Option( + "--stage2-generator", + help="Stage 2's generative objective — overrides --mode for stage 2 " + "only, e.g. combine with --stage1-generator flow for a mixed " + "flow/wgan run", + ), + ] = None, + stage2_hidden_dim: Annotated[ + Optional[int], + typer.Option("--stage2-hidden-dim", help="Stage 2 trunk width"), + ] = None, + stage2_n_res_blocks: Annotated[ + Optional[int], + typer.Option("--stage2-n-res-blocks", help="Stage 2 trunk depth"), + ] = None, + stage2_dropout: Annotated[ + Optional[float], + typer.Option("--stage2-dropout", help="Dropout inside stage 2's ResBlocks"), + ] = None, + stage2_decoder: Annotated[ + Optional[Decoder], + typer.Option( + "--stage2-decoder", + help="one_shot: predict all k_max secondary slots at once (v0.2 " + "behaviour). autoregressive: emit one secondary at a time in " + "descending-energy order (default)", + ), + ] = None, + stage2_k_max: Annotated[ + Optional[int], + typer.Option( + "--stage2-k-max", + help="Maximum secondary slots (fixed width under one_shot, a " + "generation-loop safety cap under autoregressive; default: 15)", + ), + ] = None, + stage2_context_dim: Annotated[ + Optional[int], + typer.Option( + "--stage2-context-dim", + help="Width of the projected stage-1 outcome fed into stage 2's " + "conditioning (default: 64)", + ), + ] = None, + stage2_stage1_context: Annotated[ + Optional[Stage1Context], + typer.Option( + "--stage2-stage1-context", + help="What stage 2 conditions on during training: 'truth' (the " + "ground-truth stage-1 target, detached — default) or 'sampled' " + "(stage 1's own sampled output, closing the train/inference gap " + "at the cost of an extra sampling pass per batch)", + ), + ] = None, conditioning: Annotated[ Optional[Conditioning], typer.Option( @@ -457,6 +536,52 @@ def train( "(default: same as --lr)", ), ] = None, + stage1_n_critic: Annotated[ + Optional[int], + typer.Option("--stage1-n-critic", help="Overrides --n-critic for stage 1 only"), + ] = None, + stage1_gp_weight: Annotated[ + Optional[float], + typer.Option( + "--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only" + ), + ] = None, + stage1_noise_dim: Annotated[ + Optional[int], + typer.Option( + "--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only" + ), + ] = None, + stage1_critic_lr: Annotated[ + Optional[float], + typer.Option( + "--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only" + ), + ] = None, + stage2_n_critic: Annotated[ + Optional[int], + typer.Option("--stage2-n-critic", help="Overrides --n-critic for stage 2 only"), + ] = None, + stage2_gp_weight: Annotated[ + Optional[float], + typer.Option( + "--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only" + ), + ] = None, + stage2_noise_dim: Annotated[ + Optional[int], + typer.Option( + "--stage2-noise-dim", + help="Overrides --noise-dim for stage 2 only; under " + "--stage2-decoder autoregressive a fresh draw is made per token", + ), + ] = None, + stage2_critic_lr: Annotated[ + Optional[float], + typer.Option( + "--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only" + ), + ] = None, val_fraction: Annotated[ Optional[float], typer.Option("--val-fraction", "-f") ] = None, @@ -597,9 +722,11 @@ def train( }.items() if v is not None } - # Architecture-shape flags apply to stage1_model only (decision: stage 2 - # is tuned via a hand-edited config.toml until stage-prefixed --stage2-* - # flags land — see docs/v0.3.0-design.md decision 7). + # --hidden-dim/--n-blocks/--dropout are stage-1-only shorthands kept for + # backward compatibility (they predate stage2_model having its own + # flags); --stage1-*/--stage2-* below are the explicit, discoverable + # per-stage flags docs/v0.3.0-design.md decision 7 calls for, and take + # precedence when both are given. cli_stage1_model: dict[str, object] = { k: v for k, v in { @@ -609,6 +736,32 @@ def train( }.items() if v is not None } + cli_stage1_model.update( + { + k: v + for k, v in { + "hidden_dim": stage1_hidden_dim, + "n_res_blocks": stage1_n_res_blocks, + "dropout": stage1_dropout, + }.items() + if v is not None + } + ) + cli_stage2_model: dict[str, object] = { + k: v + for k, v in { + "hidden_dim": stage2_hidden_dim, + "n_res_blocks": stage2_n_res_blocks, + "dropout": stage2_dropout, + "decoder": stage2_decoder.value if stage2_decoder is not None else None, + "k_max": stage2_k_max, + "context_dim": stage2_context_dim, + "stage1_context": stage2_stage1_context.value + if stage2_stage1_context is not None + else None, + }.items() + if v is not None + } cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) if cli_router: cli_stage1_model["router"] = cli_router @@ -628,37 +781,65 @@ def train( overrides["train"] = cli_train if cli_stage1_model: overrides["stage1_model"] = cli_stage1_model + if cli_stage2_model: + overrides["stage2_model"] = cli_stage2_model if cli_conditioning: overrides["conditioning"] = cli_conditioning - # --mode/--n-critic/--gp-weight/--critic-lr apply to BOTH stages — v0.2 - # had one global mode/wgan config shared by both (see - # giant.config.migrate_config's train.mode / - # train.{n_critic,gp_weight,critic_lr} precedent), unlike the - # architecture-shape flags above. + # --mode/--n-critic/--gp-weight/--critic-lr/--noise-dim apply to BOTH + # stages by default (v0.2 had one global mode/wgan config shared by both + # — see giant.config.migrate_config's train.mode / + # train.{n_critic,gp_weight,critic_lr} precedent); the --stage1-*/ + # --stage2-* variants below override a single stage independently (decision + # 7), which is what actually enables e.g. `--stage1-generator flow + # --stage2-generator wgan`. if mode is not None: overrides.setdefault("stage1_model", {})["generator"] = mode.value overrides.setdefault("stage2_model", {})["generator"] = mode.value - if noise_dim is not None: - overrides.setdefault("stage1_model", {}).setdefault("wgan", {})["noise_dim"] = ( - noise_dim - ) - wgan_overrides = { + if stage1_generator is not None: + overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value + if stage2_generator is not None: + overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value + + shared_wgan_overrides = { k: v for k, v in { "n_critic": n_critic, "gp_weight": gp_weight, + "noise_dim": noise_dim, "critic_lr": critic_lr, }.items() if v is not None } - if wgan_overrides: - overrides.setdefault("stage1_model", {}).setdefault("wgan", {}).update( - wgan_overrides - ) - overrides.setdefault("stage2_model", {}).setdefault("wgan", {}).update( - wgan_overrides - ) + stage1_wgan_overrides = { + k: v + for k, v in { + "n_critic": stage1_n_critic, + "gp_weight": stage1_gp_weight, + "noise_dim": stage1_noise_dim, + "critic_lr": stage1_critic_lr, + }.items() + if v is not None + } + stage2_wgan_overrides = { + k: v + for k, v in { + "n_critic": stage2_n_critic, + "gp_weight": stage2_gp_weight, + "noise_dim": stage2_noise_dim, + "critic_lr": stage2_critic_lr, + }.items() + if v is not None + } + for stage_name, stage_specific in ( + ("stage1_model", stage1_wgan_overrides), + ("stage2_model", stage2_wgan_overrides), + ): + stage_wgan = {**shared_wgan_overrides, **stage_specific} + if stage_wgan: + overrides.setdefault(stage_name, {}).setdefault("wgan", {}).update( + stage_wgan + ) cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides) gconfig.validate_config(cfg) @@ -735,6 +916,36 @@ def new_run( n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None, emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None, dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None, + stage1_generator: Annotated[ + Optional[Mode], typer.Option("--stage1-generator") + ] = None, + stage1_hidden_dim: Annotated[ + Optional[int], typer.Option("--stage1-hidden-dim") + ] = None, + stage1_n_res_blocks: Annotated[ + Optional[int], typer.Option("--stage1-n-res-blocks") + ] = None, + stage1_dropout: Annotated[Optional[float], typer.Option("--stage1-dropout")] = None, + stage2_generator: Annotated[ + Optional[Mode], typer.Option("--stage2-generator") + ] = None, + stage2_hidden_dim: Annotated[ + Optional[int], typer.Option("--stage2-hidden-dim") + ] = None, + stage2_n_res_blocks: Annotated[ + Optional[int], typer.Option("--stage2-n-res-blocks") + ] = None, + stage2_dropout: Annotated[Optional[float], typer.Option("--stage2-dropout")] = None, + stage2_decoder: Annotated[ + Optional[Decoder], typer.Option("--stage2-decoder") + ] = None, + stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None, + stage2_context_dim: Annotated[ + Optional[int], typer.Option("--stage2-context-dim") + ] = None, + stage2_stage1_context: Annotated[ + Optional[Stage1Context], typer.Option("--stage2-stage1-context") + ] = None, conditioning: Annotated[ Optional[Conditioning], typer.Option("--conditioning") ] = None, @@ -802,6 +1013,32 @@ def new_run( }.items() if v is not None } + cli_stage1_model.update( + { + k: v + for k, v in { + "hidden_dim": stage1_hidden_dim, + "n_res_blocks": stage1_n_res_blocks, + "dropout": stage1_dropout, + }.items() + if v is not None + } + ) + cli_stage2_model: dict[str, object] = { + k: v + for k, v in { + "hidden_dim": stage2_hidden_dim, + "n_res_blocks": stage2_n_res_blocks, + "dropout": stage2_dropout, + "decoder": stage2_decoder.value if stage2_decoder is not None else None, + "k_max": stage2_k_max, + "context_dim": stage2_context_dim, + "stage1_context": stage2_stage1_context.value + if stage2_stage1_context is not None + else None, + }.items() + if v is not None + } cli_router = _router_cli_overrides(router, router_type, n_experts, router_axis) if cli_router: cli_stage1_model["router"] = cli_router @@ -818,11 +1055,17 @@ def new_run( overrides["train"] = cli_train if cli_stage1_model: overrides["stage1_model"] = cli_stage1_model + if cli_stage2_model: + overrides["stage2_model"] = cli_stage2_model if cli_conditioning: overrides["conditioning"] = cli_conditioning if mode is not None: overrides.setdefault("stage1_model", {})["generator"] = mode.value overrides.setdefault("stage2_model", {})["generator"] = mode.value + if stage1_generator is not None: + overrides.setdefault("stage1_model", {})["generator"] = stage1_generator.value + if stage2_generator is not None: + overrides.setdefault("stage2_model", {})["generator"] = stage2_generator.value cfg = gconfig.merge_cli_overrides(gconfig.DEFAULT_CONFIG, config, overrides) gconfig.validate_config(cfg) @@ -988,10 +1231,11 @@ def predict( raise typer.Exit(1) model_cfg = ckpt["model_config"] - _check_conditioning_onehot_support(model_cfg, "predict") pdg_topn_map = _load_pdg_topn_map(ckpt) + mat_topn_map = _load_mat_topn_map(ckpt) other_policy = _particle_type_other_policy(model_cfg) stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1") + stage2_k_max = _stage_cfg(model_cfg, "stage2").get("k_max", K_MAX) if batch_size_auto: est_hidden_dim, est_n_blocks = _batch_size_estimate_dims( @@ -1013,7 +1257,21 @@ def predict( assert batch_size_value is not None bs = batch_size_value - conditioning = _conditioning_str(model_cfg) + particle_conditioning, material_conditioning = _conditioning_axes(model_cfg) + if particle_conditioning == "onehot" and pdg_topn_map is None: + typer.echo( + "error: checkpoint's conditioning.particle.type='onehot' but has " + "no pdg_topn_map — retrain with the current code", + err=True, + ) + raise typer.Exit(1) + if material_conditioning == "onehot" and mat_topn_map is None: + typer.echo( + "error: checkpoint's conditioning.material.type='onehot' but has " + "no mat_topn_map — retrain with the current code", + err=True, + ) + raise typer.Exit(1) 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"]) @@ -1050,7 +1308,18 @@ def predict( skipped = 0 unknown_pdg_counts: Counter[int] = Counter() total_rows = sum(pq.ParquetFile(path).metadata.num_rows for path in files) - chunk_iter = iter_file_chunks if coord == Coord.local else iter_cond_chunks + + def chunk_iter(path: Path, offset: int): + if coord == Coord.local: + return iter_file_chunks(path, offset=offset, k_max=stage2_k_max) + return iter_cond_chunks(path, offset=offset) + + cond_pdg_topn = ( + pdg_topn_map.class_map if particle_conditioning == "onehot" else None + ) + cond_mat_topn = ( + mat_topn_map.class_map if material_conditioning == "onehot" else None + ) def _concat( a: dict[str, np.ndarray], b: dict[str, np.ndarray] @@ -1062,12 +1331,26 @@ def predict( if coord == Coord.local: cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features( - piece, pdg_map, mat_map, conditioning=conditioning + piece, + pdg_map, + mat_map, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, + pdg_topn_map=cond_pdg_topn, + mat_topn_map=cond_mat_topn, + k_max=stage2_k_max, ) cond_cont = cond_norm.transform(cond_cont) else: cond_cont, cond_cat = build_cond_features( - piece, pdg_map, mat_map, cond_norm, conditioning=conditioning + piece, + pdg_map, + mat_map, + cond_norm, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, + pdg_topn_map=cond_pdg_topn, + mat_topn_map=cond_mat_topn, ) cc = torch.from_numpy(cond_cont).float().to(_device) @@ -1409,9 +1692,23 @@ def rollout( training_cfg = gconfig.load_checkpoint_config(checkpoint) model_cfg = ckpt["model_config"] - _check_conditioning_onehot_support(model_cfg, "rollout") - conditioning = _conditioning_str(model_cfg) + particle_conditioning, material_conditioning = _conditioning_axes(model_cfg) pdg_topn_map = _load_pdg_topn_map(ckpt) + mat_topn_map = _load_mat_topn_map(ckpt) + if particle_conditioning == "onehot" and pdg_topn_map is None: + typer.echo( + "error: checkpoint's conditioning.particle.type='onehot' but has " + "no pdg_topn_map — retrain with the current code", + err=True, + ) + raise typer.Exit(1) + if material_conditioning == "onehot" and mat_topn_map is None: + typer.echo( + "error: checkpoint's conditioning.material.type='onehot' but has " + "no mat_topn_map — retrain with the current code", + err=True, + ) + raise typer.Exit(1) other_policy = _particle_type_other_policy(model_cfg) stage1_ddpm_steps = _ddpm_steps(model_cfg, "stage1") stage2_ddpm_steps = _ddpm_steps(model_cfg, "stage2") @@ -1489,8 +1786,10 @@ def rollout( max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, on_chunk=_write_chunk, - conditioning=conditioning, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, pdg_topn_map=pdg_topn_map, + mat_topn_map=mat_topn_map, other_policy=other_policy, seed=seed, stage1_ddpm_steps=stage1_ddpm_steps, diff --git a/giant/config.py b/giant/config.py index 4357c70..28ed354 100644 --- a/giant/config.py +++ b/giant/config.py @@ -680,18 +680,6 @@ def validate_config(cfg: dict) -> None: single block's defaults. """ particle_type = _get_path(cfg, "conditioning.particle.type") - material_type = _get_path(cfg, "conditioning.material.type") - - if particle_type != material_type: - raise ValueError( - "conditioning.particle.type " - f"({particle_type!r}) != conditioning.material.type " - f"({material_type!r}) — the data pipeline (giant/data/transforms.py's " - "build_features/build_cond_features) doesn't support mixed " - "particle/material conditioning types yet, even though " - "ConditionEncoder itself (giant/model/network.py) is already able " - "to. Not implemented in v0.3.0 — use the same type for both axes." - ) pt_target = _get_path(cfg, "stage2_model.particle_type.target") if pt_target == "embedding" and particle_type != "embedding": @@ -733,6 +721,21 @@ def validate_config(cfg: dict) -> None: "(standalone stage-2 evaluation only, never for rollout)" ) + if ( + _get_path(cfg, "stage2_model.n_sec.mode") == "truth" + and _get_path(cfg, "stage1_model.active") + and _get_path(cfg, "stage2_model.active") + ): + raise ValueError( + "stage2_model.n_sec.mode = 'truth' is invalid for a " + "rollout-capable checkpoint (both stage1_model.active and " + "stage2_model.active = true — see docs/v0.3.0-design.md §9): " + "'truth' takes n_sec from ground truth, which giant rollout " + "doesn't have. 'truth' is for standalone stage-2 evaluation " + "only — set stage1_model.active = false for that, or use " + "'head' (default) for a rollout-capable checkpoint." + ) + if _get_path(cfg, "stage2_model.decoder") == "autoregressive": history = _get_path(cfg, "stage2_model.autoregressive.history") if history not in ("markov", "attention"): diff --git a/giant/data/dataset.py b/giant/data/dataset.py index 3736b92..6f5f216 100644 --- a/giant/data/dataset.py +++ b/giant/data/dataset.py @@ -6,6 +6,7 @@ import numpy as np import torch from torch.utils.data import IterableDataset +from giant.constants import K_MAX from giant.data.loader import event_id_offset, iter_file_chunks from giant.data.transforms import Normalizer, build_features, sorted_membership @@ -46,7 +47,7 @@ class StreamingStepsDataset(IterableDataset): (see docs/v0.3.0-design.md decision 4) target_s1: (B, 9) float32 — normalised Stage-1 primary target n_sec: (B,) int64 — true secondary count per step - sec_cont: (B, K_MAX, SEC_SLOT_DIM) float32 — [stick_logit, + sec_cont: (B, k_max, SEC_SLOT_DIM) float32 — [stick_logit, local_dir, log_mass, charge] per slot (mass/charge normalised iff `sec_phys_normalizer` was given); always computed the same way regardless of @@ -54,9 +55,13 @@ class StreamingStepsDataset(IterableDataset): downstream under target="physical" proc_idx: (B,) int64 — process-class label (ProcessRouter supervision only; zeros when `proc_map` is None) - sec_type_idx: (B, K_MAX) int64 — per-slot class index into + sec_type_idx: (B, k_max) int64 — per-slot class index into `sec_type_class_map`, for particle_type.target in ("onehot", "embedding"); zeros (unused) otherwise + + `k_max` (constructor arg, default the module constant) should match + `stage2_model.k_max` (docs/v0.3.0-design.md §8) — it sets the padded + width of `sec_cont`/`sec_type_idx` above. """ def __init__( @@ -71,11 +76,13 @@ class StreamingStepsDataset(IterableDataset): shuffle_buffer: int = 65536, shuffle: bool = True, proc_map: dict[str, int] | None = None, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", + material_conditioning: str = "embedding", sec_phys_normalizer: Normalizer | None = None, pdg_topn_map: dict[int, int] | None = None, mat_topn_map: dict[str, int] | None = None, sec_type_class_map: dict | None = None, + k_max: int = K_MAX, ) -> None: self.files = list(files) self._offsets = {path: event_id_offset(i) for i, path in enumerate(self.files)} @@ -89,11 +96,13 @@ class StreamingStepsDataset(IterableDataset): self.shuffle_buffer = max(shuffle_buffer, batch_size) self.shuffle = shuffle self.proc_map = proc_map - self.conditioning = conditioning + self.particle_conditioning = particle_conditioning + self.material_conditioning = material_conditioning self.sec_phys_normalizer = sec_phys_normalizer self.pdg_topn_map = pdg_topn_map self.mat_topn_map = mat_topn_map self.sec_type_class_map = sec_type_class_map + self.k_max = k_max def __iter__(self): worker_info = torch.utils.data.get_worker_info() @@ -115,7 +124,9 @@ class StreamingStepsDataset(IterableDataset): buf_n = 0 for path in files: - for chunk in iter_file_chunks(path, offset=self._offsets[path]): + for chunk in iter_file_chunks( + path, offset=self._offsets[path], k_max=self.k_max + ): mask = sorted_membership(chunk["event_id"], self._events_arr) if not mask.any(): continue @@ -140,10 +151,12 @@ class StreamingStepsDataset(IterableDataset): sec_phys_normalizer=self.sec_phys_normalizer, proc_map=self.proc_map, require_secondaries=True, - conditioning=self.conditioning, + particle_conditioning=self.particle_conditioning, + material_conditioning=self.material_conditioning, pdg_topn_map=self.pdg_topn_map, mat_topn_map=self.mat_topn_map, sec_type_class_map=self.sec_type_class_map, + k_max=self.k_max, ) buf_cont.append(cond_cont) buf_cat.append(cond_cat) diff --git a/giant/data/loader.py b/giant/data/loader.py index 40f6a6b..e5a77bf 100644 --- a/giant/data/loader.py +++ b/giant/data/loader.py @@ -6,6 +6,8 @@ import numpy as np import pandas as pd import pyarrow.parquet as pq +from giant.constants import K_MAX + # A manifest is a plain text file listing one parquet path per line, used to # name a curated subset of files (e.g. a train/holdout pool) without copying # or symlinking the underlying parquet files. Lines are resolved relative to @@ -113,9 +115,9 @@ def _pad_dir_col(dx: pd.Series, dy: pd.Series, dz: pd.Series, K: int) -> np.ndar return out -def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: - from giant.constants import K_MAX - +def _df_to_dict( + df: pd.DataFrame, offset: int = 0, k_max: int = K_MAX +) -> dict[str, np.ndarray]: has_sec_lists = "sec_E_list" in df.columns d: dict[str, np.ndarray] = { @@ -147,17 +149,19 @@ def _df_to_dict(df: pd.DataFrame, offset: int = 0) -> dict[str, np.ndarray]: } if has_sec_lists: - d["sec_E_list"] = _pad_list_col(df["sec_E_list"], K_MAX) - d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], K_MAX) + d["sec_E_list"] = _pad_list_col(df["sec_E_list"], k_max) + d["sec_pdg_list"] = _pad_list_col_int(df["sec_pdg_list"], k_max) d["sec_dir_list"] = _pad_dir_col( - df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], K_MAX + df["sec_dx_list"], df["sec_dy_list"], df["sec_dz_list"], k_max ) return d -def load_steps(path: str | Path, offset: int = 0) -> dict[str, np.ndarray]: - return _df_to_dict(pd.read_parquet(path), offset=offset) +def load_steps( + path: str | Path, offset: int = 0, k_max: int = K_MAX +) -> dict[str, np.ndarray]: + return _df_to_dict(pd.read_parquet(path), offset=offset, k_max=k_max) def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray: @@ -167,12 +171,16 @@ def load_event_ids(path: str | Path, offset: int = 0) -> np.ndarray: def iter_file_chunks( - path: str | Path, offset: int = 0 + path: str | Path, offset: int = 0, k_max: int = K_MAX ) -> Iterator[dict[str, np.ndarray]]: - """Yield one parquet row-group at a time so a large file never fully loads.""" + """Yield one parquet row-group at a time so a large file never fully loads. + + `k_max` sets the padded width of the sec_*_list columns (should match + `stage2_model.k_max` — see docs/v0.3.0-design.md §8); defaults to the + module constant for callers that don't care (e.g. Stage-1-only reads).""" pf = pq.ParquetFile(path) for i in range(pf.num_row_groups): - yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset) + yield _df_to_dict(pf.read_row_group(i).to_pandas(), offset=offset, k_max=k_max) _COND_COLS = [ diff --git a/giant/data/setup_cache.py b/giant/data/setup_cache.py index a20d0b6..ce38def 100644 --- a/giant/data/setup_cache.py +++ b/giant/data/setup_cache.py @@ -96,10 +96,22 @@ def fingerprint_files(files: list[Path]) -> list[list]: return out -def normalizer_key(val_fraction: float, seed: int, conditioning: str) -> str: +def normalizer_key( + val_fraction: float, + seed: int, + particle_conditioning: str, + material_conditioning: str, +) -> str: # .6g avoids float-repr drift (e.g. 0.1 vs 0.10000000000000002) causing - # spurious cache misses between runs with the "same" val_fraction. - return f"valfrac={val_fraction:.6g}_seed={seed}_cond={conditioning}" + # spurious cache misses between runs with the "same" val_fraction. The two + # conditioning axes are independent (docs/v0.3.0-design.md §3.1) and both + # affect which cond_cont columns are computed for real vs. zero-filled + # (giant.data.transforms._physical_cond_columns), so both must be part of + # the key or two mixed-axis runs could collide on the same cache entry. + return ( + f"valfrac={val_fraction:.6g}_seed={seed}_pcond={particle_conditioning}" + f"_mcond={material_conditioning}" + ) # Top-N-map axes (docs/v0.3.0-design.md §8): "pdg" keys match pdg_map's int diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 562abf0..555ed22 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -2,6 +2,8 @@ import warnings import numpy as np +from giant.constants import K_MAX + _EPS = 1e-8 # Floor added to each energy fraction before taking log-ratios so the simplex @@ -719,14 +721,23 @@ def decode_secondaries( def _physical_cond_columns( - data: dict[str, np.ndarray], conditioning: str + data: dict[str, np.ndarray], + particle_conditioning: str, + material_conditioning: str, ) -> np.ndarray: """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns. - "embedding"/"onehot" modes zero-fill (cheap, and ConditionEncoder never - reads these columns in either mode — so an unfilled giant.materials table - can never crash an "embedding"/"onehot"-mode run). "physical" mode - computes them for real: particle columns come from + The particle and material blocks are gated independently + (docs/v0.3.0-design.md §3.1: "configured independently and may mix + freely — e.g. material `physical` with particle `embedding`"), so e.g. + `particle_conditioning="embedding"` + `material_conditioning="physical"` + zero-fills only the particle columns and computes the material ones for + real. + + "embedding"/"onehot" zero-fill their block (cheap, and ConditionEncoder + never reads these columns in either mode — so an unfilled + giant.materials table can never crash an "embedding"/"onehot"-mode run). + "physical" computes it for real: particle columns come from `data["mass"]`/`data["charge"]` when the caller already knows them directly (rollout.py, for a track descended from a model-predicted secondary — see giant/rollout.py's "no snapping" design), else derived @@ -736,36 +747,47 @@ def _physical_cond_columns( """ from giant.constants import MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM - if conditioning in ("embedding", "onehot"): - n = len(next(iter(data.values()))) - return np.zeros((n, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM), dtype=np.float32) - if conditioning != "physical": - raise ValueError(f"unknown conditioning mode {conditioning!r}") + n = len(next(iter(data.values()))) - from giant.materials import material_properties_array - from giant.particles import particle_phys_array + if particle_conditioning == "physical": + from giant.particles import particle_phys_array - if "mass" in data and "charge" in data: - mass = np.asarray(data["mass"], dtype=np.float32) - charge = np.asarray(data["charge"], dtype=np.float32) + if "mass" in data and "charge" in data: + mass = np.asarray(data["mass"], dtype=np.float32) + charge = np.asarray(data["charge"], dtype=np.float32) + else: + mass, charge = particle_phys_array(data["pdg"]).T + particle_cols = np.column_stack([log_transform(mass), charge]) + elif particle_conditioning in ("embedding", "onehot"): + particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32) else: - mass, charge = particle_phys_array(data["pdg"]).T + raise ValueError( + f"unknown conditioning.particle.type {particle_conditioning!r}" + ) - z_eff, a_eff, density, x0, lambda_int = material_properties_array( - data["material"] - ).T + if material_conditioning == "physical": + from giant.materials import material_properties_array - return np.column_stack( - [ - log_transform(mass), - charge, - z_eff, - a_eff, - log_transform(density), - log_transform(x0), - log_transform(lambda_int), - ] - ).astype(np.float32) + z_eff, a_eff, density, x0, lambda_int = material_properties_array( + data["material"] + ).T + material_cols = np.column_stack( + [ + z_eff, + a_eff, + log_transform(density), + log_transform(x0), + log_transform(lambda_int), + ] + ) + elif material_conditioning in ("embedding", "onehot"): + material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32) + else: + raise ValueError( + f"unknown conditioning.material.type {material_conditioning!r}" + ) + + return np.column_stack([particle_cols, material_cols]).astype(np.float32) def build_cond_features( @@ -773,19 +795,24 @@ def build_cond_features( pdg_map: dict[int, int], mat_map: dict[str, int], cond_normalizer: "Normalizer | None" = None, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", + material_conditioning: str = "embedding", pdg_topn_map: dict[int, int] | None = None, mat_topn_map: dict[str, int] | None = None, ) -> tuple[np.ndarray, np.ndarray]: """Build conditioning arrays only — no target, no post-step variables. + `particle_conditioning`/`material_conditioning` are independent + (docs/v0.3.0-design.md §3.1) — e.g. `particle_conditioning="embedding"` + + `material_conditioning="physical"` is a valid mix. + `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat` columns read by `ConditionEncoder`'s `"onehot"` mode (docs/v0.3.0-design.md decision 4): pdg topN index at column 2 (iff `pdg_topn_map` given), material topN index at column 3 (iff `mat_topn_map` given, after column 2 if both are). Only ever given when - `conditioning == "onehot"`; `cond_cat` stays `(N, 2)` otherwise. + the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise. """ cond_cont = np.column_stack( [ @@ -796,7 +823,10 @@ def build_cond_features( ] ).astype(np.float32) cond_cont = np.column_stack( - [cond_cont, _physical_cond_columns(data, conditioning)] + [ + cond_cont, + _physical_cond_columns(data, particle_conditioning, material_conditioning), + ] ).astype(np.float32) # In "physical" mode cond_cat's first two columns are only a @@ -807,10 +837,11 @@ def build_cond_features( # conditioning signal, so an unmapped value must still raise loudly # rather than silently misassign. In "onehot" mode they again go unread # (the topN columns below are the real signal), so they're as permissive - # as "physical". - strict = conditioning == "embedding" - pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict) - mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict) + # as "physical". Each axis's strictness is independent. + pdg_strict = particle_conditioning == "embedding" + mat_strict = material_conditioning == "embedding" + pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=pdg_strict) + mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=mat_strict) cat_cols = [pdg_idx, mat_idx] if pdg_topn_map is not None: cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map)) @@ -819,34 +850,44 @@ def build_cond_features( cond_cat = np.column_stack(cat_cols) if cond_normalizer is not None: - cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning) + cond_cont = _cond_normalizer_transform( + cond_cont, cond_normalizer, particle_conditioning, material_conditioning + ) return cond_cont, cond_cat def _cond_normalizer_transform( - cond_cont: np.ndarray, cond_normalizer: "Normalizer", conditioning: str + cond_cont: np.ndarray, + cond_normalizer: "Normalizer", + particle_conditioning: str, + material_conditioning: str, ) -> np.ndarray: """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed. Checkpoints trained before physical-property conditioning (``COND_DIM`` 8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond normalizer, fit before ``build_cond_features`` grew the extra physical - columns. In "embedding" mode those columns are never read by + columns. When NEITHER axis is "physical" those columns are never read by ``ConditionEncoder`` (``giant/model/network.py``), so padding the missing entries with mean=0/std=1 is a safe no-op that keeps such checkpoints - usable under the current, always-``COND_DIM``-wide contract. In - "physical" mode the physical columns are load-bearing, so a mismatch - there is a real incompatibility, not something to paper over. + usable under the current, always-``COND_DIM``-wide contract. If EITHER + axis is "physical" its columns are load-bearing, so a mismatch there is a + real incompatibility, not something to paper over. """ mean, std = cond_normalizer.mean, cond_normalizer.std assert mean is not None and std is not None, "Normalizer not fitted" width = cond_cont.shape[-1] if mean.shape[-1] < width: - if conditioning != "embedding": + physical_load_bearing = "physical" in ( + particle_conditioning, + material_conditioning, + ) + if physical_load_bearing: raise ValueError( f"cond normalizer has {mean.shape[-1]} columns, expected " - f"{width}, and conditioning={conditioning!r} reads the " + f"{width}, and particle_conditioning={particle_conditioning!r}/" + f"material_conditioning={material_conditioning!r} reads the " "physical columns directly — this checkpoint predates " "physical-property conditioning and can't be safely padded; " "retrain it under the current code." @@ -867,11 +908,13 @@ def build_features( fit: bool = False, proc_map: dict[str, int] | None = None, require_secondaries: bool = False, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", + material_conditioning: str = "embedding", sec_phys_only: bool = False, pdg_topn_map: dict[int, int] | None = None, mat_topn_map: dict[str, int] | None = None, sec_type_class_map: dict | None = None, + k_max: int = K_MAX, ) -> tuple[ np.ndarray, np.ndarray, @@ -922,8 +965,14 @@ def build_features( top-N-plus-other map's `class_map` for `target = "onehot"`, or the dense `pdg_map` for `target = "embedding"` (pass `pdg_map` itself). `None` for `target = "physical"`. + + k_max: should match `stage2_model.k_max` (docs/v0.3.0-design.md §8) — + overridden internally by `data["sec_E_list"]`'s own padded width when + present (the loader already padded it to some k_max; that width is + authoritative), so this only actually matters when secondary list + columns are absent (Stage-1-only reads, or a pre-secondary-join + file), where it sets `sec_cont`/`sec_type_idx`'s zero-filled width. """ - from giant.constants import K_MAX post_dir_local = local_frame_rotation(data["pre_dir"], data["post_dir"]) travel_dir_local = local_frame_rotation( @@ -953,7 +1002,10 @@ def build_features( ] ).astype(np.float32) # (N, COND_DIM_BASE=8) cond_cont = np.column_stack( - [cond_cont, _physical_cond_columns(data, conditioning)] + [ + cond_cont, + _physical_cond_columns(data, particle_conditioning, material_conditioning), + ] ).astype(np.float32) # (N, COND_DIM=15) pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map) @@ -968,20 +1020,27 @@ def build_features( n_sec_raw = data["n_sec"].astype( np.int64 ) # (N,) unclamped, for the valid-slot mask - # Clamp the classification label to K_MAX: the head only has K_MAX+1 classes - # (0..K_MAX), and truncating here mirrors the K_MAX-slot truncation already - # applied to sec_cont by the loader's list padding. Without this, a rare - # high-multiplicity step (real data goes up to ~37) hands cross_entropy - # an out-of-range target and CUDA asserts. - n_sec = np.minimum(n_sec_raw, K_MAX).astype(np.int64) # (N,) # Secondary continuous targets sec_E_list = data.get("sec_E_list") sec_dir_list = data.get("sec_dir_list") sec_pdg_list = data.get("sec_pdg_list") + if sec_E_list is not None: + # The loader already padded sec_*_list to some k_max (see + # giant.data.loader.iter_file_chunks); that padded width is + # authoritative over whatever this call happened to pass in, so the + # two can never drift apart (docs/v0.3.0-design.md §8). + k_max = sec_E_list.shape[1] + + # Clamp the classification label to k_max: the head only has k_max+1 + # classes (0..k_max), and truncating here mirrors the k_max-slot + # truncation already applied to sec_cont by the loader's list padding. + # Without this, a rare high-multiplicity step (real data goes up to ~37) + # hands cross_entropy an out-of-range target and CUDA asserts. + n_sec = np.minimum(n_sec_raw, k_max).astype(np.int64) # (N,) if sec_E_list is not None and sec_dir_list is not None and sec_pdg_list is not None: - sec_valid = np.arange(K_MAX)[None, :] < n_sec_raw[:, None] # (N, K_MAX) + sec_valid = np.arange(k_max)[None, :] < n_sec_raw[:, None] # (N, k_max) sec_cont = encode_secondaries( sec_E_list, sec_dir_list, @@ -990,11 +1049,11 @@ def build_features( data["pre_dir"], sec_pdg_list=sec_pdg_list, phys_only=sec_phys_only, - ) # (N, K_MAX, 6) + ) # (N, k_max, 6) sec_type_idx = ( encode_secondary_type_idx(sec_pdg_list, sec_valid, sec_type_class_map) if sec_type_class_map is not None - else np.zeros((len(n_sec), K_MAX), dtype=np.int64) + else np.zeros((len(n_sec), k_max), dtype=np.int64) ) else: # Guard against silently training Stage 2 on zeroed targets: if any step @@ -1017,8 +1076,8 @@ def build_features( "require_secondaries=False for Stage-1-only use." ) N = len(n_sec) - sec_cont = np.zeros((N, K_MAX, 6), dtype=np.float32) - sec_type_idx = np.zeros((N, K_MAX), dtype=np.int64) + sec_cont = np.zeros((N, k_max, 6), dtype=np.float32) + sec_type_idx = np.zeros((N, k_max), dtype=np.int64) if fit: cond_normalizer = Normalizer().fit(cond_cont) diff --git a/giant/model/network.py b/giant/model/network.py index a058af6..f7161aa 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -990,7 +990,12 @@ class Stage1Model(nn.Module): (docs/v0.3.0-design.md §2) moves it to stage 2, except for a migrated v0.2 checkpoint (`n_sec_head_k_max` given), where it stays attached here since that's where its weights live and what conditioning it was trained - against (see `_migrate_legacy_model_config`).""" + against (see `_migrate_legacy_model_config`). + + `cond_enc`, if given, is used in place of building a fresh + `ConditionEncoder` — `conditioning.share_stages = true` (§3.1): `build_models` + constructs one shared instance and passes it to both stages, halving the + conditioning parameter count and forcing a common representation.""" def __init__( self, @@ -1008,12 +1013,17 @@ class Stage1Model(nn.Module): noise_dim: int = 64, router: Router | None = None, n_sec_head_k_max: int | None = None, + cond_enc: ConditionEncoder | None = None, ) -> None: super().__init__() self.generator_kind = generator self.noise_dim = noise_dim - self.cond_enc = ConditionEncoder( - pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder( + pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim + ) ) has_time = generator in ("flow", "ddpm") self.time_emb = SinusoidalEmbedding(time_dim) if has_time else None @@ -1083,6 +1093,9 @@ class Stage2OneShot(nn.Module): Under `generator == "wgan"` the type slice stays folded into `sec_dim` (just `emb_dim` instead of `PARTICLE_PHYS_DIM` wide) and `type_head` is unused (`None`) — the WGAN trainer handles the ST-Gumbel relaxation. + + `cond_enc`, if given, is used in place of building a fresh + `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). """ def __init__( @@ -1105,6 +1118,7 @@ class Stage2OneShot(nn.Module): router: Router | None = None, build_n_sec_head: bool = True, particle_type_cfg: dict | None = None, + cond_enc: ConditionEncoder | None = None, ) -> None: super().__init__() self.generator_kind = generator @@ -1112,8 +1126,12 @@ class Stage2OneShot(nn.Module): 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 + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder( + pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim + ) ) self.context_adapter = ContextAdapter(x_dim, context_dim) self.fuse = nn.Sequential( @@ -1227,6 +1245,9 @@ class Stage2Autoregressive(nn.Module): on token position; `_token_cond` additionally fuses in the history encoding and two running scalars (remaining energy-budget fraction, normalized slot index), and feeds `forward`/`predict_type`/the trunk. + + `cond_enc`, if given, is used in place of building a fresh + `ConditionEncoder` — see `Stage1Model`'s docstring (`conditioning.share_stages`). """ def __init__( @@ -1251,6 +1272,7 @@ class Stage2Autoregressive(nn.Module): history: str = "markov", attn_n_heads: int = 4, attn_n_layers: int = 2, + cond_enc: ConditionEncoder | None = None, ) -> None: super().__init__() if history not in ("markov", "attention"): @@ -1266,8 +1288,12 @@ class Stage2Autoregressive(nn.Module): emb_dim = particle_cfg["emb_dim"] self.type_dim = stage2_type_dim(self.particle_type_cfg, emb_dim) - self.cond_enc = ConditionEncoder( - pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim + self.cond_enc = ( + cond_enc + if cond_enc is not None + else ConditionEncoder( + pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim + ) ) self.context_adapter = ContextAdapter(x_dim, context_dim) self.base_fuse = nn.Sequential( @@ -1562,6 +1588,20 @@ def _migrate_legacy_model_config(model_config: dict) -> dict: k_max = m.get("k_max", K_MAX) noise_dim = m.get("noise_dim", 64) router_cfg = dict(m.get("router") or {}) + expert_hidden_dim = router_cfg.pop("expert_hidden_dim", 0) + expert_n_blocks = router_cfg.pop("expert_n_blocks", 0) + if expert_hidden_dim or expert_n_blocks: + raise ValueError( + "this checkpoint's model_config.router sets expert_hidden_dim/" + f"expert_n_blocks to a non-default value " + f"({expert_hidden_dim!r}, {expert_n_blocks!r}); v0.3.0 removed " + "per-expert sizing (experts always inherit the stage's " + "hidden_dim/n_res_blocks), so this checkpoint's routed experts " + "have a different width/depth than the monolith — silently " + "dropping these keys would resize the experts instead of " + "refusing (docs/v0.3.0-design.md §4.2). This checkpoint can " + "only be loaded by v0.2 code." + ) router_cfg.setdefault("enabled", False) return { @@ -1657,6 +1697,13 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: instance rather than building a second, independently-parameterized one (v0.2's actual — probably accidental — behaviour: two routers built from one config with no semantic relationship between them). + + `conditioning.share_stages = true` (§3.1) builds one `ConditionEncoder` + instance here and passes it to both stages (`Stage1Model`/`Stage2OneShot`/ + `Stage2Autoregressive`'s `cond_enc` param), instead of each stage + building its own — halving the conditioning parameter count and forcing a + common representation. `false` (default) keeps v0.2 behaviour: + independent instances with identical config but independent weights. """ cfg = ( model_config @@ -1666,17 +1713,17 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: pdg_vocab = cfg["pdg_vocab"] mat_vocab = cfg["mat_vocab"] conditioning = cfg["conditioning"] - if conditioning.get("share_stages"): - raise NotImplementedError( - "conditioning.share_stages = true is not implemented yet — each " - "stage always builds its own ConditionEncoder for now" - ) particle_cfg = conditioning["particle"] material_cfg = conditioning["material"] particle_conditioning = particle_cfg["type"] s1cfg = cfg["stage1_model"] s2cfg = cfg["stage2_model"] cond_out_dim = conditioning.get("out_dim", 128) + shared_cond_enc: ConditionEncoder | None = None + if conditioning.get("share_stages"): + shared_cond_enc = ConditionEncoder( + pdg_vocab, mat_vocab, particle_cfg, material_cfg, out_dim=cond_out_dim + ) result: dict[str, nn.Module | None] = {"stage1": None, "stage2": None} @@ -1707,6 +1754,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: noise_dim=(s1cfg.get("wgan") or {}).get("noise_dim", 64), router=stage1_router, n_sec_head_k_max=n_sec_head_k_max, + cond_enc=shared_cond_enc, ) if s2cfg.get("active", True): @@ -1748,6 +1796,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: history=ar_cfg.get("history", "markov"), attn_n_heads=ar_cfg.get("attn_n_heads", 4), attn_n_layers=ar_cfg.get("attn_n_layers", 2), + cond_enc=shared_cond_enc, ) else: sec_dim = stage2_trunk_sec_dim( @@ -1771,6 +1820,7 @@ def build_models(model_config: dict) -> dict[str, nn.Module | None]: router=stage2_router, build_n_sec_head=legacy_owner != "stage1", particle_type_cfg=particle_type_cfg, + cond_enc=shared_cond_enc, ) return result diff --git a/giant/model/schedule.py b/giant/model/schedule.py index 5c1e019..25aff5e 100644 --- a/giant/model/schedule.py +++ b/giant/model/schedule.py @@ -108,12 +108,13 @@ def flow_matching_loss_secondary( averaged over their own width first and then combined with equal weight — this stays correct if type_dim/CONT_SLOT_DIM change. """ - from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM + from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM if type_dim is None: type_dim = PARTICLE_PHYS_DIM B = x1.size(0) + k_max = sec_mask.size(1) t = torch.rand(B, device=x1.device) x0 = torch.randn_like(x1) x_t = (1.0 - t.view(-1, 1)) * x0 + t.view(-1, 1) * x1 @@ -121,8 +122,8 @@ def flow_matching_loss_secondary( v_t = model(x_t, cond_cont, cond_cat, stage1_out, t=t) slot_dim = CONT_SLOT_DIM + type_dim - err = ((v_t - u_t) ** 2).view(B, K_MAX, slot_dim) - cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1) # (B, K_MAX) + err = ((v_t - u_t) ** 2).view(B, k_max, slot_dim) + cont_err = err[..., :CONT_SLOT_DIM].mean(dim=-1) # (B, k_max) mask = sec_mask.float() denom = mask.sum().clamp(min=1) diff --git a/giant/pipeline.py b/giant/pipeline.py index 5584996..a10eea8 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -9,7 +9,6 @@ from torch.utils.data import DataLoader from giant import config from giant.constants import ( COND_DIM, - K_MAX, PARTICLE_PHYS_DIM, X_DIM, ) @@ -106,11 +105,9 @@ def run_setup_stage( Stage-2 normalizers. `cfg` is the full merged v0.3 config (`conditioning`/`stage1_model`/ - `stage2_model`), already passed through `giant.config.validate_config` — - in particular, `conditioning.particle.type == conditioning.material.type` - is assumed (the data pipeline doesn't support mixed conditioning types - yet, see `validate_config`), so a single shared conditioning string - drives `giant.data.transforms.build_features`. + `stage2_model`), already passed through `giant.config.validate_config`. + `conditioning.particle.type` and `conditioning.material.type` are + independent (docs/v0.3.0-design.md §3.1) and may differ. Reads from and writes to the `giant.data.setup_cache` sidecar when `cache_setup` is set (`rebuild_setup_cache` ignores — but still @@ -118,7 +115,9 @@ def run_setup_stage( is mutated in place: an active `EnergyRouter` (`router.type == "energy"`) gets its `centers_init` seeded from real data quantiles here. """ - conditioning = cfg["conditioning"]["particle"]["type"] + particle_conditioning = cfg["conditioning"]["particle"]["type"] + material_conditioning = cfg["conditioning"]["material"]["type"] + k_max = cfg["stage2_model"]["k_max"] stage1_router = cfg["stage1_model"].get("router") or {} stage2_router = cfg["stage2_model"].get("router") or {} @@ -255,7 +254,9 @@ def run_setup_stage( for r in (stage1_router, stage2_router) ) energy_idx = 3 - norm_key = setup_cache.normalizer_key(val_fraction, seed, conditioning) + norm_key = setup_cache.normalizer_key( + val_fraction, seed, particle_conditioning, material_conditioning + ) entry = cache.normalizers.get(norm_key) if cache is not None else None if entry is not None: @@ -284,7 +285,7 @@ def run_setup_stage( _ReservoirSampler(capacity=100_000) if collect_energy_sample else None ) for i, path in enumerate(files): - for chunk in iter_file_chunks(path, offset=event_id_offset(i)): + for chunk in iter_file_chunks(path, offset=event_id_offset(i), k_max=k_max): mask = sorted_membership(chunk["event_id"], events_arr) if not mask.any(): continue @@ -296,15 +297,17 @@ def run_setup_stage( mat_map, proc_map=proc_map, require_secondaries=True, - conditioning=conditioning, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, sec_phys_only=True, + k_max=k_max, ) ) cond_acc.update(cond_cont) tgt_acc.update(target_s1) if energy_sampler is not None: energy_sampler.update(cond_cont[:, energy_idx]) - sec_valid = np.arange(K_MAX)[None, :] < n_sec[:, None] + sec_valid = np.arange(sec_cont.shape[1])[None, :] < n_sec[:, None] sec_phys = sec_cont[:, :, 4:6][sec_valid] if len(sec_phys) > 0: sec_phys_acc.update(sec_phys) @@ -375,7 +378,9 @@ def run_train_job( ) config.validate_config(cfg) - conditioning = cfg["conditioning"]["particle"]["type"] + particle_conditioning = cfg["conditioning"]["particle"]["type"] + material_conditioning = cfg["conditioning"]["material"]["type"] + k_max = cfg["stage2_model"]["k_max"] setup = run_setup_stage( data, val_fraction=t["val_fraction"], @@ -399,18 +404,20 @@ def run_train_job( ) # cond_cat's onehot columns (docs/v0.3.0-design.md decision 4) are - # present only under conditioning="onehot" — validate_config enforces - # conditioning.particle.type == conditioning.material.type, so both - # axes' topN columns are always present or absent together. run_setup_stage - # builds both maps whenever conditioning=="onehot" (see its own - # particle_cfg["type"] == "onehot" check), so they're guaranteed non-None - # here — asserted, not just assumed, so a future wiring bug fails loudly - # instead of silently dropping the onehot columns. + # present per-axis, independently, under that axis's own + # conditioning.{particle,material}.type == "onehot" (§3.1: the two axes + # may mix freely). run_setup_stage builds each map whenever its own axis + # is "onehot" (see its own particle_cfg["type"]/material_cfg["type"] + # checks), so they're guaranteed non-None here — asserted, not just + # assumed, so a future wiring bug fails loudly instead of silently + # dropping the onehot columns. cond_pdg_topn = None cond_mat_topn = None - if conditioning == "onehot": - assert setup.pdg_topn_map is not None and setup.mat_topn_map is not None + if particle_conditioning == "onehot": + assert setup.pdg_topn_map is not None cond_pdg_topn = setup.pdg_topn_map.class_map + if material_conditioning == "onehot": + assert setup.mat_topn_map is not None cond_mat_topn = setup.mat_topn_map.class_map # The secondary type-index map depends on stage2_model.particle_type.target, @@ -441,11 +448,13 @@ def run_train_job( shuffle_buffer=shuffle_buffer, shuffle=True, proc_map=proc_map, - conditioning=conditioning, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, sec_phys_normalizer=sec_phys_norm, pdg_topn_map=cond_pdg_topn, mat_topn_map=cond_mat_topn, sec_type_class_map=sec_type_class_map, + k_max=k_max, ) val_ds = StreamingStepsDataset( files=files, @@ -457,11 +466,13 @@ def run_train_job( batch_size=t["batch_size"], shuffle=False, proc_map=proc_map, - conditioning=conditioning, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, sec_phys_normalizer=sec_phys_norm, pdg_topn_map=cond_pdg_topn, mat_topn_map=cond_mat_topn, sec_type_class_map=sec_type_class_map, + k_max=k_max, ) pin = device.type == "cuda" diff --git a/giant/rollout.py b/giant/rollout.py index 3bca136..c38edd4 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -352,7 +352,7 @@ def make_seed_frontier( pre_pos: np.ndarray, pre_E: np.ndarray, pre_dir: np.ndarray, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", ) -> tuple[dict[str, np.ndarray], dict[int, int]]: """Build the initial frontier from primary entry states. @@ -372,17 +372,17 @@ def make_seed_frontier( dir_ = dir_ / np.clip(np.linalg.norm(dir_, axis=1, keepdims=True), 1e-12, None) pdg_arr = np.asarray(pdg, dtype=np.int64) - if conditioning == "physical": + if particle_conditioning == "physical": # Real primaries always have a genuine ground-truth PDG code, looked # up once here and carried forward unchanged for the track's lifetime # (its species never changes mid-track) — same lifecycle as "pdg" # itself. mass, charge = particle_phys_array(pdg_arr).T else: - # "embedding" mode never reads mass/charge (see + # "embedding"/"onehot" never read mass/charge (see # _physical_cond_columns), so resolving them here would only risk - # crashing an embedding-mode rollout on a PDG code giant.particles - # can't resolve, for a value that's never used. + # crashing a rollout on a PDG code giant.particles can't resolve, for + # a value that's never used. mass = np.zeros(n, dtype=np.float64) charge = np.zeros(n, dtype=np.float64) @@ -457,8 +457,10 @@ def rollout( max_tracks_per_event: int | None = None, escape_threshold: float | None = None, on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None, - conditioning: str = "embedding", + particle_conditioning: str = "embedding", + material_conditioning: str = "embedding", pdg_topn_map: "TopNMap | None" = None, + mat_topn_map: "TopNMap | None" = None, other_policy: str = "sample", seed: int | None = None, stage1_ddpm_steps: int = 1000, @@ -485,8 +487,13 @@ def rollout( 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 + `pdg_topn_map`/`mat_topn_map` serve two independent purposes that happen + to share `pdg_topn_map` (docs/v0.3.0-design.md §8 — one PDG map, not + two): they're required whenever `particle_conditioning`/ + `material_conditioning` is `"onehot"` (feeds `build_cond_features`'s + extra `cond_cat` top-N columns — §3.1), and `pdg_topn_map`/`other_policy` + are additionally read under `stage2_model.particle_type.target = + "onehot"` (§3.3, secondary-species decode). `seed` seeds the `other_policy = "sample"` draw only (torch/numpy sampling itself is seeded by the caller, same as today). @@ -494,6 +501,16 @@ def rollout( diagnostic across the whole run — see `L1DistCollector`. Only populated under `particle_type.target = "embedding"`; a no-op otherwise. """ + if particle_conditioning == "onehot" and pdg_topn_map is None: + raise RuntimeError( + "conditioning.particle.type='onehot' rollout needs pdg_topn_map " + "(the checkpoint's saved top-N map) — see ckpt['pdg_topn_map']" + ) + if material_conditioning == "onehot" and mat_topn_map is None: + raise RuntimeError( + "conditioning.material.type='onehot' rollout needs mat_topn_map " + "(the checkpoint's saved top-N map) — see ckpt['mat_topn_map']" + ) device = device or torch.device("cpu") stage1_model.eval() sec_decoder.eval() @@ -507,7 +524,7 @@ def rollout( seeds["pre_pos"], seeds["pre_E"], seeds["pre_dir"], - conditioning=conditioning, + particle_conditioning=particle_conditioning, ) rec = _Recorder(sink=on_chunk) @@ -534,8 +551,10 @@ def rollout( steps, device, max_tracks_per_event, - conditioning, + particle_conditioning, + material_conditioning, pdg_topn_map, + mat_topn_map, other_policy, rng, stage1_ddpm_steps, @@ -570,8 +589,10 @@ def _step_chunk( steps, device, max_tracks_per_event, - conditioning, + particle_conditioning, + material_conditioning, pdg_topn_map, + mat_topn_map, other_policy, rng, stage1_ddpm_steps, @@ -587,7 +608,7 @@ def _step_chunk( tr["_material"] = material tr["_layer_id"] = layer_id - if conditioning == "physical": + if particle_conditioning == "physical": # Under physical-property conditioning, mass/charge (already resolved # on every track — see the cond_dict comment below) drive the model, # not a training-vocab PDG embedding — build_cond_features passes @@ -655,7 +676,18 @@ def _step_chunk( "charge": tr["charge"], } cond_cont, cond_cat = build_cond_features( - cond_dict, pdg_map, mat_map, cond_norm, conditioning=conditioning + cond_dict, + pdg_map, + mat_map, + cond_norm, + particle_conditioning=particle_conditioning, + material_conditioning=material_conditioning, + pdg_topn_map=pdg_topn_map.class_map + if particle_conditioning == "onehot" + else None, + mat_topn_map=mat_topn_map.class_map + if material_conditioning == "onehot" + else None, ) cc = torch.from_numpy(cond_cont).float().to(device) ck = torch.from_numpy(cond_cat).long().to(device) diff --git a/giant/train.py b/giant/train.py index 6a3689c..774638c 100644 --- a/giant/train.py +++ b/giant/train.py @@ -4,7 +4,6 @@ import math import os import signal import time -import warnings from pathlib import Path from types import FrameType from typing import Callable @@ -16,7 +15,7 @@ import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm -from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM +from giant.constants import CONT_SLOT_DIM, PARTICLE_PHYS_DIM from giant.data.loader import TopNMap from giant.data.setup_cache import topnmap_to_json from giant.model.network import Router, stage2_type_dim @@ -483,9 +482,10 @@ class FlowDDPMStageTrainer(StageTrainer): ) -> None: if is_stage2 and generator not in ("flow",): raise NotImplementedError( - f"stage2_model.generator={generator!r} is not implemented for " - "stage 2 (only 'flow' and 'wgan' — 'ddpm' has no stage-2 " - "secondary-decoder loss yet, see docs/v0.3.0-design.md)" + f"stage2_model.generator={generator!r} is accepted by the " + "schema but not implemented in v0.3.0 for stage 2 (only " + "'flow' and 'wgan' have a stage-2 secondary-decoder loss — " + "see docs/v0.3.0-design.md §11.2)" ) self.name = name self.is_stage2 = is_stage2 @@ -653,7 +653,9 @@ class FlowDDPMStageTrainer(StageTrainer): proc_idx, sec_type_idx, ) = _batch_to_device(batch, device) - sec_mask = torch.arange(K_MAX, device=device).unsqueeze(0) < n_sec.unsqueeze(1) + sec_mask = torch.arange(sec_cont.size(1), device=device).unsqueeze( + 0 + ) < n_sec.unsqueeze(1) stage1_ctx = x1_s1.detach() x1_s2 = None @@ -957,8 +959,9 @@ class WGANStageTrainer(StageTrainer): self.particle_type_cfg, self.particle_type_emb_dim ) slot_width = CONT_SLOT_DIM + type_dim + k_max = sec_cont.size(1) - sec_mask = torch.arange(K_MAX, device=device).unsqueeze( + sec_mask = torch.arange(k_max, device=device).unsqueeze( 0 ) < n_sec.unsqueeze(1) mask = ( @@ -1002,7 +1005,7 @@ class WGANStageTrainer(StageTrainer): ).reshape(B, -1) * mask ) - z = torch.randn(B, K_MAX, self.model.noise_dim, device=device) + z = torch.randn(B, k_max, self.model.noise_dim, device=device) fake_raw = self.model( z, cond_cont, @@ -1042,7 +1045,7 @@ class WGANStageTrainer(StageTrainer): self.type_gumbel_tau_end, ) fake_raw = _relax_onehot_type_slice( - fake_raw, K_MAX, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe + fake_raw, k_max, CONT_SLOT_DIM, type_dim, tau, grad_probe=grad_probe ) fake = fake_raw * mask @@ -1161,39 +1164,14 @@ class WGANStageTrainer(StageTrainer): group["lr"] = resumed_critic_lr -_WARNED_MARGINAL_VALIDATION_BROKEN = False - - def _try_validate_marginals(trainer: StageTrainer, val_loader, device, **kwargs): - """`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 + """Runs `validate_marginals` on `trainer`'s sampling model (EMA model if + present, else the raw model). `validate_marginals` itself dispatches + through `giant.sample.sample_stage1`/`sample_stage2`/`resolve_n_sec` + (docs/v0.3.0-design.md §10), so this is generator- and + one-shot-vs-autoregressive-agnostic.""" 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/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/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 - return None + return validate_marginals(model, val_loader, device=device, **kwargs) def _build_stage_trainers( @@ -1688,7 +1666,6 @@ def train( stage1_tr, val_loader, device, - mode=cfg["stage1_model"]["generator"], sec_decoder=trainers["stage2"].sampling_model() if "stage2" in trainers else None, @@ -1716,11 +1693,11 @@ def train( stage1_tr, val_loader, device, - mode=cfg["stage1_model"]["generator"], - schedule=stage1_tr.ddpm_schedule - if isinstance(stage1_tr, FlowDDPMStageTrainer) - else None, steps=validate_steps, + ddpm_steps=stage1_tr.ddpm_schedule.T + if isinstance(stage1_tr, FlowDDPMStageTrainer) + and stage1_tr.ddpm_schedule is not None + else 1000, sec_decoder=trainers["stage2"].sampling_model() if "stage2" in trainers else None, diff --git a/giant/validate.py b/giant/validate.py index 62a88df..ea8a34e 100644 --- a/giant/validate.py +++ b/giant/validate.py @@ -2,23 +2,12 @@ import numpy as np import torch from torch.utils.data import DataLoader -from giant.constants import K_MAX, LOCAL_TARGET_NAMES -from giant.sample import ( - sample_flow, - sample_ddpm, - sample_ddim, - sample_secondaries, - sample_wgan, - sample_secondaries_wgan, -) +from giant.constants import LOCAL_TARGET_NAMES +from giant.sample import resolve_n_sec, sample_stage1, sample_stage2 _SEC_PHYS_NAMES = ["log_mass", "charge"] -def _kw(steps: int | None) -> dict[str, int]: - return {} if steps is None else {"steps": steps} - - def _histogram_kl( p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8 ) -> float: @@ -44,15 +33,44 @@ def _bincount_frac(x: np.ndarray, minlength: int) -> np.ndarray: return counts / total if total > 0 else counts +def _categorical_kl( + real_idx: np.ndarray, gen_idx: np.ndarray, n_classes: int, eps: float = 1e-8 +) -> float: + """KL(P_real || Q_gen) between two class-index samples over `n_classes` + categories, estimated from bincount fractions. NaN if either side has no + valid samples (mirrors `_histogram_kl`'s empty-input handling).""" + if len(real_idx) == 0 or len(gen_idx) == 0: + return float("nan") + p = _bincount_frac(real_idx, n_classes) + eps + q = _bincount_frac(gen_idx, n_classes) + eps + p /= p.sum() + q /= q.sum() + return float(np.sum(p * np.log(p / q))) + + +def _embedding_nearest_class( + vectors: torch.Tensor, emb_weight: torch.Tensor +) -> np.ndarray: + """Nearest row index (L1) of `vectors` (..., emb_dim) against `emb_weight` + (vocab, emb_dim) — same computation as + `giant.particles.decode_embedding_nearest`, but returning the raw class + index instead of a decoded PDG code: validate.py only needs a + real-vs-generated class-distribution comparison, not a rollout-usable + identity, so there's no need for the pdg_map inversion here.""" + flat = vectors.reshape(-1, vectors.size(-1)) + dist = (flat.unsqueeze(1) - emb_weight.detach().unsqueeze(0)).abs().sum(-1) + nearest = dist.argmin(dim=1) + return nearest.reshape(vectors.shape[:-1]).cpu().numpy() + + def validate_marginals( - model: torch.nn.Module, + stage1_model: torch.nn.Module, val_loader: DataLoader, - mode: str = "flow", - schedule=None, device: torch.device | None = None, n_batches: int | None = None, kl_bins: int = 50, - steps: int | None = None, + steps: int = 10, + ddpm_steps: int = 1000, sec_decoder: torch.nn.Module | None = None, ) -> dict[str, np.ndarray | float]: """Compare per-dimension marginals of generated vs. real steps. @@ -61,51 +79,61 @@ def validate_marginals( normalised space. `kl_divergence[j]` is KL(real || generated) for dimension j, estimated from a shared histogram over both samples. - `steps` overrides the number of sampler steps (flow ODE steps or DDIM - substeps); `None` keeps each sampler's own default. Unused in "ddpm" - mode, which always runs the full schedule. + Stage 1 is sampled via `giant.sample.sample_stage1`, which dispatches on + `stage1_model.generator_kind` — `steps`/`ddpm_steps` are forwarded but + only one of them is actually read, depending on that dispatch. - When `sec_decoder` is given, also validates Stage 2: n_sec distribution - (+ classification accuracy), predicted secondary physical-identity - (log_mass, charge) marginals, and per-slot energy-fraction marginals — - restricted to each side's own valid slots (real: `n_sec`; generated: the - Stage-1 head's argmax), since the two need not agree on how many slots - are valid. Compared directly in normalised space (no denormalising — - KL estimated from a shared per-sample histogram is invariant to a shared - affine rescaling of both sides). Adds {"n_sec_real", "n_sec_pred", - "n_sec_accuracy", "phys_real", "phys_generated", "phys_kl", - "energy_fraction_kl"} to the returned dict. + When `sec_decoder` is given, also validates Stage 2 via + `giant.sample.sample_stage2`/`resolve_n_sec` (generator- and + one-shot-vs-autoregressive-agnostic, docs/v0.3.0-design.md §10): n_sec + distribution (+ classification accuracy), per-slot energy-fraction + marginals, and a particle-type marginal whose shape depends on + `sec_decoder.particle_type_cfg["target"]` — restricted to each side's own + valid slots (real: `n_sec`; generated: the resolved `n_sec_pred`), since + the two need not agree on how many slots are valid. Adds {"n_sec_real", + "n_sec_pred", "n_sec_accuracy", "energy_fraction_kl"} plus, under + `target = "physical"`, {"phys_real", "phys_generated", "phys_kl"} + (continuous log_mass/charge marginals — v0.2 behaviour), or under + `target` in `("onehot", "embedding")`, {"type_class_real", + "type_class_gen", "type_class_kl"} (categorical class-index marginal: + argmax for "onehot", L1-nearest conditioning-embedding row for + "embedding" — see `_embedding_nearest_class`). Compared directly in + normalised space (no denormalising — KL estimated from a shared + per-sample histogram/bincount is invariant to a shared affine rescaling + of both sides). """ if device is None: - device = next(model.parameters()).device - model.eval() + device = next(stage1_model.parameters()).device + stage1_model.eval() if sec_decoder is not None: sec_decoder.eval() + k_max = sec_decoder.k_max if sec_decoder is not None else 0 + target = ( + sec_decoder.particle_type_cfg.get("target", "physical") + if sec_decoder is not None + else "physical" + ) + all_real, all_gen = [], [] all_n_sec_real, all_n_sec_pred = [], [] all_phys_real, all_phys_gen = [], [] - all_frac_real: list[list[np.ndarray]] = [[] for _ in range(K_MAX)] - all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(K_MAX)] + all_type_class_real, all_type_class_gen = [], [] + all_frac_real: list[list[np.ndarray]] = [[] for _ in range(k_max)] + all_frac_gen: list[list[np.ndarray]] = [[] for _ in range(k_max)] for i, batch in enumerate(val_loader): if n_batches is not None and i >= n_batches: break - # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx). - cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx = batch + # Batch is (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, + # sec_type_idx). + cond_cont, cond_cat, x1, n_sec, sec_cont, _proc_idx, sec_type_idx = batch cond_cont = cond_cont.to(device) cond_cat = cond_cat.to(device) - if mode == "flow": - gen, n_sec_pred = sample_flow(model, cond_cont, cond_cat, **_kw(steps)) - elif mode == "ddpm": - gen, n_sec_pred = sample_ddpm(model, cond_cont, cond_cat, schedule) - elif mode == "wgan": - gen, n_sec_pred = sample_wgan(model, cond_cont, cond_cat) - else: - gen, n_sec_pred = sample_ddim( - model, cond_cont, cond_cat, schedule, **_kw(steps) - ) + gen, n_sec_pred = sample_stage1( + stage1_model, cond_cont, cond_cat, steps=steps, ddpm_steps=ddpm_steps + ) all_real.append(x1.numpy()) all_gen.append(gen.cpu().numpy()) @@ -113,37 +141,41 @@ def validate_marginals( if sec_decoder is None: continue + n_sec_pred = resolve_n_sec( + stage1_model, sec_decoder, cond_cont, cond_cat, gen, n_sec_pred + ) n_sec_pred_np = n_sec_pred.cpu().numpy() n_sec_np = n_sec.numpy() all_n_sec_real.append(n_sec_np) all_n_sec_pred.append(n_sec_pred_np) - real_valid = np.arange(K_MAX)[None, :] < n_sec_np[:, None] # (B, K_MAX) + real_valid = np.arange(k_max)[None, :] < n_sec_np[:, None] # (B, k_max) real_frac = 1.0 / (1.0 + np.exp(-sec_cont[:, :, 0].numpy().astype(np.float64))) - real_phys = sec_cont[:, :, 4:6].numpy() # (B, K_MAX, 2) [log_mass, charge] - if mode == "wgan": - sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries_wgan( - sec_decoder, cond_cont, cond_cat, gen, n_sec_pred - ) - else: - sec_cont_pred, sec_phys_pred, sec_valid_pred = sample_secondaries( - sec_decoder, - cond_cont, - cond_cat, - gen, - n_sec_pred, - steps=steps if steps is not None else 10, - ) + sec_cont_pred, sec_type_pred, sec_valid_pred = sample_stage2( + sec_decoder, cond_cont, cond_cat, gen, n_sec_pred, steps=steps + ) gen_frac = 1.0 / ( 1.0 + np.exp(-sec_cont_pred[:, :, 0].cpu().numpy().astype(np.float64)) ) - gen_phys = sec_phys_pred.cpu().numpy() gen_valid = sec_valid_pred.cpu().numpy() - all_phys_real.append(real_phys[real_valid]) - all_phys_gen.append(gen_phys[gen_valid]) - for j in range(K_MAX): + if target == "physical": + real_phys = sec_cont[:, :, 4:6].numpy() # (B, k_max, 2) [log_mass, charge] + gen_phys = sec_type_pred.cpu().numpy() + all_phys_real.append(real_phys[real_valid]) + all_phys_gen.append(gen_phys[gen_valid]) + else: + sec_type_idx_np = sec_type_idx.numpy() + all_type_class_real.append(sec_type_idx_np[real_valid]) + if target == "onehot": + gen_class = sec_type_pred.argmax(dim=-1).cpu().numpy() + else: # "embedding" + emb_weight = sec_decoder.cond_enc.pdg_emb.weight + gen_class = _embedding_nearest_class(sec_type_pred, emb_weight) + all_type_class_gen.append(gen_class[gen_valid]) + + for j in range(k_max): all_frac_real[j].append(real_frac[real_valid[:, j], j]) all_frac_gen[j].append(gen_frac[gen_valid[:, j], j]) @@ -181,20 +213,8 @@ def validate_marginals( n_sec_real = np.concatenate(all_n_sec_real, axis=0) n_sec_pred_all = np.concatenate(all_n_sec_pred, axis=0) n_sec_accuracy = float((n_sec_real == n_sec_pred_all).mean()) - phys_real = np.concatenate(all_phys_real, axis=0) # (M, 2) - phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2) - if len(phys_real) > 0 and len(phys_gen) > 0: - phys_kl = np.array( - [ - _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) - for j in range(2) - ] - ) - else: - phys_kl = np.full(2, np.nan) - - energy_fraction_kl = np.full(K_MAX, np.nan) + energy_fraction_kl = np.full(k_max, np.nan) print( f"\n{'n_sec':<20} accuracy={n_sec_accuracy:.4f} " f"mean|Δ|={np.abs(n_sec_real - n_sec_pred_all).mean():.4f}" @@ -208,26 +228,12 @@ def validate_marginals( for v in range(max_n_sec): print(f"{v:<20} {real_n_sec_dist[v]:>10.4f} {gen_n_sec_dist[v]:>10.4f}") - print( - f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} " - f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}" - ) - print("-" * 68) - for j, name in enumerate(_SEC_PHYS_NAMES): - r, g = phys_real[:, j], phys_gen[:, j] - if len(r) == 0 or len(g) == 0: - continue - print( - f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} " - f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}" - ) - print( f"\n{'sec slot (energy frac.)':<24} {'real_mean':>10} {'gen_mean':>10} " f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}" ) print("-" * 90) - for j in range(K_MAX): + for j in range(k_max): r = np.concatenate(all_frac_real[j]) if all_frac_real[j] else np.array([]) g = np.concatenate(all_frac_gen[j]) if all_frac_gen[j] else np.array([]) if len(r) == 0 or len(g) == 0: @@ -244,10 +250,63 @@ def validate_marginals( "n_sec_real": n_sec_real, "n_sec_pred": n_sec_pred_all, "n_sec_accuracy": n_sec_accuracy, - "phys_real": phys_real, - "phys_generated": phys_gen, - "phys_kl": phys_kl, "energy_fraction_kl": energy_fraction_kl, } ) + + if target == "physical": + phys_real = np.concatenate(all_phys_real, axis=0) # (M, 2) + phys_gen = np.concatenate(all_phys_gen, axis=0) # (M, 2) + + if len(phys_real) > 0 and len(phys_gen) > 0: + phys_kl = np.array( + [ + _histogram_kl(phys_real[:, j], phys_gen[:, j], bins=kl_bins) + for j in range(2) + ] + ) + else: + phys_kl = np.full(2, np.nan) + + print( + f"\n{'sec phys (normalised)':<20} {'real_mean':>10} {'gen_mean':>10} " + f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}" + ) + print("-" * 68) + for j, name in enumerate(_SEC_PHYS_NAMES): + r, g = phys_real[:, j], phys_gen[:, j] + if len(r) == 0 or len(g) == 0: + continue + print( + f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} " + f"{r.std():>10.4f} {g.std():>10.4f} {phys_kl[j]:>14.4f}" + ) + + result.update( + {"phys_real": phys_real, "phys_generated": phys_gen, "phys_kl": phys_kl} + ) + else: + type_class_real = np.concatenate(all_type_class_real, axis=0) + type_class_gen = np.concatenate(all_type_class_gen, axis=0) + n_classes = ( + sec_decoder.type_dim + if target == "onehot" + else sec_decoder.cond_enc.pdg_emb.weight.size(0) + ) + type_class_kl = _categorical_kl(type_class_real, type_class_gen, n_classes) + + print( + f"\n{'sec type class (' + target + ')':<24} " + f"n={len(type_class_real)}/{len(type_class_gen)} " + f"KL(real||gen)={type_class_kl:.4f}" + ) + + result.update( + { + "type_class_real": type_class_real, + "type_class_gen": type_class_gen, + "type_class_kl": type_class_kl, + } + ) + return result diff --git a/scripts/dwarf.py b/scripts/dwarf.py index 200b1ec..311e384 100644 --- a/scripts/dwarf.py +++ b/scripts/dwarf.py @@ -522,10 +522,21 @@ def warm_cache( "--seed", "-s", help="Must match the `giant train` run(s) to warm for" ), ] = 0, - conditioning: Annotated[ + particle_conditioning: Annotated[ Conditioning, typer.Option( - "--conditioning", help="Must match the `giant train` run(s) to warm for" + "--particle-conditioning", + help="Must match the `giant train` run(s)' conditioning.particle.type " + "to warm for", + ), + ] = Conditioning.physical, + material_conditioning: Annotated[ + Conditioning, + typer.Option( + "--material-conditioning", + help="Must match the `giant train` run(s)' conditioning.material.type " + "to warm for — independent of --particle-conditioning " + "(docs/v0.3.0-design.md §3.1: the two axes may differ)", ), ] = Conditioning.physical, router: Annotated[ @@ -552,15 +563,17 @@ def warm_cache( """Precompute `giant train`'s setup-stage sidecar for `data` ahead of time. Warms the vocab maps, event-id split index, and the normalizer entry for - the given --val-fraction/--seed/--conditioning, so a later `giant train` - run (or a `dwarf hparam-scan` sweep, which shares one such entry across - every run) skips straight to training. See giant/data/setup_cache.py. + the given --val-fraction/--seed/--particle-conditioning/ + --material-conditioning, so a later `giant train` run (or a `dwarf + hparam-scan` sweep, which shares one such entry across every run) skips + straight to training. See giant/data/setup_cache.py. """ run_warm_setup_cache( data=str(data), val_fraction=val_fraction, seed=seed, - conditioning=conditioning.value, + particle_conditioning=particle_conditioning.value, + material_conditioning=material_conditioning.value, router_enabled=router, router_type=router_type, n_experts=n_experts, diff --git a/scripts/warm_setup_cache.py b/scripts/warm_setup_cache.py index 030529b..115463c 100644 --- a/scripts/warm_setup_cache.py +++ b/scripts/warm_setup_cache.py @@ -9,6 +9,7 @@ for the sidecar itself. from pathlib import Path +from giant.constants import K_MAX from giant.pipeline import run_setup_stage @@ -16,7 +17,8 @@ def run_warm_setup_cache( data: str, val_fraction: float = 0.1, seed: int = 0, - conditioning: str = "physical", + particle_conditioning: str = "physical", + material_conditioning: str = "physical", router_enabled: bool = False, router_type: str = "energy", n_experts: int = 4, @@ -25,10 +27,12 @@ def run_warm_setup_cache( ) -> None: """Populate (or refresh) the setup cache sidecar for `data`. - `val_fraction`/`seed`/`conditioning` select the normalizer cache entry + `val_fraction`/`seed`/`particle_conditioning`/`material_conditioning` + select the normalizer cache entry (`giant.data.setup_cache.normalizer_key`) — pass the same values a later - `giant train` invocation will use so it hits this warmed entry. - `router_enabled`/`router_type`/`n_experts` only matter for + `giant train` invocation will use so it hits this warmed entry. The two + conditioning axes are independent (docs/v0.3.0-design.md §3.1) and may + differ. `router_enabled`/`router_type`/`n_experts` only matter for `router_type == "process"` (warms that `n_experts`'s process map); the energy-router quantile summary is always collected regardless, so a later `--router-type energy` run never needs to rescan just to seed @@ -40,16 +44,17 @@ def run_warm_setup_cache( "n_experts": n_experts, } # A minimal v0.3 cfg — only the keys run_setup_stage actually reads - # (conditioning.particle.type, stage{1,2}_model.router). This CLI only - # ever configures one router (matching today's single --router-type - # flag), so it's placed on stage1_model; stage2_model's stays disabled. + # (conditioning.{particle,material}.type, stage{1,2}_model.router). This + # CLI only ever configures one router (matching today's single + # --router-type flag), so it's placed on stage1_model; stage2_model's + # stays disabled. cfg = { "conditioning": { - "particle": {"type": conditioning}, - "material": {"type": conditioning}, + "particle": {"type": particle_conditioning}, + "material": {"type": material_conditioning}, }, "stage1_model": {"router": router_cfg}, - "stage2_model": {"router": {"enabled": False}}, + "stage2_model": {"router": {"enabled": False}, "k_max": K_MAX}, } run_setup_stage( Path(data), diff --git a/tests/test_cli_predict.py b/tests/test_cli_predict.py index 1ee8af6..a6f74fb 100644 --- a/tests/test_cli_predict.py +++ b/tests/test_cli_predict.py @@ -1,73 +1,14 @@ import uuid -import pytest -import typer import yaml from giant.cli import ( _CEPH_PREDICTIONS, - _check_conditioning_onehot_support, _resolve_prediction_output, _write_prediction_ref, ) -# --------------------------------------------------------------------------- -# _check_conditioning_onehot_support -# --------------------------------------------------------------------------- - - -def _nested_model_cfg( - particle_type="physical", material_type="physical", target="physical" -): - return { - "conditioning": { - "particle": {"type": particle_type, "emb_dim": 8}, - "material": {"type": material_type, "emb_dim": 8}, - }, - "stage2_model": {"particle_type": {"target": target}}, - } - - -def test_check_conditioning_onehot_support_allows_physical(): - _check_conditioning_onehot_support(_nested_model_cfg(), "predict") # no raise - - -def test_check_conditioning_onehot_support_rejects_onehot_particle_conditioning(): - cfg = _nested_model_cfg(particle_type="onehot") - with pytest.raises(typer.Exit): - _check_conditioning_onehot_support(cfg, "predict") - - -def test_check_conditioning_onehot_support_rejects_onehot_material_conditioning(): - cfg = _nested_model_cfg(material_type="onehot") - with pytest.raises(typer.Exit): - _check_conditioning_onehot_support(cfg, "rollout") - - -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") - _check_conditioning_onehot_support(cfg, "predict") # no raise - - -def test_check_conditioning_onehot_support_allows_embedding_particle_type_target(): - cfg = _nested_model_cfg( - particle_type="embedding", material_type="embedding", target="embedding" - ) - _check_conditioning_onehot_support(cfg, "predict") # no raise - - -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, so this must be a silent no-op rather - than crash on `.get("particle")` against a string.""" - cfg = {"conditioning": "embedding", "mode": "flow"} - _check_conditioning_onehot_support(cfg, "predict") # no raise - - # --------------------------------------------------------------------------- # _resolve_prediction_output # --------------------------------------------------------------------------- diff --git a/tests/test_cli_train_overrides.py b/tests/test_cli_train_overrides.py new file mode 100644 index 0000000..2341b69 --- /dev/null +++ b/tests/test_cli_train_overrides.py @@ -0,0 +1,97 @@ +"""Tests for `giant train`'s stage-prefixed CLI flags (docs/v0.3.0-design.md +decision 7 / docs/v0.3.0-followups.md item 2): --stage1-*/--stage2-* must +independently override each stage's config block, and must take precedence +over the older shared flags (--mode/--hidden-dim/--n-critic/... ) that still +apply the same value to both stages for backward compatibility.""" + +from __future__ import annotations + +from pathlib import Path + +from typer.testing import CliRunner + +import giant.cli as cli + +runner = CliRunner() + + +def _invoke_and_capture_cfg(monkeypatch, tmp_path: Path, args: list[str]) -> dict: + captured: dict = {} + + def _fake_run_train_job(*, data, cfg, out_dir, **kwargs): + captured["cfg"] = cfg + + monkeypatch.setattr(cli, "run_train_job", _fake_run_train_job) + + result = runner.invoke( + cli.app, + ["train", "dummy.parquet", "--out", str(tmp_path / "run")] + args, + ) + assert result.exit_code == 0, result.output + return captured["cfg"] + + +def test_stage_prefixed_generator_overrides_shared_mode(monkeypatch, tmp_path): + cfg = _invoke_and_capture_cfg( + monkeypatch, + tmp_path, + ["--mode", "wgan", "--stage1-generator", "flow"], + ) + assert cfg["stage1_model"]["generator"] == "flow" + assert cfg["stage2_model"]["generator"] == "wgan" + + +def test_stage2_only_knobs(monkeypatch, tmp_path): + cfg = _invoke_and_capture_cfg( + monkeypatch, + tmp_path, + [ + "--stage2-decoder", + "one_shot", + "--stage2-k-max", + "8", + "--stage2-hidden-dim", + "32", + "--stage2-context-dim", + "16", + "--stage2-stage1-context", + "sampled", + ], + ) + assert cfg["stage2_model"]["decoder"] == "one_shot" + assert cfg["stage2_model"]["k_max"] == 8 + assert cfg["stage2_model"]["hidden_dim"] == 32 + assert cfg["stage2_model"]["context_dim"] == 16 + assert cfg["stage2_model"]["stage1_context"] == "sampled" + # untouched stage1 defaults + assert cfg["stage1_model"]["hidden_dim"] == 256 + + +def test_stage1_hidden_dim_flag_overrides_legacy_hidden_dim_flag(monkeypatch, tmp_path): + cfg = _invoke_and_capture_cfg( + monkeypatch, + tmp_path, + ["--hidden-dim", "64", "--stage1-hidden-dim", "128"], + ) + assert cfg["stage1_model"]["hidden_dim"] == 128 + + +def test_wgan_knobs_split_per_stage(monkeypatch, tmp_path): + cfg = _invoke_and_capture_cfg( + monkeypatch, + tmp_path, + [ + "--mode", + "wgan", + "--n-critic", + "5", + "--stage1-n-critic", + "3", + "--stage2-gp-weight", + "2.5", + ], + ) + assert cfg["stage1_model"]["wgan"]["n_critic"] == 3 + assert cfg["stage1_model"]["wgan"]["gp_weight"] == 10.0 + assert cfg["stage2_model"]["wgan"]["n_critic"] == 5 + assert cfg["stage2_model"]["wgan"]["gp_weight"] == 2.5 diff --git a/tests/test_config.py b/tests/test_config.py index b57b5cf..ca50de6 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -584,21 +584,19 @@ def test_validate_config_embedding_target_passes_with_embedding_conditioning(): gconfig.validate_config(cfg) # must not raise -def test_validate_config_mixed_particle_material_conditioning_not_supported(): - """The data pipeline doesn't support mixed conditioning types yet, even - though ConditionEncoder itself already can (docs/v0.3.0-design.md §3.1 - vs. giant/data/transforms.py's still-single conditioning param).""" +def test_validate_config_mixed_particle_material_conditioning_is_valid(): + """docs/v0.3.0-design.md §3.1: the particle and material conditioning + axes are configured independently and may mix freely — e.g. material + "physical" with particle "embedding" — and the data pipeline + (giant/data/transforms.py) now implements that end-to-end, so + validate_config must not reject it.""" cfg = _cfg_with( **{ "conditioning.particle.type": "physical", "conditioning.material.type": "embedding", } ) - try: - gconfig.validate_config(cfg) - assert False, "expected ValueError" - except ValueError as e: - assert "mixed" in str(e) or "conditioning.material.type" in str(e) + gconfig.validate_config(cfg) # must not raise def test_validate_config_pdg_router_incompatible_with_physical_conditioning(): @@ -639,6 +637,49 @@ def test_validate_config_stop_token_not_implemented(): assert "stop_token" in str(e) +def test_validate_config_n_sec_truth_rejected_for_rollout_capable_checkpoint(): + """docs/v0.3.0-design.md §9: 'n_sec.mode = "truth" is invalid for a + rollout-capable checkpoint' — both stages active means giant rollout + could load this checkpoint, but 'truth' has no ground truth to draw + n_sec from at rollout time.""" + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "truth", + "stage1_model.active": True, + "stage2_model.active": True, + } + ) + try: + gconfig.validate_config(cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "n_sec.mode" in str(e) and "truth" in str(e) + + +def test_validate_config_n_sec_truth_allowed_for_stage2_only_checkpoint(): + """'truth' is exactly the standalone stage-2 evaluation mode the design + doc carves out — stage1_model.active = false must still pass.""" + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "truth", + "stage1_model.active": False, + "stage2_model.active": True, + } + ) + gconfig.validate_config(cfg) # must not raise + + +def test_validate_config_n_sec_truth_allowed_when_stage2_inactive(): + cfg = _cfg_with( + **{ + "stage2_model.n_sec.mode": "truth", + "stage1_model.active": True, + "stage2_model.active": False, + } + ) + gconfig.validate_config(cfg) # must not raise + + def test_validate_config_ar_default_markov_always_passes(): """DEFAULT_CONFIG already has decoder='autoregressive', history='markov', teacher_forcing='always' — must not raise (v0.3.0 diff --git a/tests/test_dataset.py b/tests/test_dataset.py index e311dae..5d2e71e 100644 --- a/tests/test_dataset.py +++ b/tests/test_dataset.py @@ -126,7 +126,8 @@ def test_streaming_dataset_offsets_colliding_event_ids_across_files(tmp_path): target_normalizer=tgt_norm, batch_size=4, shuffle=False, - conditioning="embedding", + particle_conditioning="embedding", + material_conditioning="embedding", ) return sum(len(batch[0]) for batch in ds) diff --git a/tests/test_dwarf.py b/tests/test_dwarf.py index 14b2673..767828e 100644 --- a/tests/test_dwarf.py +++ b/tests/test_dwarf.py @@ -103,7 +103,7 @@ def test_warm_cache_writes_sidecar(tmp_path): assert loaded is not None assert loaded.vocab is not None assert loaded.event_index is not None - assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers + assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers def test_warm_cache_second_run_hits_cache(tmp_path): @@ -167,5 +167,5 @@ def test_warm_cache_different_val_fraction_is_separate_entry(tmp_path): assert "fitting normalizer (streaming)" in result.output loaded = setup_cache.load(data, [data]) assert loaded is not None - assert "valfrac=0.1_seed=0_cond=physical" in loaded.normalizers - assert "valfrac=0.3_seed=0_cond=physical" in loaded.normalizers + assert "valfrac=0.1_seed=0_pcond=physical_mcond=physical" in loaded.normalizers + assert "valfrac=0.3_seed=0_pcond=physical_mcond=physical" in loaded.normalizers diff --git a/tests/test_flow.py b/tests/test_flow.py index f1276f0..aae4a93 100644 --- a/tests/test_flow.py +++ b/tests/test_flow.py @@ -52,7 +52,7 @@ def test_sample_flow_shape(): cond_cat = torch.zeros(B, 2, dtype=torch.long) sample, n_sec = sample_flow(_small_model(), cond_cont, cond_cat, steps=5) assert sample.shape == (B, 9) - assert n_sec.shape == (B,) + assert n_sec is not None and n_sec.shape == (B,) def test_ddpm_loss_nonneg(): @@ -69,4 +69,4 @@ def test_sample_ddim_shape(): cond_cat = torch.zeros(B, 2, dtype=torch.long) sample, n_sec = sample_ddim(_small_model(), cond_cont, cond_cat, schedule, steps=5) assert sample.shape == (B, 9) - assert n_sec.shape == (B,) + assert n_sec is not None and n_sec.shape == (B,) diff --git a/tests/test_migration_v02_v03.py b/tests/test_migration_v02_v03.py index 488bf37..6c5d10c 100644 --- a/tests/test_migration_v02_v03.py +++ b/tests/test_migration_v02_v03.py @@ -198,6 +198,56 @@ def test_migrate_legacy_model_config_shape(): assert migrated["stage2_model"]["decoder"] == "one_shot" +def test_migrate_legacy_model_config_nonzero_expert_dims_raises(): + """docs/v0.3.0-followups.md item 8 regression: a v0.2 checkpoint's + model_config carrying a non-default expert_hidden_dim/expert_n_blocks + must fail loudly through this path too (§4.2) — not just + giant.config.migrate_config's parallel TOML-load path. Silently dropping + these keys (build_router's kwarg filtering) would resize the experts + instead of refusing.""" + legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical") + legacy_cfg["router"] = { + "enabled": True, + "expert_hidden_dim": 128, + "expert_n_blocks": 0, + } + try: + net._migrate_legacy_model_config(legacy_cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "expert_hidden_dim" in str(e) + + +def test_migrate_legacy_model_config_zero_expert_dims_dropped_silently(): + legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical") + legacy_cfg["router"] = { + "enabled": True, + "expert_hidden_dim": 0, + "expert_n_blocks": 0, + } + migrated = net._migrate_legacy_model_config(legacy_cfg) + assert "expert_hidden_dim" not in migrated["stage1_model"]["router"] + assert "expert_n_blocks" not in migrated["stage1_model"]["router"] + assert "expert_hidden_dim" not in migrated["stage2_model"]["router"] + assert "expert_n_blocks" not in migrated["stage2_model"]["router"] + + +def test_build_models_with_legacy_config_nonzero_expert_dims_raises(): + """The same check must also fire through the actual caller, + build_models, not just the internal helper directly.""" + legacy_cfg = _legacy_model_config(mode="flow", conditioning="physical") + legacy_cfg["router"] = { + "enabled": True, + "expert_hidden_dim": 128, + "expert_n_blocks": 0, + } + try: + net.build_models(legacy_cfg) + assert False, "expected ValueError" + except ValueError as e: + assert "expert_hidden_dim" in str(e) + + def test_build_models_accepts_new_nested_shape_unchanged(): """A dict that already has a 'stage1_model' key (the new shape) is passed through build_models without going through the legacy migration diff --git a/tests/test_network.py b/tests/test_network.py index f3c5293..0a254f6 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -1,5 +1,8 @@ +import copy + import pytest import torch +from giant import config as gconfig from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM from giant.model.network import ( AttentionHistory, @@ -9,6 +12,7 @@ from giant.model.network import ( Stage1Model, Stage2Autoregressive, Stage2OneShot, + build_models, cat_col_layout, stage2_trunk_sec_dim, stage2_type_dim, @@ -628,3 +632,52 @@ def test_stage2_autoregressive_history_step_matches_parallel_history_encoder(): def test_stage2_autoregressive_init_history_cache_is_none_for_markov(): model = _build_stage2_ar("physical", "wgan", history="markov") assert model.init_history_cache() is None + + +# ── build_models: conditioning.share_stages ───────────────────────────────── + + +def _minimal_model_config(share_stages: bool) -> dict: + cfg = copy.deepcopy(gconfig.DEFAULT_CONFIG) + cfg["conditioning"]["share_stages"] = share_stages + cfg["conditioning"]["particle"]["emb_dim"] = 4 + cfg["conditioning"]["material"]["emb_dim"] = 4 + cfg["stage1_model"].update({"hidden_dim": 8, "n_res_blocks": 1}) + cfg["stage2_model"].update({"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3}) + return { + "pdg_vocab": 3, + "mat_vocab": 2, + "conditioning": cfg["conditioning"], + "stage1_model": cfg["stage1_model"], + "stage2_model": cfg["stage2_model"], + } + + +def test_build_models_share_stages_true_shares_condition_encoder_instance(): + built = build_models(_minimal_model_config(share_stages=True)) + stage1, stage2 = built["stage1"], built["stage2"] + assert stage1 is not None and stage2 is not None + assert stage1.cond_enc is stage2.cond_enc + + +def test_build_models_share_stages_false_builds_independent_condition_encoders(): + built = build_models(_minimal_model_config(share_stages=False)) + stage1, stage2 = built["stage1"], built["stage2"] + assert stage1 is not None and stage2 is not None + assert stage1.cond_enc is not stage2.cond_enc + + +def test_build_models_share_stages_true_shared_params_are_in_both_stage_parameter_lists(): + """The shared encoder's parameters must actually appear in both stages' + own `.parameters()` — that's what makes each stage's independent + optimizer include (and update) them, which is the actual mechanism behind + "shared weights, forced common representation" (docs/v0.3.0-design.md + §3.1), not just object identity on `.cond_enc`.""" + built = build_models(_minimal_model_config(share_stages=True)) + stage1, stage2 = built["stage1"], built["stage2"] + assert stage1 is not None and stage2 is not None + + shared_ids = {id(p) for p in stage1.cond_enc.parameters()} + assert shared_ids + assert shared_ids <= {id(p) for p in stage1.parameters()} + assert shared_ids <= {id(p) for p in stage2.parameters()} diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index f4c093c..f85beca 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -7,6 +7,7 @@ import torch from giant import config as gconfig from giant.data import setup_cache +from giant.data.transforms import Normalizer from giant.pipeline import run_train_job @@ -235,6 +236,74 @@ def test_run_train_job_new_val_fraction_is_partial_hit(tmp_path, data, monkeypat assert "fitting normalizer (streaming)" in joined +def test_run_train_job_custom_k_max_end_to_end(tmp_path, data): + """docs/v0.3.0-followups.md item 3 regression: stage2_model.k_max other + than the K_MAX module constant's default (15) must not produce a shape + mismatch between the data pipeline (loader.py/transforms.py padding) and + the model (network.py's trunks, sized from this same config value).""" + cfg = _tiny_cfg() + cfg["stage2_model"]["k_max"] = 3 + _run(data, tmp_path / "out", cfg=cfg) + ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False) + assert ckpt["model_config"]["stage2_model"]["k_max"] == 3 + + +def test_run_train_job_mixed_particle_material_conditioning_end_to_end(tmp_path, data): + """docs/v0.3.0-followups.md item 4 regression: conditioning.particle.type + and conditioning.material.type are configured independently and may mix + freely (docs/v0.3.0-design.md §3.1) — e.g. particle "embedding" with + material "physical" — end-to-end through the real data pipeline, not + just accepted by validate_config.""" + cfg = _tiny_cfg() + cfg["conditioning"]["particle"]["type"] = "embedding" + cfg["conditioning"]["material"]["type"] = "physical" + _run(data, tmp_path / "out", cfg=cfg) + ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False) + cond_cfg = ckpt["model_config"]["conditioning"] + assert cond_cfg["particle"]["type"] == "embedding" + assert cond_cfg["material"]["type"] == "physical" + + cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"]) + # Particle block ([COND_DIM_BASE:COND_DIM_BASE+PARTICLE_PHYS_DIM]) stays + # unfitted (mean=0/std=1) since "embedding" never computes real values + # for it; the material block is fit for real under "physical". + from giant.constants import COND_DIM_BASE, PARTICLE_PHYS_DIM + + assert cond_norm.mean is not None and cond_norm.std is not None + np.testing.assert_allclose( + cond_norm.mean[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 0.0 + ) + np.testing.assert_allclose( + cond_norm.std[COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM], 1.0 + ) + material_std = cond_norm.std[COND_DIM_BASE + PARTICLE_PHYS_DIM :] + assert np.all(material_std > 0) and not np.allclose(material_std, 1.0) + + +def test_run_train_job_share_stages_end_to_end(tmp_path, data): + """docs/v0.3.0-followups.md item 5 regression: conditioning.share_stages + = true must actually train (not raise NotImplementedError), and the + resulting checkpoint's two stages must reload into a single shared + ConditionEncoder instance rather than two independent ones.""" + from giant.model.network import build_models + + cfg = _tiny_cfg() + cfg["conditioning"]["share_stages"] = True + _run(data, tmp_path / "out", cfg=cfg) + ckpt = torch.load(tmp_path / "out" / "last.pt", weights_only=False) + assert ckpt["model_config"]["conditioning"]["share_stages"] is True + + built = build_models(ckpt["model_config"]) + stage1, stage2 = built["stage1"], built["stage2"] + assert stage1 is not None and stage2 is not None + assert stage1.cond_enc is stage2.cond_enc + + stage1.load_state_dict(ckpt["model"]) + stage2.load_state_dict(ckpt["sec_decoder"]) + for p1, p2 in zip(stage1.cond_enc.parameters(), stage2.cond_enc.parameters()): + assert torch.equal(p1, p2) + + def test_run_train_job_matches_uncached_output(tmp_path, data): _run(data, tmp_path / "uncached", cache_setup=False) _run(data, tmp_path / "cached1", cache_setup=True) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index f03eab8..7c230b2 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -111,7 +111,8 @@ def _run( batch_size=128, max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, - conditioning=conditioning, + particle_conditioning=conditioning, + material_conditioning=conditioning, ) @@ -172,7 +173,7 @@ def test_seed_frontier_embedding_mode_skips_unresolvable_pdg_lookup(): frontier construction — mass/charge are simply zero-filled, unused.""" seeds = _seeds(3) seeds["pdg"] = np.full(3, 999999999, dtype=np.int64) - fr, _counts = make_seed_frontier(**seeds, conditioning="embedding") + fr, _counts = make_seed_frontier(**seeds, particle_conditioning="embedding") np.testing.assert_array_equal(fr["mass"], 0.0) np.testing.assert_array_equal(fr["charge"], 0.0) @@ -384,25 +385,43 @@ def _models_v3( 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, - ) + # Explicit kwargs rather than a shared **common dict: a dict() call whose + # values have heterogeneous types (str/int/dict/bool) widens under static + # analysis to dict[str, ], which then makes every constructor + # keyword not itself part of that union (router, cond_enc, ...) look like + # a type mismatch to `ty` even though every actual value passed is fine. if decoder == "one_shot": sec_dim = stage2_trunk_sec_dim(particle_type_cfg, generator2, k_max, emb_dim) - s2 = Stage2OneShot(sec_dim=sec_dim, **common) + s2 = Stage2OneShot( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=32, + n_res_blocks=2, + generator=generator2, + time_dim=16, + noise_dim=8, + k_max=k_max, + particle_type_cfg=particle_type_cfg, + build_n_sec_head=stage2_has_n_sec_head, + sec_dim=sec_dim, + ) else: - s2 = Stage2Autoregressive(**common) + s2 = Stage2Autoregressive( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=32, + n_res_blocks=2, + generator=generator2, + time_dim=16, + noise_dim=8, + k_max=k_max, + particle_type_cfg=particle_type_cfg, + build_n_sec_head=stage2_has_n_sec_head, + ) return s1.eval(), s2.eval() @@ -440,7 +459,8 @@ def _run_v3( batch_size=128, max_tracks_per_event=max_tracks_per_event, escape_threshold=escape_threshold, - conditioning=conditioning, + particle_conditioning=conditioning, + material_conditioning=conditioning, pdg_topn_map=pdg_topn_map, other_policy=other_policy, seed=seed, @@ -535,6 +555,88 @@ def test_rollout_onehot_target_missing_topn_map_raises(fake_material_props): _run_v3(s1, s2, pdg_topn_map=None) +# --- conditioning.{particle,material}.type = "onehot" (docs/v0.3.0-followups.md +# item 7) — a separate axis from stage2_model.particle_type.target above: this +# is what feeds cond_cat's extra top-N columns for ConditionEncoder's own +# "onehot" mode, not the secondary-species decode. --------------------------- + +COND_PDG_TOPN_MAP = TopNMap(class_map=dict(PDG_MAP), other_members={}) +COND_MAT_TOPN_MAP = TopNMap(class_map={"G4_AIR": 0, "G4_PbWO4": 1}, other_members={}) + + +def _onehot_conditioning_models(): + particle_cfg = {"type": "onehot", "emb_dim": len(PDG_MAP), "n_layers": 1} + material_cfg = {"type": "onehot", "emb_dim": len(MAT_MAP), "n_layers": 1} + s1 = Stage1Model( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=32, + n_res_blocks=2, + ) + s2 = Stage2OneShot( + pdg_vocab=3, + mat_vocab=2, + particle_cfg=particle_cfg, + material_cfg=material_cfg, + hidden_dim=32, + n_res_blocks=2, + sec_dim=stage2_trunk_sec_dim({"target": "physical"}, "flow", K_MAX, 3), + generator="flow", + time_dim=16, + ) + return s1.eval(), s2.eval() + + +def _run_onehot_conditioning( + pdg_topn_map=COND_PDG_TOPN_MAP, mat_topn_map=COND_MAT_TOPN_MAP +): + s1, s2 = _onehot_conditioning_models() + cond, tgt, sec_phys = _norms() + return rollout( + s1, + s2, + _oracle(), + _seeds(), + cond, + tgt, + sec_phys, + PDG_MAP, + MAT_MAP, + energy_cutoff=1.0, + max_steps=30, + steps=4, + batch_size=128, + max_tracks_per_event=300, + escape_threshold=1e9, + particle_conditioning="onehot", + material_conditioning="onehot", + pdg_topn_map=pdg_topn_map, + mat_topn_map=mat_topn_map, + ) + + +def test_rollout_conditioning_onehot_end_to_end(fake_material_props): + rec = _run_onehot_conditioning() + assert len(rec["event_id"]) > 0 + assert set(rec["event_id"].tolist()) == set(range(6)) + + +def test_rollout_conditioning_onehot_particle_missing_topn_map_raises( + fake_material_props, +): + with pytest.raises(RuntimeError, match="pdg_topn_map"): + _run_onehot_conditioning(pdg_topn_map=None) + + +def test_rollout_conditioning_onehot_material_missing_topn_map_raises( + fake_material_props, +): + with pytest.raises(RuntimeError, match="mat_topn_map"): + _run_onehot_conditioning(mat_topn_map=None) + + @pytest.mark.parametrize("decoder", ["one_shot", "autoregressive"]) def test_rollout_embedding_target_end_to_end(decoder): """particle_type.target="embedding" L1-snaps to the nearest row of the diff --git a/tests/test_sample.py b/tests/test_sample.py index cb8bdfd..ff52610 100644 --- a/tests/test_sample.py +++ b/tests/test_sample.py @@ -150,7 +150,7 @@ def test_sample_flow_returns_n_sec_for_legacy_stage1(): ) cond_cont, cond_cat = _cond(5) _, n_sec = sample_flow(model, cond_cont, cond_cat, steps=2) - assert n_sec.shape == (5,) + assert n_sec is not None and n_sec.shape == (5,) # ── Stage2OneShot: non-"physical" particle_type.target ────────────────────── diff --git a/tests/test_transforms.py b/tests/test_transforms.py index a2e19c2..04ce58d 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -426,7 +426,13 @@ def test_build_features_embedding_mode_zero_fills_physical_columns(): data = _minimal_step_data(3) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} - cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="embedding") + cond_cont, *_ = build_features( + data, + pdg_map, + mat_map, + particle_conditioning="embedding", + material_conditioning="embedding", + ) assert cond_cont.shape[1] == COND_DIM np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0) @@ -438,7 +444,13 @@ def test_build_features_physical_mode_shape_and_values(fake_material_props): data = _minimal_step_data(3) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} - cond_cont, *_ = build_features(data, pdg_map, mat_map, conditioning="physical") + cond_cont, *_ = build_features( + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_conditioning="physical", + ) assert cond_cont.shape[1] == COND_DIM mass, charge = particle_mass_charge(11) @@ -461,7 +473,13 @@ def test_build_features_physical_mode_unfilled_material_raises(): pdg_map, mat_map = {11: 0}, {"G4_LYSO": 0} with pytest.raises(MaterialPropertiesNotFilledError): - build_features(data, pdg_map, mat_map, conditioning="physical") + build_features( + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_conditioning="physical", + ) def test_build_cond_features_mass_charge_override(fake_material_props): @@ -474,7 +492,13 @@ def test_build_cond_features_mass_charge_override(fake_material_props): data["charge"] = np.array([2.0, -2.0], dtype=np.float32) pdg_map, mat_map = {11: 0}, {"PbWO4": 0} - cond_cont, _ = build_cond_features(data, pdg_map, mat_map, conditioning="physical") + cond_cont, _ = build_cond_features( + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_conditioning="physical", + ) np.testing.assert_allclose( cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])) @@ -494,7 +518,12 @@ def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode(): legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32) cond_cont, _ = build_cond_features( - data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding" + data, + pdg_map, + mat_map, + cond_normalizer=legacy_norm, + particle_conditioning="embedding", + material_conditioning="embedding", ) assert cond_cont.shape[-1] == COND_DIM @@ -521,7 +550,8 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode( pdg_map, mat_map, cond_normalizer=legacy_norm, - conditioning="physical", + particle_conditioning="physical", + material_conditioning="physical", ) @@ -610,13 +640,23 @@ def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_materi } cond_cont, cond_cat = build_cond_features( - data, pdg_map, mat_map, conditioning="physical" + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_conditioning="physical", ) assert cond_cont.shape[-1] == COND_DIM np.testing.assert_array_equal(cond_cat, [[0, 0]]) # dummy indices, no raise with pytest.raises(KeyError): - build_cond_features(data, pdg_map, mat_map, conditioning="embedding") + build_cond_features( + data, + pdg_map, + mat_map, + particle_conditioning="embedding", + material_conditioning="embedding", + ) # ── _WelfordAccumulator ────────────────────────────────────────────────────── diff --git a/tests/test_validate.py b/tests/test_validate.py index 61fa72f..533c933 100644 --- a/tests/test_validate.py +++ b/tests/test_validate.py @@ -1,15 +1,18 @@ import numpy as np import torch -from giant.constants import COND_DIM, K_MAX, SEC_SLOT_DIM, X_DIM -from giant.model.network import Stage1Model, Stage2OneShot +from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM +from giant.model.network import Stage1Model, Stage2OneShot, stage2_trunk_sec_dim from giant.validate import validate_marginals _PARTICLE_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} _MATERIAL_CFG = {"type": "physical", "emb_dim": 8, "n_layers": 1} +_K_MAX = 5 -def _tiny_models(): +def _tiny_models(particle_type_cfg: dict | None = None): + """A fresh v0.3.0 pair: Stage1Model owns no n_sec_head (decision 1), so + n_sec always comes from Stage2OneShot.""" s1 = Stage1Model( pdg_vocab=3, mat_vocab=2, @@ -17,7 +20,13 @@ def _tiny_models(): material_cfg=_MATERIAL_CFG, hidden_dim=16, n_res_blocks=1, - n_sec_head_k_max=K_MAX, + ) + target = (particle_type_cfg or {}).get("target", "physical") + sec_dim = stage2_trunk_sec_dim( + particle_type_cfg or {"target": "physical"}, + "flow", + _K_MAX, + int(_PARTICLE_CFG["emb_dim"]), ) s2 = Stage2OneShot( pdg_vocab=3, @@ -28,23 +37,29 @@ def _tiny_models(): n_res_blocks=1, generator="flow", time_dim=16, + k_max=_K_MAX, + sec_dim=sec_dim, + particle_type_cfg=particle_type_cfg, ) + assert s2.particle_type_cfg.get("target", "physical") == target return s1.eval(), s2.eval() -def _zero_secondaries_loader(B=4, n_batches=2): - """A val_loader whose every batch has n_sec=0 (real side) — matches the - (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx) tuple shape - StreamingStepsDataset yields.""" +def _loader(B: int = 4, n_batches: int = 2, n_sec_value: int = 0, n_classes: int = 8): + """A val_loader matching StreamingStepsDataset's 7-tuple batch shape: + (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx, sec_type_idx).""" batches = [] for _ in range(n_batches): cond_cont = torch.randn(B, COND_DIM) cond_cat = torch.zeros(B, 2, dtype=torch.long) x1 = torch.randn(B, X_DIM) - n_sec = torch.zeros(B, dtype=torch.long) - sec_cont = torch.zeros(B, K_MAX, SEC_SLOT_DIM) + n_sec = torch.full((B,), n_sec_value, dtype=torch.long) + sec_cont = torch.randn(B, _K_MAX, SEC_SLOT_DIM) proc_idx = torch.zeros(B, dtype=torch.long) - batches.append((cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx)) + sec_type_idx = torch.randint(0, n_classes, (B, _K_MAX), dtype=torch.long) + batches.append( + (cond_cont, cond_cat, x1, n_sec, sec_cont, proc_idx, sec_type_idx) + ) return batches @@ -52,22 +67,56 @@ def test_validate_marginals_all_zero_secondaries_returns_nan_phys_kl(monkeypatch """If n_sec_pred collapses to 0 across the whole validated set (realistic during early/unstable training), phys_kl must degrade to NaN instead of crashing on the empty-array .min()/.max() reduction inside - _histogram_kl -- a regression the old species/bincount code this - replaced explicitly guarded against.""" + _histogram_kl.""" s1, s2 = _tiny_models() - loader = _zero_secondaries_loader() + loader = _loader(n_sec_value=0) - # Force the Stage-1 n_sec head's prediction to 0 for every sample too, so - # the generated side's valid-slot mask is also empty (real side is - # already all n_sec=0 by construction of the fake loader above). - def _fake_sample_flow(model, cond_cont, cond_cat, **kw): - B = cond_cont.size(0) - return torch.randn(B, X_DIM), torch.zeros(B, dtype=torch.long) + def _fake_resolve_n_sec( + stage1_model, sec_decoder, cond_cont, cond_cat, stage1_out, n_sec_pred + ): + return torch.zeros(cond_cont.size(0), dtype=torch.long) - monkeypatch.setattr("giant.validate.sample_flow", _fake_sample_flow) + monkeypatch.setattr("giant.validate.resolve_n_sec", _fake_resolve_n_sec) - result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2) + result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=2, steps=2) assert np.asarray(result["phys_real"]).shape == (0, 2) assert np.asarray(result["phys_generated"]).shape == (0, 2) assert np.isnan(np.asarray(result["phys_kl"])).all() + + +def test_validate_marginals_physical_target_shapes(): + s1, s2 = _tiny_models() + loader = _loader(n_sec_value=2) + + result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2) + + assert np.asarray(result["real"]).shape == (4, X_DIM) + assert np.asarray(result["generated"]).shape == (4, X_DIM) + assert np.asarray(result["kl_divergence"]).shape == (X_DIM,) + assert "phys_real" in result and "phys_generated" in result and "phys_kl" in result + assert "type_class_real" not in result + + +def test_validate_marginals_onehot_type_class_marginal(): + particle_type_cfg = {"target": "onehot"} + s1, s2 = _tiny_models(particle_type_cfg) + loader = _loader(n_sec_value=2, n_classes=s2.type_dim) + + result = validate_marginals(s1, loader, sec_decoder=s2, n_batches=1, steps=2) + + assert "phys_real" not in result + # Real/generated valid-slot counts need not agree (real: ground-truth + # n_sec=2 always; generated: the untrained n_sec_head's own prediction). + assert np.asarray(result["type_class_real"]).ndim == 1 + assert np.asarray(result["type_class_gen"]).ndim == 1 + assert np.asarray(result["type_class_real"]).shape[0] > 0 + + +def test_validate_marginals_without_sec_decoder_returns_stage1_only(): + s1, _ = _tiny_models() + loader = _loader(n_sec_value=0) + + result = validate_marginals(s1, loader, n_batches=1, steps=2) + + assert set(result) == {"real", "generated", "kl_divergence"} diff --git a/tests/test_wgan.py b/tests/test_wgan.py index 07840b7..851527b 100644 --- a/tests/test_wgan.py +++ b/tests/test_wgan.py @@ -120,7 +120,7 @@ def test_sample_wgan_shape(): cond_cont, cond_cat = _cond(B) sample, n_sec = sample_wgan(model, cond_cont, cond_cat) assert sample.shape == (B, X_DIM) - assert n_sec.shape == (B,) + assert n_sec is not None and n_sec.shape == (B,) # --- Stage-2 generator/critic ---