Allow exogenous variables in BayesianSARIMAX, rename SARIMA -> SARIMAX

Author: jessegrabowski

notebooks/SARMA Example.ipynb notebooks/SARIMAX Example.ipynbnotebooks/SARMA Example.ipynb renamed to notebooks/SARIMAX Example.ipynb

@@ -1,13 +1,13 @@
 from pymc_extras.statespace.core.compile import compile_statespace
 from pymc_extras.statespace.models import structural
 from pymc_extras.statespace.models.ETS import BayesianETS
-from pymc_extras.statespace.models.SARIMAX import BayesianSARIMA
+from pymc_extras.statespace.models.SARIMAX import BayesianSARIMAX
 from pymc_extras.statespace.models.VARMAX import BayesianVARMAX
 
 __all__ = [
     "compile_statespace",
     "structural",
     "BayesianETS",
-    "BayesianSARIMA",
+    "BayesianSARIMAX",
     "BayesianVARMAX",
 ]

pymc_extras/statespace/__init__.py

@@ -22,6 +22,7 @@
     SARIMAX_STATE_STRUCTURES,
     SEASONAL_AR_PARAM_DIM,
     SEASONAL_MA_PARAM_DIM,
+    TIME_DIM,
 )
 
 
@@ -38,13 +39,12 @@ def _verify_order(p, d, q, P, D, Q, S):
             )
 
 
-class BayesianSARIMA(PyMCStateSpace):
+class BayesianSARIMAX(PyMCStateSpace):
     r"""
     Seasonal AutoRegressive Integrated Moving Average with eXogenous regressors.
 
-    This class implements a Bayesian approach to SARIMA models, which are used for
-    modeling univariate time series data with seasonal and non-seasonal components.
-    The model supports exogenous regressors and can be represented in state-space form.
+    This class implements a Bayesian approach to SARIMAX models, which are used for modeling univariate time series data
+    with seasonal and non-seasonal components. The model supports exogenous regressors.
 
     Notes
     -----
@@ -98,14 +98,14 @@ class BayesianSARIMA(PyMCStateSpace):
 
     Examples
     --------
-    The following example shows how to build an ARMA(1, 1) model -- ARIMA(1, 0, 1) -- using the BayesianSARIMA class:
+    The following example shows how to build an ARMA(1, 1) model -- ARIMA(1, 0, 1) -- using the BayesianSARIMAX class:
 
     .. code:: python
 
         import pymc_extras.statespace as pmss
         import pymc as pm
 
-        ss_mod = pmss.BayesianSARIMA(order=(1, 0, 1), verbose=True)
+        ss_mod = pmss.BayesianSARIMAX(order=(1, 0, 1), verbose=True)
 
         with pm.Model(coords=ss_mod.coords) as arma_model:
             state_sigmas = pm.HalfNormal("sigma_state", sigma=1.0, dims=ss_mod.param_dims["sigma_state"])
@@ -140,11 +140,11 @@ def __init__(
         mode: str | Mode | None = None,
     ):
         """
-        Initialize a BayesianSARIMA model.
+        Initialize a BayesianSARIMAX model.
 
         Parameters
         ----------
-        order : tuple(int, int, int)
+        order : tuple of int, int, int
             Order of the ARIMA process. The order has the notation (p, d, q), where p is the number of autoregressive
             lags, q is the number of moving average components, and d is order of integration -- the number of
             differences needed to render the data stationary.
