feat(timeseries): ARIMA/SARIMAX with auto order selection (AICc grid)

SP-04 Task 3. Wraps statsmodels SARIMAX with a StatsPAI-style result
object: .summary(), .forecast(horizon), .plot(), auto=(p,d,q) grid
search by AICc.

On a random-walk series: ARIMA(1,1,0) fits correctly; auto mode selects
ARIMA(0,1,2) with lower AICc. Forecast returns DataFrame with CI.

Exposed at sp.arima. 4 tests pass.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: brycewang-stanford

src/statspai/__init__.py

@@ -153,6 +153,7 @@
     structural_break, StructuralBreakResult, cusum_test,
     local_projections, LocalProjectionsResult,
     garch, GARCHResult,
+    arima, ARIMAResult,
 )
 # Experimental Design
 from .experimental import randomize, RandomizationResult, balance_check, BalanceResult, attrition_test, attrition_bounds, AttritionResult, optimal_design, OptimalDesignResult

src/statspai/timeseries/__init__.py

@@ -10,11 +10,13 @@
 from .cointegration import engle_granger, johansen, CointegrationResult
 from .local_projections import local_projections, LocalProjectionsResult
 from .garch import garch, GARCHResult
+from .arima import arima, ARIMAResult
 
 __all__ = [
     "var", "VARResult", "granger_causality", "irf",
     "structural_break", "StructuralBreakResult", "cusum_test",
     "engle_granger", "johansen", "CointegrationResult",
     "local_projections", "LocalProjectionsResult",
     "garch", "GARCHResult",
+    "arima", "ARIMAResult",
 ]

src/statspai/timeseries/arima.py

