Add shared_state argument to AutoregressiveComponent

Author: jessegrabowski

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

@@ -23,6 +23,11 @@ class AutoregressiveComponent(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
     -----
     An autoregressive component can be thought of as a way o introducing serially correlated errors into the model.
@@ -73,45 +78,58 @@ def __init__(
         order: int = 1,
         name: str = "auto_regressive",
         observed_state_names: list[str] | None = None,
+        share_states: bool = False,
     ):
         if observed_state_names is None:
             observed_state_names = ["data"]
 
-        k_posdef = k_endog = len(observed_state_names)
+        k_endog = len(observed_state_names)
+        k_endog_effective = k_posdef = 1 if share_states else k_endog
 
         order = order_to_mask(order)
         ar_lags = np.flatnonzero(order).ravel().astype(int) + 1
         k_states = len(order)
 
+        self.share_states = share_states
         self.order = order
         self.ar_lags = ar_lags
 
         super().__init__(
             name=name,
             k_endog=k_endog,
-            k_states=k_states * k_endog,
+            k_states=k_states * k_endog_effective,
             k_posdef=k_posdef,
             measurement_error=True,
             combine_hidden_states=True,
             observed_state_names=observed_state_names,
-            obs_state_idxs=np.tile(np.r_[[1.0], np.zeros(k_states - 1)], k_endog),
+            obs_state_idxs=np.tile(np.r_[[1.0], np.zeros(k_states - 1)], k_endog_effective),
         )
 
     def populate_component_properties(self):
-        k_states = self.k_states // self.k_endog  # this is also the number of AR lags
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
 
-        self.state_names = [
-            f"L{i + 1}[{state_name}]"
-            for state_name in self.observed_state_names
-            for i in range(k_states)
-        ]
+        k_states = self.k_states // k_endog_effective  # this is also the number of AR lags
+        base_names = [f"L{i + 1}_{self.name}" for i in range(k_states)]
+
+        if self.share_states:
+            self.state_names = [f"{name}[shared]" for name in base_names]
+            self.shock_names = [f"{self.name}[shared]"]
+        else:
+            self.state_names = [
+                f"{name}[{state_name}]"
+                for state_name in self.observed_state_names
+                for name in base_names
+            ]
+            self.shock_names = [
+                f"{self.name}[{obs_name}]" for obs_name in self.observed_state_names
+            ]
 
-        self.shock_names = [f"{self.name}[{obs_name}]" for obs_name in self.observed_state_names]
         self.param_names = [f"params_{self.name}", f"sigma_{self.name}"]
         self.param_dims = {f"params_{self.name}": (f"lag_{self.name}",)}
         self.coords = {f"lag_{self.name}": self.ar_lags.tolist()}
 
-        if self.k_endog > 1:
+        if k_endog_effective > 1:
             self.param_dims[f"params_{self.name}"] = (
                 f"endog_{self.name}",
                 f"lag_{self.name}",
@@ -140,26 +158,29 @@ def populate_component_properties(self):
 
     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
         k_posdef = self.k_posdef
 
         k_nonzero = int(sum(self.order))
         ar_params = self.make_and_register_variable(
-            f"params_{self.name}", shape=(k_nonzero,) if k_endog == 1 else (k_endog, k_nonzero)
+            f"params_{self.name}",
+            shape=(k_nonzero,) if k_endog_effective == 1 else (k_endog_effective, k_nonzero),
         )
         sigma_ar = self.make_and_register_variable(
-            f"sigma_{self.name}", shape=() if k_endog == 1 else (k_endog,)
+            f"sigma_{self.name}", shape=() if k_endog_effective == 1 else (k_endog_effective,)
         )
 
-        if k_endog == 1:
+        if k_endog_effective == 1:
             T = pt.eye(k_states, k=-1)
             ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
             T = T[ar_idx].set(ar_params)
 
         else:
             transition_matrices = []
 
-            for i in range(k_endog):
+            for i in range(k_endog_effective):
                 T = pt.eye(k_states, k=-1)
                 ar_idx = (np.zeros(k_nonzero, dtype="int"), np.nonzero(self.order)[0])
                 T = T[ar_idx].set(ar_params[i])
@@ -171,18 +192,21 @@ def make_symbolic_graph(self) -> None:
         self.ssm["transition", :, :] = T
 
         R = np.eye(k_states)
-        R_mask = np.full((k_states), False)
+        R_mask = np.full((k_states,), False)
         R_mask[0] = True
         R = R[:, R_mask]
 
         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, k_posdef)
         )
 