@@ -153,7 +153,7 @@ def __init__(
             This is only possible if state_structure = 'fast'. For interpretable states, the user must manually
             difference the data prior to calling the `build_statespace_graph` method.
 
-        seasonal_order : tuple(int, int, int, int), optional
+        seasonal_order : tuple of int, int, int, int, optional
             Seasonal order of the SARIMA process. The order has the notation (P, D, Q, S), where P is the number of seasonal
             lags to include, Q is the number of seasonal innovation lags to include, and D is the number of seasonal
             differences to perform. S is the length of the season.
@@ -232,6 +232,12 @@ def __init__(
 
         self.stationary_initialization = stationary_initialization
 
+        if (self.d or self.D) and self.stationary_initialization:
+            raise ValueError(
+                "Cannot use stationary initialization with differencing. "
+                "Set stationary_initialization=False."
+            )
+
         self.state_structure = state_structure
 
         self._p_max = max(1, self.p + self.P * self.S)
@@ -271,6 +277,7 @@ def __init__(
             measurement_error=measurement_error,
             mode=mode,
         )
+        self._needs_exog_data = self.k_exog > 0
 
     @property
     def param_names(self):
@@ -303,6 +310,17 @@ def param_names(self):
 
         return names
 
+    @property
+    def data_info(self) -> dict[str, dict[str, Any]]:
+        info = {
+            "exogenous_data": {
+                "dims": (TIME_DIM, "exogenous"),
+                "shape": (None, self.k_exog),
+            }
+        }
+
+        return {name: info[name] for name in self.data_names}
+
     @property
     def param_info(self) -> dict[str, dict[str, Any]]:
         info = {
@@ -332,7 +350,7 @@ def param_info(self) -> dict[str, dict[str, Any]]:
             },
             "seasonal_ar_params": {"shape": (self.P,), "constraints": "None"},
             "seasonal_ma_params": {"shape": (self.Q,), "constraints": "None"},
-            "beta_exog": {"shape": (self.k_exog,), "constraints": "None"},
+            "beta_exog": {"shape": (self.k_exog,), "constraints": "None", "dims": ("exogenous",)},
         }
 
         for name in self.param_names:
@@ -357,11 +375,14 @@ def state_names(self):
         else:
             raise NotImplementedError()
 
-        if self.k_exog > 0:
-            states += ["exogenous"]
-
         return states
 
+    @property
+    def data_names(self) -> list[str]:
+        if self.k_exog > 0:
+            return ["exogenous_data"]
+        return []
+
     @property
     def observed_states(self):
         return [self.state_names[0]]
@@ -396,7 +417,7 @@ def param_dims(self):
             del coord_map["seasonal_ar_params"]
         if self.Q == 0:
             del coord_map["seasonal_ma_params"]
-        if not self.k_exog == 0:
+        if self.k_exog == 0:
             del coord_map["beta_exog"]
         if self.stationary_initialization:
             del coord_map["P0"]
@@ -426,7 +447,7 @@ def _stationary_initialization(self):
         Q = self.ssm["state_cov"]
         c = self.ssm["state_intercept"]
 
-        x0 = pt.linalg.solve(pt.identity_like(T) - T, c, assume_a="gen", check_finite=True)
+        x0 = pt.linalg.solve(pt.identity_like(T) - T, c, assume_a="gen", check_finite=False)
         P0 = solve_discrete_lyapunov(T, pt.linalg.matrix_dot(R, Q, R.T), method="bilinear")
 
         return x0, P0
@@ -575,7 +596,7 @@ def make_symbolic_graph(self) -> None:
                 "beta_exog", shape=(self.k_exog,), dtype=floatX
             )
 
-            self.ssm["obs_intercept", :, :] = exog_data @ exog_beta
+            self.ssm["obs_intercept"] = (exog_data @ exog_beta)[:, None]
 
         # Set up the state covariance matrix
         state_cov_idx = ("state_cov", *np.diag_indices(self.k_posdef))

pymc_extras/statespace/models/SARIMAX.py

@@ -1,6 +1,6 @@
 from pymc_extras.statespace.models import structural
 from pymc_extras.statespace.models.ETS import BayesianETS
-from pymc_extras.statespace.models.SARIMAX import BayesianSARIMA
+from pymc_extras.statespace.models.SARIMAX import BayesianSARIMAX
 from pymc_extras.statespace.models.VARMAX import BayesianVARMAX
 
-__all__ = ["structural", "BayesianSARIMA", "BayesianVARMAX", "BayesianETS"]
+__all__ = ["structural", "BayesianSARIMAX", "BayesianVARMAX", "BayesianETS"]

pymc_extras/statespace/models/__init__.py

