Regression component bugfix and tests

Author: jessegrabowski

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

@@ -10,7 +10,7 @@ class RegressionComponent(Component):
     def __init__(
         self,
         k_exog: int | None = None,
-        name: str | None = "Exogenous",
+        name: str | None = "regression",
         state_names: list[str] | None = None,
         observed_state_names: list[str] | None = None,
         innovations=False,
@@ -61,38 +61,38 @@ def _handle_input_data(self, k_exog: int, state_names: list[str] | None, name) -
     def make_symbolic_graph(self) -> None:
         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))
+        betas = self.make_and_register_variable(
+            f"beta_{self.name}", shape=(k_endog, k_states) if k_endog > 1 else (k_states,)
+        )
         regression_data = self.make_and_register_data(f"data_{self.name}", shape=(None, 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["initial_state", :] = betas.ravel()
+        self.ssm["transition", :, :] = pt.eye(self.k_states)
+        self.ssm["selection", :, :] = pt.eye(self.k_states)
+
         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(
-                f"sigma_beta_{self.name}", (self.k_states,)
+                f"sigma_beta_{self.name}", (k_states,) if k_endog == 1 else (k_endog, k_states)
             )
             row_idx, col_idx = np.diag_indices(self.k_states)
-            self.ssm["state_cov", row_idx, col_idx] = sigma_beta**2
+            self.ssm["state_cov", row_idx, col_idx] = sigma_beta.ravel() ** 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_endog", "exog_state"),
+            f"beta_{self.name}": (f"{self.name}_endog", f"{self.name}_state"),
         }
 
         base_names = self.state_names
@@ -102,9 +102,11 @@ def populate_component_properties(self) -> None:
 
         self.param_info = {
             f"beta_{self.name}": {
-                "shape": (k_endog, k_states),
+                "shape": (k_endog, k_states) if k_endog > 1 else (k_states,),
                 "constraints": None,
-                "dims": ("exog_endog", "exog_state"),
+                "dims": (f"{self.name}_endog", f"{self.name}_state")
+                if k_endog > 1
+                else (f"{self.name}_state",),
             },
         }
 
@@ -114,13 +116,18 @@ def populate_component_properties(self) -> None:
                 "dims": (TIME_DIM, "exog_state"),
             },
         }
-        self.coords = {"exog_state": base_names, "exog_endog": self.observed_state_names}
+        self.coords = {
+            f"{self.name}_state": self.state_names,
+            f"{self.name}_endog": self.observed_state_names,
+        }
 
         if self.innovations:
             self.param_names += [f"sigma_beta_{self.name}"]
-            self.param_dims[f"sigma_beta_{self.name}"] = "exog_state"
+            self.param_dims[f"sigma_beta_{self.name}"] = f"{self.name}_state"
             self.param_info[f"sigma_beta_{self.name}"] = {
                 "shape": (),
                 "constraints": "Positive",
-                "dims": ("exog_state",),
+                "dims": (f"{self.name}_state",)
+                if k_endog == 1
+                else (f"{self.name}_endog", f"{self.name}_state"),
             }

tests/statespace/models/structural/components/test_regression.py

@@ -5,6 +5,7 @@
 from numpy.testing import assert_allclose
 from pytensor import config
 from pytensor import tensor as pt
+from pytensor.graph.basic import explicit_graph_inputs
 
 from pymc_extras.statespace.models import structural as st
 from tests.statespace.models.structural.conftest import _assert_basic_coords_correct
@@ -25,7 +26,7 @@ def test_exogenous_component(rng):
     # Check that the generated data is just a linear regression
     assert_allclose(y, data @ params["beta_exog"], atol=ATOL, rtol=RTOL)
 
-    mod.build(verbose=False)
+    mod = mod.build(verbose=False)
     _assert_basic_coords_correct(mod)
     assert mod.coords["exog_state"] == ["feature_1", "feature_2"]
 
@@ -42,6 +43,73 @@ def test_adding_exogenous_component(rng):
     assert_allclose(mod.ssm["design", 5, 0, :2].eval({"data_exog": data}), data[5])
 
 
