diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 50bd1c8..e41041e 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -339,12 +339,22 @@ def sorted_membership(values: np.ndarray, sorted_arr: np.ndarray) -> np.ndarray: return sorted_arr[idx] == values -def _vectorized_map_lookup(values: np.ndarray, mapping: dict) -> np.ndarray: +def _vectorized_map_lookup( + values: np.ndarray, mapping: dict, strict: bool = True +) -> np.ndarray: """Vectorized equivalent of `np.array([mapping[v] for v in values], dtype=np.int64)`. Replaces a per-element Python dict lookup with one `searchsorted` call. Raises `KeyError` if any value in `values` isn't a key of `mapping`, - matching the dict-comprehension it replaces (never silently misassigns). + matching the dict-comprehension it replaces (never silently misassigns) + — unless `strict=False`, in which case unmapped values get a dummy index + of 0 instead. Only pass `strict=False` where the caller has independently + verified the resulting index is never actually read (e.g. + `build_cond_features` under `conditioning="physical"`, where + `ConditionEncoder` ignores `cond_cat` entirely); it exists so a rollout + can be seeded with a species/material outside the training vocab without + a spurious `KeyError`, which is the entire point of physical-property + conditioning. """ keys = np.asarray(list(mapping.keys())) vals = np.asarray(list(mapping.values()), dtype=np.int64) @@ -355,6 +365,10 @@ def _vectorized_map_lookup(values: np.ndarray, mapping: dict) -> np.ndarray: pos = np.clip(pos, 0, len(keys_sorted) - 1) found = keys_sorted[pos] == values if not found.all(): + if not strict: + out = np.zeros(values.shape, dtype=np.int64) + out[found] = vals_sorted[pos[found]] + return out missing = np.unique(values[~found]) raise KeyError(f"value(s) not in mapping: {missing[:10].tolist()}") return vals_sorted[pos] @@ -693,8 +707,15 @@ def build_cond_features( [cond_cont, _physical_cond_columns(data, conditioning)] ).astype(np.float32) - pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map) - mat_idx = _vectorized_map_lookup(data["material"], mat_map) + # In "physical" mode cond_cat is only a reporting/router convenience — + # ConditionEncoder never reads it (giant/model/network.py) — so a + # species/material outside the training vocab (the whole point of + # physical-property conditioning) gets a dummy index instead of raising. + # In "embedding" mode cond_cat IS the conditioning signal, so an unmapped + # value must still raise loudly rather than silently misassign. + strict = conditioning == "embedding" + pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=strict) + mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=strict) cond_cat = np.column_stack([pdg_idx, mat_idx]) if cond_normalizer is not None: diff --git a/giant/model/network.py b/giant/model/network.py index d04e4c4..9964d92 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -1228,7 +1228,40 @@ def _parse_composed_axes(router_cfg: dict) -> list[dict]: return [axes[i] for i in range(len(axes))] -def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> Router: +# Router types that read cond_cat's pdg index through their own +# nn.Embedding(pdg_vocab, ...), regardless of the trunk's `conditioning` +# mode — see _check_router_conditioning_compat. +_VOCAB_SCOPED_ROUTER_TYPES = ("pdg", "process") + + +def _check_router_conditioning_compat(router_types: list[str], conditioning: str) -> None: + """Reject a router axis that reintroduces a training-vocab PDG lookup + under `conditioning="physical"`. + + `PdgRouter`/`ProcessRouter` always build their own dataset-scoped + `nn.Embedding(pdg_vocab, ...)` (network.py's PdgRouter/ProcessRouter), + independent of `ConditionEncoder`'s `conditioning` mode. Pairing either + with `conditioning="physical"` would silently reintroduce a + training-menu-scoped lookup at the routing layer — defeating the entire + point of physical-property conditioning, which is to generalize to a + species/material outside that menu (see giant/rollout.py's + `build_cond_features(strict=...)` gate for the same concern on the + trunk side). Raised loudly at model-build time rather than left to + surface as a confusing rollout/generalization-benchmark result. + """ + bad = sorted(set(router_types) & set(_VOCAB_SCOPED_ROUTER_TYPES)) + if bad and conditioning == "physical": + raise ValueError( + f"router type(s) {bad} always use a training-vocab PDG embedding, " + "which is incompatible with conditioning='physical' (whose whole " + "point is generalizing beyond that vocab) — pick a different " + "router type (e.g. 'energy') or use conditioning='embedding'." + ) + + +def _build_router_from_cfg( + router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding" +) -> Router: """Resolve one `model.router` config into a `Router`, single-axis or composed. `router_cfg["type"] == "composed"` reads `axis{i}_{field}` flat keys @@ -1243,9 +1276,12 @@ def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int) -> """ shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab) if router_cfg["type"] == "composed": - router = build_composed_router(_parse_composed_axes(router_cfg), **shared_vocab) + axes = _parse_composed_axes(router_cfg) + _check_router_conditioning_compat([a["type"] for a in axes], conditioning) + router = build_composed_router(axes, **shared_vocab) router.gumbel = bool(router_cfg.get("gumbel", False)) return router + _check_router_conditioning_compat([router_cfg["type"]], conditioning) router_kwargs = { k: v for k, v in router_cfg.items() if k not in ("enabled", "type", "n_experts") } @@ -1303,13 +1339,18 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]: dropout=model_config.get("dropout", 0.1), conditioning=model_config.get("conditioning", "embedding"), ) + conditioning = shared["conditioning"] stage1 = RoutedDenoisingMLP( - router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab), + router=_build_router_from_cfg( + router_cfg, pdg_vocab, mat_vocab, conditioning + ), k_max=model_config.get("k_max", K_MAX), **shared, ) sec_decoder = RoutedSecondaryDecoder( - router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab), + router=_build_router_from_cfg( + router_cfg, pdg_vocab, mat_vocab, conditioning + ), **shared, ) return stage1, sec_decoder diff --git a/giant/rollout.py b/giant/rollout.py index 5fde841..8270a5e 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -407,7 +407,16 @@ def _step_chunk( tr["_material"] = material tr["_layer_id"] = layer_id - known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool) + if conditioning == "physical": + # Under physical-property conditioning, mass/charge (already resolved + # on every track — see the cond_dict comment below) drive the model, + # not a training-vocab PDG embedding — build_cond_features passes + # strict=False for exactly this mode, so an out-of-vocab species no + # longer raises. Terminating on it here would defeat the entire + # point of physical conditioning: generalizing to a held-out species. + known_pdg = np.ones(n, dtype=bool) + else: + known_pdg = np.array([int(p) in pdg_map for p in tr["pdg"]], dtype=bool) # --- Pre-step termination gates (in priority order; each track picks one) --- stop = np.zeros(n, dtype=bool) diff --git a/tests/test_rollout.py b/tests/test_rollout.py index ec96b17..5227109 100644 --- a/tests/test_rollout.py +++ b/tests/test_rollout.py @@ -119,6 +119,21 @@ def test_rollout_physical_conditioning_end_to_end(fake_material_props): assert set(np.unique(rec["pdg"]).tolist()) <= set(PDG_MAP.keys()) +def test_rollout_physical_conditioning_generalizes_to_out_of_vocab_pdg( + fake_material_props, +): + """A real, giant.particles-resolvable species outside the training PDG + vocab (muon, 13) must run through physical-property conditioning rather + than terminate via TERM_UNKNOWN_PDG — that generalization is the entire + point of "physical" mode (see build_cond_features(strict=...)).""" + seeds = _seeds(6) + seeds["pdg"] = np.full(6, 13, dtype=np.int64) + assert 13 not in PDG_MAP + rec = _run(seeds=seeds, conditioning="physical") + assert len(rec["event_id"]) > 0 + assert TERM_UNKNOWN_PDG not in set(rec["termination_reason"].tolist()) + + def test_seed_frontier_track_ids(): seeds = _seeds(3) fr, counts = make_seed_frontier(**seeds) diff --git a/tests/test_router.py b/tests/test_router.py index 34c669e..a8e9174 100644 --- a/tests/test_router.py +++ b/tests/test_router.py @@ -1,5 +1,6 @@ """Tests for the mixture-of-experts routing prototype (giant/model/network.py).""" +import pytest import torch from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM @@ -484,6 +485,49 @@ def test_build_models_routed_with_pdg_router(): assert stage1.router.pdg_emb.num_embeddings == 4 +def test_build_models_rejects_pdg_router_with_physical_conditioning(): + """conditioning="physical" is meant to generalize beyond the training PDG + vocab; PdgRouter always uses a training-vocab nn.Embedding regardless of + conditioning, so the combination must raise rather than silently building + a model that can't actually generalize the way it claims to.""" + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + conditioning="physical", + router={"enabled": True, "type": "pdg", "n_experts": 3}, + ) + with pytest.raises(ValueError, match="physical"): + build_models(model_config) + + +def test_build_models_rejects_composed_router_with_pdg_axis_and_physical_conditioning(): + model_config = dict( + pdg_vocab=4, + mat_vocab=2, + emb_dim=16, + dropout=0.1, + k_max=K_MAX, + expert_hidden_dim=16, + expert_n_blocks=2, + conditioning="physical", + router={ + "enabled": True, + "type": "composed", + "axis0_type": "energy", + "axis0_n_experts": 2, + "axis1_type": "pdg", + "axis1_n_experts": 3, + }, + ) + with pytest.raises(ValueError, match="physical"): + build_models(model_config) + + # ── ProcessRouter ──────────────────────────────────────────────────────────── diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 899e976..dcda475 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -574,6 +574,51 @@ def test_vectorized_map_lookup_raises_keyerror_on_missing_value(): _vectorized_map_lookup(values, mapping) +def test_vectorized_map_lookup_strict_false_dummy_indexes_unmapped_values(): + """strict=False must leave found values untouched and only dummy-index + (0) the unmapped ones — never raise, and never disturb a value that IS + in the mapping (e.g. one that happens to map to a nonzero index).""" + mapping = {1: 5, 2: 7} + values = np.array([1, 99, 2, 100]) + result = _vectorized_map_lookup(values, mapping, strict=False) + np.testing.assert_array_equal(result, [5, 0, 7, 0]) + + +def test_build_cond_features_physical_mode_tolerates_out_of_vocab_pdg_and_material(): + """conditioning="physical" must not KeyError on a pdg/material outside + the training-dataset vocab (mat_map/pdg_map) — that's the entire point + of the mode (see giant.rollout's known_pdg gate for the paired fix). + "embedding" mode must still raise, since cond_cat IS the conditioning + signal there. Note this is specifically about the dataset-scoped + vocab index, not giant.materials' physical-properties table — a + material must still be a real, known Geant4 material (e.g. "G4_Pb", + just not one *this* mat_map happened to include) for "physical" mode + to derive its Z_eff/A_eff/density/X0/λ_int; a genuinely unknown + material name correctly still raises via giant.materials, same as the + documented G4_LYSO precedent — that's a separate, intentional guard.""" + pdg_map = {11: 0, 22: 1} + mat_map = {"G4_AIR": 0} + data = { + "pre_pos": np.zeros((1, 3), dtype=np.float32), + "pre_E": np.array([10.0], dtype=np.float32), + "pre_dir": np.array([[0.0, 0.0, 1.0]], dtype=np.float32), + "layer_id": np.array([0], dtype=np.int32), + "pdg": np.array([13], dtype=np.int64), # not in pdg_map + "material": np.array(["G4_Pb"], dtype=object), # not in mat_map + "mass": np.array([105.7], dtype=np.float32), + "charge": np.array([-1.0], dtype=np.float32), + } + + cond_cont, cond_cat = build_cond_features( + data, pdg_map, mat_map, conditioning="physical" + ) + assert cond_cont.shape[-1] == COND_DIM + np.testing.assert_array_equal(cond_cat, [[0, 0]]) # dummy indices, no raise + + with pytest.raises(KeyError): + build_cond_features(data, pdg_map, mat_map, conditioning="embedding") + + # ── _WelfordAccumulator ──────────────────────────────────────────────────────