@@ -46,7 +46,7 @@ class AutoregressiveComponent(Component):
     The coefficient :math:`\rho_3` has been constrained to zero.
 
     .. warning:: This class is meant to be used as a component in a structural time series model. For modeling of
-              stationary processes with ARIMA, use ``statespace.BayesianSARIMA``.
+              stationary processes with ARIMA, use ``statespace.BayesianSARIMAX``.
 
     Examples
     --------

pymc_extras/statespace/models/structural/components/autoregressive.py

@@ -6,8 +6,6 @@
 from numpy.testing import assert_allclose
 from pytensor import config
 from pytensor import tensor as pt
-from pytensor.graph.basic import explicit_graph_inputs
-from scipy.linalg import block_diag
 
 from pymc_extras.statespace.models import structural as st
 from tests.statespace.models.structural.conftest import _assert_basic_coords_correct
@@ -327,13 +325,14 @@ def test_regression_mixed_shared_and_not_shared():
         Z,
         np.concat(
             (
-                block_diag(*[data_individual[:, np.newaxis] for _ in range(mod.k_endog)]),
-                np.concat((data_joint[:, np.newaxis], data_joint[:, np.newaxis]), axis=1),
+                pt.linalg.block_diag(
+                    *[data_individual[:, None] for _ in range(mod.k_endog)]
+                ).eval(),
+                np.concat((data_joint[:, None], data_joint[:, None]), axis=1),
             ),
             axis=2,
         ),
     )
 
     np.testing.assert_allclose(T, np.eye(mod.k_states))
-
     np.testing.assert_allclose(R, np.eye(mod.k_states))

tests/statespace/models/structural/components/test_regression.py

@@ -1,15 +1,18 @@
 from itertools import combinations
 
 import numpy as np
+import pandas as pd
 import pymc as pm
 import pytensor
 import pytensor.tensor as pt
 import pytest
 import statsmodels.api as sm
 
 from numpy.testing import assert_allclose, assert_array_less
+from pymc.testing import mock_sample_setup_and_teardown
+from pytensor.graph.basic import explicit_graph_inputs
 
-from pymc_extras.statespace import BayesianSARIMA
+from pymc_extras.statespace import BayesianSARIMAX
 from pymc_extras.statespace.models.utilities import (
     make_harvey_state_names,
     make_SARIMA_transition_matrix,
@@ -27,6 +30,8 @@
     simulate_from_numpy_model,
 )
 
+mock_sample = pytest.fixture()(mock_sample_setup_and_teardown)
+
 floatX = pytensor.config.floatX
 ATOL = 1e-8 if floatX.endswith("64") else 1e-6
 RTOL = 0 if floatX.endswith("64") else 1e-6
@@ -167,7 +172,7 @@ def data():
 
 @pytest.fixture(scope="session")
 def arima_mod():
-    return BayesianSARIMA(order=(2, 0, 1), stationary_initialization=True, verbose=False)
+    return BayesianSARIMAX(order=(2, 0, 1), stationary_initialization=True, verbose=False)
 
 
 @pytest.fixture(scope="session")
@@ -188,7 +193,7 @@ def pymc_mod(arima_mod):
 
 @pytest.fixture(scope="session")
 def arima_mod_interp():
