Error on missing secondary lists instead of silently zeroing Stage-2 targets
A parquet that carries child_track_ids/e_sec but was never run through the parent->child join lacks the per-secondary columns (sec_E_list/sec_pdg_list/ sec_dir_list). build_features would fall back to all-zero sec_cont/sec_pdg_idx, collapsing every secondary to PDG index 0 and a constant energy fraction — a broken Stage 2 that trained with no error (single-species validation tables). Add an opt-in require_secondaries flag that raises when n_sec > 0 but the lists are absent, and enable it on the training paths (StreamingStepsDataset and the normalizer-fit pass). giant predict keeps the default False for Stage-1-only use. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -239,3 +239,51 @@ def test_build_features_clamps_n_sec_label_to_k_max():
|
||||
|
||||
assert n_sec.max() <= K_MAX
|
||||
np.testing.assert_array_equal(n_sec, [0, 5, K_MAX])
|
||||
|
||||
|
||||
def _step_data_no_sec_lists(n_sec: np.ndarray) -> dict:
|
||||
"""Minimal build_features input with n_sec but no per-secondary list columns
|
||||
(mimics a parquet that skipped the parent->child join)."""
|
||||
N = len(n_sec)
|
||||
rng = np.random.default_rng(0)
|
||||
return {
|
||||
"pdg": np.full(N, 11, dtype=np.int32),
|
||||
"material": np.full(N, "PbWO4", dtype=object),
|
||||
"pre_pos": rng.standard_normal((N, 3)).astype(np.float32),
|
||||
"pre_E": np.full(N, 10.0, dtype=np.float32),
|
||||
"pre_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
|
||||
"layer_id": np.zeros(N, dtype=np.int32),
|
||||
"n_sec": np.asarray(n_sec, dtype=np.int32),
|
||||
"e_sec": np.full(N, 1.0, dtype=np.float32),
|
||||
"step_length": np.full(N, 1.0, dtype=np.float32),
|
||||
"post_E": np.full(N, 9.0, dtype=np.float32),
|
||||
"edep": np.full(N, 1.0, dtype=np.float32),
|
||||
"post_dir": np.tile(np.array([0, 0, 1], dtype=np.float32), (N, 1)),
|
||||
"post_pos": rng.standard_normal((N, 3)).astype(np.float32),
|
||||
}
|
||||
|
||||
|
||||
def test_build_features_require_secondaries_raises_when_lists_missing():
|
||||
"""A parquet with n_sec > 0 but no per-secondary list columns was never run
|
||||
through the parent->child join; require_secondaries must catch it instead of
|
||||
silently zeroing every Stage-2 target (regression: this collapsed the
|
||||
secondary species to a single PDG index during training)."""
|
||||
data = _step_data_no_sec_lists(np.array([0, 2, 1]))
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
with pytest.raises(ValueError, match="per-secondary columns"):
|
||||
build_features(data, pdg_map, mat_map, require_secondaries=True)
|
||||
|
||||
|
||||
def test_build_features_require_secondaries_ok_when_no_secondaries():
|
||||
"""require_secondaries only fires when secondaries actually exist; a file
|
||||
with n_sec == 0 everywhere (e.g. Stage-1-only) must still load."""
|
||||
data = _step_data_no_sec_lists(np.zeros(3, dtype=np.int32))
|
||||
pdg_map, mat_map = {11: 0}, {"PbWO4": 0}
|
||||
|
||||
_, _, _, _, sec_cont, sec_pdg_idx, _, _ = build_features(
|
||||
data, pdg_map, mat_map, require_secondaries=True
|
||||
)
|
||||
|
||||
assert not sec_cont.any()
|
||||
assert not sec_pdg_idx.any()
|
||||
|
||||
Reference in New Issue
Block a user