Bug fixes for statespace

pymc_experimental/statespace/core/statespace.py

@@ -284,6 +284,8 @@ 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():
@@ -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):
             [

pymc_experimental/statespace/models/structural.py

@@ -20,7 +20,6 @@
     ALL_STATE_DIM,
     AR_PARAM_DIM,
     LONG_MATRIX_NAMES,
-    OBS_STATE_DIM,
     POSITION_DERIVATIVE_NAMES,
     TIME_DIM,
 )
@@ -40,8 +39,7 @@ def order_to_mask(order):
 def _frequency_transition_block(s, j):
     lam = 2 * np.pi * j / s
 
-    # Squeeze because otherwise if lamb has shape (1,), T will have shape (2, 2, 1)
-    return pt.stack([[pt.cos(lam), pt.sin(lam)], [-pt.sin(lam), pt.cos(lam)]]).squeeze()
+    return pt.stack([[pt.cos(lam), pt.sin(lam)], [-pt.sin(lam), pt.cos(lam)]])
 
 
 class StructuralTimeSeries(PyMCStateSpace):
@@ -911,17 +909,17 @@ def __init__(self, name: str = "MeasurementError"):
 
     def populate_component_properties(self):
         self.param_names = [f"sigma_{self.name}"]
-        self.param_dims = {f"sigma_{self.name}": (OBS_STATE_DIM,)}
+        self.param_dims = {}
         self.param_info = {
             f"sigma_{self.name}": {
-                "shape": (1,),
+                "shape": (),
                 "constraints": "Positive",
-                "dims": (OBS_STATE_DIM,),
+                "dims": None,
             }
         }
 
     def make_symbolic_graph(self) -> None:
-        sigma_shape = () if self.k_endog == 1 else (self.k_endog,)
+        sigma_shape = ()
         error_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=sigma_shape)
         diag_idx = np.diag_indices(self.k_endog)
         idx = np.s_["obs_cov", diag_idx[0], diag_idx[1]]
@@ -1015,13 +1013,13 @@ def populate_component_properties(self):
                 "constraints": None,
                 "dims": (AR_PARAM_DIM,),
             },
-            "sigma_ar": {"shape": (1,), "constraints": "Positive", "dims": None},
+            "sigma_ar": {"shape": (), "constraints": "Positive", "dims": None},
         }
 
     def make_symbolic_graph(self) -> None:
         k_nonzero = int(sum(self.order))
         ar_params = self.make_and_register_variable("ar_params", shape=(k_nonzero,))
-        sigma_ar = self.make_and_register_variable("sigma_ar", shape=(1,))
+        sigma_ar = self.make_and_register_variable("sigma_ar", shape=())
 
         T = np.eye(self.k_states, k=-1)
         self.ssm["transition", :, :] = T
@@ -1150,6 +1148,7 @@ def __init__(
         innovations: bool = True,
         name: Optional[str] = None,
         state_names: Optional[list] = None,
+        pop_state: bool = True,
     ):
         if name is None:
             name = f"Seasonal[s={season_length}]"
@@ -1162,11 +1161,14 @@ def __init__(
                 )
             state_names = state_names.copy()
         self.innovations = innovations
+        self.pop_state = pop_state
 
-        # The first state doesn't get a coefficient, it is defined as -sum(state_coefs)
-        # TODO: Can I stash that information in the model somewhere so users don't have to know that?
-        state_0 = state_names.pop(0)
-        k_states = season_length - 1
+        if self.pop_state:
+            # In traditional models, the first state isn't identified, so we can help out the user by automatically
+            # discarding it.
+            # TODO: Can this be stashed and reconstructed automatically somehow?
+            state_names.pop(0)
+            k_states = season_length - 1
 
         super().__init__(
             name=name,
@@ -1194,7 +1196,7 @@ def populate_component_properties(self):
         if self.innovations:
             self.param_names += [f"sigma_{self.name}"]
             self.param_info[f"sigma_{self.name}"] = {
-                "shape": (1,),
+                "shape": (),
                 "constraints": "Positive",
                 "dims": None,
             }
@@ -1214,7 +1216,7 @@ def make_symbolic_graph(self) -> None:
 
         if self.innovations:
             self.ssm["selection", 0, 0] = 1
-            season_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=(1,))
+            season_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=())
             cov_idx = ("state_cov", *np.diag_indices(1))
             self.ssm[cov_idx] = season_sigma**2
 
@@ -1313,7 +1315,7 @@ def make_symbolic_graph(self) -> None:
         self.ssm["transition", :, :] = T
 
         if self.innovations:
