updated plurality for some of the constants in constants.py

Author: Dekermanjian

pymc_extras/statespace/core/statespace.py

@@ -805,16 +805,16 @@ def _register_kalman_filter_outputs_with_pymc_model(outputs: tuple[pt.TensorVari
         states, covs = outputs[:4], outputs[4:]
 
         state_names = [
-            "filtered_state",
-            "predicted_state",
-            "predicted_observed_state",
-            "smoothed_state",
+            "filtered_states",
+            "predicted_states",
+            "predicted_observed_states",
+            "smoothed_states",
         ]
         cov_names = [
-            "filtered_covariance",
-            "predicted_covariance",
-            "predicted_observed_covariance",
-            "smoothed_covariance",
+            "filtered_covariances",
+            "predicted_covariances",
+            "predicted_observed_covariances",
+            "smoothed_covariances",
         ]
 
         with mod:
@@ -939,7 +939,7 @@ def build_statespace_graph(
             all_kf_outputs = [*states, smooth_states, *covs, smooth_covariances]
             self._register_kalman_filter_outputs_with_pymc_model(all_kf_outputs)
 
-        obs_dims = FILTER_OUTPUT_DIMS["predicted_observed_state"]
+        obs_dims = FILTER_OUTPUT_DIMS["predicted_observed_states"]
         obs_dims = obs_dims if all([dim in pm_mod.coords.keys() for dim in obs_dims]) else None
 
         SequenceMvNormal(
@@ -1727,14 +1727,14 @@ def sample_filter_outputs(
             )
 
             # Filter output names are singular in constants.py but are returned as plural from kalman_.build_graph()
-            filter_output_dims_mapping = {}
-            for k in FILTER_OUTPUT_DIMS.keys():
-                filter_output_dims_mapping[k + "s"] = FILTER_OUTPUT_DIMS[k]
+            # filter_output_dims_mapping = {}
+            # for k in FILTER_OUTPUT_DIMS.keys():
+            #     filter_output_dims_mapping[k + "s"] = FILTER_OUTPUT_DIMS[k]
 
             all_filter_outputs = filter_outputs[:-1] + list(smoother_outputs)
             # This excludes observed states and observed covariances from the filter outputs
             all_filter_outputs = [
-                output for output in all_filter_outputs if output.name in filter_output_dims_mapping
+                output for output in all_filter_outputs if output.name in FILTER_OUTPUT_DIMS
             ]
 
             if filter_output_names is None:
@@ -1752,7 +1752,7 @@ def sample_filter_outputs(
                 ]
 
             for output in filter_output_names:
-                dims = filter_output_dims_mapping[output.name]
+                dims = FILTER_OUTPUT_DIMS[output.name]
                 pm.Deterministic(output.name, output, dims=dims)
 
         frozen_model = freeze_dims_and_data(m)

pymc_extras/statespace/utils/constants.py

@@ -38,14 +38,14 @@
 LONG_NAME_TO_SHORT = dict(zip(LONG_MATRIX_NAMES, MATRIX_NAMES))
 
 FILTER_OUTPUT_NAMES = [
-    "filtered_state",
-    "predicted_state",
-    "filtered_covariance",
-    "predicted_covariance",
+    "filtered_states",
+    "predicted_states",
+    "filtered_covariances",
+    "predicted_covariances",
 ]
 
-SMOOTHER_OUTPUT_NAMES = ["smoothed_state", "smoothed_covariance"]
-OBSERVED_OUTPUT_NAMES = ["predicted_observed_state", "predicted_observed_covariance"]
+SMOOTHER_OUTPUT_NAMES = ["smoothed_states", "smoothed_covariances"]
+OBSERVED_OUTPUT_NAMES = ["predicted_observed_states", "predicted_observed_covariances"]
 
 MATRIX_DIMS = {
     "x0": (ALL_STATE_DIM,),
@@ -60,14 +60,14 @@
 }
 
 FILTER_OUTPUT_DIMS = {
-    "filtered_state": (TIME_DIM, ALL_STATE_DIM),
-    "smoothed_state": (TIME_DIM, ALL_STATE_DIM),
-    "predicted_state": (TIME_DIM, ALL_STATE_DIM),
-    "filtered_covariance": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
-    "smoothed_covariance": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
-    "predicted_covariance": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
-    "predicted_observed_state": (TIME_DIM, OBS_STATE_DIM),
-    "predicted_observed_covariance": (TIME_DIM, OBS_STATE_DIM, OBS_STATE_AUX_DIM),
+    "filtered_states": (TIME_DIM, ALL_STATE_DIM),
+    "smoothed_states": (TIME_DIM, ALL_STATE_DIM),
+    "predicted_states": (TIME_DIM, ALL_STATE_DIM),
+    "filtered_covariances": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
+    "smoothed_covariances": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
+    "predicted_covariances": (TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM),
+    "predicted_observed_states": (TIME_DIM, OBS_STATE_DIM),
+    "predicted_observed_covariances": (TIME_DIM, OBS_STATE_DIM, OBS_STATE_AUX_DIM),
 }
 
 POSITION_DERIVATIVE_NAMES = ["level", "trend", "acceleration", "jerk", "snap", "crackle", "pop"]

tests/statespace/core/test_statespace.py

@@ -487,16 +487,16 @@ def test_build_statespace_graph_raises_if_data_has_missing_fill():
 
 def test_build_statespace_graph(pymc_mod):
     for name in [
-        "filtered_state",
-        "predicted_state",
-        "predicted_covariance",
-        "filtered_covariance",
+        "filtered_states",
+        "predicted_states",
+        "predicted_covariances",
+        "filtered_covariances",
     ]:
         assert name in [x.name for x in pymc_mod.deterministics]
 
 
 def test_build_smoother_graph(ss_mod, pymc_mod):
-    names = ["smoothed_state", "smoothed_covariance"]
+    names = ["smoothed_states", "smoothed_covariances"]
     for name in names:
         assert name in [x.name for x in pymc_mod.deterministics]
 
@@ -1193,11 +1193,11 @@ def test_build_forecast_model(rng, exog_ss_mod, exog_pymc_mod, exog_data, idata_
 
     # Check that the frozen states and covariances correctly match the sliced index
     np.testing.assert_allclose(
-        idata_exog.posterior["predicted_covariance"].sel(time=t0).mean(("chain", "draw")).values,
+        idata_exog.posterior["predicted_covariances"].sel(time=t0).mean(("chain", "draw")).values,
         idata_forecast.posterior_predictive["P0_slice"].mean(("chain", "draw")).values,
     )
     np.testing.assert_allclose(
-        idata_exog.posterior["predicted_state"].sel(time=t0).mean(("chain", "draw")).values,
+        idata_exog.posterior["predicted_states"].sel(time=t0).mean(("chain", "draw")).values,
         idata_forecast.posterior_predictive["x0_slice"].mean(("chain", "draw")).values,
     )
 

tests/statespace/models/test_SARIMAX.py

@@ -321,7 +321,7 @@ def test_SARIMAX_update_matches_statsmodels(p, d, q, P, D, Q, S, data, rng):
 
 @pytest.mark.parametrize("filter_output", ["filtered", "predicted", "smoothed"])
 def test_all_prior_covariances_are_PSD(filter_output, pymc_mod, rng):
-    rv = pymc_mod[f"{filter_output}_covariance"]
+    rv = pymc_mod[f"{filter_output}_covariances"]
     cov_mats = pm.draw(rv, 100, random_seed=rng)
     w, v = np.linalg.eig(cov_mats)
     assert_array_less(0, w, err_msg=f"Smallest eigenvalue: {min(w.ravel())}")

tests/statespace/models/test_VARMAX.py

@@ -156,7 +156,7 @@ def test_VARMAX_update_matches_statsmodels(data, order, rng):
 
 @pytest.mark.parametrize("filter_output", ["filtered", "predicted", "smoothed"])
 def test_all_prior_covariances_are_PSD(filter_output, pymc_mod, rng):
-    rv = pymc_mod[f"{filter_output}_covariance"]
+    rv = pymc_mod[f"{filter_output}_covariances"]
     cov_mats = pm.draw(rv, 100, random_seed=rng)
     w, v = np.linalg.eig(cov_mats)
     assert_array_less(0, w, err_msg=f"Smallest eigenvalue: {min(w.ravel())}")

tests/statespace/utils/test_coord_assignment.py

@@ -93,7 +93,7 @@ def test_filter_output_coord_assignment(f, warning, create_model):
     with warning:
         pymc_model = create_model(f)
 
-    for output in FILTER_OUTPUT_NAMES + SMOOTHER_OUTPUT_NAMES + ["predicted_observed_state"]:
+    for output in FILTER_OUTPUT_NAMES + SMOOTHER_OUTPUT_NAMES + ["predicted_observed_states"]:
         assert pymc_model.named_vars_to_dims[output] == FILTER_OUTPUT_DIMS[output]