Add shared_state argument to LevelTrendComponent

Author: jessegrabowski

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

@@ -13,13 +13,11 @@ class LevelTrendComponent(Component):
     Parameters
     ----------
     order : int
-
         Number of time derivatives of the trend to include in the model. For example, when order=3, the trend will
         be of the form ``y = a + b * t + c * t ** 2``, where the coefficients ``a, b, c`` come from the initial
         state values.
 
     innovations_order : int or sequence of int, optional
-
         The number of stochastic innovations to include in the model. By default, ``innovations_order = order``
 
     name : str, default "level_trend"
@@ -28,6 +26,11 @@ class LevelTrendComponent(Component):
     observed_state_names : list[str] | None, default None
         List of strings for observed state labels. If None, defaults to ["data"].
 
+    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. This argument has no effect if `k_endog` is 1.
+
     Notes
     -----
     This class implements the level and trend components of the general structural time series model. In the most
@@ -120,7 +123,10 @@ def __init__(
         innovations_order: int | list[int] | None = None,
         name: str = "level_trend",
         observed_state_names: list[str] | None = None,
+        share_states: bool = False,
     ):
+        self.share_states = share_states
+
         if innovations_order is None:
             innovations_order = order
 
@@ -156,37 +162,50 @@ def __init__(
         super().__init__(
             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,
             observed_state_names=observed_state_names,
             measurement_error=False,
             combine_hidden_states=False,
-            obs_state_idxs=np.tile(np.array([1.0] + [0.0] * (k_states - 1)), k_endog),
+            obs_state_idxs=np.tile(
+                np.array([1.0] + [0.0] * (k_states - 1)), k_endog if not share_states else 1
+            ),
         )
 
     def populate_component_properties(self):
         k_endog = self.k_endog
-        k_states = self.k_states // k_endog
-        k_posdef = self.k_posdef // k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
+        k_posdef = self.k_posdef // k_endog_effective
 
         name_slice = POSITION_DERIVATIVE_NAMES[:k_states]
         self.param_names = [f"initial_{self.name}"]
         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
-        ]
+
+        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_dims = {f"initial_{self.name}": (f"state_{self.name}",)}
         self.coords = {f"state_{self.name}": base_names}
 
         if k_endog > 1:
+            self.coords[f"endog_{self.name}"] = self.observed_state_names
+
+        if k_endog_effective > 1:
             self.param_dims[f"state_{self.name}"] = (
                 f"endog_{self.name}",
                 f"state_{self.name}",
             )
             self.param_dims = {f"initial_{self.name}": (f"endog_{self.name}", f"state_{self.name}")}
-            self.coords[f"endog_{self.name}"] = self.observed_state_names
 
-        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,)
         self.param_info = {f"initial_{self.name}": {"shape": shape, "constraints": None}}
 
         if self.k_posdef > 0:
@@ -196,20 +215,23 @@ def populate_component_properties(self):
                 name for name, mask in zip(name_slice, self.innovations_order) if mask
             ]
 