-            sigma_season = self.make_and_register_variable(f"sigma_{self.name}", shape=(1,))
+            sigma_season = self.make_and_register_variable(f"sigma_{self.name}", shape=())
             self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_season**2
             self.ssm["selection", :, :] = np.eye(self.k_states)
 
@@ -1339,7 +1341,7 @@ def populate_component_properties(self):
             self.shock_names = self.state_names.copy()
             self.param_names += [f"sigma_{self.name}"]
             self.param_info[f"sigma_{self.name}"] = {
-                "shape": (1,),
+                "shape": (),
                 "constraints": "Positive",
                 "dims": None,
             }
@@ -1480,24 +1482,24 @@ def make_symbolic_graph(self) -> None:
         self.ssm["initial_state", :] = init_state
 
         if self.estimate_cycle_length:
-            lamb = self.make_and_register_variable(f"{self.name}_length", shape=(1,))
+            lamb = self.make_and_register_variable(f"{self.name}_length", shape=())
         else:
             lamb = self.cycle_length
 
         if self.dampen:
-            rho = self.make_and_register_variable(f"{self.name}_dampening_factor", shape=(1,))
+            rho = self.make_and_register_variable(f"{self.name}_dampening_factor", shape=())
         else:
             rho = 1
 
         T = rho * _frequency_transition_block(lamb, j=1)
         self.ssm["transition", :, :] = T
 
         if self.innovations:
-            sigma_cycle = self.make_and_register_variable(f"sigma_{self.name}", shape=(1,))
+            sigma_cycle = self.make_and_register_variable(f"sigma_{self.name}", shape=())
             self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_cycle**2
 
     def populate_component_properties(self):
-        self.state_names = [f"{self.name}_{f}" for f in ["Sin", "Cos"]]
+        self.state_names = [f"{self.name}_{f}" for f in ["Cos", "Sin"]]
         self.param_names = [f"{self.name}"]
 
         self.param_info = {
@@ -1511,23 +1513,23 @@ def populate_component_properties(self):
         if self.estimate_cycle_length:
             self.param_names += [f"{self.name}_length"]
             self.param_info[f"{self.name}_length"] = {
-                "shape": (1,),
+                "shape": (),
                 "constraints": "Positive, non-zero",
                 "dims": None,
             }
 
         if self.dampen:
             self.param_names += [f"{self.name}_dampening_factor"]
             self.param_info[f"{self.name}_dampening_factor"] = {
-                "shape": (1,),
+                "shape": (),
                 "constraints": "0 < x ≤ 1",
                 "dims": None,
             }
 
         if self.innovations:
             self.param_names += [f"sigma_{self.name}"]
             self.param_info[f"sigma_{self.name}"] = {
-                "shape": (1,),
+                "shape": (),
                 "constraints": "Positive",
                 "dims": None,
             }
@@ -1609,7 +1611,11 @@ def populate_component_properties(self) -> None:
         }
 
         self.param_info = {
-            f"beta_{self.name}": {"shape": (1,), "constraints": None, "dims": ("exog_state",)},
+            f"beta_{self.name}": {
+                "shape": (self.k_states,),
+                "constraints": None,
+                "dims": ("exog_state",),
+            },
         }
 
         self.data_info = {
@@ -1624,7 +1630,7 @@ def populate_component_properties(self) -> None:
             self.param_names += [f"sigma_beta_{self.name}"]
             self.param_dims[f"sigma_beta_{self.name}"] = "exog_state"
             self.param_info[f"sigma_beta_{self.name}"] = {
-                "shape": (1,),
+                "shape": (),
                 "constraints": "Positive",
                 "dims": ("exog_state",),
             }

pymc_experimental/statespace/utils/data_tools.py

@@ -112,7 +112,18 @@ def add_data_to_active_model(values, index):
     if OBS_STATE_DIM in pymc_mod.coords:
         data_dims = [TIME_DIM, OBS_STATE_DIM]
 
-    pymc_mod.add_coord(TIME_DIM, index)
+    if TIME_DIM not in pymc_mod.coords:
+        pymc_mod.add_coord(TIME_DIM, index)
+    else:
+        found_time = pymc_mod.coords[TIME_DIM]
+        if found_time is None:
+            pymc_mod.coords.update({TIME_DIM: index})
+        elif not np.array_equal(found_time, index):
+            raise ValueError(
+                "Provided data has a different time index than the model. Please ensure that the time values "
+                "set on coords matches that of the exogenous data."
+            )
+
     data = pm.Data("data", values, dims=data_dims)
 
     return data