[ENH] Add prediction intervals to ARIMA iterative_forecast (#3100)

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)

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)