-        Z = pt.zeros((1, k_states))[0, 0].set(1.0)
-        self.ssm["design", :, :] = pt.specify_shape(
-            pt.linalg.block_diag(*[Z for _ in range(k_endog)]), (self.k_endog, self.k_states)
-        )
+        Zs = [pt.zeros((1, k_states))[0, 0].set(1.0) for _ in range(k_endog)]
+
+        if self.share_states:
+            Z = pt.join(0, *Zs)
+        else:
+            Z = pt.linalg.block_diag(*Zs)
+        self.ssm["design", :, :] = pt.specify_shape(Z, (k_endog, self.k_states))
 
         cov_idx = ("state_cov", *np.diag_indices(k_posdef))
         self.ssm[cov_idx] = sigma_ar**2

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

@@ -15,7 +15,6 @@
 def test_autoregressive_model(order, rng):
     ar = st.AutoregressiveComponent(order=order).build(verbose=False)
 
-    # Check coords
     _assert_basic_coords_correct(ar)
 
     lags = np.arange(len(order) if isinstance(order, list) else order, dtype="int") + 1
@@ -25,34 +24,34 @@ def test_autoregressive_model(order, rng):
 
 
 def test_autoregressive_multiple_observed_build(rng):
-    ar = st.AutoregressiveComponent(order=3, observed_state_names=["data_1", "data_2"])
+    ar = st.AutoregressiveComponent(order=3, name="ar", observed_state_names=["data_1", "data_2"])
     mod = ar.build(verbose=False)
 
     assert mod.k_endog == 2
     assert mod.k_states == 6
     assert mod.k_posdef == 2
 
     assert mod.state_names == [
-        "L1[data_1]",
-        "L2[data_1]",
-        "L3[data_1]",
-        "L1[data_2]",
-        "L2[data_2]",
-        "L3[data_2]",
+        "L1_ar[data_1]",
+        "L2_ar[data_1]",
+        "L3_ar[data_1]",
+        "L1_ar[data_2]",
+        "L2_ar[data_2]",
+        "L3_ar[data_2]",
     ]
 
