giant/analysis/context.py

@@ -85,7 +85,11 @@ _LO_Q, _HI_Q = 0.001, 0.999
 
 
 def _row_subsample(lf: pl.LazyFrame, sample_rows: int, seed: int) -> pl.LazyFrame:
-    """Hash-subsample ~``sample_rows`` rows (for range estimation only)."""
+    """Hash-subsample ~``sample_rows`` rows (for range estimation only).
+
+    ``lf`` must still carry ``pre_E`` (the hash key) — subsample before
+    projecting away columns, not after.
+    """
     n_total = lf.select(pl.len()).collect(engine="streaming").item()
     if n_total <= sample_rows:
         return lf
@@ -195,20 +199,29 @@ def build_context(
     if predictions:
         sides = {ps.name: open_prediction(ps.source) for ps in predictions}
         present_vars = sorted(set().union(*(paired_vars_for_coord(s.coord) for s in sides.values())))
-        for var in present_vars:
-            true_samples, pred_samples, residual_samples = [], [], []
-            for s in sides.values():
-                if var not in paired_vars_for_coord(s.coord):
-                    continue
-                cols = [f"pred_{var}"] + ([f"true_{var}"] if s.has_truth else [])
-                sample = _row_subsample(s.paired.select(cols), sample_rows, seed).collect(engine="streaming")
-                pred_samples.append(sample[f"pred_{var}"].to_numpy())
+        true_samples: dict[str, list[np.ndarray]] = {v: [] for v in present_vars}
+        pred_samples: dict[str, list[np.ndarray]] = {v: [] for v in present_vars}
+        residual_samples: dict[str, list[np.ndarray]] = {v: [] for v in present_vars}
+
+        # One subsample+collect per side (not per variable) — `s.paired`
+        # still carries `pre_E`, which `_row_subsample`'s hash needs, so
+        # subsample before projecting down to the pred/true columns.
+        for s in sides.values():
+            vars_here = paired_vars_for_coord(s.coord)
+            cols = [f"pred_{v}" for v in vars_here] + ([f"true_{v}" for v in vars_here] if s.has_truth else [])
+            sample = _row_subsample(s.paired, sample_rows, seed).select(cols).collect(engine="streaming")
+            for var in vars_here:
+                p = sample[f"pred_{var}"].to_numpy()
+                pred_samples[var].append(p)
                 if s.has_truth:
-                    true_samples.append(sample[f"true_{var}"].to_numpy())
-                    residual_samples.append(sample[f"pred_{var}"].to_numpy() - sample[f"true_{var}"].to_numpy())
-            pred_var_ranges[var] = _combined_quantiles([*true_samples, *pred_samples], _LO_Q, _HI_Q)
-            if residual_samples:
-                pred_residual_ranges[var] = _combined_quantiles(residual_samples, _LO_Q, _HI_Q)
+                    t = sample[f"true_{var}"].to_numpy()
+                    true_samples[var].append(t)
+                    residual_samples[var].append(p - t)
+
+        for var in present_vars:
+            pred_var_ranges[var] = _combined_quantiles([*true_samples[var], *pred_samples[var]], _LO_Q, _HI_Q)
+            if residual_samples[var]:
+                pred_residual_ranges[var] = _combined_quantiles(residual_samples[var], _LO_Q, _HI_Q)
 
         sec_pdg_counts: dict[int, int] = {}
         for s in sides.values():

tests/test_analysis_prediction.py

@@ -193,6 +193,28 @@ def test_build_context_resolves_prediction_ranges():
     assert ctx.pred_top_sec_pdgs  # secondaries present in the fixture
 
 
+def test_build_context_resolves_prediction_ranges_when_subsampled():
+    """Regression test: `sample_rows` smaller than the prediction row count
+    must still work — `_row_subsample`'s hash key (`pre_E`) has to survive
+    into the prediction branch's subsample call, not be projected away first
+    (see giant/analysis/context.py's predictions loop)."""
+    ctx = build_context(
+        [RolloutSpec("rollout", _rollout_frame())],
+        _reference_frame(),
+        predictions=[PredictionSpec("pred", _global_prediction_frame())],
+        n_energy_bins=2,
+        n_marginal_bins=10,
+        top_k_pdg=3,
+        sample_rows=2,  # < the 3-row prediction fixture: forces the hash-filter branch
+        seed=0,
+    )
+    for var in PAIRED_SCALARS + ("cos_scatter", "cos_travel"):
+        lo, hi = ctx.pred_var_ranges[var]
+        assert np.isfinite(lo) and np.isfinite(hi) and lo < hi
+        lo, hi = ctx.pred_residual_ranges[var]
+        assert np.isfinite(lo) and np.isfinite(hi) and lo < hi
+
+
 def test_prediction_specs_compute_valid_reduced():
     ctx = _ctx_with_predictions()
     bundle = Bundle.open(