diff --git a/giant/analysis/reduce.py b/giant/analysis/reduce.py index b050563..ce65c89 100644 --- a/giant/analysis/reduce.py +++ b/giant/analysis/reduce.py @@ -29,8 +29,17 @@ from giant.constants import TERM_ESCAPED def _bin_expr(value: pl.Expr, lo: float, hi: float, nbins: int) -> pl.Expr: - """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins.""" - return ((value - lo) / (hi - lo) * nbins).floor().cast(pl.Int32).clip(0, nbins - 1) + """Uniform bin index of ``value`` over ``[lo, hi]`` into ``nbins`` bins. + + Out-of-range values clamp into the edge bins, and the clamp deliberately + happens in f64 *before* the integer cast: a rollout is free to emit a wildly + out-of-range outlier (a step_length of 1e10 mm, say) or an inf, whose + unclamped bin index overflows i32 and makes the cast fail outright. NaN has + no edge to clamp to, so it becomes null and is dropped by the callers below + — the same thing ``np.histogram`` does with it. + """ + idx = ((value - lo) / (hi - lo) * nbins).floor().clip(0, nbins - 1) + return pl.when(idx.is_nan()).then(None).otherwise(idx).cast(pl.Int32) def hist1d( @@ -50,6 +59,7 @@ def hist1d( group = pl.lit(0, dtype=pl.Int64) if group is None else group res = ( lf.select(group.alias("_g"), _bin_expr(value, lo, hi, nbins).alias("_b")) + .drop_nulls("_b") .group_by("_g", "_b") .agg(pl.len().alias("_n")) .collect(engine="streaming") @@ -188,6 +198,7 @@ def profile_partial( _bin_expr(coord, lo, hi, nbins).alias("_b"), weight.alias("_w"), ) + .drop_nulls("_b") .group_by("event_id", "_b") .agg(pl.col("_w").sum().alias("_ws")) .collect(engine="streaming") diff --git a/tests/test_analysis_reduce.py b/tests/test_analysis_reduce.py index 877f274..614233c 100644 --- a/tests/test_analysis_reduce.py +++ b/tests/test_analysis_reduce.py @@ -102,6 +102,31 @@ def test_hist1d_overall_and_grouped(): assert hg[11].sum() == 4 +def test_hist1d_clamps_extreme_values_and_drops_nan(): + # A rollout can emit a wildly out-of-range step_length (or an inf/NaN); the + # fixed-edge binning must clamp rather than overflow the i32 bin cast. + lf = pl.DataFrame({"x": [5.0, 1.0725e10, float("inf"), -float("inf"), float("nan"), None]}).lazy() + edges = np.linspace(0.0, 50.0, 6) # width 10 + h = R.hist1d(lf, pl.col("x"), edges) + # 5 -> bin 0; 1e10 and +inf -> top bin; -inf -> bin 0; NaN/null dropped + assert h[0].tolist() == [2, 0, 0, 0, 2] + + +def test_profile_partial_clamps_extreme_values_and_drops_nan(): + lf = pl.DataFrame( + { + "event_id": [1, 1, 1, 1], + "z": [5.0, 1.0725e10, float("nan"), 45.0], + "w": [1.0, 2.0, 4.0, 8.0], + } + ).lazy() + edges = np.linspace(0.0, 50.0, 6) + ev, mat = R.profile_partial(lf, pl.col("z"), edges, pl.col("w")) + assert ev.tolist() == [1] + # 1e10 clamps into the top bin alongside 45; the NaN row's weight is dropped + assert mat[0].tolist() == [1.0, 0.0, 0.0, 0.0, 10.0] + + def test_physical_steps_drops_synthetic_rollout_rows_only(): lf = _rollout_frame() phys = physical_steps(lf, Side.rollout).collect()