Fix coefs_ to params_ renaming in broader tests

Author: AlexAndorra

pymc_extras/statespace/models/structural/core.py

@@ -120,7 +120,7 @@ class StructuralTimeSeries(PyMCStateSpace):
             initial_trend = pm.Normal('initial_trend', sigma=10, dims=ss_mod.param_dims['initial_trend'])
             sigma_trend = pm.HalfNormal('sigma_trend', sigma=1, dims=ss_mod.param_dims['sigma_trend'])
 
-            seasonal_coefs = pm.Normal('seasonal_coefs', sigma=1, dims=ss_mod.param_dims['seasonal_coefs'])
+            seasonal_coefs = pm.Normal('params_seasonal', sigma=1, dims=ss_mod.param_dims['params_seasonal'])
             sigma_seasonal = pm.HalfNormal('sigma_seasonal', sigma=1)
 
             sigma_obs = pm.Exponential('sigma_obs', 1, dims=ss_mod.param_dims['sigma_obs'])

tests/statespace/models/structural/test_against_statsmodels.py

@@ -67,6 +67,11 @@ def _assert_coord_shapes_match_matrices(mod, params):
     n_shocks = max(1, len(mod.coords[SHOCK_DIM]))
     n_obs = len(mod.coords[OBS_STATE_DIM])
 
+    print(f"{mod.coords[ALL_STATE_DIM] = }")
+    print(f"{mod.coords[SHOCK_DIM] = }")
+    print(f"{mod.coords[OBS_STATE_DIM] = }")
+    print(f"{R = }")
+
     assert x0.shape[-1:] == (
         n_states,
     ), f"x0 expected to have shape (n_states, ), found {x0.shape[-1:]}"
@@ -104,12 +109,12 @@ def _assert_keys_match(test_dict, expected_dict):
     expected_keys = list(expected_dict.keys())
     param_keys = list(test_dict.keys())
     key_diff = set(expected_keys) - set(param_keys)
-    assert len(key_diff) == 0, f'{", ".join(key_diff)} were not found in the test_dict keys.'
+    assert len(key_diff) == 0, f"{', '.join(key_diff)} were not found in the test_dict keys."
 
     key_diff = set(param_keys) - set(expected_keys)
     assert (
         len(key_diff) == 0
-    ), f'{", ".join(key_diff)} were keys of the tests_dict not in expected_dict.'
+    ), f"{', '.join(key_diff)} were keys of the tests_dict not in expected_dict."
 
 
 def _assert_param_dims_correct(param_dims, expected_dims):
@@ -296,8 +301,8 @@ def create_structural_model_and_equivalent_statsmodel(
     if seasonal is not None:
         state_names = [f"seasonal_{i}" for i in range(seasonal)][1:]
         seasonal_coefs = rng.normal(size=(seasonal - 1,)).astype(floatX)
-        params["coefs_seasonal"] = seasonal_coefs
-        expected_param_dims["coefs_seasonal"] += ("state_seasonal",)
+        params["params_seasonal"] = seasonal_coefs
+        expected_param_dims["params_seasonal"] += ("state_seasonal",)
 
         expected_coords["state_seasonal"] += tuple(state_names)
         expected_coords[ALL_STATE_DIM] += state_names
@@ -335,8 +340,8 @@ def create_structural_model_and_equivalent_statsmodel(
 
             seasonal_params = rng.normal(size=n_states).astype(floatX)
 
-            params[f"seasonal_{s}"] = seasonal_params
-            expected_param_dims[f"seasonal_{s}"] += (f"state_seasonal_{s}",)
+            params[f"params_seasonal_{s}"] = seasonal_params
+            expected_param_dims[f"params_seasonal_{s}"] += (f"state_seasonal_{s}",)
             expected_coords[ALL_STATE_DIM] += state_names
             expected_coords[ALL_STATE_AUX_DIM] += state_names
             expected_coords[f"state_seasonal_{s}"] += (
@@ -404,7 +409,7 @@ def create_structural_model_and_equivalent_statsmodel(
         components.append(comp)
 
     if autoregressive is not None:
-        ar_names = [f"L{i+1}" for i in range(autoregressive)]
+        ar_names = [f"L{i + 1}" for i in range(autoregressive)]
         params_ar = rng.normal(size=(autoregressive,)).astype(floatX)
         if autoregressive == 1:
             params_ar = params_ar.item()
@@ -421,8 +426,8 @@ def create_structural_model_and_equivalent_statsmodel(
 
         sm_params["sigma2.ar"] = sigma2
         for i, rho in enumerate(params_ar):
-            sm_init[f"ar.L{i+1}"] = 0
-            sm_params[f"ar.L{i+1}"] = rho
+            sm_init[f"ar.L{i + 1}"] = 0
+            sm_params[f"ar.L{i + 1}"] = rho
 
         comp = st.AutoregressiveComponent(name="ar", order=autoregressive)
         components.append(comp)
@@ -439,7 +444,7 @@ def create_structural_model_and_equivalent_statsmodel(
 
         for i, beta in enumerate(betas):
             sm_params[f"beta.x{i + 1}"] = beta
-            sm_init[f"beta.x{i+1}"] = beta
+            sm_init[f"beta.x{i + 1}"] = beta
         comp = st.RegressionComponent(name="exog", state_names=names)
         components.append(comp)
 

tests/statespace/models/structural/test_core.py

@@ -28,7 +28,7 @@ def test_add_components():
         "sigma_level_trend": np.ones(2, dtype=floatX),
     }
     se_params = {
-        "coefs_seasonal": np.ones(11, dtype=floatX),
+        "params_seasonal": np.ones(11, dtype=floatX),
         "sigma_seasonal": 1.0,
     }
     all_params = ll_params.copy()
@@ -97,7 +97,7 @@ def test_extract_components_from_idata(rng):
         beta_exog = pm.Normal("beta_exog", dims=["state_exog"])
         initial_trend = pm.Normal("initial_level_trend", dims=["state_level_trend"])
         sigma_trend = pm.Exponential("sigma_level_trend", 1, dims=["shock_level_trend"])
-        seasonal_coefs = pm.Normal("seasonal", dims=["state_seasonal"])
+        seasonal_coefs = pm.Normal("params_seasonal", dims=["state_seasonal"])
         sigma_obs = pm.Exponential("sigma_obs", 1)
 
         mod.build_statespace_graph(y)
@@ -144,7 +144,7 @@ def test_extract_multiple_observed(rng):
         sigma_auto_regressive = pm.Normal("sigma_auto_regressive", dims=["endog_auto_regressive"])
         initial_trend = pm.Normal("initial_trend", dims=["endog_trend", "state_trend"])
         sigma_trend = pm.Exponential("sigma_trend", 1, dims=["endog_trend", "shock_trend"])
-        seasonal_coefs = pm.Normal("seasonal", dims=["state_seasonal"])
+        seasonal_coefs = pm.Normal("params_seasonal", dims=["state_seasonal"])
         sigma_obs = pm.Exponential("sigma_obs", 1, dims=["endog_obs"])
 
         mod.build_statespace_graph(y)