+def test_regression_with_multiple_observed_states(rng):
+    from scipy.linalg import block_diag
+
+    data = rng.normal(size=(100, 2)).astype(config.floatX)
+    mod = st.RegressionComponent(
+        state_names=["feature_1", "feature_2"],
+        name="exog",
+        observed_state_names=["data_1", "data_2"],
+    )
+
+    params = {"beta_exog": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=config.floatX)}
+    exog_data = {"data_exog": data}
+    x, y = simulate_from_numpy_model(mod, rng, params, exog_data)
+
+    assert x.shape == (100, 4)  # 2 features, 2 states
+    assert y.shape == (100, 2)
+
+    # Check that the generated data are two independent linear regressions
+    assert_allclose(y[:, 0], data @ params["beta_exog"][0], atol=ATOL, rtol=RTOL)
+    assert_allclose(y[:, 1], data @ params["beta_exog"][1], atol=ATOL, rtol=RTOL)
+
+    mod = mod.build(verbose=False)
+    assert mod.coords["exog_state"] == [
+        "feature_1[data_1]",
+        "feature_2[data_1]",
+        "feature_1[data_2]",
+        "feature_2[data_2]",
+    ]
+
+    Z = mod.ssm["design"].eval({"data_exog": data})
+    vec_block_diag = np.vectorize(block_diag, signature="(n,m),(o,p)->(q,r)")
+    assert Z.shape == (100, 2, 4)
+    assert np.allclose(Z, vec_block_diag(data[:, None, :], data[:, None, :]))
+
+
+def test_add_regression_components_with_multiple_observed_states(rng):
+    from scipy.linalg import block_diag
+
+    data_1 = rng.normal(size=(100, 2)).astype(config.floatX)
+    data_2 = rng.normal(size=(100, 1)).astype(config.floatX)
+
+    reg1 = st.RegressionComponent(
+        state_names=["a", "b"], name="exog1", observed_state_names=["data_1", "data_2"]
+    )
+    reg2 = st.RegressionComponent(state_names=["c"], name="exog2", observed_state_names=["data_3"])
+
+    mod = (reg1 + reg2).build(verbose=False)
+    assert mod.coords["exog1_state"] == ["a[data_1]", "b[data_1]", "a[data_2]", "b[data_2]"]
+    assert mod.coords["exog2_state"] == ["c[data_3]"]
+
+    Z = mod.ssm["design"].eval({"data_exog1": data_1, "data_exog2": data_2})
+    vec_block_diag = np.vectorize(block_diag, signature="(n,m),(o,p)->(q,r)")
+    assert Z.shape == (100, 3, 5)
+    assert np.allclose(
+        Z,
+        vec_block_diag(vec_block_diag(data_1[:, None, :], data_1[:, None, :]), data_2[:, None, :]),
+    )
+
+    x0 = mod.ssm["initial_state"].eval(
+        {
+            "beta_exog1": np.array([[1.0, 2.0], [3.0, 4.0]], dtype=config.floatX),
+            "beta_exog2": np.array([5.0], dtype=config.floatX),
+        }
+    )
+    np.testing.assert_allclose(x0, np.array([1.0, 2.0, 3.0, 4.0, 5.0], dtype=config.floatX))
+
+
 def test_filter_scans_time_varying_design_matrix(rng):
     time_idx = pd.date_range(start="2000-01-01", freq="D", periods=100)
     data = pd.DataFrame(rng.normal(size=(100, 2)), columns=["a", "b"], index=time_idx)

tests/statespace/models/structural/test_core.py

@@ -106,7 +106,7 @@ def test_extract_components_from_idata(rng):
     filter_prior = mod.sample_conditional_prior(prior)
     comp_prior = mod.extract_components_from_idata(filter_prior)
     comp_states = comp_prior.filtered_prior.coords["state"].values
-    expected_states = ["LevelTrend[level]", "LevelTrend[trend]", "seasonal", "exog[a]", "exog[b]"]
+    expected_states = ["level_trend[level]", "level_trend[trend]", "seasonal", "exog[a]", "exog[b]"]
     missing = set(comp_states) - set(expected_states)
 
     assert len(missing) == 0, missing