Fix bug in Statespace model state names (#634)

* Fix ETS states

* Fix DFM states

* Add workflow tests for SARIMAX and VARMAX

Author: jessegrabowski

pymc_extras/statespace/models/DFM.py

@@ -486,7 +486,12 @@ def set_states(self) -> State | tuple[State, ...] | None:
                     for endog_name in self.endog_names
                 )
 
-        return tuple(State(name=name, observed=False, shared=False) for name in names)
+        hidden_states = [State(name=name, observed=False, shared=False) for name in names]
+        observed_states = [
+            State(name=name, observed=True, shared=False) for name in self.endog_names
+        ]
+
+        return *hidden_states, *observed_states
 
     def set_shocks(self) -> Shock | tuple[Shock, ...] | None:
         shock_names = [f"factor_shock_{i}" for i in range(self.k_factors)]
@@ -549,13 +554,6 @@ def set_coords(self) -> Coord | tuple[Coord, ...] | None:
 
         return tuple(coords)
 
-    @property
-    def observed_states(self) -> tuple[str, ...]:
-        """
-        Returns the names of the observed states (i.e., the endogenous variables).
-        """
-        return self.endog_names
-
     def make_symbolic_graph(self):
         if not self.exog_flag:
             x0 = self.make_and_register_variable("x0", shape=(self.k_states,), dtype=floatX)

pymc_extras/statespace/models/ETS.py

@@ -446,15 +446,13 @@ def set_states(self) -> State | tuple[State, ...] | None:
         else:
             state_names = base_states
 
-        # First state for each endog is the innovation (observed), rest are hidden
-        states = []
-        states_per_endog = len(base_states)
-        for i, name in enumerate(state_names):
-            # innovation states are "observed" in the sense they directly affect the observation
-            is_observed = (i % states_per_endog) == 0
-            states.append(State(name=name, observed=is_observed, shared=False))
-
-        return tuple(states)
+        hidden_states = [State(name=name, observed=False, shared=False) for name in state_names]
+
+        observed_states = [
+            State(name=name, observed=True, shared=False) for name in self.endog_names
+        ]
+
+        return *hidden_states, *observed_states
 
     def set_shocks(self) -> Shock | tuple[Shock, ...] | None:
         k_endog = self.k_endog

tests/statespace/models/test_DFM.py

@@ -9,6 +9,7 @@
 import statsmodels.api as sm
 
 from numpy.testing import assert_allclose
+from pymc.testing import mock_sample_setup_and_teardown
 from pytensor.graph.traversal import explicit_graph_inputs
 from statsmodels.tsa.statespace.dynamic_factor import DynamicFactor
 
@@ -28,6 +29,8 @@
 )
 from tests.statespace.shared_fixtures import rng
 
+mock_sample = pytest.fixture(scope="function")(mock_sample_setup_and_teardown)
+
 floatX = pytensor.config.floatX
 
 
@@ -718,3 +721,50 @@ def test_exog_not_shared_no_exog_innovations(self):
         assert len(mod.shock_names) == k_factors + k_endog + (
             k_exog if shared_exog_states else k_exog * k_endog
         )
