Nixtla/statsforecast · #1246

fix: use the finite-sample conformal quantile for conformal intervals

nicoloangileri · merged Oct 14, 20265 files · 180 + / 46 −
python/statsforecast/models.py15 + / 22 −
@@ -87,6 +87,7 @@ def _warn_distribution_mismatch(error_distribution, fitted_dist):     _VALID_DISTRIBUTIONS,     _calculate_intervals,     _calculate_sigma,+    _conformal_quantile,     _ensure_float,     _naive,     _quantiles,@@ -164,11 +165,9 @@ def _add_conformal_distribution_intervals(     cuts.extend(1 - alpha / 200 for alpha in alphas)     mean = fcst["mean"].reshape(1, -1)     scores = np.vstack([mean - cs, mean + cs])-    quantiles = np.quantile(-        scores,-        cuts,-        axis=0,-    )+    # lower cuts come first, then upper ones, as in `out_cols` below+    upper = [False] * len(level) + [True] * len(level)+    quantiles = _conformal_quantile(scores, cuts, upper, axis=0)     quantiles = quantiles.reshape(len(cuts), -1)     lo_cols = [f"lo-{lv}" for lv in reversed(level)]     hi_cols = [f"hi-{lv}" for lv in level]@@ -185,11 +184,15 @@ def _add_conformal_error_intervals( ) -> Dict:     r"""     Adds conformal intervals to the `fcst` dict based on conformal scores `cs`.-    `level` should be already sorted. Uses the `lv / 100` quantile of the-    absolute conformity scores as the error margin around `fcst["mean"]`.+    `level` should be already sorted. Uses the finite-sample conformal quantile+    of the absolute conformity scores (see `_conformal_quantile`) as the error+    margin around `fcst["mean"]`.     Returns the modified `fcst` dict.     """-    quantiles = {lv: np.quantile(cs, lv / 100, axis=0) for lv in level}+    # the margin on absolute scores is an upper bound at every level, even below 50+    quantiles = {+        lv: _conformal_quantile(cs, lv / 100, upper=True, axis=0) for lv in level+    }     for lv in reversed(level):         fcst[f"lo-{lv}"] = fcst["mean"] - quantiles[lv]     for lv in level:@@ -4398,22 +4401,13 @@ def __init__(             alias = "CSP-Adaptive" if variant == "adaptive" else "CSP-Fixed"         self.alias = alias -    @staticmethod-    def _oriented_index(q: float, n: int) -> float:-        if n <= 0:-            return float(q)-        if q < 0.5:-            return max(0.0, float(np.floor((n + 1.0) * q)) / n)-        return min(1.0, float(np.ceil((n + 1.0) * q)) / n)-     def _intervals_from_samples(         self, samples: np.ndarray, level: Sequence[Union[int, float]]     ) -> Dict[str, np.ndarray]:         cuts = [(1 - lv / 100) / 2 for lv in reversed(level)]         cuts += [1 - (1 - lv / 100) / 2 for lv in level]-        n = samples.shape[0]-        oriented = [self._oriented_index(c, n) for c in cuts]-        quantiles = np.quantile(samples, oriented, axis=0)+        upper = [False] * len(level) + [True] * len(level)+        quantiles = _conformal_quantile(samples, cuts, upper, axis=0)         lo_cols = [f"lo-{lv}" for lv in reversed(level)]         hi_cols = [f"hi-{lv}" for lv in level]         return {col: quantiles[i] for i, col in enumerate(lo_cols + hi_cols)}@@ -4512,9 +4506,8 @@ def predict_in_sample(self, level: Optional[List[int]] = None):                 lo_q = (1 - lv / 100) / 2                 hi_q = 1 - lo_q                 if R.size > 0:-                    n_r = R.size-                    lo_offset = float(np.quantile(R, self._oriented_index(lo_q, n_r)))-                    hi_offset = float(np.quantile(R, self._oriented_index(hi_q, n_r)))+                    lo_offset = float(_conformal_quantile(R, lo_q, upper=False))+                    hi_offset = float(_conformal_quantile(R, hi_q, upper=True))                 else:                     lo_offset = hi_offset = np.nan                 res[f"fitted-lo-{lv}"] = fitted + lo_offset
python/statsforecast/utils.py37 + / 0 −
@@ -328,6 +328,43 @@ def _calculate_sigma(residuals, n):     return sigma  +def _conformal_rank(n: int, q: float, upper: bool) -> int:+    r"""1-based rank of the finite-sample conformal quantile among `n` scores.++    Split conformal prediction takes the `ceil((n + 1) * q)`-th smallest score+    for an upper bound and the `floor((n + 1) * q)`-th for a lower bound, so+    each bound covers at least its nominal probability when the scores are+    exchangeable. The tail is explicit rather than inferred from `q`: an+    absolute-score margin is an upper bound even when `q < 0.5`. The rank is+    clipped to `[1, n]`: beyond that only an infinite bound would reach the+    nominal level.+    """+    # rounding guards against `(n + 1) * q` landing just off an integer+    pos = round((n + 1) * q, 8)+    k = math.ceil(pos) if upper else math.floor(pos)+    return min(max(k, 1), n)+++def _conformal_quantile(scores, q, upper, axis: int = 0) -> np.ndarray:+    r"""Finite-sample conformal quantile(s) of `scores` along `axis`.++    Returns the exact order statistic of `_conformal_rank` rather than an+    interpolated quantile: linear interpolation lands between two scores and+    under-covers the lower tail. `upper` is a bool, or one bool per level in+    `q`. Like `np.quantile`, a sequence of levels adds a leading axis to the+    result.+    """+    scores = np.asarray(scores)+    n = scores.shape[axis]+    if n == 0:+        raise ValueError("Cannot compute a conformal quantile of zero scores")+    qs = np.atleast_1d(q)+    uppers = np.broadcast_to(np.asarray(upper, dtype=bool), qs.shape)+    ranks = [_conformal_rank(n, float(qi), bool(ui)) - 1 for qi, ui in zip(qs, uppers)]+    out = np.moveaxis(np.take(np.sort(scores, axis=axis), ranks, axis=axis), axis, 0)+    return out[0] if np.ndim(q) == 0 else out++ class ConformalIntervals:     """Class for storing conformal intervals metadata information. 
tests/test_csp.py19 + / 20 −
@@ -4,6 +4,7 @@ import pytest  from statsforecast.models import ConformalSeasonalPool, SeasonalNaive+from statsforecast.utils import _conformal_rank   @pytest.fixture@@ -213,15 +214,12 @@ def test_predict_in_sample_empty_R_returns_nan_intervals():     assert np.isnan(result["fitted-hi-90"]).all(), "offsets should be NaN when R is empty"  -def test_oriented_index_values():-    # alpha=0.05, n=100: floor(101*0.025)/100 = 2/100 = 0.02  (differs from plain 0.025)-    assert ConformalSeasonalPool._oriented_index(0.025, 100) == pytest.approx(0.02)-    # upper tail: ceil(101*0.975)/100 = 99/100 = 0.99  (differs from plain 0.975)-    assert ConformalSeasonalPool._oriented_index(0.975, 100) == pytest.approx(0.99)-    # alpha=0.10, n=100: floor(101*0.05)/100 = 5/100 = 0.05  (exact — no change)-    assert ConformalSeasonalPool._oriented_index(0.05, 100) == pytest.approx(0.05)-    # degenerate: n=0 returns q unchanged-    assert ConformalSeasonalPool._oriented_index(0.025, 0) == pytest.approx(0.025)+def test_csp_uses_shared_conformal_ranks():+    # alpha=0.05, n=100: floor(101*0.025) = 2 and ceil(101*0.975) = 99+    assert _conformal_rank(100, 0.025, upper=False) == 2+    assert _conformal_rank(100, 0.975, upper=True) == 99+    # alpha=0.10, n=100: floor(101*0.05) = 5+    assert _conformal_rank(100, 0.05, upper=False) == 5   def test_intervals_use_orientation_correction():@@ -232,13 +230,14 @@ def test_intervals_use_orientation_correction():     model = ConformalSeasonalPool(season_length=12)     result = model._intervals_from_samples(samples, level=[95]) -    # At n=100, alpha=0.05: oriented lo = floor(101*0.025)/100 = 0.02 (not 0.025)-    oriented_lo = np.quantile(samples, 0.02, axis=0).item()-    plain_lo = np.quantile(samples, 0.025, axis=0).item()--    np.testing.assert_allclose(result["lo-95"], [oriented_lo])+    # At n=100, alpha=0.05 the bounds are the floor(101*0.025) = 2nd and the+    # ceil(101*0.975) = 99th smallest samples, not interpolated quantiles+    srt = np.sort(samples, axis=0)+    np.testing.assert_allclose(result["lo-95"], srt[1])+    np.testing.assert_allclose(result["hi-95"], srt[98])+    plain_lo = np.quantile(samples, 0.025, axis=0)     # The two values must actually differ for this test to be meaningful-    assert abs(oriented_lo - plain_lo) > 1e-10+    assert abs(srt[1] - plain_lo).item() > 1e-10   def test_predict_in_sample_uses_orientation_correction():@@ -250,11 +249,11 @@ def test_predict_in_sample_uses_orientation_correction():     model = ConformalSeasonalPool(season_length=m, calib_frac=0.5).fit(y)     R = model.model_["calib_residuals"]     n_r = R.size  # 30-    # level=95: lo_q=0.025, oriented = floor(31*0.025)/30 = 0/30 = 0.0 != 0.025-    oriented_lo = ConformalSeasonalPool._oriented_index(0.025, n_r)-    plain_lo = 0.025-    assert oriented_lo != plain_lo, "oriented and plain must differ for this test to be meaningful"+    # level=95: lo_q=0.025, floor(31*0.025) = 0 is clipped to rank 1, the+    # smallest residual; the plain 0.025 quantile would interpolate above it+    assert _conformal_rank(n_r, 0.025, upper=False) == 1+    assert np.quantile(R, 0.025) > R.min(), "plain quantile must differ for this test to be meaningful"     result = model.predict_in_sample(level=[95])     fitted = model.model_["fitted"]-    expected_lo = fitted + float(np.quantile(R, oriented_lo))+    expected_lo = fitted + float(R.min())     np.testing.assert_allclose(result["fitted-lo-95"], expected_lo)
tests/test_models.py34 + / 3 −
@@ -164,16 +164,47 @@ def test_conformal_error_intervals():     assert np.all(fcst["hi-90"] >= fcst["hi-80"])     # Hi and lo are anti-symmetric (implied by the exact-value checks below).     # Exact-value check: for ZeroModel (mean=0), conformal_error intervals-    # should match the lv/100 quantile of the conformity scores.+    # should match the finite-sample conformal quantile of the conformity+    # scores. With 2 windows, ceil(3 * 0.8) and ceil(3 * 0.9) both exceed 2,+    # so both levels use the largest score.     cs = zero_model._conformity_scores(ap)-    q90 = np.quantile(cs, 0.9, axis=0)-    q80 = np.quantile(cs, 0.8, axis=0)+    q90 = cs.max(axis=0)+    q80 = cs.max(axis=0)     np.testing.assert_allclose(fcst["hi-90"], q90)     np.testing.assert_allclose(fcst["hi-80"], q80)     np.testing.assert_allclose(fcst["lo-90"], -q90)     np.testing.assert_allclose(fcst["lo-80"], -q80)  +@pytest.mark.parametrize("method", ["conformal_error", "conformal_distribution"])+def test_conformal_intervals_finite_sample_rank(method):+    # With n windows both methods put the bounds exactly at -/+ the+    # ceil((n + 1) * lv / 100)-th smallest conformity score, which covers lv+    # percent of exchangeable errors.+    n_windows = 10+    conf_intervals = ConformalIntervals(h=12, n_windows=n_windows, method=method)+    zero_model = ZeroModel(conf_intervals)+    fcst = zero_model.forecast(ap, h=12, level=[70, 80])+    cs_sorted = np.sort(zero_model._conformity_scores(ap), axis=0)+    # ceil(11 * 0.7) = 8 and ceil(11 * 0.8) = 9+    for lv, rank in [(70, 8), (80, 9)]:+        np.testing.assert_allclose(fcst[f"hi-{lv}"], cs_sorted[rank - 1])+        np.testing.assert_allclose(fcst[f"lo-{lv}"], -cs_sorted[rank - 1])+++@pytest.mark.parametrize("method", ["conformal_error", "conformal_distribution"])+def test_conformal_intervals_levels_below_50(method):+    # Below 50% the margin is still an upper bound, so the rank rounds up:+    # ceil(11 * 0.3) = 4, ceil(11 * 0.4) = 5 and ceil(11 * 0.49) = 6.+    conf_intervals = ConformalIntervals(h=12, n_windows=10, method=method)+    zero_model = ZeroModel(conf_intervals)+    fcst = zero_model.forecast(ap, h=12, level=[30, 40, 49])+    cs_sorted = np.sort(zero_model._conformity_scores(ap), axis=0)+    for lv, rank in [(30, 4), (40, 5), (49, 6)]:+        np.testing.assert_allclose(fcst[f"hi-{lv}"], cs_sorted[rank - 1])+        np.testing.assert_allclose(fcst[f"lo-{lv}"], -cs_sorted[rank - 1])++ def assert_class(     cls_,     x,
tests/test_utils.py75 + / 1 −
@@ -1,7 +1,8 @@ import warnings  import numpy as np-from statsforecast.utils import _seasonal_naive+import pytest+from statsforecast.utils import _conformal_quantile, _conformal_rank, _seasonal_naive   def test_seasonal_naive():@@ -71,3 +72,76 @@ def test_seasonal_naive_partial_season():      assert len(w) >= 1     assert "shorter than season_length" in str(w[-1].message)+++@pytest.mark.parametrize(+    "n, q, upper, expected",+    [+        # upper bound: ceil((n + 1) * q), capped at n+        (2, 0.8, True, 2),+        (9, 0.9, True, 9),+        (19, 0.9, True, 18),+        (20, 0.8, True, 17),+        (99, 0.95, True, 95),+        (100, 0.975, True, 99),+        # an upper bound below the median still rounds up+        (10, 0.4, True, 5),+        (10, 0.3, True, 4),+        # lower bound: floor((n + 1) * q), floored at 1+        (100, 0.025, False, 2),+        (100, 0.05, False, 5),+        (10, 0.1, False, 1),+        (30, 0.025, False, 1),+    ],+)+def test_conformal_rank(n, q, upper, expected):+    assert _conformal_rank(n, q, upper) == expected+++def test_conformal_quantile_is_an_exact_order_statistic():+    rng = np.random.default_rng(0)+    scores = rng.standard_normal((10, 3))+    srt = np.sort(scores, axis=0)+    # scalar level: the ceil(11 * 0.8) = 9th smallest, no interpolation+    np.testing.assert_array_equal(_conformal_quantile(scores, 0.8, upper=True), srt[8])+    # sequence of levels adds a leading axis, as np.quantile does+    out = _conformal_quantile(scores, [0.1, 0.9], [False, True])+    assert out.shape == (2, 3)+    np.testing.assert_array_equal(out, srt[[0, 9]])+    # other axes work the same way+    np.testing.assert_array_equal(+        _conformal_quantile(scores.T, 0.8, upper=True, axis=1), srt[8]+    )+++@pytest.mark.parametrize("n, level", [(10, 80), (20, 90), (50, 95)])+def test_conformal_quantile_tails_cover(n, level):+    # Each tail of a two-sided interval should miss at most (1 - level) / 2 of+    # new exchangeable draws. Linear interpolation misses the lower tail.+    rng = np.random.default_rng(1)+    reps = 20_000+    alpha = 1 - level / 100+    draws = rng.standard_normal((reps, n + 1))+    lo, hi = _conformal_quantile(+        draws[:, :n], [alpha / 2, 1 - alpha / 2], [False, True], axis=1+    )+    se = np.sqrt(alpha / 2 * (1 - alpha / 2) / reps)+    assert np.mean(draws[:, n] < lo) <= alpha / 2 + 3 * se+    assert np.mean(draws[:, n] > hi) <= alpha / 2 + 3 * se+++def test_conformal_quantile_rejects_empty_scores():+    with pytest.raises(ValueError, match="zero scores"):+        _conformal_quantile(np.empty((0, 2)), 0.9, upper=True)+++@pytest.mark.parametrize("level", [30, 40, 49])+def test_conformal_quantile_upper_bound_below_the_median_covers(level):+    # An absolute-score margin is an upper bound even when level < 50, so it+    # must round the rank up to cover at least `level` percent.+    rng = np.random.default_rng(2)+    n, reps = 10, 20_000+    scores = np.abs(rng.standard_normal((reps, n + 1)))+    margin = _conformal_quantile(scores[:, :n], level / 100, upper=True, axis=1)+    se = np.sqrt(0.25 / reps)+    assert np.mean(scores[:, n] <= margin) >= level / 100 - 3 * se