From 024b07b15877332e65816a7540bfafdb07e51e93 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Sat, 15 Nov 2025 14:46:20 +0530 Subject: [PATCH] [ENH] Add prediction intervals to ARIMA iterative_forecast (#3100) --- aeon/forecasting/stats/_arima.py | 59 +++++++++++++++++++--- aeon/forecasting/stats/tests/test_arima.py | 23 +++++++++ 2 files changed, 74 insertions(+), 8 deletions(-) diff --git a/aeon/forecasting/stats/_arima.py b/aeon/forecasting/stats/_arima.py index f1d62eb7f1..6e5b2783ba 100644 --- a/aeon/forecasting/stats/_arima.py +++ b/aeon/forecasting/stats/_arima.py @@ -8,7 +8,9 @@ __all__ = ["ARIMA", "AutoARIMA"] import numpy as np +import pandas as pd from numba import njit +from scipy.stats import norm from aeon.forecasting.base import BaseForecaster, IterativeForecastingMixin from aeon.forecasting.utils._extract_paras import _extract_arma_params @@ -207,7 +209,26 @@ def _forecast(self, y, exog=None): self._fit(y, exog) return float(self.forecast_) - def iterative_forecast(self, y, prediction_horizon): + def iterative_forecast(self, y, prediction_horizon, alpha=None): + """Perform iterative (recursive) multi-step forecast. + + Parameters + ---------- + y : np.ndarray + The time series to make forecasts about. + prediction_horizon : int + Number of future time steps to forecast. + alpha : float, optional (default=None) + Significance level for prediction intervals. If provided, returns + lower and upper prediction bounds in addition to the mean forecast. + + Returns + ------- + np.ndarray or pandas.DataFrame + If alpha is None, returns a NumPy array of shape `(prediction_horizon,)`. + If alpha is provided, returns a DataFrame + with columns ["mean", "lower", "upper"]. + """ self.fit(y) n = len(self._differenced_series) p, q = self.p, self.q @@ -240,9 +261,26 @@ def iterative_forecast(self, y, prediction_horizon): # Correct differencing using forecast values y_forecast_diff = forecast_series[n : n + h] if self.d == 0: - return y_forecast_diff + preds = y_forecast_diff else: - return _undifference(y_forecast_diff, self._series[-self.d :])[self.d :] + preds = _undifference(y_forecast_diff, self._series[-self.d :])[self.d :] + + # If alpha is not given, return exactly as before + if alpha is None: + return preds + + # Residual standard deviation + sigma = np.std(self.residuals_) if len(self.residuals_) > 1 else 0.0 + + # z multiplier + z = norm.ppf(1 - alpha / 2) + + # SE grows with sqrt(horizon) + se = sigma * np.sqrt(np.arange(1, h + 1)) + lower = preds - z * se + upper = preds + z * se + + return pd.DataFrame({"mean": preds, "lower": lower, "upper": upper}) class AutoARIMA(BaseForecaster, IterativeForecastingMixin): @@ -364,7 +402,7 @@ def _forecast(self, y, exog=None): self._fit(y, exog) return float(self.final_model_.forecast_) - def iterative_forecast(self, y, prediction_horizon): + def iterative_forecast(self, y, prediction_horizon, alpha=None): """Forecast ``prediction_horizon`` prediction using a single model fit on `y`. This function implements the iterative forecasting strategy (also called @@ -380,19 +418,24 @@ def iterative_forecast(self, y, prediction_horizon): ``(n_channels, n_timepoints)`` if a multivariate time series. prediction_horizon : int The number of future time steps to forecast. + alpha : float, optional (default=None) + Significance level for prediction intervals. If provided, returns + lower and upper prediction bounds in addition to the mean forecast. Returns ------- - np.ndarray - An array of shape `(prediction_horizon,)` containing the forecasts for - each horizon. + np.ndarray or pandas.DataFrame + If alpha is None, returns a NumPy array of shape `(prediction_horizon,)` + containing point forecasts. + If alpha is provided, returns a DataFrame + with columns "mean", "lower", and "upper". Raises ------ ValueError if prediction_horizon` less than 1. """ - return self.final_model_.iterative_forecast(y, prediction_horizon) + return self.final_model_.iterative_forecast(y, prediction_horizon, alpha=alpha) @njit(cache=True, fastmath=True) diff --git a/aeon/forecasting/stats/tests/test_arima.py b/aeon/forecasting/stats/tests/test_arima.py index e1bbf9825b..2809addb15 100644 --- a/aeon/forecasting/stats/tests/test_arima.py +++ b/aeon/forecasting/stats/tests/test_arima.py @@ -1,6 +1,7 @@ """Test the ARIMA forecaster.""" import numpy as np +import pandas as pd import pytest from aeon.forecasting.stats._arima import ARIMA, AutoARIMA @@ -48,6 +49,18 @@ def test_arima_iterative_forecast(): assert preds.shape == (horizon,) +def test_arima_iterative_forecast_with_alpha(): + """Test multi-step forecasting using iterative_forecast method with alpha.""" + model = ARIMA(p=1, d=1, q=1) + horizon = 3 + out = model.iterative_forecast(y, prediction_horizon=horizon, alpha=0.1) + assert isinstance(out, pd.DataFrame) + assert list(out.columns) == ["mean", "lower", "upper"] + assert out.shape == (horizon, 3) + assert (out["lower"] <= out["mean"]).all() + assert (out["upper"] >= out["mean"]).all() + + @pytest.mark.parametrize( "y_input, error_match", [ @@ -164,6 +177,16 @@ def test_autoarima_iterative_forecast_shape_and_validity(): assert np.all(np.isfinite(preds)) +def test_autoarima_iterative_forecast_with_alpha(): + """AutoARIMA should forward alpha to the underlying ARIMA model.""" + forecaster = AutoARIMA() + forecaster.fit(y) + out = forecaster.iterative_forecast(y, prediction_horizon=2, alpha=0.2) + assert isinstance(out, pd.DataFrame) + assert list(out.columns) == ["mean", "lower", "upper"] + assert out.shape == (2, 3) + + def test_autoarima_respects_small_max_orders(): """With small max orders, ensure discovered orders don’t exceed those limits.""" forecaster = AutoARIMA(max_p=1, max_d=1, max_q=1)