-            self.shock_names = [
-                f"{name}[{obs_name}]"
-                for obs_name in self.observed_state_names
-                for name in base_shock_names
-            ]
+            if self.share_states:
+                self.shock_names = [f"{name}[{self.name}_shared]" for name in base_shock_names]
+            else:
+                self.shock_names = [
+                    f"{name}[{obs_name}]"
+                    for obs_name in self.observed_state_names
+                    for name in base_shock_names
+                ]
 
             self.param_dims[f"sigma_{self.name}"] = (
                 (f"shock_{self.name}",)
-                if k_endog == 1
+                if k_endog_effective == 1
                 else (f"endog_{self.name}", f"shock_{self.name}")
             )
             self.coords[f"shock_{self.name}"] = base_shock_names
             self.param_info[f"sigma_{self.name}"] = {
-                "shape": (k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
+                "shape": (k_posdef,) if k_endog_effective == 1 else (k_endog_effective, k_posdef),
                 "constraints": "Positive",
             }
 
@@ -218,40 +240,49 @@ def populate_component_properties(self):
 
     def make_symbolic_graph(self) -> None:
         k_endog = self.k_endog
-        k_states = self.k_states // k_endog
-        k_posdef = self.k_posdef // k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+
+        k_states = self.k_states // k_endog_effective
+        k_posdef = self.k_posdef // k_endog_effective
 
         initial_trend = self.make_and_register_variable(
             f"initial_{self.name}",
-            shape=(k_states,) if k_endog == 1 else (k_endog, k_states),
+            shape=(k_states,) if k_endog_effective == 1 else (k_endog, k_states),
         )
         self.ssm["initial_state", :] = initial_trend.ravel()
 
         triu_idx = pt.triu_indices(k_states)
         T = pt.zeros((k_states, k_states))[triu_idx[0], triu_idx[1]].set(1)
 
         self.ssm["transition", :, :] = pt.specify_shape(
-            pt.linalg.block_diag(*[T for _ in range(k_endog)]), (self.k_states, self.k_states)
+            pt.linalg.block_diag(*[T for _ in range(k_endog_effective)]),
+            (self.k_states, self.k_states),
         )
 
         R = np.eye(k_states)
         R = R[:, self.innovations_order]
 
         self.ssm["selection", :, :] = pt.specify_shape(
-            pt.linalg.block_diag(*[R for _ in range(k_endog)]), (self.k_states, self.k_posdef)
+            pt.linalg.block_diag(*[R for _ in range(k_endog_effective)]),
+            (self.k_states, self.k_posdef),
         )
 
         Z = np.array([1.0] + [0.0] * (k_states - 1)).reshape((1, -1))
 
-        self.ssm["design", :, :] = pt.specify_shape(
-            pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
-        )
+        if self.share_states:
+            self.ssm["design", :, :] = pt.specify_shape(
+                pt.join(0, *[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+            )
+        else:
+            self.ssm["design", :, :] = pt.specify_shape(
+                pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
+            )
 
         if k_posdef > 0:
             sigma_trend = self.make_and_register_variable(
                 f"sigma_{self.name}",
-                shape=(k_posdef,) if k_endog == 1 else (k_endog, k_posdef),
+                shape=(k_posdef,) if k_endog_effective == 1 else (k_endog, k_posdef),
             )
-            diag_idx = np.diag_indices(k_posdef * k_endog)
+            diag_idx = np.diag_indices(k_posdef * k_endog_effective)
             idx = np.s_["state_cov", diag_idx[0], diag_idx[1]]
             self.ssm[idx] = (sigma_trend**2).ravel()

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

@@ -91,6 +91,44 @@ def test_level_trend_multiple_observed_construction():
     )
 
 
+def test_level_trend_multiple_shared_construction():
+    mod = st.LevelTrendComponent(
+        order=2, innovations_order=1, observed_state_names=["data_1", "data_2"], share_states=True
+    )
+    mod = mod.build(verbose=False)
+
+    assert mod.k_endog == 2
+    assert mod.k_states == 2
+    assert mod.k_posdef == 1
+
+    assert mod.coords["state_level_trend"] == ["level", "trend"]
+    assert mod.coords["endog_level_trend"] == ["data_1", "data_2"]
+
+    assert mod.state_names == [
+        "level[level_trend_shared]",
+        "trend[level_trend_shared]",
+    ]
+    assert mod.shock_names == ["level[level_trend_shared]"]
+
+    Z, T, R = pytensor.function(
+        [], [mod.ssm["design"], mod.ssm["transition"], mod.ssm["selection"]], mode="FAST_COMPILE"
+    )()
+
+    np.testing.assert_allclose(
+        Z,
+        np.array(
+            [
+                [1.0, 0.0],
+                [1.0, 0.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(T, np.array([[1.0, 1.0], [0.0, 1.0]]))
+
+    np.testing.assert_allclose(R, np.array([[1.0], [0.0]]))
+
+
 def test_level_trend_multiple_observed(rng):
     mod = st.LevelTrendComponent(
         order=2, innovations_order=0, observed_state_names=["data_1", "data_2", "data_3"]
@@ -102,6 +140,19 @@ def test_level_trend_multiple_observed(rng):
     assert (np.diff(x, axis=0) == np.array([[1.0, 0.0, 2.0, 0.0, 3.0, 0.0]])).all().all()
 
 
+def test_level_trend_multiple_shared_observed(rng):
+    mod = st.LevelTrendComponent(
+        order=2,
+        innovations_order=0,
+        observed_state_names=["data_1", "data_2", "data_3"],
+        share_states=True,
+    )
+    params = {"initial_level_trend": np.array([10.0, 0.1])}
+    x, y = simulate_from_numpy_model(mod, rng, params)
+    np.testing.assert_allclose(y[:, 0], y[:, 1])
+    np.testing.assert_allclose(y[:, 0], y[:, 2])
+
+
 def test_add_level_trend_with_different_observed():
     mod_1 = st.LevelTrendComponent(
         name="ll", order=2, innovations_order=[0, 1], observed_state_names=["data_1"]
@@ -156,3 +207,77 @@ def test_add_level_trend_with_different_observed():
             ]
         ),
     )
+
+
+def test_mixed_shared_and_not_shared():
+    mod_1 = st.LevelTrendComponent(
+        name="individual",
+        order=2,
+        innovations_order=[0, 1],
+        observed_state_names=["data_1", "data_2"],
+    )
+    mod_2 = st.LevelTrendComponent(
+        name="joint",
+        order=2,
+        innovations_order=[1, 1],
+        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 == 6
+    assert mod.k_posdef == 4
+
+    assert mod.state_names == [
+        "level[data_1]",
+        "trend[data_1]",
+        "level[data_2]",
+        "trend[data_2]",
+        "level[joint_shared]",
+        "trend[joint_shared]",
+    ]
+
+    assert mod.shock_names == [
+        "trend[data_1]",
+        "trend[data_2]",
+        "level[joint_shared]",
+        "trend[joint_shared]",
+    ]
+
+    Z, T, R = pytensor.function(
+        [], [mod.ssm["design"], mod.ssm["transition"], mod.ssm["selection"]], mode="FAST_COMPILE"
+    )()
+
+    np.testing.assert_allclose(
+        Z, np.array([[1.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 1.0, 0.0, 1.0, 0.0]])
+    )
+
+    np.testing.assert_allclose(
+        T,
+        np.array(
+            [
+                [1.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0, 1.0, 1.0],
+                [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(
+        R,
+        np.array(
+            [
+                [0.0, 0.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0, 0.0],
+                [0.0, 0.0, 1.0, 0.0],
+                [0.0, 0.0, 0.0, 1.0],
+            ]
+        ),
+    )