Nixtla/statsforecast · #1246
fix: use the finite-sample conformal quantile for conformal intervals
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_offsetpython/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