From 670f57c3098a3e713f631efb58ae7d0f25926142 Mon Sep 17 00:00:00 2001 From: Lars Bogner Date: Mon, 13 Jul 2026 13:45:13 +0200 Subject: [PATCH] Rescale secondary energies to exactly consume the e_sec budget MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit decode_secondaries's stick-breaking only guarantees valid secondary slots sum to <= e_sec, leaving a shortfall that rollout.py silently dumped into that step's edep. Rescale the valid slots by one common per-row factor instead, so they sum to exactly e_sec whenever n_sec > 0: this spreads any shortfall proportionally across all secondaries rather than concentrating it in whichever slot is last by energy rank (which would let that one low-energy secondary balloon and distort the shower's topology). Rows where every valid slot decodes to ~zero fall back to an even split. n_sec == 0 rows are unchanged (still nothing to carry the budget, so rollout.py's edep top-up still applies there) — narrowed the related caveat in load_rollout_vs_truth's docstring to just that case. Co-Authored-By: Claude Sonnet 5 --- giant/analysis.py | 16 ++--- giant/data/transforms.py | 26 +++++++- giant/rollout.py | 8 ++- tests/test_phase2.py | 127 +++++++++++++++++++++++++++++++++++++++ 4 files changed, 164 insertions(+), 13 deletions(-) diff --git a/giant/analysis.py b/giant/analysis.py index 0453ac9..6184850 100644 --- a/giant/analysis.py +++ b/giant/analysis.py @@ -2293,13 +2293,15 @@ def load_rollout_vs_truth( The rollout side drops synthetic termination-bookkeeping rows (see `_SYNTHETIC_ROLLOUT_TERMINATION_REASONS`) before decoding, since those - aren't real generated steps. Even among the real steps that remain, - `edep` isn't perfectly analogous between the two files: `rollout.py` - tops up a step's `edep` with any secondary-energy budget Stage 2 didn't - allocate to an actual spawned secondary (so every step still conserves - energy exactly), which truth's Geant4-recorded `edep` never does. A - generated `edep` that runs a bit high relative to truth can be this - bookkeeping, not necessarily a Stage-1/Stage-2 miscalibration. + aren't real generated steps. One remaining case where `edep` still isn't + perfectly analogous between the two files: `decode_secondaries` rescales + the valid secondary slots to sum to exactly `e_sec` whenever `n_sec > 0` + (see that function's docstring), but when Stage 1 predicts a nonzero + `e_sec` while Stage 2's `n_sec` head predicts 0 secondaries, there's no + slot to carry that budget at all — `rollout.py` deposits it into that + step's `edep` instead, which truth's Geant4-recorded `edep` never does. + That specific disagreement between the two Stage-1/Stage-2 heads is rare + but not otherwise fixable at decode time. """ if not (0 < sample_frac <= 1): raise ValueError(f"sample_frac must be in (0, 1], got {sample_frac}") diff --git a/giant/data/transforms.py b/giant/data/transforms.py index 2994393..ee859eb 100644 --- a/giant/data/transforms.py +++ b/giant/data/transforms.py @@ -360,6 +360,8 @@ def decode_secondaries( pdg_map_inv: maps model index → PDG code Returns (sec_E, sec_dir_world, sec_pdg_code, sec_valid) each shape (N, K_MAX). + The valid slots' energies (`sec_E[sec_valid]`, per row) always sum to + exactly `e_sec` — see the rescaling below. """ N, K, _ = sec_cont.shape stick_logits = sec_cont[:, :, 0] # (N, K) @@ -372,15 +374,33 @@ def decode_secondaries( fractions = 1.0 / (1.0 + np.exp(-stick_logits.astype(np.float64))) - sec_E = np.zeros((N, K), dtype=np.float32) + sec_E = np.zeros((N, K), dtype=np.float64) e_sec = np.asarray(e_sec, dtype=np.float64) remaining = e_sec.copy() for i in range(K): - sec_E[:, i] = (fractions[:, i] * remaining).astype(np.float32) - remaining = np.maximum(remaining - sec_E[:, i].astype(np.float64), 0.0) + sec_E[:, i] = fractions[:, i] * remaining + remaining = np.maximum(remaining - sec_E[:, i], 0.0) sec_valid = np.arange(K)[None, :] < n_sec[:, None] # (N, K) + # Stick-breaking guarantees sum(sec_E[valid]) <= e_sec (each fraction is in + # [0,1] of an already-shrinking remainder) but rarely hits it exactly, so + # rescale the valid slots by one common per-row factor to close that gap — + # rather than dumping the shortfall into whichever slot happens to be last + # by energy rank, which would let one low-energy secondary balloon and + # distort the shower's topology. This preserves each row's relative split + # across its secondaries and only ever scales up (valid_sum <= e_sec). + # Rows where every valid slot decoded to ~zero (scale undefined) fall back + # to an even split of e_sec across the n_sec valid slots. + sec_E = sec_E * sec_valid + valid_sum = sec_E.sum(axis=1) + degenerate = (valid_sum <= _EPS) & (n_sec > 0) + scale = np.where(valid_sum > _EPS, e_sec / np.maximum(valid_sum, _EPS), 0.0) + sec_E = sec_E * scale[:, None] + even_share = e_sec / np.maximum(n_sec, 1).astype(np.float64) + sec_E = np.where(degenerate[:, None] & sec_valid, even_share[:, None], sec_E) + sec_E = sec_E.astype(np.float32) + sec_dir_world = np.zeros((N, K, 3), dtype=np.float32) for i in range(K): valid = sec_valid[:, i] diff --git a/giant/rollout.py b/giant/rollout.py index a0fe071..109faa0 100644 --- a/giant/rollout.py +++ b/giant/rollout.py @@ -394,9 +394,11 @@ def _step_chunk( max_tracks_per_event, ) # Energy bookkeeping so each step conserves exactly (edep + carried + post_E - # == pre_E): the primary lost `e_sec` to secondaries, but the decoded - # secondaries only carry `sec_E[valid].sum()`. Deposit the unallocated - # residual locally, plus the energy of any sub-cap secondaries we dropped. + # == pre_E): `decode_secondaries` already rescales valid slots to sum to + # exactly `e_sec` whenever n_sec > 0, so `residual` here is ~0 except when + # n_sec == 0 (no secondary to carry the budget at all — the whole `e_sec` + # becomes residual). Also deposit the energy of any sub-cap secondaries + # we dropped for hitting `max_tracks_per_event`. sec_E_valid_sum = (sec_E * sec_valid).sum(axis=1) residual = np.maximum(e_sec - sec_E_valid_sum, 0.0) edep = edep + residual + dropped_edep diff --git a/tests/test_phase2.py b/tests/test_phase2.py index 8e46740..e3dec95 100644 --- a/tests/test_phase2.py +++ b/tests/test_phase2.py @@ -231,3 +231,130 @@ def test_encode_secondaries_direction_encoding(): local_dirs = sec_cont[:, :2, 1:4] # (N, 2, 3) — valid slots only norms_out = np.linalg.norm(local_dirs, axis=-1) np.testing.assert_allclose(norms_out, 1.0, atol=1e-5) + + +# ── decode_secondaries: exact energy conservation ──────────────────────────── + + +def _random_sec_cont(rng, N, stick_logit_scale=1.0): + sec_cont = rng.standard_normal((N, K_MAX, 4)).astype(np.float32) + sec_cont[:, :, 0] *= stick_logit_scale + dirs = sec_cont[:, :, 1:] + dirs /= np.linalg.norm(dirs, axis=-1, keepdims=True) + return sec_cont + + +def test_decode_secondaries_valid_slots_sum_to_e_sec(): + """The valid slots' energies must sum to exactly e_sec, not just <= e_sec. + + Rows with n_sec=0 are excluded: there's no slot to put the budget in, so + valid_sum is correctly 0 regardless of e_sec there (see + test_decode_secondaries_zero_n_sec_has_zero_energy) — the shortfall in + that case is handled downstream (e.g. rollout.py dumps it into edep). + """ + from giant.data.transforms import decode_secondaries + + rng = np.random.default_rng(0) + N = 200 + sec_cont = _random_sec_cont(rng, N) + sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) + n_sec = rng.integers(0, K_MAX + 1, size=N) + e_sec = rng.uniform(0.0, 50.0, size=N).astype(np.float32) + pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) + + sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries( + sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22} + ) + + valid_sum = (sec_E * sec_valid).sum(axis=1) + has_secondaries = n_sec > 0 + np.testing.assert_allclose( + valid_sum[has_secondaries], + e_sec[has_secondaries], + atol=1e-3, + rtol=1e-5, + ) + + +def test_decode_secondaries_zero_n_sec_has_zero_energy(): + """n_sec=0 rows get no secondaries and no forced energy assignment.""" + from giant.data.transforms import decode_secondaries + + rng = np.random.default_rng(1) + N = 10 + sec_cont = _random_sec_cont(rng, N) + sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) + n_sec = np.zeros(N, dtype=np.int64) + e_sec = rng.uniform(1.0, 10.0, size=N).astype(np.float32) + pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) + + sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries( + sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22} + ) + + assert not sec_valid.any() + np.testing.assert_allclose(sec_E, 0.0) + + +def test_decode_secondaries_degenerate_row_falls_back_to_even_split(): + """All-zero stick fractions for the valid slots fall back to an even split.""" + from giant.data.transforms import decode_secondaries + + rng = np.random.default_rng(2) + N = 4 + sec_cont = _random_sec_cont(rng, N) + # Drive every valid slot's stick-breaking fraction to ~0 (huge negative logit). + n_sec = np.array([0, 1, 3, K_MAX]) + for i, k in enumerate(n_sec): + sec_cont[i, :k, 0] = -80.0 + sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) + e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32) + pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32) + + sec_E, _sec_dir, _sec_pdg, sec_valid = decode_secondaries( + sec_cont, sec_pdg_pred, n_sec, e_sec, pre_dir, {0: 22} + ) + + for i, k in enumerate(n_sec): + if k == 0: + continue + np.testing.assert_allclose(sec_E[i, :k], e_sec[i] / k, atol=1e-4) + np.testing.assert_allclose(sec_E[i, :k].sum(), e_sec[i], atol=1e-3) + + +def test_decode_secondaries_rescale_preserves_relative_shares(): + """Rescaling should keep each valid slot's *share* of the budget unchanged. + + A shortfall shouldn't get dumped into whichever slot is last by energy + rank — it should be spread proportionally, i.e. sec_E[i] / sec_E[j] for + two valid slots must match before and after the e_sec rescale. + """ + from giant.data.transforms import decode_secondaries + + rng = np.random.default_rng(3) + N = 1 + sec_cont = _random_sec_cont(rng, N) + n_sec = np.array([4]) + pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32) + sec_pdg_pred = np.zeros((N, K_MAX), dtype=np.int64) + + sec_E_small, _, _, sec_valid = decode_secondaries( + sec_cont, + sec_pdg_pred, + n_sec, + np.array([5.0], dtype=np.float32), + pre_dir, + {0: 22}, + ) + sec_E_large, _, _, _ = decode_secondaries( + sec_cont, + sec_pdg_pred, + n_sec, + np.array([50.0], dtype=np.float32), + pre_dir, + {0: 22}, + ) + + ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0] + ratio_large = sec_E_large[0, :4] / sec_E_large[0, 0] + np.testing.assert_allclose(ratio_small, ratio_large, rtol=1e-4)