-    return BayesianSARIMA(
+    return BayesianSARIMAX(
         order=(3, 0, 3),
         stationary_initialization=False,
         verbose=False,
@@ -219,7 +224,7 @@ def pymc_mod_interp(arima_mod_interp):
 
 def test_mode_argument():
     # Mode argument should be passed to the parent class
-    mod = BayesianSARIMA(order=(0, 0, 3), mode="FAST_RUN", verbose=False)
+    mod = BayesianSARIMAX(order=(0, 0, 3), mode="FAST_RUN", verbose=False)
     assert mod.mode == "FAST_RUN"
 
 
@@ -269,7 +274,7 @@ def test_SARIMAX_update_matches_statsmodels(p, d, q, P, D, Q, S, data, rng):
     param_d = {name: getattr(np, floatX)(rng.normal(scale=0.1) ** 2) for name in param_names}
 
     res = sm_sarimax.fit_constrained(param_d)
-    mod = BayesianSARIMA(
+    mod = BayesianSARIMAX(
         order=(p, d, q), seasonal_order=(P, D, Q, S), verbose=False, stationary_initialization=False
     )
 
@@ -331,9 +336,9 @@ def test_interpretable_raises_if_d_nonzero():
     with pytest.raises(
         ValueError, match="Cannot use interpretable state structure with statespace differencing"
     ):
-        BayesianSARIMA(
+        BayesianSARIMAX(
             order=(2, 1, 1),
-            stationary_initialization=True,
+            stationary_initialization=False,
             verbose=False,
             state_structure="interpretable",
         )
@@ -383,7 +388,7 @@ def test_representations_are_equivalent(p, d, q, P, D, Q, S, data, rng):
 
     for representation in SARIMAX_STATE_STRUCTURES:
         rng = np.random.default_rng(sum(map(ord, "representation test")))
-        mod = BayesianSARIMA(
+        mod = BayesianSARIMAX(
             order=(p, d, q),
             seasonal_order=(P, D, Q, S),
             stationary_initialization=False,
@@ -414,4 +419,52 @@ def test_representations_are_equivalent(p, d, q, P, D, Q, S, data, rng):
 def test_invalid_order_raises(order, name):
     p, P, q, Q = order
     with pytest.raises(ValueError, match=f"The following {name} and seasonal {name} terms overlap"):
-        BayesianSARIMA(order=(p, 0, q), seasonal_order=(P, 0, Q, 4))
+        BayesianSARIMAX(order=(p, 0, q), seasonal_order=(P, 0, Q, 4))
+
+
+def test_SARIMA_with_exogenous(rng, mock_sample):
+    ss_mod = BayesianSARIMAX(order=(3, 0, 1), seasonal_order=(1, 0, 0, 12), k_exog=2)
+
+    assert ss_mod.param_dims["beta_exog"] == ("exogenous",)
+    assert ss_mod.data_names == ["exogenous_data"]
+    assert ss_mod.coords["exogenous"] == ["exogenous_0", "exogenous_1"]
+
+    obs_intercept = ss_mod.ssm["obs_intercept"]
+    assert obs_intercept.type.shape == (None, ss_mod.k_endog)
+
+    intercept_fn = pytensor.function(
+        inputs=list(explicit_graph_inputs(obs_intercept)), outputs=obs_intercept
+    )
+    data_val = rng.normal(size=(100, 2)).astype(floatX)
+    beta_val = rng.normal(size=(2,)).astype(floatX)
+
+    intercept_val = intercept_fn(data_val, beta_val)
+    np.testing.assert_allclose(intercept_val, intercept_fn(data_val, beta_val))
+
+    data_df = pd.DataFrame(
+        rng.normal(size=(100, 1)),
+        index=pd.date_range(start="2020-01-01", periods=100, freq="D"),
+        columns=["endog"],
+    )
+
+    with pm.Model(coords=ss_mod.coords) as pymc_mod:
+        pm.Data("exogenous_data", data_val, dims=["time", "exogenous"])
+
+        ar_params = pm.Normal("ar_params", dims=["lag_ar"])
+        ma_params = pm.Normal("ma_params", dims=["lag_ma"])
+        seasonal_ar_params = pm.Normal("seasonal_ar_params", dims=["seasonal_lag_ar"])
+
+        beta_exog = pm.Normal("beta_exog", dims=["exogenous"])
+
+        sigma_state = pm.Exponential("sigma_state", 1.0)
+
+        ss_mod.build_statespace_graph(data=data_df, save_kalman_filter_outputs_in_idata=True)
+        idata = pm.sample(chains=2, draws=100)
+
+    assert "exogenous_data" in idata.constant_data
+    assert idata.posterior.beta_exog.shape == (
+        2,
+        100,
+        2,
+    )
+    np.testing.assert_allclose(ss_mod._fit_exog_data["exogenous_data"]["value"], data_val)