Allow multiple observed in TimeSeasonality component

Author: jessegrabowski

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

@@ -154,27 +154,41 @@ def __init__(
             # TODO: Can this be stashed and reconstructed automatically somehow?
             state_names.pop(0)
 
+        self.provided_state_names = state_names
+
         k_states = season_length - int(self.remove_first_state)
+        k_endog = len(observed_state_names)
+        k_posdef = int(innovations)
 
         super().__init__(
             name=name,
-            k_endog=len(observed_state_names),
-            k_states=k_states,
-            k_posdef=int(innovations),
-            state_names=state_names,
+            k_endog=k_endog,
+            k_states=k_states * k_endog,
+            k_posdef=k_posdef * k_endog,
             observed_state_names=observed_state_names,
             measurement_error=False,
             combine_hidden_states=True,
-            obs_state_idxs=np.r_[[1.0], np.zeros(k_states - 1)],
+            obs_state_idxs=np.tile(np.array([1.0] + [0.0] * (k_states - 1)), k_endog),
         )
 
     def populate_component_properties(self):
+        k_states = self.k_states // self.k_endog
+        k_endog = self.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
+        ]
         self.param_names = [f"{self.name}_coefs"]
+
         self.param_info = {
             f"{self.name}_coefs": {
-                "shape": (self.k_states,),
+                "shape": (k_states,) if k_endog == 1 else (k_endog, k_states),
                 "constraints": None,
-                "dims": (f"{self.name}_state",),
+                "dims": (f"{self.name}_state",)
+                if k_endog == 1
+                else (f"{self.name}_endog", f"{self.name}_state"),
             }
         }
         self.param_dims = {f"{self.name}_coefs": (f"{self.name}_state",)}
@@ -187,32 +201,41 @@ def populate_component_properties(self):
                 "constraints": "Positive",
                 "dims": None,
             }
-            self.shock_names = [f"{self.name}"]
+            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_posdef = self.k_posdef // self.k_endog
+        k_endog = self.k_endog
+
         if self.remove_first_state:
             # In this case, parameters are normalized to sum to zero, so the current state is the negative sum of
             # all previous states.
-            T = np.eye(self.k_states, k=-1)
+            T = np.eye(k_states, k=-1)
             T[0, :] = -1
         else:
             # In this case we assume the user to be responsible for ensuring the states sum to zero, so T is just a
             # circulant matrix that cycles between the states.
-            T = np.eye(self.k_states, k=1)
+            T = np.eye(k_states, k=1)
             T[-1, 0] = 1
 
-        self.ssm["transition", :, :] = T
-        self.ssm["design", 0, 0] = 1
+        self.ssm["transition", :, :] = pt.linalg.block_diag(*[T for _ in range(k_endog)])
+
+        Z = pt.zeros((1, k_states))[0, 0].set(1)
+        self.ssm["design", :, :] = pt.linalg.block_diag(*[Z for _ in range(k_endog)])
 
         initial_states = self.make_and_register_variable(
-            f"{self.name}_coefs", shape=(self.k_states,)
+            f"{self.name}_coefs", shape=(k_states,) if k_endog == 1 else (k_endog, k_states)
         )
-        self.ssm["initial_state", np.arange(self.k_states, dtype=int)] = initial_states
+        self.ssm["initial_state", :] = initial_states.ravel()
 
         if self.innovations:
-            self.ssm["selection", 0, 0] = 1
-            season_sigma = self.make_and_register_variable(f"sigma_{self.name}", shape=())
-            cov_idx = ("state_cov", *np.diag_indices(1))
+            R = pt.zeros((k_states, k_posdef))[0, 0].set(1.0)
+            self.ssm["selection", :, :] = pt.join(0, *[R for _ in range(k_endog)])
+            season_sigma = self.make_and_register_variable(
+                f"sigma_{self.name}", shape=() if k_endog == 1 else (k_endog,)
+            )
+            cov_idx = ("state_cov", *np.diag_indices(k_posdef * k_endog))
             self.ssm[cov_idx] = season_sigma**2
 
 

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

@@ -1,7 +1,9 @@
 import numpy as np
+import pytensor
 import pytest
 
 from pytensor import config
+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
@@ -35,7 +37,7 @@ def random_word(rng):
     x0[0] = 1
 
     params = {"season_coefs": x0}
-    if mod.innovations:
+    if innovations:
         params["sigma_season"] = 0.0
 
     x, y = simulate_from_numpy_model(mod, rng, params)
@@ -44,12 +46,175 @@ def random_word(rng):
         assert_pattern_repeats(y, s, atol=ATOL, rtol=RTOL)
 
     # Check coords
-    mod.build(verbose=False)
+    mod = mod.build(verbose=False)
     _assert_basic_coords_correct(mod)
     test_slice = slice(1, None) if remove_first_state else slice(None)
     assert mod.coords["season_state"] == state_names[test_slice]
 
 
