Be more consistent with the names of filter outputs

Author: jessegrabowski

pymc_experimental/statespace/core/statespace.py

@@ -758,22 +758,28 @@ def _register_matrices_with_pymc_model(self) -> list[pt.TensorVariable]:
     @staticmethod
     def _register_kalman_filter_outputs_with_pymc_model(outputs: tuple[pt.TensorVariable]) -> None:
         mod = modelcontext(None)
+        coords = mod.coords
+
         states, covs = outputs[:4], outputs[4:]
 
-        state_names = ["filtered_state", "predicted_state", "observed_state", "smoothed_state"]
+        state_names = [
+            "filtered_state",
+            "predicted_state",
+            "predicted_observed_state",
+            "smoothed_state",
+        ]
         cov_names = [
             "filtered_covariance",
             "predicted_covariance",
-            "observed_covariance",
+            "predicted_observed_covariance",
             "smoothed_covariance",
         ]
 
         with mod:
-            for state, name in zip(states, state_names):
-                pm.Deterministic(name, state, dims=FILTER_OUTPUT_DIMS.get(name, None))
-
-            for cov, name in zip(covs, cov_names):
-                pm.Deterministic(name, cov, dims=FILTER_OUTPUT_DIMS.get(name, None))
+            for var, name in zip(states + covs, state_names + cov_names):
+                dim_names = FILTER_OUTPUT_DIMS.get(name, None)
+                dims = tuple([dim if dim in coords.keys() else None for dim in dim_names])
+                pm.Deterministic(name, var, dims=dims)
 
     def build_statespace_graph(
         self,
@@ -866,7 +872,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["obs"]
+        obs_dims = FILTER_OUTPUT_DIMS["predicted_observed_state"]
         obs_dims = obs_dims if all([dim in pm_mod.coords.keys() for dim in obs_dims]) else None
 
         SequenceMvNormal(

pymc_experimental/statespace/utils/constants.py

@@ -44,7 +44,7 @@
 ]
 
 SMOOTHER_OUTPUT_NAMES = ["smoothed_state", "smoothed_covariance"]
-OBSERVED_OUTPUT_NAMES = ["observed_state", "observed_covariance"]
+OBSERVED_OUTPUT_NAMES = ["predicted_observed_state", "predicted_observed_covariance"]
 
 MATRIX_DIMS = {
     "x0": (ALL_STATE_DIM,),
@@ -65,7 +65,8 @@
     "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),
-    "obs": (TIME_DIM, OBS_STATE_DIM),
+    "predicted_observed_state": (TIME_DIM, OBS_STATE_DIM),
+    "predicted_observed_covariance": (TIME_DIM, OBS_STATE_DIM, OBS_STATE_AUX_DIM),
 }
 
 POSITION_DERIVATIVE_NAMES = ["level", "trend", "acceleration", "jerk", "snap", "crackle", "pop"]

pymc_experimental/tests/statespace/test_coord_assignment.py

@@ -94,7 +94,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 + ["obs"]:
+    for output in FILTER_OUTPUT_NAMES + SMOOTHER_OUTPUT_NAMES + ["predicted_observed_state"]:
         assert pymc_model.named_vars_to_dims[output] == FILTER_OUTPUT_DIMS[output]