Add shared_state argument to TimeSeasonality

Author: jessegrabowski

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

@@ -44,6 +44,11 @@ class TimeSeasonality(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
     -----
     A seasonal effect is any pattern that repeats at fixed intervals. There are several ways to model such effects;
@@ -235,6 +240,7 @@ def __init__(
         state_names: list | None = None,
         remove_first_state: bool = True,
         observed_state_names: list[str] | None = None,
+        share_states: bool = False,
     ):
         if observed_state_names is None:
             observed_state_names = ["data"]
@@ -261,6 +267,7 @@ def __init__(
                 )
             state_names = state_names.copy()
 
+        self.share_states = share_states
         self.innovations = innovations
         self.duration = duration
         self.remove_first_state = remove_first_state
@@ -281,44 +288,53 @@ 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 if share_states else k_states * k_endog,
+            k_posdef=k_posdef if share_states else k_posdef * k_endog,
             observed_state_names=observed_state_names,
             measurement_error=False,
             combine_hidden_states=True,
-            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)), 1 if share_states else k_endog
+            ),
         )
 
     def populate_component_properties(self):
-        k_states = self.k_states // self.k_endog
         k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
 
-        self.state_names = [
-            f"{state_name}[{endog_name}]"
-            for endog_name in self.observed_state_names
-            for state_name in self.provided_state_names
-        ]
+        k_states = self.k_states // k_endog_effective
+
+        if self.share_states:
+            self.state_names = [
+                f"{state_name}[{self.name}_shared]" for state_name in self.provided_state_names
+            ]
+        else:
+            self.state_names = [
+                f"{state_name}[{endog_name}]"
+                for endog_name in self.observed_state_names
+                for state_name in self.provided_state_names
+            ]
         self.param_names = [f"coefs_{self.name}"]
 
         self.param_info = {
             f"coefs_{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_effective, k_states),
                 "constraints": None,
                 "dims": (f"state_{self.name}",)
-                if k_endog == 1
+                if k_endog_effective == 1
                 else (f"endog_{self.name}", f"state_{self.name}"),
             }
         }
 
         self.param_dims = {
             f"coefs_{self.name}": (f"state_{self.name}",)
-            if k_endog == 1
+            if k_endog_effective == 1
             else (f"endog_{self.name}", f"state_{self.name}")
         }
 
         self.coords = (
             {f"state_{self.name}": self.provided_state_names}
-            if k_endog == 1
+            if k_endog_effective == 1
             else {
                 f"endog_{self.name}": self.observed_state_names,
                 f"state_{self.name}": self.provided_state_names,
@@ -332,14 +348,19 @@ def populate_component_properties(self):
                 "constraints": "Positive",
                 "dims": None,
             }
-            self.shock_names = [f"{self.name}[{name}]" for name in self.observed_state_names]
+            if self.share_states:
+                self.shock_names = [f"{self.name}[shared]"]
+            else:
+                self.shock_names = [f"{self.name}[{name}]" for name in self.observed_state_names]
 
     def make_symbolic_graph(self) -> None:
-        k_states = self.k_states // self.k_endog
+        k_endog = self.k_endog
+        k_endog_effective = 1 if self.share_states else k_endog
+        k_states = self.k_states // k_endog_effective
         duration = self.duration
+
         k_unique_states = k_states // duration
-        k_posdef = self.k_posdef // self.k_endog
-        k_endog = self.k_endog
+        k_posdef = self.k_posdef // k_endog_effective
 
         if self.remove_first_state:
             # In this case, parameters are normalized to sum to zero, so the current state is the negative sum of
@@ -371,16 +392,18 @@ def make_symbolic_graph(self) -> None:
             T = pt.eye(k_states, k=1)
             T = pt.set_subtensor(T[-1, 0], 1)
 
-        self.ssm["transition", :, :] = pt.linalg.block_diag(*[T for _ in range(k_endog)])
+        self.ssm["transition", :, :] = pt.linalg.block_diag(*[T for _ in range(k_endog_effective)])
 
         Z = pt.zeros((1, k_states))[0, 0].set(1)
-        self.ssm["design", :, :] = pt.linalg.block_diag(*[Z for _ in range(k_endog)])
+        self.ssm["design", :, :] = pt.linalg.block_diag(*[Z for _ in range(k_endog_effective)])
 
         initial_states = self.make_and_register_variable(
             f"coefs_{self.name}",
-            shape=(k_unique_states,) if k_endog == 1 else (k_endog, k_unique_states),
+            shape=(k_unique_states,)
+            if k_endog_effective == 1
+            else (k_endog_effective, k_unique_states),
         )
-        if k_endog == 1:
+        if k_endog_effective == 1:
             self.ssm["initial_state", :] = pt.extra_ops.repeat(initial_states, duration, axis=0)
         else:
             self.ssm["initial_state", :] = pt.extra_ops.repeat(
@@ -389,11 +412,11 @@ def make_symbolic_graph(self) -> None:
 
         if self.innovations:
             R = pt.zeros((k_states, k_posdef))[0, 0].set(1.0)
-            self.ssm["selection", :, :] = pt.join(0, *[R for _ in range(k_endog)])
+            self.ssm["selection", :, :] = pt.join(0, *[R for _ in range(k_endog_effective)])
             season_sigma = 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,)
             )
-            cov_idx = ("state_cov", *np.diag_indices(k_posdef * k_endog))
+            cov_idx = ("state_cov", *np.diag_indices(k_posdef * k_endog_effective))
             self.ssm[cov_idx] = season_sigma**2
 
 

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

@@ -146,6 +146,42 @@ def test_time_seasonality_multiple_observed(rng, d, remove_first_state):
         np.testing.assert_allclose(matrix, expected)
 
 
+def test_time_seasonality_shared_states():
+    mod = st.TimeSeasonality(
+        season_length=3,
+        duration=1,
+        innovations=True,
+        name="season",
+        state_names=["season_1", "season_2", "season_3"],
+        observed_state_names=["data_1", "data_2"],
+        remove_first_state=False,
+        share_states=True,
+    )
+
+    assert mod.k_endog == 2
+    assert mod.k_states == 3
+    assert mod.k_posdef == 1
+
+    assert mod.coords["state_season"] == ["season_1", "season_2", "season_3"]
+
+    assert mod.state_names == [
+        "season_1[season_shared]",
+        "season_2[season_shared]",
+        "season_3[season_shared]",
+    ]
+    assert mod.shock_names == ["season[shared]"]
+
+    Z, T, R = pytensor.function(
+        [], [mod.ssm["design"], mod.ssm["transition"], mod.ssm["selection"]], mode="FAST_COMPILE"
+    )()
+
+    np.testing.assert_allclose(np.array([[1.0, 0.0, 0.0], [1.0, 0.0, 0.0]]), Z)
+
+    np.testing.assert_allclose(np.array([[0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [1.0, 0.0, 0.0]]), T)
+
+    np.testing.assert_allclose(np.array([[1.0], [0.0], [0.0]]), R)
+
+
 @pytest.mark.parametrize("d1, d2", [(1, 1), (1, 3), (3, 1), (3, 3)])
 def test_add_two_time_seasonality_different_observed(rng, d1, d2):
     mod1 = st.TimeSeasonality(