diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 9b74ea7..d83148c 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -530,11 +530,44 @@ def build_cond_features( cond_cat = np.column_stack([pdg_idx, mat_idx]) if cond_normalizer is not None: - cond_cont = cond_normalizer.transform(cond_cont) + cond_cont = _cond_normalizer_transform(cond_cont, cond_normalizer, conditioning) return cond_cont, cond_cat +def _cond_normalizer_transform( + cond_cont: np.ndarray, cond_normalizer: "Normalizer", 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 + ``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. + """ + 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": + raise ValueError( + f"cond normalizer has {mean.shape[-1]} columns, expected " + f"{width}, and conditioning={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." + ) + pad = width - mean.shape[-1] + mean = np.concatenate([mean, np.zeros(pad, dtype=mean.dtype)]) + std = np.concatenate([std, np.ones(pad, dtype=std.dtype)]) + return ((cond_cont - mean) / std).astype(np.float32) + + def build_features( data: dict[str, np.ndarray], pdg_map: dict[int, int], diff --git a/tests/test_transforms.py b/tests/test_transforms.py index a7af975..c9bc7ac 100644 --- a/tests/test_transforms.py +++ b/tests/test_transforms.py @@ -397,3 +397,46 @@ def test_build_cond_features_mass_charge_override(fake_material_props): cond_cont[:, COND_DIM_BASE], log_transform(np.array([123.0, 456.0])) ) np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE + 1], [2.0, -2.0]) + + +def test_build_cond_features_pads_legacy_normalizer_in_embedding_mode(): + """A pre-physical-conditioning checkpoint's cond normalizer is COND_DIM_BASE + (8) wide, fit before build_cond_features grew the extra physical columns. + In "embedding" mode those columns are never read downstream, so a legacy + normalizer should be usable as-is (padded, not rejected).""" + 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_cond_features( + data, pdg_map, mat_map, cond_normalizer=legacy_norm, conditioning="embedding" + ) + + assert cond_cont.shape[-1] == COND_DIM + # padded physical columns are zero-filled pre-normalization and + # mean=0/std=1 post-normalization, so they should come out as zero + np.testing.assert_allclose(cond_cont[:, COND_DIM_BASE:], 0.0) + + +def test_build_cond_features_rejects_legacy_normalizer_in_physical_mode( + fake_material_props, +): + """Unlike "embedding" mode, "physical" mode actually reads the physical + columns, so a legacy 8-wide normalizer can't be silently padded — that + would silently feed the network un-normalized physical properties.""" + 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_cond_features( + data, + pdg_map, + mat_map, + cond_normalizer=legacy_norm, + conditioning="physical", + )