Added multivariate regression-component

pymc_extras/statespace/models/structural/components/level_trend.py

@@ -167,15 +167,19 @@ def populate_component_properties(self):
 
         name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
         self.param_names = ["initial_trend"]
-        self.state_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
+        base_names = [name for name, mask in zip(name_slice, self._order_mask) if mask]
+        self.state_names = [
+            f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
+        ]
         self.param_dims = {"initial_trend": ("trend_state",)}
-        self.coords = {"trend_state": self.state_names}
+        self.coords = {"trend_state": base_names}
 
         if k_endog > 1:
             self.param_dims["trend_state"] = (
                 "trend_endog",
                 "trend_state",
             )
+            self.param_dims = {"initial_trend": ("trend_endog", "trend_state")}
             self.coords["trend_endog"] = self.observed_state_names
 
         shape = (k_endog, k_states) if k_endog > 1 else (k_states,)

pymc_extras/statespace/models/structural/components/regression.py

@@ -28,8 +28,8 @@ def __init__(
         super().__init__(
             name=name,
             k_endog=k_endog,
-            k_states=k_states,
-            k_posdef=k_posdef,
+            k_states=k_states * k_endog,
+            k_posdef=k_posdef * k_endog,
             state_names=self.state_names,
             observed_state_names=observed_state_names,
             measurement_error=False,
@@ -59,15 +59,21 @@ def _handle_input_data(self, k_exog: int, state_names: list[str] | None, name) -
         return k_exog
 
     def make_symbolic_graph(self) -> None:
-        betas = self.make_and_register_variable(f"beta_{self.name}", shape=(self.k_states,))
-        regression_data = self.make_and_register_data(
-            f"data_{self.name}", shape=(None, self.k_states)
-        )
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+        self.k_posdef // k_endog
+
+        betas = self.make_and_register_variable(f"beta_{self.name}", shape=(k_endog, k_states))
+        regression_data = self.make_and_register_data(f"data_{self.name}", shape=(None, k_states))
 
-        self.ssm["initial_state", :] = betas
-        self.ssm["transition", :, :] = np.eye(self.k_states)
+        self.ssm["initial_state", :] = betas.reshape((1, -1)).squeeze()
+        T = np.eye(k_states)
+        self.ssm["transition", :, :] = pt.linalg.block_diag(*[T for _ in range(k_endog)])
         self.ssm["selection", :, :] = np.eye(self.k_states)
-        self.ssm["design"] = pt.expand_dims(regression_data, 1)
+        Z = pt.linalg.block_diag(*[pt.expand_dims(regression_data, 1) for _ in range(k_endog)])
+        self.ssm["design"] = pt.specify_shape(
+            Z, (None, k_endog, regression_data.type.shape[1] * k_endog)
+        )
 
         if self.innovations:
             sigma_beta = self.make_and_register_variable(
@@ -77,29 +83,38 @@ def make_symbolic_graph(self) -> None:
             self.ssm["state_cov", row_idx, col_idx] = sigma_beta**2
 
     def populate_component_properties(self) -> None:
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+        self.k_posdef // k_endog
+
         self.shock_names = self.state_names
 
         self.param_names = [f"beta_{self.name}"]
         self.data_names = [f"data_{self.name}"]
         self.param_dims = {
-            f"beta_{self.name}": ("exog_state",),
+            f"beta_{self.name}": ("exog_endog", "exog_state"),
         }
 
+        base_names = self.state_names
+        self.state_names = [
+            f"{name}[{obs_name}]" for obs_name in self.observed_state_names for name in base_names
+        ]
+
         self.param_info = {
             f"beta_{self.name}": {
-                "shape": (self.k_states,),
+                "shape": (k_endog, k_states),
                 "constraints": None,
-                "dims": ("exog_state",),
+                "dims": ("exog_endog", "exog_state"),
             },
         }
 
         self.data_info = {
             f"data_{self.name}": {
-                "shape": (None, self.k_states),
+                "shape": (None, k_states),
                 "dims": (TIME_DIM, "exog_state"),
             },
         }
-        self.coords = {"exog_state": self.state_names}
+        self.coords = {"exog_state": base_names, "exog_endog": self.observed_state_names}
 
         if self.innovations:
             self.param_names += [f"sigma_beta_{self.name}"]