+
+
+def test_dfm_workflow(rng, mock_sample):
+    df = pd.read_csv(
+        "tests/statespace/_data/statsmodels_macrodata_processed.csv",
+        index_col=0,
+        parse_dates=True,
+    ).astype(floatX)
+    df.index.freq = df.index.inferred_freq
+
+    ss_mod = BayesianDynamicFactor(
+        endog_names=df.columns.tolist(),
+        k_factors=1,
+        factor_order=1,
+        error_order=0,
+        measurement_error=True,
+        verbose=False,
+    )
+
+    with pm.Model(coords=ss_mod.coords) as m:
+        pm.Normal("x0", dims=["state"])
+        P0_diag = pm.Exponential("P0_diag", 1, dims=["state"])
+        pm.Deterministic("P0", pt.diag(P0_diag), dims=["state", "state_aux"])
+
+        pm.Normal("factor_loadings", dims=["observed_state", "factor"])
+        pm.Normal("factor_ar", dims=["factor", "lag_ar"])
+        pm.Exponential("error_sigma", 1, dims=["observed_state"])
+        pm.Exponential("sigma_obs", 1, dims=["observed_state"])
+
+        ss_mod.build_statespace_graph(df)
+
+        idata = pm.sample()
+
+    post = ss_mod.sample_conditional_posterior(idata, mvn_method="svd")
+    assert "filtered_posterior" in post
+    assert "smoothed_posterior" in post
+    assert "predicted_posterior" in post
+
+    forecast = ss_mod.forecast(idata, periods=10, random_seed=rng)
+    assert "forecast_latent" in forecast
+    assert "forecast_observed" in forecast
+    assert np.isfinite(forecast.forecast_latent.values).all()
+    assert np.isfinite(forecast.forecast_observed.values).all()
+
+    irf = ss_mod.impulse_response_function(idata, n_steps=10, random_seed=rng)
+    assert "irf" in irf
+    assert np.isfinite(irf.irf.values).all()

tests/statespace/models/test_ETS.py

@@ -1,9 +1,11 @@
 import numpy as np
+import pymc as pm
 import pytensor
 import pytest
 import statsmodels.api as sm
 
 from numpy.testing import assert_allclose
+from pymc.testing import mock_sample_setup_and_teardown
 from pytensor.graph.traversal import explicit_graph_inputs
 from scipy import linalg
 
@@ -12,6 +14,8 @@
 from tests.statespace.shared_fixtures import rng
 from tests.statespace.test_utilities import load_nile_test_data
 
+mock_sample = pytest.fixture(scope="function")(mock_sample_setup_and_teardown)
+
 
 @pytest.fixture(scope="session")
 def data():
@@ -419,3 +423,44 @@ def test_ETS_stationary_initialization():
     P0_expected = linalg.solve_discrete_lyapunov(T_stationary, R @ Q @ R.T)
 
     assert_allclose(outputs["initial_state_cov"], P0_expected, rtol=1e-8, atol=1e-8)
+
+
+def test_ets_workflow(mock_sample):
+    data = load_nile_test_data()
+
+    ss_mod = BayesianETS(
+        order=("A", "Ad", "N"),
+        endog_names=["height"],
+        stationary_initialization=True,
+        measurement_error=True,
+        initialization_dampening=0.8,
+    )
+
+    with pm.Model(coords=ss_mod.coords) as m:
+        pm.Normal("initial_level", 0, 1)
+        pm.Normal("initial_trend", 0, 1)
+        pm.Beta("alpha", 1, 1)
+        pm.Beta("beta", 1, 1)
+        pm.Beta("phi", 1, 1)
+
+        pm.Exponential("sigma_state", 1)
+        pm.Exponential("sigma_obs", 1)
+
+        ss_mod.build_statespace_graph(data)
+
+        idata = pm.sample()
+
+    post = ss_mod.sample_conditional_posterior(idata, mvn_method="cholesky")
+    assert "filtered_posterior" in post
+    assert "smoothed_posterior" in post
+    assert "predicted_posterior" in post
+
+    forecast = ss_mod.forecast(idata, periods=10, random_seed=42)
+    assert "forecast_latent" in forecast
+    assert "forecast_observed" in forecast
+    assert np.isfinite(forecast.forecast_latent.values).all()
+    assert np.isfinite(forecast.forecast_observed.values).all()
+
+    irf = ss_mod.impulse_response_function(idata, n_steps=10, random_seed=42)
+    assert "irf" in irf
+    assert np.isfinite(irf.irf.values).all()

tests/statespace/models/test_SARIMAX.py

@@ -469,3 +469,39 @@ def test_SARIMA_with_exogenous(rng, mock_sample):
         2,
     )
     np.testing.assert_allclose(ss_mod._fit_exog_data["exogenous_data"]["value"], data_val)