+@pytest.mark.parametrize(
+    "remove_first_state", [True, False], ids=["remove_first_state", "keep_first_state"]
+)
+def test_time_seasonality_multiple_observed(rng, remove_first_state):
+    s = 3
+    state_names = [f"state_{i}" for i in range(s)]
+    mod = st.TimeSeasonality(
+        season_length=s,
+        innovations=True,
+        name="season",
+        state_names=state_names,
+        observed_state_names=["data_1", "data_2"],
+        remove_first_state=remove_first_state,
+    )
+    x0 = np.zeros((mod.k_endog, mod.k_states // mod.k_endog), dtype=config.floatX)
+
+    expected_states = [
+        f"state_{i}[data_{j}]" for j in range(1, 3) for i in range(int(remove_first_state), s)
+    ]
+    assert mod.state_names == expected_states
+    assert mod.shock_names == ["season[data_1]", "season[data_2]"]
+
+    x0[0, 0] = 1
+    x0[1, 0] = 2.0
+
+    params = {"season_coefs": x0, "sigma_season": np.array([0.0, 0.0], dtype=config.floatX)}
+
+    x, y = simulate_from_numpy_model(mod, rng, params, steps=123)
+    assert_pattern_repeats(y[:, 0], s, atol=ATOL, rtol=RTOL)
+    assert_pattern_repeats(y[:, 1], s, atol=ATOL, rtol=RTOL)
+
+    mod = mod.build(verbose=False)
+    x0, *_, T, Z, R, _, Q = mod._unpack_statespace_with_placeholders()
+
+    input_vars = explicit_graph_inputs([x0, T, Z, R, Q])
+
+    fn = pytensor.function(
+        inputs=list(input_vars),
+        outputs=[x0, T, Z, R, Q],
+        mode="FAST_COMPILE",
+    )
+
+    params["sigma_season"] = np.array([0.1, 0.8], dtype=config.floatX)
+    x0, T, Z, R, Q = fn(**params)
+
+    if remove_first_state:
+        expected_x0 = np.array([1.0, 0.0, 2.0, 0.0])
+
+        expected_T = np.array(
+            [
+                [-1.0, -1.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, -1.0, -1.0],
+                [0.0, 0.0, 1.0, 0.0],
+            ]
+        )
+        expected_R = np.array([[1.0, 1.0], [0.0, 0.0], [1.0, 1.0], [0.0, 0.0]])
+        expected_Z = np.array([[1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0]])
+
+    else:
+        expected_x0 = np.array([1.0, 0.0, 0.0, 2.0, 0.0, 0.0])
+        expected_T = np.array(
+            [
+                [0.0, 1.0, 0.0, 0.0, 0.0, 0.0],
+                [0.0, 0.0, 1.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, 1.0, 0.0],
+                [0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+                [0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+            ]
+        )
+        expected_R = np.array(
+            [[1.0, 1.0], [0.0, 0.0], [0.0, 0.0], [1.0, 1.0], [0.0, 0.0], [0.0, 0.0]]
+        )
+        expected_Z = np.array([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0, 0.0]])
+
+    expected_Q = np.array([[0.1**2, 0.0], [0.0, 0.8**2]])
+
+    for matrix, expected in zip(
+        [x0, T, Z, R, Q],
+        [expected_x0, expected_T, expected_Z, expected_R, expected_Q],
+    ):
+        np.testing.assert_allclose(matrix, expected)
+
+
+def test_add_two_time_seasonality_different_observed(rng):
+    mod1 = st.TimeSeasonality(
+        season_length=3,
+        innovations=True,
+        name="season1",
+        state_names=[f"state_{i}" for i in range(3)],
+        observed_state_names=["data_1"],
+        remove_first_state=False,
+    )
+    mod2 = st.TimeSeasonality(
+        season_length=5,
+        innovations=True,
+        name="season2",
+        state_names=[f"state_{i}" for i in range(5)],
+        observed_state_names=["data_2"],
+    )
+
+    mod = (mod1 + mod2).build(verbose=False)
+
+    params = {
+        "season1_coefs": np.array([1.0, 0.0, 0.0], dtype=config.floatX),
+        "season2_coefs": np.array([3.0, 0.0, 0.0, 0.0], dtype=config.floatX),
+        "sigma_season1": np.array(0.0, dtype=config.floatX),
+        "sigma_season2": np.array(0.0, dtype=config.floatX),
+        "initial_state_cov": np.eye(mod.k_states, dtype=config.floatX),
+    }
+
+    x, y = simulate_from_numpy_model(mod, rng, params, steps=3 * 5 * 5)
+    assert_pattern_repeats(y[:, 0], 3, atol=ATOL, rtol=RTOL)
+    assert_pattern_repeats(y[:, 1], 5, atol=ATOL, rtol=RTOL)
+
+    assert mod.state_names == [
+        "state_0[data_1]",
+        "state_1[data_1]",
+        "state_2[data_1]",
+        "state_1[data_2]",
+        "state_2[data_2]",
+        "state_3[data_2]",
+        "state_4[data_2]",
+    ]
+
+    assert mod.shock_names == ["season1[data_1]", "season2[data_2]"]
+
+    x0, *_, T = mod._unpack_statespace_with_placeholders()[:5]
+    input_vars = explicit_graph_inputs([x0, T])
+    fn = pytensor.function(
+        inputs=list(input_vars),
+        outputs=[x0, T],
+        mode="FAST_COMPILE",
+    )
+
+    x0, T = fn(
+        season1_coefs=np.array([1.0, 0.0, 0.0], dtype=config.floatX),
+        season2_coefs=np.array([3.0, 0.0, 0.0, 1.2], dtype=config.floatX),
+    )
+
+    np.testing.assert_allclose(
+        np.array([1.0, 0.0, 0.0, 3.0, 0.0, 0.0, 1.2]), x0, atol=ATOL, rtol=RTOL
+    )
+
+    np.testing.assert_allclose(
+        np.array(
+            [
+                [0.0, 1.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, 0.0, 0.0],
+                [0.0, 0.0, 0.0, -1.0, -1.0, -1.0, -1.0],
+                [0.0, 0.0, 0.0, 1.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, 1.0, 0.0],
+            ]
+        ),
+        T,
+        atol=ATOL,
+        rtol=RTOL,
+    )
+
+
 def get_shift_factor(s):
     s_str = str(s)
     if "." not in s_str: