Update tests to respect new naming convention

Author: jessegrabowski

tests/statespace/filters/test_distributions.py

@@ -52,8 +52,8 @@ def pymc_model(data):
         data = pm.Data("data", data.values)
         P0_diag = pm.Exponential("P0_diag", 1, shape=(2,))
         P0 = pm.Deterministic("P0", pt.diag(P0_diag))
-        initial_trend = pm.Normal("level_trend_initial", shape=(2,))
-        sigma_trend = pm.Exponential("level_trend_sigma", 1, shape=(2,))
+        initial_trend = pm.Normal("initial_level_trend", shape=(2,))
+        sigma_trend = pm.Exponential("sigma_level_trend", 1, shape=(2,))
 
     return mod
 
@@ -69,8 +69,8 @@ def pymc_model_2(data):
     with pm.Model(coords=coords) as mod:
         P0_diag = pm.Exponential("P0_diag", 1, shape=(2,))
         P0 = pm.Deterministic("P0", pt.diag(P0_diag))
-        initial_trend = pm.Normal("level_trend_initial", shape=(2,))
-        sigma_trend = pm.Exponential("level_trend_sigma", 1, shape=(2,))
+        initial_trend = pm.Normal("initial_level_trend", shape=(2,))
+        sigma_trend = pm.Exponential("sigma_level_trend", 1, shape=(2,))
         sigma_me = pm.Exponential("sigma_error", 1)
 
     return mod
@@ -207,8 +207,8 @@ def test_lgss_with_time_varying_inputs(output_name, rng):
         exog_data = pm.Data("data_exog", X)
         P0_diag = pm.Exponential("P0_diag", 1, shape=(mod.k_states,))
         P0 = pm.Deterministic("P0", pt.diag(P0_diag))
-        initial_trend = pm.Normal("level_trend_initial", shape=(2,))
-        sigma_trend = pm.Exponential("level_trend_sigma", 1, shape=(2,))
+        initial_trend = pm.Normal("initial_level_trend", shape=(2,))
+        sigma_trend = pm.Exponential("sigma_level_trend", 1, shape=(2,))
         beta_exog = pm.Normal("beta_exog", shape=(3,))
 
         mod._insert_random_variables()

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

@@ -14,15 +14,15 @@
 
 def test_level_trend_model(rng):
     mod = st.LevelTrendComponent(order=2, innovations_order=0)
-    params = {"level_trend_initial": [0.0, 1.0]}
+    params = {"initial_level_trend": [0.0, 1.0]}
     x, y = simulate_from_numpy_model(mod, rng, params)
 
     assert_allclose(np.diff(y), 1, atol=ATOL, rtol=RTOL)
 
     # Check coords
     mod = mod.build(verbose=False)
     _assert_basic_coords_correct(mod)
-    assert mod.coords["level_trend_state"] == ["level", "trend"]
+    assert mod.coords["state_level_trend"] == ["level", "trend"]
 
 
 def test_level_trend_multiple_observed_construction():
@@ -34,8 +34,8 @@ def test_level_trend_multiple_observed_construction():
     assert mod.k_states == 6
     assert mod.k_posdef == 3
 
-    assert mod.coords["level_trend_state"] == ["level", "trend"]
-    assert mod.coords["level_trend_endog"] == ["data_1", "data_2", "data_3"]
+    assert mod.coords["state_level_trend"] == ["level", "trend"]
+    assert mod.coords["endog_level_trend"] == ["data_1", "data_2", "data_3"]
 
     assert mod.state_names == [
         "level[data_1]",
@@ -95,7 +95,7 @@ def test_level_trend_multiple_observed(rng):
     mod = st.LevelTrendComponent(
         order=2, innovations_order=0, observed_state_names=["data_1", "data_2", "data_3"]
     )
-    params = {"level_trend_initial": np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])}
+    params = {"initial_level_trend": np.array([[0.0, 1.0], [0.0, 2.0], [0.0, 3.0]])}
 
     x, y = simulate_from_numpy_model(mod, rng, params)
     assert (np.diff(y, axis=0) == np.array([[1.0, 2.0, 3.0]])).all().all()
@@ -115,8 +115,8 @@ def test_add_level_trend_with_different_observed():
     assert mod.k_states == 3
     assert mod.k_posdef == 2
 
-    assert mod.coords["ll_state"] == ["level", "trend"]
-    assert mod.coords["grw_state"] == ["level"]
+    assert mod.coords["state_ll"] == ["level", "trend"]
+    assert mod.coords["state_grw"] == ["level"]
 
     assert mod.state_names == ["level[data_1]", "trend[data_1]", "level[data_2]"]
     assert mod.shock_names == ["trend[data_1]", "level[data_2]"]

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

@@ -28,7 +28,7 @@ def test_exogenous_component(rng):
 
     mod = mod.build(verbose=False)
     _assert_basic_coords_correct(mod)
-    assert mod.coords["exog_state"] == ["feature_1", "feature_2"]
+    assert mod.coords["state_exog"] == ["feature_1", "feature_2"]
 
 
 def test_adding_exogenous_component(rng):
@@ -65,11 +65,9 @@ def test_regression_with_multiple_observed_states(rng):
     assert_allclose(y[:, 1], data @ params["beta_exog"][1], atol=ATOL, rtol=RTOL)
 
     mod = mod.build(verbose=False)
-    assert mod.coords["exog_state"] == [
-        "feature_1[data_1]",
-        "feature_2[data_1]",
-        "feature_1[data_2]",
-        "feature_2[data_2]",
+    assert mod.coords["state_exog"] == [
+        "feature_1",
+        "feature_2",
     ]
 
     Z = mod.ssm["design"].eval({"data_exog": data})
@@ -90,8 +88,8 @@ def test_add_regression_components_with_multiple_observed_states(rng):
     reg2 = st.RegressionComponent(state_names=["c"], name="exog2", observed_state_names=["data_3"])
 
     mod = (reg1 + reg2).build(verbose=False)
-    assert mod.coords["exog1_state"] == ["a[data_1]", "b[data_1]", "a[data_2]", "b[data_2]"]
-    assert mod.coords["exog2_state"] == ["c[data_3]"]
+    assert mod.coords["state_exog1"] == ["a", "b"]
+    assert mod.coords["state_exog2"] == ["c"]
 
     Z = mod.ssm["design"].eval({"data_exog1": data_1, "data_exog2": data_2})
     vec_block_diag = np.vectorize(block_diag, signature="(n,m),(o,p)->(q,r)")
@@ -124,7 +122,7 @@ def test_filter_scans_time_varying_design_matrix(rng):
 
         x0 = pm.Normal("x0", dims=["state"])
         P0 = pm.Deterministic("P0", pt.eye(mod.k_states), dims=["state", "state_aux"])
-        beta_exog = pm.Normal("beta_exog", dims=["exog_state"])
+        beta_exog = pm.Normal("beta_exog", dims=["state_exog"])
 
         mod.build_statespace_graph(y)
         x0, P0, c, d, T, Z, R, H, Q = mod.unpack_statespace()

tests/statespace/models/structural/test_against_statsmodels.py

@@ -220,7 +220,7 @@ def create_structural_model_and_equivalent_statsmodel(
 
     if level:
         level_trend_order[0] = 1
-        expected_coords["level_state"] += [
+        expected_coords["state_level"] += [
             "level",
         ]
         expected_coords[ALL_STATE_DIM] += [
@@ -241,7 +241,7 @@ def create_structural_model_and_equivalent_statsmodel(
 
     if trend:
         level_trend_order[1] = 1
-        expected_coords["level_state"] += [
+        expected_coords["state_level"] += [
             "trend",
         ]
         expected_coords[ALL_STATE_DIM] += [
@@ -258,7 +258,7 @@ def create_structural_model_and_equivalent_statsmodel(
             expected_coords[SHOCK_AUX_DIM] += ["trend"]
 
     if level or trend:
-        expected_param_dims["level_initial"] += ("level_state",)
+        expected_param_dims["initial_level"] += ("state_level",)
         level_value = np.where(
             level_trend_order,
             rng.normal(
@@ -272,13 +272,13 @@ def create_structural_model_and_equivalent_statsmodel(
         max_order = np.flatnonzero(level_value)[-1].item() + 1
         level_trend_order = level_trend_order[:max_order]
 
-        params["level_initial"] = level_value[:max_order]
+        params["initial_level"] = level_value[:max_order]
         sm_init["level"] = level_value[0]
         sm_init["trend"] = level_value[1]
 
         if sum(level_trend_innov_order) > 0:
-            expected_param_dims["level_sigma"] += ("level_shock",)
-            params["level_sigma"] = np.sqrt(sigma_level_value2)
+            expected_param_dims["sigma_level"] += ("level_shock",)
+            params["sigma_level"] = np.sqrt(sigma_level_value2)
 
         sigma_level_value = sigma_level_value2.tolist()
         if stochastic_level:

tests/statespace/models/structural/test_core.py

@@ -22,8 +22,8 @@ def test_add_components():
     mod = ll + se
 
     ll_params = {
-        "level_trend_initial": np.zeros(2, dtype=floatX),
-        "level_trend_sigma": np.ones(2, dtype=floatX),
+        "initial_level_trend": np.zeros(2, dtype=floatX),
+        "sigma_level_trend": np.ones(2, dtype=floatX),
     }
     se_params = {
         "seasonal_coefs": np.ones(11, dtype=floatX),
@@ -92,9 +92,9 @@ def test_extract_components_from_idata(rng):
 
         x0 = pm.Normal("x0", dims=["state"])
         P0 = pm.Deterministic("P0", pt.eye(mod.k_states), dims=["state", "state_aux"])
-        beta_exog = pm.Normal("beta_exog", dims=["exog_state"])
-        initial_trend = pm.Normal("level_trend_initial", dims=["level_trend_state"])
-        sigma_trend = pm.Exponential("level_trend_sigma", 1, dims=["level_trend_shock"])
+        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=["level_trend_shock"])
         seasonal_coefs = pm.Normal("seasonal", dims=["seasonal_state"])
         sigma_obs = pm.Exponential("sigma_obs", 1)
 

tests/statespace/utils/test_coord_assignment.py

@@ -80,8 +80,8 @@ def _create_model(f):
                 dims="state",
             )
             P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=("state", "state_aux"))
-            initial_trend = pm.Normal("level_trend_initial", dims="level_trend_state")
-            sigma_trend = pm.Exponential("level_trend_sigma", 1, dims="level_trend_shock")
+            initial_trend = pm.Normal("initial_level_trend", dims="state_level_trend")
+            sigma_trend = pm.Exponential("sigma_level_trend", 1, dims="level_trend_shock")
             ss_mod.build_statespace_graph(data, save_kalman_filter_outputs_in_idata=True)
         return mod
 
@@ -103,8 +103,8 @@ def test_model_build_without_coords(load_dataset):
     with pm.Model() as mod:
         P0_diag = pm.Exponential("P0_diag", 1, shape=(2,))
         P0 = pm.Deterministic("P0", pt.diag(P0_diag))
-        initial_trend = pm.Normal("level_trend_initial", shape=(2,))
-        sigma_trend = pm.Exponential("level_trend_sigma", 1, shape=(2,))
+        initial_trend = pm.Normal("initial_level_trend", shape=(2,))
+        sigma_trend = pm.Exponential("sigma_level_trend", 1, shape=(2,))
         ss_mod.build_statespace_graph(data, register_data=False)
 
     assert mod.coords == {}
@@ -131,8 +131,8 @@ def make_model(index):
         P0_diag = pm.Gamma("P0_diag", alpha=5, beta=5)
         P0 = pm.Deterministic("P0", pt.eye(ss_mod.k_states) * P0_diag, dims=P0_dims)
 
-        initial_trend = pm.Normal("level_trend_initial", dims=initial_trend_dims)
-        sigma_trend = pm.Gamma("level_trend_sigma", alpha=2, beta=50, dims=sigma_trend_dims)
+        initial_trend = pm.Normal("initial_level_trend", dims=initial_trend_dims)
+        sigma_trend = pm.Gamma("sigma_level_trend", alpha=2, beta=50, dims=sigma_trend_dims)
 
         with pytest.warns(UserWarning, match="No time index found on the supplied data"):
             ss_mod.build_statespace_graph(