diff --git a/giant/cond_layout.py b/giant/cond_layout.py new file mode 100644 index 0000000..cf4a0d4 --- /dev/null +++ b/giant/cond_layout.py @@ -0,0 +1,103 @@ +"""Single source of truth for the conditioning arrays' column layout (gitea #37). + +`cond_cont` and `cond_cat` are built in `giant.data.transforms` and consumed in +`giant.model.encoders` / `giant.model.routers`. Their column order used to be +written down independently on each side, kept in sync only by parallel comments +— so getting it wrong produced silently mis-indexed columns rather than an +exception, and adding a conditioning axis meant a coordinated multi-file edit. + +`CondLayout` owns that order. Both sides construct one from the same +`conditioning.particle.type` / `conditioning.material.type` pair and read named +slices off it, so the layout is stated exactly once. This module depends only on +`giant.constants`, so both the data and model packages can import it. +""" + +from dataclasses import dataclass +from typing import ClassVar + +from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM + +# The three per-axis conditioning modes. Mirrors giant.config.Conditioning, +# which this module deliberately does not import (giant.config pulls in the +# whole model package). +AXIS_TYPES = ("physical", "embedding", "onehot") + + +@dataclass(frozen=True) +class CondLayout: + """Column layout of `cond_cont`/`cond_cat` for one (particle, material) mode pair. + + `cond_cont` is unconditionally `COND_DIM` wide regardless of mode: the base + block, then the particle physical block, then the material physical block. + An axis that isn't `"physical"` gets its block zero-filled and never reads + it (see `giant.data.transforms._physical_cond_columns`), so the widths are + mode-independent and only the *meaning* of a block changes. + + `cond_cat` is 2 to 4 wide. Columns `PDG_COL`/`MAT_COL` are always the dense + training-vocab index; an axis in `"onehot"` mode appends one more column + holding its top-N-plus-other class index, particle before material. + """ + + particle_type: str + material_type: str + + # cond_cat's dense-vocab columns, present in every mode. Under + # "physical"/"onehot" they are a reporting/router convenience the + # ConditionEncoder never reads; under "embedding" they are the signal. + PDG_COL: ClassVar[int] = 0 + MAT_COL: ClassVar[int] = 1 + + def __post_init__(self) -> None: + if self.particle_type not in AXIS_TYPES: + raise ValueError(f"unknown conditioning.particle.type {self.particle_type!r}") + if self.material_type not in AXIS_TYPES: + raise ValueError(f"unknown conditioning.material.type {self.material_type!r}") + + @classmethod + def from_types(cls, particle_type: str, material_type: str) -> "CondLayout": + """Named constructor — the entry point both sides use.""" + return cls(particle_type=particle_type, material_type=material_type) + + # --- cond_cont --------------------------------------------------------- + + @property + def base(self) -> slice: + """pre_pos(3), log(pre_E)(1), pre_dir(3), layer_id(1).""" + return slice(0, COND_DIM_BASE) + + @property + def particle_phys(self) -> slice: + """log(mass), charge — see `giant.particles`.""" + return slice(COND_DIM_BASE, COND_DIM_BASE + PARTICLE_PHYS_DIM) + + @property + def material_phys(self) -> slice: + """Z_eff, A_eff, log(density), log(X0), log(lambda_int) — see `giant.materials`.""" + start = COND_DIM_BASE + PARTICLE_PHYS_DIM + return slice(start, start + MATERIAL_PHYS_DIM) + + @property + def cont_dim(self) -> int: + return COND_DIM + + # --- cond_cat ---------------------------------------------------------- + + @property + def particle_topn_col(self) -> int | None: + """Column of the particle top-N class index, or `None` if not `"onehot"`.""" + return self.MAT_COL + 1 if self.particle_type == "onehot" else None + + @property + def material_topn_col(self) -> int | None: + """Column of the material top-N class index, or `None` if not `"onehot"`. + + Comes after the particle top-N column when both axes are `"onehot"`. + """ + if self.material_type != "onehot": + return None + return self.MAT_COL + (2 if self.particle_type == "onehot" else 1) + + @property + def cat_dim(self) -> int: + """Total `cond_cat` width: 2, plus one column per `"onehot"` axis.""" + return self.MAT_COL + 1 + (self.particle_type == "onehot") + (self.material_type == "onehot") diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 5f36fcf..e436211 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -3,6 +3,7 @@ from typing import NamedTuple import numpy as np +from giant.cond_layout import CondLayout from giant.constants import K_MAX _EPS = 1e-8 @@ -693,18 +694,13 @@ def decode_secondaries( return sec_E, sec_dir_world, sec_mass, sec_charge, sec_valid -def _physical_cond_columns( - data: dict[str, np.ndarray], - particle_conditioning: str, - material_conditioning: str, -) -> np.ndarray: +def _physical_cond_columns(data: dict[str, np.ndarray], layout: CondLayout) -> np.ndarray: """(N, PARTICLE_PHYS_DIM + MATERIAL_PHYS_DIM) physical conditioning columns. The particle and material blocks are gated independently and may mix freely — e.g. material `physical` with particle `embedding` — so e.g. - `particle_conditioning="embedding"` + `material_conditioning="physical"` - zero-fills only the particle columns and computes the material ones for - real. + `particle_type="embedding"` + `material_type="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 @@ -721,7 +717,7 @@ def _physical_cond_columns( n = len(next(iter(data.values()))) - if particle_conditioning == "physical": + if layout.particle_type == "physical": from giant.particles import particle_phys_array if "mass" in data and "charge" in data: @@ -730,12 +726,10 @@ def _physical_cond_columns( 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: - raise ValueError(f"unknown conditioning.particle.type {particle_conditioning!r}") + particle_cols = np.zeros((n, PARTICLE_PHYS_DIM), dtype=np.float32) - if material_conditioning == "physical": + if layout.material_type == "physical": from giant.materials import material_properties_array z_eff, a_eff, density, x0, lambda_int = material_properties_array(data["material"]).T @@ -748,14 +742,66 @@ def _physical_cond_columns( 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}") + material_cols = np.zeros((n, MATERIAL_PHYS_DIM), dtype=np.float32) return np.column_stack([particle_cols, material_cols]).astype(np.float32) +def _build_cond_arrays( + data: dict[str, np.ndarray], + pdg_map: dict[int, int], + mat_map: dict[str, int], + layout: CondLayout, + pdg_topn_map: dict[int, int] | None, + mat_topn_map: dict[str, int] | None, +) -> tuple[np.ndarray, np.ndarray]: + """The un-normalized `(cond_cont, cond_cat)` pair, in `layout`'s column order. + + Both `build_cond_features` and `build_features` go through here, so the + column order — and everything that depends on it — is stated once. See + `giant.cond_layout.CondLayout` for the layout itself. + """ + cond_cont = np.column_stack( + [ + data["pre_pos"], + log_transform(data["pre_E"]), + data["pre_dir"], + data["layer_id"].astype(np.float32), + ] + ).astype(np.float32) # (N, COND_DIM_BASE=8) + cond_cont = np.column_stack([cond_cont, _physical_cond_columns(data, layout)]).astype( + np.float32 + ) # (N, COND_DIM=15) + + # In "physical" mode cond_cat's first two columns are only a + # reporting/router convenience — ConditionEncoder never reads them + # (giant/model/encoders.py) — so a species/material outside the training + # vocab (the whole point of physical-property conditioning) gets a dummy + # index instead of raising. In "embedding" mode those columns ARE the + # conditioning signal, so an unmapped value must still raise loudly + # rather than silently misassign. In "onehot" mode they again go unread + # (the topN columns below are the real signal), so they're as permissive + # as "physical". Each axis's strictness is independent. + pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=layout.particle_type == "embedding") + mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=layout.material_type == "embedding") + # Which extra columns exist is the layout's call, not "did the caller + # happen to pass a map" — that's what used to let the producer and + # ConditionEncoder disagree. A map for a non-"onehot" axis is unused. + cat_cols = [pdg_idx, mat_idx] + if layout.particle_topn_col is not None: + if pdg_topn_map is None: + raise ValueError("conditioning.particle.type='onehot' needs pdg_topn_map") + cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map)) + if layout.material_topn_col is not None: + if mat_topn_map is None: + raise ValueError("conditioning.material.type='onehot' needs mat_topn_map") + cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map)) + cond_cat = np.column_stack(cat_cols) # (N, layout.cat_dim) + + return cond_cont, cond_cat + + def build_cond_features( data: dict[str, np.ndarray], pdg_map: dict[int, int], @@ -773,49 +819,16 @@ def build_cond_features( `material_conditioning="physical"` is a valid mix. `pdg_topn_map`/`mat_topn_map` (a top-N-plus-other `class_map`, see - `giant.data.loader.build_topn_map_from_files`) append extra `cond_cat` - columns read by `ConditionEncoder`'s `"onehot"` mode: pdg topN index at - column 2 (iff `pdg_topn_map` given), material topN index at column 3 - (iff `mat_topn_map` given, after column 2 if both are). Only ever given when - the corresponding axis is `"onehot"`; `cond_cat` stays `(N, 2)` otherwise. + `giant.data.loader.build_topn_map_from_files`) supply the extra `cond_cat` + columns read by `ConditionEncoder`'s `"onehot"` mode, and are required + whenever the corresponding axis is `"onehot"`. See + `giant.cond_layout.CondLayout` for which columns exist where. """ - cond_cont = np.column_stack( - [ - data["pre_pos"], - log_transform(data["pre_E"]), - data["pre_dir"], - data["layer_id"].astype(np.float32), - ] - ).astype(np.float32) - cond_cont = np.column_stack( - [ - 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 - # reporting/router convenience — ConditionEncoder never reads them - # (giant/model/network.py) — so a species/material outside the training - # vocab (the whole point of physical-property conditioning) gets a dummy - # index instead of raising. In "embedding" mode those columns ARE the - # conditioning signal, so an unmapped value must still raise loudly - # rather than silently misassign. In "onehot" mode they again go unread - # (the topN columns below are the real signal), so they're as permissive - # as "physical". Each axis's strictness is independent. - pdg_strict = particle_conditioning == "embedding" - mat_strict = material_conditioning == "embedding" - pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map, strict=pdg_strict) - mat_idx = _vectorized_map_lookup(data["material"], mat_map, strict=mat_strict) - cat_cols = [pdg_idx, mat_idx] - if pdg_topn_map is not None: - cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map)) - if mat_topn_map is not None: - cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map)) - cond_cat = np.column_stack(cat_cols) + layout = CondLayout.from_types(particle_conditioning, material_conditioning) + cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map) if cond_normalizer is not None: - cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, particle_conditioning, material_conditioning) + cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, layout) return cond_cont, cond_cat @@ -823,8 +836,7 @@ def build_cond_features( def _cond_normalizer_transform( cond_cont: np.ndarray, cond_normalizer: "Normalizer", - particle_conditioning: str, - material_conditioning: str, + layout: CondLayout, ) -> np.ndarray: """Apply ``cond_normalizer``, padding a legacy narrower normalizer if needed. @@ -832,7 +844,7 @@ def _cond_normalizer_transform( 8->15, ``giant/constants.py``) saved a ``COND_DIM_BASE``-wide (8) cond normalizer, fit before ``build_cond_features`` grew the extra physical columns. When NEITHER axis is "physical" those columns are never read by - ``ConditionEncoder`` (``giant/model/network.py``), so padding the missing + ``ConditionEncoder`` (``giant/model/encoders.py``), so padding the missing entries with mean=0/std=1 is a safe no-op that keeps such checkpoints usable under the current, always-``COND_DIM``-wide contract. If EITHER axis is "physical" its columns are load-bearing, so a mismatch there is a @@ -843,14 +855,14 @@ def _cond_normalizer_transform( width = cond_cont.shape[-1] if mean.shape[-1] < width: physical_load_bearing = "physical" in ( - particle_conditioning, - material_conditioning, + layout.particle_type, + layout.material_type, ) if physical_load_bearing: raise ValueError( f"cond normalizer has {mean.shape[-1]} columns, expected " - f"{width}, and particle_conditioning={particle_conditioning!r}/" - f"material_conditioning={material_conditioning!r} reads the " + f"{width}, and particle_conditioning={layout.particle_type!r}/" + f"material_conditioning={layout.material_type!r} reads the " "physical columns directly — this checkpoint predates " "physical-property conditioning and can't be safely padded; " "retrain it under the current code." @@ -928,7 +940,7 @@ def build_features( instead) for callers (normalizer fitting) that only read `sec_cont[:, :, 4:6]` and would otherwise discard that work. - pdg_topn_map/mat_topn_map: appended `cond_cat` columns for + pdg_topn_map/mat_topn_map: source of the extra `cond_cat` columns for `ConditionEncoder`'s `"onehot"` mode — see `build_cond_features`. sec_type_class_map: the map `sec_type_idx` is looked up against — a @@ -959,29 +971,8 @@ def build_features( ).astype(np.float32) # (N, 9) # Phase 2: conditioning drops n_sec and log(e_sec) - cond_cont = np.column_stack( - [ - data["pre_pos"], - log_transform(data["pre_E"]), - data["pre_dir"], - data["layer_id"].astype(np.float32), - ] - ).astype(np.float32) # (N, COND_DIM_BASE=8) - cond_cont = np.column_stack( - [ - cond_cont, - _physical_cond_columns(data, particle_conditioning, material_conditioning), - ] - ).astype(np.float32) # (N, COND_DIM=15) - - pdg_idx = _vectorized_map_lookup(data["pdg"], pdg_map) - mat_idx = _vectorized_map_lookup(data["material"], mat_map) - cat_cols = [pdg_idx, mat_idx] - if pdg_topn_map is not None: - cat_cols.append(_vectorized_map_lookup(data["pdg"], pdg_topn_map)) - if mat_topn_map is not None: - cat_cols.append(_vectorized_map_lookup(data["material"], mat_topn_map)) - cond_cat = np.column_stack(cat_cols) # (N, 2/3/4) + layout = CondLayout.from_types(particle_conditioning, material_conditioning) + cond_cont, cond_cat = _build_cond_arrays(data, pdg_map, mat_map, layout, pdg_topn_map, mat_topn_map) n_sec_raw = data["n_sec"].astype(np.int64) # (N,) unclamped, for the valid-slot mask @@ -1048,7 +1039,7 @@ def build_features( target_normalizer = Normalizer().fit(target_s1) if cond_normalizer is not None: - cond_cont = cond_normalizer.transform(cond_cont) + cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, layout) if target_normalizer is not None: target_s1 = target_normalizer.transform(target_s1) if sec_phys_normalizer is not None: diff --git a/giant/model/encoders.py b/giant/model/encoders.py index 384903b..5aeee3f 100644 --- a/giant/model/encoders.py +++ b/giant/model/encoders.py @@ -5,32 +5,11 @@ import torch import torch.nn as nn import torch.nn.functional as F +from giant.cond_layout import CondLayout from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM from giant.model.layers import _make_axis_mlp -def cat_col_layout(particle_type: str, material_type: str) -> tuple[int | None, int | None]: - """`cond_cat` column indices for each axis's top-N-onehot index, or - `None` if that axis isn't `"onehot"`. - - Columns 0/1 are always the dense pdg/material vocab index. The particle - top-N column (if any) comes next, then the material top-N column (if - any) — `giant.data.transforms.build_cond_features`/`build_features` - append columns in this same order, so the two sides must never drift - apart. - """ - col = 2 - particle_col = None - if particle_type == "onehot": - particle_col = col - col += 1 - material_col = None - if material_type == "onehot": - material_col = col - col += 1 - return particle_col, material_col - - class ConditionEncoder(nn.Module): """Fuses continuous conditioning with particle/material identity. @@ -41,13 +20,17 @@ class ConditionEncoder(nn.Module): - "embedding": a learned `nn.Embedding` lookup, indexed by `cond_cat`'s dense training-vocab index. Memorizes the training menu. - "physical": an `n_layers`-deep MLP over the axis's raw physical - properties (already present in `cond_cont[:, COND_DIM_BASE:]` — see + properties (already present in `cond_cont`'s physical block — see giant.data.transforms.build_features), computable for any PDG code / material name rather than only ones seen in training. - "onehot": a fixed, unlearned one-hot vector over a top-N-plus-other class map (`giant.data.loader.build_topn_map_from_files`/ `build_pdg_topn_map_from_files`), read from `cond_cat`'s extra - top-N-index column(s) — see `_cat_col_layout`. + top-N-index column(s). + + Every column index/slice comes from `self.layout` + (`giant.cond_layout.CondLayout`), the same object the feature builders + lay the arrays out with, so the two sides cannot drift apart. """ def __init__( @@ -62,7 +45,8 @@ class ConditionEncoder(nn.Module): super().__init__() self.particle_cfg = dict(particle_cfg) self.material_cfg = dict(material_cfg) - self._particle_topn_col, self._material_topn_col = cat_col_layout(particle_cfg["type"], material_cfg["type"]) + # Also validates both axis types — an unknown one raises here. + self.layout = CondLayout.from_types(particle_cfg["type"], material_cfg["type"]) p_type = particle_cfg["type"] p_emb_dim = particle_cfg["emb_dim"] @@ -70,8 +54,6 @@ class ConditionEncoder(nn.Module): self.pdg_emb = nn.Embedding(pdg_vocab, p_emb_dim) elif p_type == "physical": self.particle_mlp = _make_axis_mlp(PARTICLE_PHYS_DIM, p_emb_dim, particle_cfg.get("n_layers", 1)) - elif p_type != "onehot": - raise ValueError(f"unknown conditioning.particle.type {p_type!r}") m_type = material_cfg["type"] m_emb_dim = material_cfg["emb_dim"] @@ -79,8 +61,6 @@ class ConditionEncoder(nn.Module): self.mat_emb = nn.Embedding(mat_vocab, m_emb_dim) elif m_type == "physical": self.material_mlp = _make_axis_mlp(MATERIAL_PHYS_DIM, m_emb_dim, material_cfg.get("n_layers", 1)) - elif m_type != "onehot": - raise ValueError(f"unknown conditioning.material.type {m_type!r}") in_dim = COND_DIM_BASE + p_emb_dim + m_emb_dim self.mlp = nn.Sequential( @@ -92,31 +72,29 @@ class ConditionEncoder(nn.Module): def _particle_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): p_type = self.particle_cfg["type"] if p_type == "embedding": - return self.pdg_emb(cond_cat[:, 0]) + return self.pdg_emb(cond_cat[:, self.layout.PDG_COL]) if p_type == "physical": - particle_phys = cond_cont[:, COND_DIM_BASE : COND_DIM_BASE + PARTICLE_PHYS_DIM] - return self.particle_mlp(particle_phys) - assert self._particle_topn_col is not None + return self.particle_mlp(cond_cont[:, self.layout.particle_phys]) + assert self.layout.particle_topn_col is not None return F.one_hot( - cond_cat[:, self._particle_topn_col], + cond_cat[:, self.layout.particle_topn_col], num_classes=self.particle_cfg["emb_dim"], ).float() def _material_embed(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor): m_type = self.material_cfg["type"] if m_type == "embedding": - return self.mat_emb(cond_cat[:, 1]) + return self.mat_emb(cond_cat[:, self.layout.MAT_COL]) if m_type == "physical": - material_phys = cond_cont[:, COND_DIM_BASE + PARTICLE_PHYS_DIM :] - return self.material_mlp(material_phys) - assert self._material_topn_col is not None + return self.material_mlp(cond_cont[:, self.layout.material_phys]) + assert self.layout.material_topn_col is not None return F.one_hot( - cond_cat[:, self._material_topn_col], + cond_cat[:, self.layout.material_topn_col], num_classes=self.material_cfg["emb_dim"], ).float() def forward(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: pdg_e = self._particle_embed(cond_cont, cond_cat) mat_e = self._material_embed(cond_cont, cond_cat) - x = torch.cat([cond_cont[:, :COND_DIM_BASE], pdg_e, mat_e], dim=-1) + x = torch.cat([cond_cont[:, self.layout.base], pdg_e, mat_e], dim=-1) return self.mlp(x) diff --git a/giant/model/network.py b/giant/model/network.py index 65dd3b9..547b39b 100644 --- a/giant/model/network.py +++ b/giant/model/network.py @@ -9,7 +9,7 @@ import X` call site keeps working unchanged. from giant.model._legacy import _migrate_legacy_model_config, migrate_legacy_state_dict from giant.model.builders import build_critics, build_models -from giant.model.encoders import ConditionEncoder, cat_col_layout +from giant.model.encoders import ConditionEncoder from giant.model.history import ( HISTORY_REGISTRY, AttentionHistory, @@ -122,7 +122,6 @@ __all__ = [ "build_objective", "build_router", "build_trunk", - "cat_col_layout", "migrate_legacy_state_dict", "register_block", "register_history", diff --git a/giant/model/routers.py b/giant/model/routers.py index c5c69db..12afa36 100644 --- a/giant/model/routers.py +++ b/giant/model/routers.py @@ -11,6 +11,7 @@ import torch import torch.nn as nn import torch.nn.functional as F +from giant.cond_layout import CondLayout from giant.constants import COND_DIM # --------------------------------------------------------------------------- @@ -211,7 +212,7 @@ class PdgRouter(Router): self.register_buffer("centers", centers) def gate(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - e = self.pdg_emb(cond_cat[:, 0]) # (B, emb_dim) + e = self.pdg_emb(cond_cat[:, CondLayout.PDG_COL]) # (B, emb_dim) d2 = ((e.unsqueeze(1) - self.centers.unsqueeze(0)) ** 2).sum(-1) # (B, n_experts) return torch.softmax(-d2 / self.temperature, dim=-1) @@ -243,8 +244,8 @@ class ProcessRouter(Router): ) def logits(self, cond_cont: torch.Tensor, cond_cat: torch.Tensor) -> torch.Tensor: - pdg_e = self.pdg_emb(cond_cat[:, 0]) - mat_e = self.mat_emb(cond_cat[:, 1]) + pdg_e = self.pdg_emb(cond_cat[:, CondLayout.PDG_COL]) + mat_e = self.mat_emb(cond_cat[:, CondLayout.MAT_COL]) h = torch.cat([cond_cont, pdg_e, mat_e], dim=-1) return self.classifier(h) diff --git a/giant/pipeline.py b/giant/pipeline.py index 7145062..3b567a7 100644 --- a/giant/pipeline.py +++ b/giant/pipeline.py @@ -271,6 +271,13 @@ def run_setup_stage( particle_conditioning=particle_conditioning, material_conditioning=material_conditioning, sec_phys_only=True, + # This pass reads only cond_cont/sec_cont, never cond_cat — + # but cond_cat's width is the conditioning modes' call + # (giant.cond_layout.CondLayout), so an "onehot" axis still + # has to be handed its map rather than silently yielding a + # narrower array. + pdg_topn_map=pdg_topn_map.class_map if pdg_topn_map is not None else None, + mat_topn_map=mat_topn_map.class_map if mat_topn_map is not None else None, k_max=k_max, ) cond_cont = feats.cond_cont diff --git a/tests/test_cond_layout.py b/tests/test_cond_layout.py new file mode 100644 index 0000000..f3bd1fb --- /dev/null +++ b/tests/test_cond_layout.py @@ -0,0 +1,87 @@ +import pytest +from giant.cond_layout import AXIS_TYPES, CondLayout +from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM + +# ── cond_cat column layout ─────────────────────────────────────────────────── + + +def test_topn_cols_neither_onehot(): + layout = CondLayout.from_types("physical", "embedding") + assert (layout.particle_topn_col, layout.material_topn_col) == (None, None) + assert layout.cat_dim == 2 + + +def test_topn_cols_particle_only(): + layout = CondLayout.from_types("onehot", "physical") + assert (layout.particle_topn_col, layout.material_topn_col) == (2, None) + assert layout.cat_dim == 3 + + +def test_topn_cols_material_only(): + layout = CondLayout.from_types("physical", "onehot") + assert (layout.particle_topn_col, layout.material_topn_col) == (None, 2) + assert layout.cat_dim == 3 + + +def test_topn_cols_both_onehot_particle_then_material(): + layout = CondLayout.from_types("onehot", "onehot") + assert (layout.particle_topn_col, layout.material_topn_col) == (2, 3) + assert layout.cat_dim == 4 + + +def test_dense_vocab_cols_are_mode_independent(): + """Columns 0/1 are always the dense pdg/material index — giant.model.routers + reads them without knowing the conditioning mode.""" + assert (CondLayout.PDG_COL, CondLayout.MAT_COL) == (0, 1) + for particle in AXIS_TYPES: + for material in AXIS_TYPES: + layout = CondLayout.from_types(particle, material) + assert layout.particle_topn_col not in (layout.PDG_COL, layout.MAT_COL) + assert layout.material_topn_col not in (layout.PDG_COL, layout.MAT_COL) + + +# ── cond_cont slice layout ─────────────────────────────────────────────────── + + +def test_cont_slices_tile_cond_cont_exactly(): + """base / particle_phys / material_phys must partition cond_cont with no + gap and no overlap — a gap or overlap is exactly the silent + mis-indexing this object exists to prevent.""" + layout = CondLayout.from_types("physical", "physical") + covered = list(range(*layout.base.indices(COND_DIM))) + covered += list(range(*layout.particle_phys.indices(COND_DIM))) + covered += list(range(*layout.material_phys.indices(COND_DIM))) + assert covered == list(range(COND_DIM)) + + +def test_cont_slice_widths_match_constants(): + layout = CondLayout.from_types("embedding", "embedding") + assert layout.base == slice(0, COND_DIM_BASE) + assert layout.particle_phys.stop - layout.particle_phys.start == PARTICLE_PHYS_DIM + assert layout.material_phys.stop - layout.material_phys.start == MATERIAL_PHYS_DIM + assert layout.cont_dim == COND_DIM + + +def test_cont_slices_are_mode_independent(): + """cond_cont is COND_DIM wide in every mode — a non-"physical" axis gets + its block zero-filled rather than dropped, so the slices never move.""" + physical = CondLayout.from_types("physical", "physical") + for particle in AXIS_TYPES: + for material in AXIS_TYPES: + layout = CondLayout.from_types(particle, material) + assert layout.base == physical.base + assert layout.particle_phys == physical.particle_phys + assert layout.material_phys == physical.material_phys + + +# ── validation ─────────────────────────────────────────────────────────────── + + +def test_unknown_particle_type_raises(): + with pytest.raises(ValueError, match="unknown conditioning.particle.type 'bogus'"): + CondLayout.from_types("bogus", "physical") + + +def test_unknown_material_type_raises(): + with pytest.raises(ValueError, match="unknown conditioning.material.type 'bogus'"): + CondLayout.from_types("physical", "bogus") diff --git a/tests/test_network.py b/tests/test_network.py index d55aef9..4a99ad7 100644 --- a/tests/test_network.py +++ b/tests/test_network.py @@ -17,7 +17,6 @@ from giant.model.network import ( build_critics, build_history, build_models, - cat_col_layout, stage2_trunk_sec_dim, stage2_type_dim, ) @@ -129,23 +128,8 @@ def test_stage1_model_n_sec_head_cfg_controls_hidden_width_and_depth(): assert model.n_sec_head[0].out_features == 16 -# --- cat_col_layout / stage2_type_dim / stage2_trunk_sec_dim --------------- - - -def test_cat_col_layout_neither_onehot(): - assert cat_col_layout("physical", "embedding") == (None, None) - - -def test_cat_col_layout_particle_only(): - assert cat_col_layout("onehot", "physical") == (2, None) - - -def test_cat_col_layout_material_only(): - assert cat_col_layout("physical", "onehot") == (None, 2) - - -def test_cat_col_layout_both_onehot_particle_then_material(): - assert cat_col_layout("onehot", "onehot") == (2, 3) +# --- stage2_type_dim / stage2_trunk_sec_dim -------------------------------- +# (the cond_cat column-layout tests live in tests/test_cond_layout.py) def test_stage2_type_dim_physical_is_particle_phys_dim(): diff --git a/tests/test_transforms.py b/tests/test_transforms.py index 7f19326..b1dfb43 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -2,6 +2,7 @@ import warnings import numpy as np import pytest +from giant.cond_layout import AXIS_TYPES, CondLayout from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX from giant.data.transforms import ( build_cond_features, @@ -531,6 +532,129 @@ def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode( ) +# ── build_cond_features / build_features share one column layout (gitea #37) ── + + +@pytest.mark.parametrize("particle_type", AXIS_TYPES) +@pytest.mark.parametrize("material_type", AXIS_TYPES) +def test_both_builders_agree_column_for_column(particle_type, material_type, fake_material_props): + """The two builders used to lay out cond_cont/cond_cat independently and + drift apart silently. They now share `_build_cond_arrays`, so for every + mode pair they must produce identical arrays.""" + data = _minimal_step_data(3) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + pdg_topn = {11: 0} if particle_type == "onehot" else None + mat_topn = {"PbWO4": 0} if material_type == "onehot" else None + + cond_cont, cond_cat = build_cond_features( + data, + pdg_map, + mat_map, + particle_conditioning=particle_type, + material_conditioning=material_type, + pdg_topn_map=pdg_topn, + mat_topn_map=mat_topn, + ) + feats = build_features( + data, + pdg_map, + mat_map, + particle_conditioning=particle_type, + material_conditioning=material_type, + pdg_topn_map=pdg_topn, + mat_topn_map=mat_topn, + ) + + layout = CondLayout.from_types(particle_type, material_type) + assert cond_cat.shape[1] == layout.cat_dim + np.testing.assert_array_equal(feats.cond_cont, cond_cont) + np.testing.assert_array_equal(feats.cond_cat, cond_cat) + + +def test_build_features_physical_mode_tolerates_out_of_vocab_pdg_and_material(): + """The permissive vocab lookup added for "physical"/"onehot" mode (see + build_cond_features) applies to build_features too — `giant predict` on a + file whose pdg/material aren't in the checkpoint's dense vocab must not + KeyError when nothing reads those indices.""" + pdg_map = {11: 0, 22: 1} + mat_map = {"G4_AIR": 0} + data = _minimal_step_data(2) + data["pdg"] = np.full(2, 13, dtype=np.int64) # not in pdg_map + data["material"] = np.full(2, "G4_Pb", dtype=object) # not in mat_map + + _, cond_cat, *_ = build_features( + data, + pdg_map, + mat_map, + particle_conditioning="physical", + material_conditioning="physical", + ) + np.testing.assert_array_equal(cond_cat, [[0, 0], [0, 0]]) # dummy indices, no raise + + with pytest.raises(KeyError): + build_features( + data, + pdg_map, + mat_map, + particle_conditioning="embedding", + material_conditioning="embedding", + ) + + +def test_build_features_pads_legacy_normalizer_in_embedding_mode(): + """The legacy-normalizer padding (a pre-physical-conditioning checkpoint's + cond normalizer is COND_DIM_BASE wide) applies to build_features too — + `giant predict` reaches build_features, not build_cond_features.""" + data = _minimal_step_data(3) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + legacy_norm = Normalizer() + legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32) + legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32) + + cond_cont, *_ = build_features( + data, + pdg_map, + mat_map, + cond_normalizer=legacy_norm, + particle_conditioning="embedding", + material_conditioning="embedding", + ) + + assert cond_cont.shape[-1] == COND_DIM + np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0) + + +def test_build_features_rejects_legacy_normalizer_in_physical_mode(fake_material_props): + data = _minimal_step_data(3) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + legacy_norm = Normalizer() + legacy_norm.mean = np.zeros(COND_DIM_BASE, dtype=np.float32) + legacy_norm.std = np.ones(COND_DIM_BASE, dtype=np.float32) + + with pytest.raises(ValueError, match="predates physical-property conditioning"): + build_features( + data, + pdg_map, + mat_map, + cond_normalizer=legacy_norm, + particle_conditioning="physical", + material_conditioning="physical", + ) + + +def test_onehot_axis_without_its_topn_map_raises(): + """`cond_cat`'s width is the layout's call, so a "onehot" axis with no + top-N map is a hard error rather than a silently-narrower array that + ConditionEncoder would then index out of bounds.""" + data = _minimal_step_data(2) + pdg_map, mat_map = {11: 0}, {"PbWO4": 0} + + with pytest.raises(ValueError, match="needs pdg_topn_map"): + build_cond_features(data, pdg_map, mat_map, particle_conditioning="onehot") + with pytest.raises(ValueError, match="needs mat_topn_map"): + build_cond_features(data, pdg_map, mat_map, material_conditioning="onehot") + + # ── sorted_membership / _vectorized_map_lookup ──────────────────────────────