updated kalman filter outputs to use names defined in constants.py, updated sample_filter_outputs to allow sampling any filter outputs defined in constants.py

Author: Dekermanjian

pymc_extras/statespace/core/statespace.py

@@ -1684,20 +1684,15 @@ def sample_filter_outputs(
         if isinstance(filter_output_names, str):
             filter_output_names = [filter_output_names]
 
-        drop_keys = {"predicted_observed_states", "predicted_observed_covariances"}
-        all_filter_output_dims = {k: v for k, v in FILTER_OUTPUT_DIMS.items() if k not in drop_keys}
-
         if filter_output_names is None:
-            filter_output_names = list(all_filter_output_dims.keys())
+            filter_output_names = list(FILTER_OUTPUT_DIMS.keys())
         else:
             unknown_filter_output_names = np.setdiff1d(
-                filter_output_names, list(all_filter_output_dims.keys())
+                filter_output_names, list(FILTER_OUTPUT_DIMS.keys())
             )
             if unknown_filter_output_names.size > 0:
                 raise ValueError(f"{unknown_filter_output_names} not a valid filter output name!")
-            filter_output_names = [
-                x for x in all_filter_output_dims.keys() if x in filter_output_names
-            ]
+            filter_output_names = [x for x in FILTER_OUTPUT_DIMS.keys() if x in filter_output_names]
 
         compile_kwargs = kwargs.pop("compile_kwargs", {})
         compile_kwargs.setdefault("mode", self.mode)
@@ -1744,7 +1739,7 @@ def sample_filter_outputs(
             filter_outputs = filter_outputs[:-1] + list(smoother_outputs)
             for output in filter_outputs:
                 if output.name in filter_output_names:
-                    dims = all_filter_output_dims[output.name]
+                    dims = FILTER_OUTPUT_DIMS[output.name]
                     pm.Deterministic(output.name, output, dims=dims)
 
         with freeze_dims_and_data(m):

pymc_extras/statespace/filters/kalman_filter.py

@@ -15,10 +15,15 @@
     split_vars_into_seq_and_nonseq,
     stabilize,
 )
-from pymc_extras.statespace.utils.constants import JITTER_DEFAULT, MISSING_FILL
+from pymc_extras.statespace.utils.constants import (
+    FILTER_OUTPUT_NAMES,
+    JITTER_DEFAULT,
+    MATRIX_NAMES,
+    MISSING_FILL,
+)
 
 MVN_CONST = pt.log(2 * pt.constant(np.pi, dtype="float64"))
-PARAM_NAMES = ["c", "d", "T", "Z", "R", "H", "Q"]
+PARAM_NAMES = MATRIX_NAMES[2:]
 
 assert_time_varying_dim_correct = Assert(
     "The first dimension of a time varying matrix (the time dimension) must be "
@@ -119,7 +124,7 @@ def unpack_args(self, args) -> tuple:
         # There are always two outputs_info wedged between the seqs and non_seqs
         seqs, (a0, P0), non_seqs = args[:n_seq], args[n_seq : n_seq + 2], args[n_seq + 2 :]
         return_ordered = []
-        for name in ["c", "d", "T", "Z", "R", "H", "Q"]:
+        for name in PARAM_NAMES:
             if name in self.seq_names:
                 idx = self.seq_names.index(name)
                 return_ordered.append(seqs[idx])
@@ -253,28 +258,28 @@ def _postprocess_scan_results(self, results, a0, P0, n) -> list[TensorVariable]:
         )
 
         filtered_states = pt.specify_shape(filtered_states, (n, self.n_states))
-        filtered_states.name = "filtered_states"
+        filtered_states.name = FILTER_OUTPUT_NAMES[0]
 
         predicted_states = pt.specify_shape(predicted_states, (n, self.n_states))
-        predicted_states.name = "predicted_states"
-
-        observed_states = pt.specify_shape(observed_states, (n, self.n_endog))
-        observed_states.name = "observed_states"
+        predicted_states.name = FILTER_OUTPUT_NAMES[1]
 
         filtered_covariances = pt.specify_shape(
             filtered_covariances, (n, self.n_states, self.n_states)
         )
-        filtered_covariances.name = "filtered_covariances"
+        filtered_covariances.name = FILTER_OUTPUT_NAMES[2]
 
         predicted_covariances = pt.specify_shape(
             predicted_covariances, (n, self.n_states, self.n_states)
         )
-        predicted_covariances.name = "predicted_covariances"
+        predicted_covariances.name = FILTER_OUTPUT_NAMES[3]
+
+        observed_states = pt.specify_shape(observed_states, (n, self.n_endog))
+        observed_states.name = FILTER_OUTPUT_NAMES[4]
 
         observed_covariances = pt.specify_shape(
             observed_covariances, (n, self.n_endog, self.n_endog)
         )
-        observed_covariances.name = "observed_covariances"
+        observed_covariances.name = FILTER_OUTPUT_NAMES[5]
 
         loglike_obs = pt.specify_shape(loglike_obs.squeeze(), (n,))
         loglike_obs.name = "loglike_obs"

pymc_extras/statespace/utils/constants.py

@@ -42,6 +42,8 @@
     "predicted_states",
     "filtered_covariances",
     "predicted_covariances",
+    "predicted_observed_states",
+    "predicted_observed_covariances",
 ]
 
 SMOOTHER_OUTPUT_NAMES = ["smoothed_states", "smoothed_covariances"]