-    assert mod.shock_names == ["auto_regressive[data_1]", "auto_regressive[data_2]"]
+    assert mod.shock_names == ["ar[data_1]", "ar[data_2]"]
 
     params = {
-        "params_auto_regressive": np.full(
+        "params_ar": np.full(
             (
                 2,
                 sum(ar.order),
             ),
             0.5,
             dtype=config.floatX,
         ),
-        "sigma_auto_regressive": np.array([0.05, 0.12]),
+        "sigma_ar": np.array([0.05, 0.12]),
     }
     _, _, _, _, T, Z, R, _, Q = mod._unpack_statespace_with_placeholders()
     input_vars = explicit_graph_inputs([T, Z, R, Q])
@@ -89,6 +88,33 @@ def test_autoregressive_multiple_observed_build(rng):
     np.testing.assert_allclose(Q, np.diag([0.05**2, 0.12**2]))
 
 
+def test_autoregressive_multiple_observed_shared():
+    ar = st.AutoregressiveComponent(
+        order=1,
+        name="latent",
+        observed_state_names=["data_1", "data_2", "data_3"],
+        share_states=True,
+    )
+    mod = ar.build(verbose=False)
+
+    assert mod.k_endog == 3
+    assert mod.k_states == 1
+    assert mod.k_posdef == 1
+
+    assert mod.state_names == ["L1_latent[shared]"]
+    assert mod.shock_names == ["latent[shared]"]
+    assert mod.coords["lag_latent"] == [1]
+    assert "endog_latent" not in mod.coords
+
+    outputs = [mod.ssm["transition"], mod.ssm["design"]]
+    params = {"params_latent": np.array([0.9])}
+    T, Z = pytensor.function(list(explicit_graph_inputs(outputs)), outputs)(**params)
+
+    np.testing.assert_allclose(np.array([[1.0], [1.0], [1.0]]), Z)
+
+    np.testing.assert_allclose(np.array([[0.9]]), T)
+
+
 def test_autoregressive_multiple_observed_data(rng):
     ar = st.AutoregressiveComponent(order=1, observed_state_names=["data_1", "data_2", "data_3"])
     mod = ar.build(verbose=False)
@@ -112,21 +138,130 @@ def test_add_autoregressive_different_observed():
 
     mod = (mod_1 + mod_2).build(verbose=False)
 
-    print(mod.coords)
-
     assert mod.k_endog == 2
     assert mod.k_states == 7
     assert mod.k_posdef == 2
     assert mod.state_names == [
-        "L1[data_1]",
-        "L1[data_2]",
-        "L2[data_2]",
-        "L3[data_2]",
-        "L4[data_2]",
-        "L5[data_2]",
-        "L6[data_2]",
+        f"L1_{mod_1.name}[data_1]",
+        f"L1_{mod_2.name}[data_2]",
+        f"L2_{mod_2.name}[data_2]",
+        f"L3_{mod_2.name}[data_2]",
+        f"L4_{mod_2.name}[data_2]",
+        f"L5_{mod_2.name}[data_2]",
+        f"L6_{mod_2.name}[data_2]",
     ]
 
     assert mod.shock_names == ["ar1[data_1]", "ar6[data_2]"]
     assert mod.coords["lag_ar1"] == [1]
     assert mod.coords["lag_ar6"] == [1, 2, 3, 4, 5, 6]
+
+
+def test_autoregressive_shared_and_not_shared():
+    shared = st.AutoregressiveComponent(
+        order=3,
+        name="shared_ar",
+        observed_state_names=["data_1", "data_2", "data_3"],
+        share_states=True,
+    )
+    individual = st.AutoregressiveComponent(
+        order=3,
+        name="individual_ar",
+        observed_state_names=["data_1", "data_2", "data_3"],
+        share_states=False,
+    )
+
+    mod = (shared + individual).build(verbose=False)
+
+    assert mod.k_endog == 3
+    assert mod.k_states == 3 + 3 * 3
+    assert mod.k_posdef == 4
+
+    assert mod.state_names == [
+        "L1_shared_ar[shared]",
+        "L2_shared_ar[shared]",
+        "L3_shared_ar[shared]",
+        "L1_individual_ar[data_1]",
+        "L2_individual_ar[data_1]",
+        "L3_individual_ar[data_1]",
+        "L1_individual_ar[data_2]",
+        "L2_individual_ar[data_2]",
+        "L3_individual_ar[data_2]",
+        "L1_individual_ar[data_3]",
+        "L2_individual_ar[data_3]",
+        "L3_individual_ar[data_3]",
+    ]
+
+    assert mod.shock_names == [
+        "shared_ar[shared]",
+        "individual_ar[data_1]",
+        "individual_ar[data_2]",
+        "individual_ar[data_3]",
+    ]
+    assert mod.coords["lag_shared_ar"] == [1, 2, 3]
+    assert mod.coords["lag_individual_ar"] == [1, 2, 3]
+
+    outputs = [mod.ssm["transition"], mod.ssm["design"], mod.ssm["selection"], mod.ssm["state_cov"]]
+    T, Z, R, Q = pytensor.function(
+        list(explicit_graph_inputs(outputs)),
+        outputs,
+    )(
+        **{
+            "params_shared_ar": np.array([0.9, 0.8, 0.7]),
+            "params_individual_ar": np.full((3, 3), 0.5),
+            "sigma_shared_ar": np.array(0.1),
+            "sigma_individual_ar": np.array([0.05, 0.12, 0.22]),
+        }
+    )
+
+    np.testing.assert_allclose(
+        T,
+        np.array(
+            [
+                [0.9, 0.8, 0.7, 0.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, 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, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.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, 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, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5, 0.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, 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],
+                [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.5, 0.5, 0.5],
+                [0.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, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(
+        Z,
+        np.array(
+            [
+                [1.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, 0.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, 0.0, 0.0, 1.0, 0.0, 0.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(
+        R,
+        np.array(
+            [
+                [1.0, 0.0, 0.0, 0.0],
+                [0.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, 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, 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],
+                [0.0, 0.0, 0.0, 0.0],
+            ]
+        ),
+    )
+
+    np.testing.assert_allclose(Q, np.diag([0.1, 0.05, 0.12, 0.22]) ** 2)