cleaned up commented code, moved internal checks to the top, reduced intermediate variables

Author: Dekermanjian

pymc_extras/statespace/core/statespace.py

@@ -1684,6 +1684,21 @@ 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())
+        else:
+            unknown_filter_output_names = np.setdiff1d(
+                filter_output_names, list(all_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
+            ]
+
         compile_kwargs = kwargs.pop("compile_kwargs", {})
         compile_kwargs.setdefault("mode", self.mode)
 
@@ -1726,44 +1741,19 @@ def sample_filter_outputs(
                 T, R, Q, filter_outputs[0], filter_outputs[3]
             )
 
-            # 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]
-
-            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
-            ]
+            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]
+                    pm.Deterministic(output.name, output, dims=dims)
 
-            if filter_output_names is None:
-                filter_output_names = all_filter_outputs
-            else:
-                unknown_filter_output_names = np.setdiff1d(
-                    filter_output_names, [x.name for x in all_filter_outputs]
-                )
-                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_outputs if x.name in filter_output_names
-                ]
-
-            for output in filter_output_names:
-                dims = FILTER_OUTPUT_DIMS[output.name]
-                pm.Deterministic(output.name, output, dims=dims)
-
-        frozen_model = freeze_dims_and_data(m)
-        with frozen_model:
-            idata_filter = pm.sample_posterior_predictive(
+        with freeze_dims_and_data(m):
+            return pm.sample_posterior_predictive(
                 idata if group == "posterior" else idata.prior,
-                var_names=[x.name for x in frozen_model.deterministics],
+                var_names=filter_output_names,
                 compile_kwargs=compile_kwargs,
                 **kwargs,
             )
-        return idata_filter
 
     @staticmethod
     def _validate_forecast_args(