Use nwe name order in autoregressive component

Author: jessegrabowski

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

@@ -101,31 +101,31 @@ def populate_component_properties(self):
         ]
 
         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()}
+        self.param_names = [f"params_{self.name}", f"sigma_{self.name}"]
+        self.param_dims = {f"params_{self.name}": (f"lag_{self.name}",)}
+        self.coords = {f"lag_{self.name}": self.ar_lags.tolist()}
 
         if self.k_endog > 1:
-            self.param_dims[f"{self.name}_params"] = (
-                f"{self.name}_endog",
+            self.param_dims[f"params_{self.name}"] = (
+                f"endog_{self.name}",
                 AR_PARAM_DIM,
             )
-            self.param_dims[f"{self.name}_sigma"] = (f"{self.name}_endog",)
+            self.param_dims[f"sigma_{self.name}"] = (f"endog_{self.name}",)
 
-            self.coords[f"{self.name}_endog"] = self.observed_state_names
+            self.coords[f"endog_{self.name}"] = self.observed_state_names
 
         self.param_info = {
-            f"{self.name}_params": {
+            f"params_{self.name}": {
                 "shape": (self.k_states,) if self.k_endog == 1 else (self.k_endog, self.k_states),
                 "constraints": None,
                 "dims": (AR_PARAM_DIM,)
                 if self.k_endog == 1
                 else (
-                    f"{self.name}_endog",
+                    f"endog_{self.name}",
                     AR_PARAM_DIM,
                 ),
             },
-            f"{self.name}_sigma": {
+            f"sigma_{self.name}": {
                 "shape": () if self.k_endog == 1 else (self.k_endog,),
                 "constraints": "Positive",
                 "dims": None if self.k_endog == 1 else (f"{self.name}_endog",),
@@ -139,10 +139,10 @@ def make_symbolic_graph(self) -> None:
 
         k_nonzero = int(sum(self.order))
         ar_params = self.make_and_register_variable(
-            f"{self.name}_params", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
+            f"params_{self.name}", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
         )
         sigma_ar = self.make_and_register_variable(
-            f"{self.name}_sigma", shape=() if k_endog == 1 else (k_endog,)
+            f"sigma_{self.name}", shape=() if k_endog == 1 else (k_endog,)
         )
 
         if k_endog == 1:

pymc_extras/statespace/utils/constants.py

@@ -7,10 +7,10 @@
 SHOCK_DIM = "shock"
 SHOCK_AUX_DIM = "shock_aux"
 TIME_DIM = "time"
-AR_PARAM_DIM = "ar_lag"
-MA_PARAM_DIM = "ma_lag"
-SEASONAL_AR_PARAM_DIM = "seasonal_ar_lag"
-SEASONAL_MA_PARAM_DIM = "seasonal_ma_lag"
+AR_PARAM_DIM = "lag_ar"
+MA_PARAM_DIM = "lag_ma"
+SEASONAL_AR_PARAM_DIM = "seasonal_lag_ar"
+SEASONAL_MA_PARAM_DIM = "seasonal_lag_ma"
 ETS_SEASONAL_DIM = "seasonal_lag"
 
 NEVER_TIME_VARYING = ["initial_state", "initial_state_cov", "a0", "P0"]

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

@@ -21,7 +21,7 @@ def test_autoregressive_model(order, rng):
     lags = np.arange(len(order) if isinstance(order, list) else order, dtype="int") + 1
     if isinstance(order, list):
         lags = lags[np.flatnonzero(order)]
-    assert_allclose(ar.coords["auto_regressive_lag"], lags)
+    assert_allclose(ar.coords["lag_auto_regressive"], lags)
 
 
 def test_autoregressive_multiple_observed_build(rng):
@@ -44,15 +44,15 @@ def test_autoregressive_multiple_observed_build(rng):
     assert mod.shock_names == ["auto_regressive[data_1]", "auto_regressive[data_2]"]
 
     params = {
-        "auto_regressive_params": np.full(
+        "params_auto_regressive": np.full(
             (
                 2,
                 sum(ar.order),
             ),
             0.5,
             dtype=config.floatX,
         ),
-        "auto_regressive_sigma": np.array([0.05, 0.12]),
+        "sigma_auto_regressive": np.array([0.05, 0.12]),
     }
     _, _, _, _, T, Z, R, _, Q = mod._unpack_statespace_with_placeholders()
     input_vars = explicit_graph_inputs([T, Z, R, Q])
@@ -94,16 +94,16 @@ def test_autoregressive_multiple_observed_data(rng):
     mod = ar.build(verbose=False)
 
     params = {
-        "auto_regressive_params": np.array([0.9, 0.8, 0.5]).reshape((3, 1)),
-        "auto_regressive_sigma": np.array([0.05, 0.12, 0.22]),
+        "params_auto_regressive": np.array([0.9, 0.8, 0.5]).reshape((3, 1)),
+        "sigma_auto_regressive": np.array([0.05, 0.12, 0.22]),
         "initial_state_cov": np.eye(3),
     }
 
     # Recover the AR(1) coefficients from the simulated data via OLS
     x, y = simulate_from_numpy_model(mod, rng, params, steps=2000)
     for i in range(3):
         ols_coefs = np.polyfit(y[:-1, i], y[1:, i], 1)
-        np.testing.assert_allclose(ols_coefs[0], params["auto_regressive_params"][i, 0], atol=1e-1)
+        np.testing.assert_allclose(ols_coefs[0], params["params_auto_regressive"][i, 0], atol=1e-1)
 
 
 def test_add_autoregressive_different_observed():
@@ -128,5 +128,5 @@ def test_add_autoregressive_different_observed():
     ]
 
     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]
+    assert mod.coords["lag_ar1"] == [1]
+    assert mod.coords["lag_ar6"] == [1, 2, 3, 4, 5, 6]

tests/statespace/models/structural/test_against_statsmodels.py

@@ -405,22 +405,22 @@ def create_structural_model_and_equivalent_statsmodel(
 
     if autoregressive is not None:
         ar_names = [f"L{i+1}" for i in range(autoregressive)]
-        ar_params = rng.normal(size=(autoregressive,)).astype(floatX)
+        params_ar = rng.normal(size=(autoregressive,)).astype(floatX)
         if autoregressive == 1:
-            ar_params = ar_params.item()
+            params_ar = params_ar.item()
         sigma2 = np.abs(rng.normal()).astype(floatX)
 
-        params["ar_params"] = ar_params
-        params["ar_sigma"] = np.sqrt(sigma2)
-        expected_param_dims["ar_params"] += (AR_PARAM_DIM,)
+        params["params_ar"] = params_ar
+        params["sigma_ar"] = np.sqrt(sigma2)
+        expected_param_dims["params_ar"] += (AR_PARAM_DIM,)
         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] += ["ar"]
         expected_coords[SHOCK_AUX_DIM] += ["ar"]
 
         sm_params["sigma2.ar"] = sigma2
-        for i, rho in enumerate(ar_params):
+        for i, rho in enumerate(params_ar):
             sm_init[f"ar.L{i+1}"] = 0
             sm_params[f"ar.L{i+1}"] = rho
 

tests/statespace/models/test_SARIMAX.py

@@ -178,8 +178,8 @@ def pymc_mod(arima_mod):
         # x0  = pm.Normal('x0', dims=['state'])
         # P0_diag = pm.Gamma('P0_diag', alpha=2, beta=1, dims=['state'])
         # P0 = pm.Deterministic('P0', pt.diag(P0_diag), dims=['state', 'state_aux'])
-        ar_params = pm.Normal("ar_params", sigma=0.1, dims=["ar_lag"])
-        ma_params = pm.Normal("ma_params", sigma=1, dims=["ma_lag"])
+        ar_params = pm.Normal("ar_params", sigma=0.1, dims=["lag_ar"])
+        ma_params = pm.Normal("ma_params", sigma=1, dims=["lag_ma"])
         sigma_state = pm.Exponential("sigma_state", 0.5)
         arima_mod.build_statespace_graph(data=data, save_kalman_filter_outputs_in_idata=True)
 
@@ -207,8 +207,8 @@ def pymc_mod_interp(arima_mod_interp):
         P0 = pm.Deterministic(
             "P0", pt.eye(arima_mod_interp.k_states) * P0_sigma, dims=["state", "state_aux"]
         )
-        ar_params = pm.Normal("ar_params", sigma=0.1, dims=["ar_lag"])
-        ma_params = pm.Normal("ma_params", sigma=1, dims=["ma_lag"])
+        ar_params = pm.Normal("ar_params", sigma=0.1, dims=["lag_ar"])
+        ma_params = pm.Normal("ma_params", sigma=1, dims=["lag_ma"])
         sigma_state = pm.Exponential("sigma_state", 0.5)
         sigma_obs = pm.Exponential("sigma_obs", 0.1)
 
@@ -344,8 +344,8 @@ def test_interpretable_states_are_interpretable(arima_mod_interp, pymc_mod_inter
         prior = pm.sample_prior_predictive(draws=10)
 
     prior_outputs = arima_mod_interp.sample_unconditional_prior(prior)
-    ar_lags = prior.prior.coords["ar_lag"].values - 1
-    ma_lags = prior.prior.coords["ma_lag"].values - 1
+    ar_lags = prior.prior.coords["lag_ar"].values - 1
+    ma_lags = prior.prior.coords["lag_ma"].values - 1
 
     # Check the first p states are lags of the previous state
     for t, tm1 in zip(ar_lags[1:], ar_lags[:-1]):

tests/statespace/models/test_VARMAX.py

@@ -57,7 +57,7 @@ def pymc_mod(varma_mod, data):
             "state_chol", n=varma_mod.k_posdef, eta=1, sd_dist=pm.Exponential.dist(1)
         )
         ar_params = pm.Normal(
-            "ar_params", mu=0, sigma=0.1, dims=["observed_state", "ar_lag", "observed_state_aux"]
+            "ar_params", mu=0, sigma=0.1, dims=["observed_state", "lag_ar", "observed_state_aux"]
         )
         state_cov = pm.Deterministic(
             "state_cov", state_chol @ state_chol.T, dims=["shock", "shock_aux"]