Rescale secondary energies to exactly consume the e_sec budget

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 <noreply@anthropic.com>

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}")

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]

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

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)