Add shared_state argument to RegressionComponent

Author: Dekermanjian

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

@@ -31,6 +31,11 @@ class RegressionComponent(Component):
         Whether to include stochastic innovations in the regression coefficients,
         allowing them to vary over time. If True, coefficients follow a random walk.
 
+    share_states: bool, default False
+        Whether latent states are shared across the observed states. If True, there will be only one set of latent
+        states, which are observed by all observed states. If False, each observed state has its own set of
+        latent states.
+
     Notes
     -----
     This component implements regression with exogenous variables in a structural time series
@@ -107,7 +112,10 @@ def __init__(
         state_names: list[str] | None = None,
         observed_state_names: list[str] | None = None,
         innovations=False,
+        share_states: bool = False,
     ):
+        self.share_states = share_states
+
         if observed_state_names is None:
             observed_state_names = ["data"]
 
@@ -121,8 +129,8 @@ def __init__(
         super().__init__(
             name=name,
             k_endog=k_endog,
-            k_states=k_states * k_endog,
-            k_posdef=k_posdef * k_endog,
+            k_states=k_states * k_endog if not share_states else k_states,
+            k_posdef=k_posdef * k_endog if not share_states else k_posdef,
             state_names=self.state_names,
             observed_state_names=observed_state_names,
             measurement_error=False,
@@ -153,54 +161,74 @@ 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
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
 
         betas = self.make_and_register_variable(
-            f"beta_{self.name}", shape=(k_endog, k_states) if k_endog > 1 else (k_states,)
+            f"beta_{self.name}", shape=(k_endog, k_states) if k_endog_effective > 1 else (k_states,)
         )
         regression_data = self.make_and_register_data(f"data_{self.name}", shape=(None, 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.share_states:
+            self.ssm["design"] = pt.specify_shape(
+                pt.join(1, *[pt.expand_dims(regression_data, 1) for _ in range(k_endog)]),
+                (None, k_endog, self.k_states),
+            )
+        else:
+            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}", (k_states,) if k_endog == 1 else (k_endog, k_states)
+                f"sigma_beta_{self.name}",
+                (k_states,) if k_endog_effective == 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.ravel() ** 2
 
     def populate_component_properties(self) -> None:
         k_endog = self.k_endog
-        k_states = self.k_states // k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
 
-        self.shock_names = self.state_names
+        if self.share_states:
+            self.shock_names = [f"{state_name}_shared" for state_name in self.state_names]
+        else:
+            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}": (f"endog_{self.name}", f"state_{self.name}")
-            if k_endog > 1
+            if k_endog_effective > 1
             else (f"state_{self.name}",)
         }
 
         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
-        ]
+
+        if self.share_states:
+            self.state_names = [f"{name}[{self.name}_shared]" for name in base_names]
+        else:
+            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": (k_endog, k_states) if k_endog > 1 else (k_states,),
+                "shape": (k_endog_effective, k_states) if k_endog_effective > 1 else (k_states,),
                 "constraints": None,
                 "dims": (f"endog_{self.name}", f"state_{self.name}")
-                if k_endog > 1
+                if k_endog_effective > 1
                 else (f"state_{self.name}",),
             },
         }
@@ -223,6 +251,6 @@ def populate_component_properties(self) -> None:
                 "shape": (k_states,),
                 "constraints": "Positive",
                 "dims": (f"state_{self.name}",)
-                if k_endog == 1
+                if k_endog_effective == 1
                 else (f"endog_{self.name}", f"state_{self.name}"),
             }

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

@@ -7,6 +7,7 @@
 from pytensor import config
 from pytensor import tensor as pt
 from pytensor.graph.basic import explicit_graph_inputs
+from scipy.linalg import block_diag
 
 from pymc_extras.statespace.models import structural as st
 from tests.statespace.models.structural.conftest import _assert_basic_coords_correct
@@ -235,3 +236,104 @@ def test_filter_scans_time_varying_design_matrix(self, rng, time_series_data, in
         if innovations:
             # Check that sigma_beta parameter is included in the prior
             assert "sigma_beta_exog" in prior.prior.data_vars
+
+
+def test_regression_multiple_shared_construction():
+    rc = st.RegressionComponent(
+        state_names=["A"],
+        observed_state_names=["data_1", "data_2"],
+        innovations=True,
+        share_states=True,
+    )
+    mod = rc.build(verbose=False)
+
+    assert mod.k_endog == 2
+    assert mod.k_states == 1
+    assert mod.k_posdef == 1
+
+    assert mod.coords["state_regression"] == ["A"]
+    assert mod.coords["endog_regression"] == ["data_1", "data_2"]
+
+    assert mod.state_names == [
+        "A[regression_shared]",
+    ]
+
+    assert mod.shock_names == ["A_shared"]
+
+    data = np.random.standard_normal(size=(10, 1))
+    Z = mod.ssm["design"].eval({"data_regression": data})
+    T = mod.ssm["transition"].eval()
+    R = mod.ssm["selection"].eval()
+
+    np.testing.assert_allclose(
+        Z,
+        np.hstack(
+            [
+                data,
+                data,
+            ]
+        )[:, :, np.newaxis],
+    )
+
+    np.testing.assert_allclose(T, np.array([[1.0]]))
+    np.testing.assert_allclose(R, np.array([[1.0]]))
+
+
+def test_regression_multiple_shared_observed(rng):
+    mod = st.RegressionComponent(
+        state_names=["A"],
+        observed_state_names=["data_1", "data_2", "data_3"],
+        innovations=False,
+        share_states=True,
+    )
+    data = np.random.standard_normal(size=(10, 1))
+
+    params = {"beta_regression": np.array([1.0])}
+    data_dict = {"data_regression": data}
+    x, y = simulate_from_numpy_model(mod, rng, params, data_dict, steps=data.shape[0])
+    np.testing.assert_allclose(y[:, 0], y[:, 1])
+    np.testing.assert_allclose(y[:, 0], y[:, 2])
+
+
+def test_regression_mixed_shared_and_not_shared():
+    mod_1 = st.RegressionComponent(
+        name="individual",
+        state_names=["A"],
+        observed_state_names=["data_1", "data_2"],
+    )
+    mod_2 = st.RegressionComponent(
+        name="joint",
+        state_names=["B", "C"],
+        observed_state_names=["data_1", "data_2"],
+        share_states=True,
+    )
+
+    mod = (mod_1 + mod_2).build(verbose=False)
+
+    assert mod.k_endog == 2
+    assert mod.k_states == 4
+    assert mod.k_posdef == 4
+
+    assert mod.state_names == ["A[data_1]", "A[data_2]", "B[joint_shared]", "C[joint_shared]"]
+    assert mod.shock_names == ["A", "B_shared", "C_shared"]
+
+    data_joint = np.random.standard_normal(size=(10, 2))
+    data_individual = np.random.standard_normal(size=(10, 1))
+    Z = mod.ssm["design"].eval({"data_joint": data_joint, "data_individual": data_individual})
+    T = mod.ssm["transition"].eval()
+    R = mod.ssm["selection"].eval()
+
+    np.testing.assert_allclose(
+        Z,
+        np.concat(
+            (
+                block_diag(*[data_individual[:, np.newaxis] for _ in range(mod.k_endog)]),
+                np.concat((data_joint[:, np.newaxis], data_joint[:, np.newaxis]), axis=1),
+            ),
+            axis=2,
+        ),
+    )
+
+    np.testing.assert_allclose(T, np.eye(mod.k_states))
+
+    np.testing.assert_allclose(R, np.eye(mod.k_states))