Adjust test suite to reflect API changes

Modify structural tests to accommodate deterministic models

Save kalman filter outputs to idata for statespace tests

Remove test related to `add_exogenous`

Adjust structural module tests

Author: jessegrabowski

pymc_experimental/tests/statespace/test_statespace.py

@@ -126,7 +126,7 @@ def exog_pymc_mod(exog_ss_mod, rng):
         beta_exog = pm.Normal("beta_exog", dims=["exog_state"])
 
         sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
-        exog_ss_mod.build_statespace_graph(y)
+        exog_ss_mod.build_statespace_graph(y, save_kalman_filter_outputs_in_idata=True)
 
     return m
 
@@ -299,34 +299,3 @@ def test_forecast_fails_if_exog_needed(exog_ss_mod, idata_exog):
         forecast_idata = exog_ss_mod.forecast(
             idata_exog, start=time_idx[-1], periods=10, random_seed=rng
         )
-
-
-@pytest.mark.parametrize(
-    "shape",
-    [
-        None,
-        (10,),
-        (10, 1),
-        pytest.param((10, 3), marks=[pytest.mark.xfail]),
-        pytest.param((10, 1, 1), marks=[pytest.mark.xfail]),
-    ],
-    ids=["None", "(10,)", "(10, 1)", "(10,3)", "(10, 1, 1)"],
-)
-def test_add_exogenous(rng, shape):
-    ss_mod = st.LevelTrendComponent(order=1, innovations_order=0).build(verbose=False)
-    y = rng.normal(size=(10, 1))
-
-    with pm.Model() as m:
-        initial_trend = pm.Normal("initial_trend", shape=(1,))
-        P0 = pm.Deterministic("P0", pt.eye(1, dtype=floatX))
-        constant = pm.Normal("constant", shape=shape)
-        ss_mod.add_exogenous(constant)
-        ss_mod.build_statespace_graph(data=y)
-
-        names = ["x0", "P0", "c", "d", "T", "Z", "R", "H", "Q"]
-        for name, matrix in zip(names, ss_mod.unpack_statespace()):
-            if name not in [x.name for x in m.deterministics]:
-                pm.Deterministic(name, matrix)
-
-    d, const = pm.draw([m["d"].squeeze(), m["constant"].squeeze()], 10)
-    assert_allclose(d, const)

pymc_experimental/tests/statespace/test_structural.py

@@ -64,19 +64,45 @@ def _assert_coord_shapes_match_matrices(mod, params):
         params["initial_state_cov"] = np.eye(mod.k_states)
 
     x0, P0, c, d, T, Z, R, H, Q = unpack_symbolic_matrices_with_params(mod, params)
+
     n_states = len(mod.coords[ALL_STATE_DIM])
-    n_shocks = len(mod.coords[SHOCK_DIM])
+
+    # There will always be one shock dimension -- dummies are inserted into fully deterministic models to avoid errors
+    # in the state space representation.
+    n_shocks = max(1, len(mod.coords[SHOCK_DIM]))
     n_obs = len(mod.coords[OBS_STATE_DIM])
 
-    assert x0.shape[-1:] == (n_states,)
-    assert P0.shape[-2:] == (n_states, n_states)
-    assert c.shape[-1:] == (n_states,)
-    assert d.shape[-1:] == (n_obs,)
-    assert T.shape[-2:] == (n_states, n_states)
-    assert Z.shape[-2:] == (n_obs, n_states)
-    assert R.shape[-2:] == (n_states, n_shocks)
-    assert H.shape[-2:] == (n_obs, n_obs)
-    assert Q.shape[-2:] == (n_shocks, n_shocks)
+    assert x0.shape[-1:] == (
+        n_states,
+    ), f"x0 expected to have shape (n_states, ), found {x0.shape[-1:]}"
+    assert P0.shape[-2:] == (
+        n_states,
+        n_states,
+    ), f"P0 expected to have shape (n_states, n_states), found {P0.shape[-2:]}"
+    assert c.shape[-1:] == (
+        n_states,
+    ), f"c expected to have shape (n_states, ), found {c.shape[-1:]}"
+    assert d.shape[-1:] == (n_obs,), f"d expected to have shape (n_obs, ), found {d.shape[-1:]}"
+    assert T.shape[-2:] == (
+        n_states,
+        n_states,
+    ), f"T expected to have shape (n_states, n_states), found {T.shape[-2:]}"
+    assert Z.shape[-2:] == (
+        n_obs,
+        n_states,
+    ), f"Z expected to have shape (n_obs, n_states), found {Z.shape[-2:]}"
+    assert R.shape[-2:] == (
+        n_states,
+        n_shocks,
+    ), f"R expected to have shape (n_states, n_shocks), found {R.shape[-2:]}"
+    assert H.shape[-2:] == (
+        n_obs,
+        n_obs,
+    ), f"H expected to have shape (n_obs, n_obs), found {H.shape[-2:]}"
+    assert Q.shape[-2:] == (
+        n_shocks,
+        n_shocks,
+    ), f"Q expected to have shape (n_shocks, n_shocks), found {Q.shape[-2:]}"
 
 
 def _assert_basic_coords_correct(mod):
@@ -514,12 +540,12 @@ def test_structural_model_against_statsmodels(
 
     _assert_all_statespace_matrices_match(mod, params, sm_mod)
 
-    mod.build(verbose=False)
+    built_model = mod.build(verbose=False)
 
-    _assert_coord_shapes_match_matrices(mod, params)
-    _assert_param_dims_correct(mod.param_dims, expected_dims)
-    _assert_coords_correct(mod.coords, expected_coords)
-    _assert_params_info_correct(mod.param_info, mod.coords, mod.param_dims)
+    _assert_coord_shapes_match_matrices(built_model, params)
+    _assert_param_dims_correct(built_model.param_dims, expected_dims)
+    _assert_coords_correct(built_model.coords, expected_coords)
+    _assert_params_info_correct(built_model.param_info, built_model.coords, built_model.param_dims)
 
 
 def test_level_trend_model(rng):

pymc_experimental/tests/statespace/utilities/test_helpers.py

@@ -225,7 +225,9 @@ def unpack_symbolic_matrices_with_params(mod, param_dict, data_dict=None, mode="
         on_unused_input="raise",
         mode=mode,
     )
+
     x0, P0, c, d, T, Z, R, H, Q = f_matrices(**param_dict, **data_dict)
+
     return x0, P0, c, d, T, Z, R, H, Q