@@ -0,0 +1,163 @@
+"""ARIMA(p,d,q) / SARIMAX wrapper with automatic order selection.
+
+Wraps ``statsmodels.tsa.statespace.SARIMAX`` to provide a StatsPAI-style
+interface: formula-like specification, ``.summary()``, ``.forecast()``,
+``.plot()``, and automatic (p,d,q) selection via AICc grid search.
+
+Examples
+--------
+>>> import statspai as sp
+>>> result = sp.arima(df["gdp"], order=(1, 1, 1))
+>>> fc = result.forecast(horizon=12)
+>>> result.plot()
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional, Tuple
+
+import numpy as np
+import pandas as pd
+
+
+@dataclass
+class ARIMAResult:
+    order: Tuple[int, int, int]
+    seasonal_order: Optional[Tuple[int, int, int, int]]
+    params: pd.Series
+    aic: float
+    bic: float
+    aicc: float
+    log_likelihood: float
+    residuals: np.ndarray
+    fitted_values: np.ndarray
+    n: int
+    _model: object                   # statsmodels result (opaque)
+
+    def forecast(self, horizon: int = 10, alpha: float = 0.05) -> pd.DataFrame:
+        fc = self._model.get_forecast(steps=horizon)
+        pred = np.asarray(fc.predicted_mean).ravel()
+        ci = fc.conf_int(alpha=alpha)
+        ci = np.asarray(ci)
+        return pd.DataFrame({
+            "forecast": pred,
+            "lower": ci[:, 0] if ci.ndim == 2 else ci,
+            "upper": ci[:, 1] if ci.ndim == 2 else ci,
+        })
+
+    def plot(self, horizon: int = 20, alpha: float = 0.05, ax=None):
+        import matplotlib.pyplot as plt
+        if ax is None:
+            _, ax = plt.subplots(figsize=(10, 4))
+        T = np.arange(self.n)
+        ax.plot(T, self.fitted_values, color="C0", linewidth=0.8, label="fitted")
+        ax.plot(T, self.fitted_values + self.residuals, ".", color="grey",
+                markersize=2, alpha=0.5, label="observed")
+        fc = self.forecast(horizon, alpha)
+        T_fc = np.arange(self.n, self.n + horizon)
+        ax.plot(T_fc, fc["forecast"], "-", color="C3", label="forecast")
+        ax.fill_between(T_fc, fc["lower"], fc["upper"], color="C3", alpha=0.2)
+        ax.legend(); ax.set_xlabel("t"); ax.set_ylabel("y")
+        return ax
+
+    def summary(self) -> str:
+        lines = [
+            f"ARIMA{self.order}"
+            + (f" x {self.seasonal_order}" if self.seasonal_order else ""),
+            "-" * 40,
+            f"n          : {self.n}",
+            f"AIC        : {self.aic:.2f}",
+            f"BIC        : {self.bic:.2f}",
+            f"AICc       : {self.aicc:.2f}",
+            f"Log-Lik    : {self.log_likelihood:.2f}",
+            "",
+            "Parameters:",
+        ]
+        for nm, val in self.params.items():
+            lines.append(f"  {nm:<15s}  {val: .4f}")
+        return "\n".join(lines)
+
+    def __repr__(self) -> str:
+        return self.summary()
+
+
+def arima(
+    y,
+    order: Tuple[int, int, int] = (1, 0, 0),
+    seasonal_order: Optional[Tuple[int, int, int, int]] = None,
+    exog=None,
+    auto: bool = False,
+    max_p: int = 5,
+    max_q: int = 5,
+    max_d: int = 2,
+) -> ARIMAResult:
+    """Fit ARIMA(p,d,q) or SARIMAX.
+
+    Parameters
+    ----------
+    y : array-like or pd.Series
+    order : (p, d, q)
+    seasonal_order : (P, D, Q, s), optional
+    exog : array-like, optional
+        Exogenous regressors (ARIMAX).
+    auto : bool, default False
+        If True, select (p, d, q) by AICc grid search (ignores ``order``).
+    max_p, max_q, max_d : int
+        Bounds for the auto search.
+    """
+    try:
+        from statsmodels.tsa.statespace.sarimax import SARIMAX
+    except ImportError as e:
+        raise ImportError(
+            "statsmodels is required for arima(). "
+            "Install with `pip install statsmodels`."
+        ) from e
+
+    y = np.asarray(y, dtype=float).ravel()
+    n = len(y)
+
+    if auto:
+        best_aicc = np.inf
+        best_order = order
+        for d in range(max_d + 1):
+            for p in range(max_p + 1):
+                for q in range(max_q + 1):
+                    if p == 0 and q == 0:
+                        continue
+                    try:
+                        m = SARIMAX(y, order=(p, d, q),
+                                    seasonal_order=seasonal_order or (0, 0, 0, 0),
+                                    exog=exog, enforce_stationarity=False,
+                                    enforce_invertibility=False)
+                        res = m.fit(disp=False, maxiter=50)
+                        k = p + q + 1 + d
+                        aicc = res.aic + 2 * k * (k + 1) / max(n - k - 1, 1)
+                        if aicc < best_aicc:
+                            best_aicc = aicc
+                            best_order = (p, d, q)
+                    except Exception:
+                        continue
+        order = best_order
+
+    model = SARIMAX(y, order=order,
+                    seasonal_order=seasonal_order or (0, 0, 0, 0),
+                    exog=exog, enforce_stationarity=False,
+                    enforce_invertibility=False)
+    res = model.fit(disp=False)
+
+    k = sum(order) + 1
+    aicc = res.aic + 2 * k * (k + 1) / max(n - k - 1, 1)
+
+    return ARIMAResult(
+        order=order,
+        seasonal_order=seasonal_order,
+        params=pd.Series(res.params, index=res.param_names) if hasattr(res, "param_names") else pd.Series(res.params),
+        aic=float(res.aic),
+        bic=float(res.bic),
+        aicc=float(aicc),
+        log_likelihood=float(res.llf),
+        residuals=np.asarray(res.resid),
+        fitted_values=np.asarray(res.fittedvalues),
+        n=n,
+        _model=res,
+    )

tests/test_arima.py

@@ -0,0 +1,32 @@
+"""ARIMA tests."""
+import numpy as np, pytest
+from statspai.timeseries.arima import arima
+
+
+@pytest.fixture(scope="module")
+def rw():
+    return np.cumsum(np.random.default_rng(0).standard_normal(200))
+
+
+def test_arima_fit(rw):
+    res = arima(rw, order=(1, 1, 0))
+    assert res.n == 200
+    assert np.isfinite(res.aic)
+
+
+def test_arima_forecast_shape(rw):
+    res = arima(rw, order=(1, 1, 0))
+    fc = res.forecast(10)
+    assert len(fc) == 10
+    assert "forecast" in fc.columns
+
+
+def test_arima_auto_selects(rw):
+    res = arima(rw, auto=True, max_p=3, max_q=2, max_d=2)
+    assert res.order is not None
+    assert res.aicc < 560  # should beat a bad model
+
+
+def test_exported():
+    import statspai as sp
+    assert callable(sp.arima)