Add component name to shock state names

Author: jessegrabowski

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

@@ -100,7 +100,7 @@ def populate_component_properties(self):
             for i in range(k_states)
         ]
 
-        self.shock_names = self.observed_state_names.copy()
+        self.shock_names = [f"{self.name}[{obs_name}]" for obs_name in self.observed_state_names]
         self.param_names = [f"{self.name}_params", f"{self.name}_sigma"]
         self.param_dims = {f"{self.name}_params": (f"{self.name}_lag",)}
         self.coords = {f"{self.name}_lag": self.ar_lags.tolist()}

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

@@ -13,13 +13,7 @@
 
 @pytest.mark.parametrize("order", [1, 2, [1, 0, 1]], ids=["AR1", "AR2", "AR(1,0,1)"])
 def test_autoregressive_model(order, rng):
-    k = sum(order) if isinstance(order, list) else order
     ar = st.AutoregressiveComponent(order=order).build(verbose=False)
-    params = {
-        "auto_regressive_params": np.full((k,), 0.5, dtype=config.floatX),
-        "auto_regressive_sigma": 0.1,
-        "initial_state_cov": np.eye(k),
-    }
 
     # Check coords
     _assert_basic_coords_correct(ar)
@@ -47,7 +41,7 @@ def test_autoregressive_multiple_observed_build(rng):
         "L3[data_2]",
     ]
 
-    assert mod.shock_names == ["data_1", "data_2"]
+    assert mod.shock_names == ["auto_regressive[data_1]", "auto_regressive[data_2]"]
 
     params = {
         "auto_regressive_params": np.full(
@@ -133,6 +127,6 @@ def test_add_autoregressive_different_observed():
         "L6[data_2]",
     ]
 
-    assert mod.shock_names == ["data_1", "data_2"]
+    assert mod.shock_names == ["ar1[data_1]", "ar6[data_2]"]
     assert mod.coords["ar1_lag"] == [1]
     assert mod.coords["ar6_lag"] == [1, 2, 3, 4, 5, 6]

tests/statespace/models/structural/test_against_statsmodels.py

@@ -416,8 +416,8 @@ def create_structural_model_and_equivalent_statsmodel(
         expected_coords[AR_PARAM_DIM] += tuple(list(range(1, autoregressive + 1)))
         expected_coords[ALL_STATE_DIM] += ar_names
         expected_coords[ALL_STATE_AUX_DIM] += ar_names
-        expected_coords[SHOCK_DIM] += ["data"]
-        expected_coords[SHOCK_AUX_DIM] += ["data"]
+        expected_coords[SHOCK_DIM] += ["ar"]
+        expected_coords[SHOCK_AUX_DIM] += ["ar"]
 
         sm_params["sigma2.ar"] = sigma2
         for i, rho in enumerate(ar_params):