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:
2026-07-09 09:09:14 +02:00
co-authored by Claude Opus 4.8
parent a14a4f973a
commit 339d2fb473
4 changed files with 75 additions and 1 deletions
+48
View File
@@ -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()