diff --git a/giant/analysis/context.py b/giant/analysis/context.py index e5c4be2..d65f2af 100644 --- a/giant/analysis/context.py +++ b/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(): diff --git a/tests/test_analysis_prediction.py b/tests/test_analysis_prediction.py index e07379e..f4785b3 100644 --- a/tests/test_analysis_prediction.py +++ b/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(