_insert_random_variables infers shape from coordinates

Author: jessegrabowski

pymc_experimental/statespace/core/statespace.py

@@ -284,14 +284,16 @@ def _print_data_requirements(self) -> None:
         Prints a short report to the terminal about the data needed for the model, including their names, shapes,
         and named dimensions.
         """
+        if not self.data_info:
+            return
 
         out = ""
         for data, info in self.data_info.items():
             out += f'\t{data} -- shape: {info["shape"]}, dims: {info["dims"]}\n'
         out = out.rstrip()
 
         _log.info(
-            "The following MutableData variables should be assigned to the model inside a PyMC "
+            "The following Data variables should be assigned to the model inside a PyMC "
             f"model block: \n"
             f"{out}"
         )
@@ -618,63 +620,6 @@ def _get_matrix_shape_and_dims(
 
         return shape, dims
 
-    def _get_output_shape_and_dims(
-        self, idata: InferenceData, filter_output: str
-    ) -> tuple[
-        Optional[tuple[int]], Optional[tuple[int]], Optional[tuple[str]], Optional[tuple[str]]
-    ]:
-        """
-        Get the shapes and dimensions of the output variables from the provided InferenceData.
-
-        This method extracts the shapes and dimensions of the output variables representing the state estimates
-        and covariances from the provided ArviZ InferenceData object. The state estimates are obtained from the
-        specified `filter_output` mode, which can be one of "filtered", "predicted", or "smoothed".
-
-        Parameters
-        ----------
-        idata : arviz.InferenceData
-            The ArviZ InferenceData object containing the posterior samples.
-
-        filter_output : str
-            The name of the filter output whose shape is being checked. It can be one of "filtered",
-            "predicted", or "smoothed".
-
-        Returns
-        -------
-        mu_shape: tuple(int, int) or None
-            Shape of the mean vectors returned by the Kalman filter. Should be (n_data_timesteps, k_states).
-            If named dimensions are found, this will be None.
-
-        cov_shape: tuple (int, int, int) or None
-            Shape of the hidden state covariance matrices returned by the Kalman filter. Should be
-            (n_data_timesteps, k_states, k_states).
-            If named dimensions are found, this will be None.
-
-        mu_dims: tuple(str, str) or None
-            *Default* named dimensions associated with the mean vectors returned by the Kalman filter, or None if
-            the default names are not found.
-
-        cov_dims: tuple (str, str, str) or None
-            *Default* named dimensions associated with the covariance matrices returned by the Kalman filter, or None if
-            the default names are not found.
-        """
-
-        mu_dims = None
-        cov_dims = None
-
-        mu_shape = idata[f"{filter_output}_state"].values.shape[2:]
-        cov_shape = idata[f"{filter_output}_covariance"].values.shape[2:]
-
-        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
-            time_dim = TIME_DIM
-            mu_dims = [time_dim, ALL_STATE_DIM]
-            cov_dims = [time_dim, ALL_STATE_DIM, ALL_STATE_AUX_DIM]
-
-            mu_shape = None
-            cov_shape = None
-
-        return mu_shape, cov_shape, mu_dims, cov_dims
-
     def _insert_random_variables(self):
         """
         Replace pytensor symbolic variables with PyMC random variables.
@@ -1506,11 +1451,11 @@ def forecast(
                 "Scenario-based forcasting with exogenous variables not currently supported"
             )
 
-        dims = None
         temp_coords = self._fit_coords.copy()
 
         filter_time_dim = TIME_DIM
 
+        dims = None
         if all([dim in temp_coords for dim in [filter_time_dim, ALL_STATE_DIM, OBS_STATE_DIM]]):
             dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
 
@@ -1544,14 +1489,10 @@ def forecast(
         temp_coords["data_time"] = time_index
         temp_coords[TIME_DIM] = forecast_index
 
-        mu_shape, cov_shape, mu_dims, cov_dims = self._get_output_shape_and_dims(
-            idata.posterior, filter_output
-        )
-
-        if mu_dims is not None:
-            mu_dims = ["data_time"] + mu_dims[1:]
-        if cov_dims is not None:
-            cov_dims = ["data_time"] + cov_dims[1:]
+        mu_dims, cov_dims = None, None
+        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
+            mu_dims = ["data_time", ALL_STATE_DIM]
+            cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM]
 
         with pm.Model(coords=temp_coords):
             [