+
+
+def test_sarimax_workflow(mock_sample):
+    data = load_nile_test_data()
+
+    ss_mod = BayesianSARIMAX(
+        order=(1, 0, 1),
+        stationary_initialization=True,
+        measurement_error=True,
+        verbose=False,
+    )
+
+    with pm.Model(coords=ss_mod.coords) as m:
+        pm.Normal("ar_params", dims=["lag_ar"])
+        pm.Normal("ma_params", dims=["lag_ma"])
+        pm.Exponential("sigma_state", 1)
+        pm.Exponential("sigma_obs", 1)
+
+        ss_mod.build_statespace_graph(data)
+
+        idata = pm.sample()
+
+    post = ss_mod.sample_conditional_posterior(idata, mvn_method="svd")
+    assert "filtered_posterior" in post
+    assert "smoothed_posterior" in post
+    assert "predicted_posterior" in post
+
+    forecast = ss_mod.forecast(idata, periods=10, random_seed=42)
+    assert "forecast_latent" in forecast
+    assert "forecast_observed" in forecast
+    assert np.isfinite(forecast.forecast_latent.values).all()
+    assert np.isfinite(forecast.forecast_observed.values).all()
+
+    irf = ss_mod.impulse_response_function(idata, n_steps=10, random_seed=42)
+    assert "irf" in irf
+    assert np.isfinite(irf.irf.values).all()

tests/statespace/models/test_VARMAX.py

@@ -10,13 +10,16 @@
 
 from numpy.testing import assert_allclose, assert_array_less
 from pymc.model.transform.optimization import freeze_dims_and_data
+from pymc.testing import mock_sample_setup_and_teardown
 
 from pymc_extras.statespace import BayesianVARMAX
 from pymc_extras.statespace.utils.constants import SHORT_NAME_TO_LONG
 from tests.statespace.shared_fixtures import (  # pylint: disable=unused-import
     rng,
 )
 
+mock_sample = pytest.fixture(scope="function")(mock_sample_setup_and_teardown)
+
 floatX = pytensor.config.floatX
 ps = [0, 1, 2, 3]
 qs = [0, 1, 2, 3]
@@ -200,6 +203,48 @@ def test_forecast(varma_mod, idata, rng):
     assert np.isfinite(forecast.forecast_observed.values).all()
 
 
+def test_varmax_workflow(rng, mock_sample):
+    df = pd.read_csv(
+        "tests/statespace/_data/statsmodels_macrodata_processed.csv",
+        index_col=0,
+        parse_dates=True,
+    ).astype(floatX)
+    df.index.freq = df.index.inferred_freq
+
+    ss_mod = BayesianVARMAX(
+        endog_names=df.columns,
+        order=(1, 0),
+        stationary_initialization=True,
+        measurement_error=True,
+        verbose=False,
+    )
+
+    with pm.Model(coords=ss_mod.coords) as m:
+        state_cov_diag = pm.Exponential("state_cov_diag", 1, dims=["shock"])
+        pm.Deterministic("state_cov", pt.diag(state_cov_diag), dims=["shock", "shock_aux"])
+        pm.Normal("ar_params", sigma=0.1, dims=["observed_state", "lag_ar", "observed_state_aux"])
+        pm.Exponential("sigma_obs", 1, dims=["observed_state"])
+
+        ss_mod.build_statespace_graph(df)
+
+        idata = pm.sample()
+
+    post = ss_mod.sample_conditional_posterior(idata, mvn_method="svd")
+    assert "filtered_posterior" in post
+    assert "smoothed_posterior" in post
+    assert "predicted_posterior" in post
+
+    forecast = ss_mod.forecast(idata, periods=10, random_seed=rng)
+    assert "forecast_latent" in forecast
+    assert "forecast_observed" in forecast
+    assert np.isfinite(forecast.forecast_latent.values).all()
+    assert np.isfinite(forecast.forecast_observed.values).all()
+
+    irf = ss_mod.impulse_response_function(idata, n_steps=10, random_seed=rng)
+    assert "irf" in irf
+    assert np.isfinite(irf.irf.values).all()
+
+
 class TestVARMAXWithExogenous:
     def test_create_varmax_with_exogenous_list_of_names(self, data):
         mod = BayesianVARMAX(