Allow multiple observed in FrequencySeasonality component

Author: jessegrabowski

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

@@ -301,8 +301,10 @@ def __init__(
         if observed_state_names is None:
             observed_state_names = ["data"]
 
+        k_endog = len(observed_state_names)
+
         if n is None:
-            n = int(season_length // 2)
+            n = int(season_length / 2)
         if name is None:
             name = f"Frequency[s={season_length}, n={n}]"
 
@@ -319,58 +321,97 @@ def __init__(
 
         obs_state_idx = np.zeros(k_states)
         obs_state_idx[slice(0, k_states, 2)] = 1
+        obs_state_idx = np.tile(obs_state_idx, k_endog)
 
         super().__init__(
             name=name,
-            k_endog=1,
-            k_states=k_states,
-            k_posdef=k_states * int(self.innovations),
+            k_endog=k_endog,
+            k_states=k_states * k_endog,
+            k_posdef=k_states * int(self.innovations) * k_endog,
             observed_state_names=observed_state_names,
             measurement_error=False,
             combine_hidden_states=True,
             obs_state_idxs=obs_state_idx,
         )
 
     def make_symbolic_graph(self) -> None:
-        self.ssm["design", 0, slice(0, self.k_states, 2)] = 1
+        k_endog = self.k_endog
+        k_states = self.k_states // k_endog
+        k_posdef = self.k_posdef // k_endog
+        n_coefs = self.n_coefs
 
-        init_state = self.make_and_register_variable(f"{self.name}", shape=(self.n_coefs,))
+        Z = pt.zeros((1, k_states))[0, slice(0, k_states, 2)].set(1.0)
 
-        init_state_idx = np.arange(self.n_coefs, dtype=int)
-        self.ssm["initial_state", init_state_idx] = init_state
+        self.ssm["design", :, :] = pt.linalg.block_diag(*[Z for _ in range(k_endog)])
+
+        init_state = self.make_and_register_variable(
+            f"{self.name}", shape=(n_coefs,) if k_endog == 1 else (k_endog, n_coefs)
+        )
+
+        init_state_idx = np.concatenate(
+            [
+                np.arange(k_states * i, (i + 1) * k_states, dtype=int)[:n_coefs]
+                for i in range(k_endog)
+            ],
+            axis=0,
+        )
+
+        self.ssm["initial_state", init_state_idx] = init_state.ravel()
 
         T_mats = [_frequency_transition_block(self.season_length, j + 1) for j in range(self.n)]
         T = pt.linalg.block_diag(*T_mats)
-        self.ssm["transition", :, :] = T
+        self.ssm["transition", :, :] = pt.linalg.block_diag(*[T for _ in range(k_endog)])
 
         if self.innovations:
-            sigma_season = self.make_and_register_variable(f"sigma_{self.name}", shape=())
-            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * sigma_season**2
-            self.ssm["selection", :, :] = np.eye(self.k_states)
+            sigma_season = self.make_and_register_variable(
+                f"sigma_{self.name}", shape=() if k_endog == 1 else (k_endog,)
+            )
+            self.ssm["selection", :, :] = pt.eye(self.k_states)
+            self.ssm["state_cov", :, :] = pt.eye(self.k_posdef) * pt.repeat(
+                sigma_season**2, k_posdef
+            )
 
     def populate_component_properties(self):
-        self.state_names = [f"{self.name}_{f}_{i}" for i in range(self.n) for f in ["Cos", "Sin"]]
+        k_endog = self.k_endog
+        n_coefs = self.n_coefs
+        k_states = self.k_states // k_endog
+
+        self.state_names = [
+            f"{self.name}_{f}_{i}[{obs_state_name}]"
+            for obs_state_name in self.observed_state_names
+            for i in range(self.n)
+            for f in ["Cos", "Sin"]
+        ]
         self.param_names = [f"{self.name}"]
 
         self.param_dims = {self.name: (f"{self.name}_state",)}
         self.param_info = {
             f"{self.name}": {
-                "shape": (self.k_states - int(self.last_state_not_identified),),
+                "shape": (n_coefs,) if k_endog == 1 else (k_endog, n_coefs),
                 "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"),
             }
         }
 
-        init_state_idx = np.arange(self.k_states, dtype=int)
-        if self.last_state_not_identified:
-            init_state_idx = init_state_idx[:-1]
+        # Regardless of whether the fourier basis are saturated, there will always be one symbolic state per basis.
+        # That's why the self.states is just a simple loop over everything. But when saturated, one of those states
+        # doesn't have an associated **parameter**, so the coords need to be adjusted to reflect this.
+        init_state_idx = np.concatenate(
+            [
+                np.arange(k_states * i, (i + 1) * k_states, dtype=int)[:n_coefs]
+                for i in range(k_endog)
+            ],
+            axis=0,
+        )
         self.coords = {f"{self.name}_state": [self.state_names[i] for i in init_state_idx]}
 
         if self.innovations:
             self.shock_names = self.state_names.copy()
             self.param_names += [f"sigma_{self.name}"]
             self.param_info[f"sigma_{self.name}"] = {
-                "shape": (),
+                "shape": () if k_endog == 1 else (k_endog, n_coefs),
                 "constraints": "Positive",
-                "dims": None,
+                "dims": None if k_endog == 1 else (f"{self.name}_endog",),
             }

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

@@ -236,7 +236,7 @@ def test_frequency_seasonality(n, s, rng):
     assert_pattern_repeats(y, T, atol=ATOL, rtol=RTOL)
 
     # Check coords
-    mod.build(verbose=False)
+    mod = mod.build(verbose=False)
     _assert_basic_coords_correct(mod)
     if n is None:
         n = int(s // 2)
@@ -246,3 +246,194 @@ def test_frequency_seasonality(n, s, rng):
     if s / n == 2.0:
         states.pop()
     assert mod.coords["season_state"] == states
+
+
+def test_frequency_seasonality_multiple_observed(rng):
+    observed_state_names = ["data_1", "data_2"]
+    season_length = 4
+    mod = st.FrequencySeasonality(
+        season_length=season_length,
+        n=None,
+        name="season",
+        innovations=True,
+        observed_state_names=observed_state_names,
+    )
+    expected_state_names = [
+        "season_Cos_0[data_1]",
+        "season_Sin_0[data_1]",
+        "season_Cos_1[data_1]",
+        "season_Sin_1[data_1]",
+        "season_Cos_0[data_2]",
+        "season_Sin_0[data_2]",
+        "season_Cos_1[data_2]",
+        "season_Sin_1[data_2]",
+    ]
+    assert mod.state_names == expected_state_names
+    assert mod.shock_names == [
+        "season_Cos_0[data_1]",
+        "season_Sin_0[data_1]",
+        "season_Cos_1[data_1]",
+        "season_Sin_1[data_1]",
+        "season_Cos_0[data_2]",
+        "season_Sin_0[data_2]",
+        "season_Cos_1[data_2]",
+        "season_Sin_1[data_2]",
+    ]
+
+    # Simulate
+    x0 = np.zeros((2, 3), dtype=config.floatX)
+    x0[0, 0] = 1.0
+    x0[1, 0] = 2.0
+    params = {"season": x0, "sigma_season": np.zeros(2, dtype=config.floatX)}
+    x, y = simulate_from_numpy_model(mod, rng, params, steps=12)
+
+    # Check periodicity for each observed series
+    assert_pattern_repeats(y[:, 0], 4, atol=ATOL, rtol=RTOL)
+    assert_pattern_repeats(y[:, 1], 4, atol=ATOL, rtol=RTOL)
+
+    mod = mod.build(verbose=False)
+    assert list(mod.coords["season_state"]) == [
+        "season_Cos_0[data_1]",
+        "season_Sin_0[data_1]",
+        "season_Cos_1[data_1]",
+        "season_Cos_0[data_2]",
+        "season_Sin_0[data_2]",
+        "season_Cos_1[data_2]",
+    ]
+
+    x0_sym, *_, T_sym, Z_sym, R_sym, _, Q_sym = mod._unpack_statespace_with_placeholders()
+    input_vars = explicit_graph_inputs([x0_sym, T_sym, Z_sym, R_sym, Q_sym])
+    fn = pytensor.function(
+        inputs=list(input_vars),
+        outputs=[x0_sym, T_sym, Z_sym, R_sym, Q_sym],
+        mode="FAST_COMPILE",
+    )
+    params["sigma_season"] = np.array([0.1, 0.8], dtype=config.floatX)
+    x0_v, T_v, Z_v, R_v, Q_v = fn(**params)
+
+    # x0 should be raveled into a single vector, with data_1 states first, then data_2 states
+    np.testing.assert_allclose(
+        x0_v, np.array([1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0]), atol=ATOL, rtol=RTOL
+    )
+
+    # T_v shape: (8, 8) (k_endog * k_states)
+    # The transition matrix is block diagonal, each block is:
+    # For n=2, season_length=4:
+    # lambda_1 = 2*pi*1/4 = pi/2, cos(pi/2)=0, sin(pi/2)=1
+    # lambda_2 = 2*pi*2/4 = pi,   cos(pi)=-1, sin(pi)=0
+    # Block 1 (Cos_0, Sin_0):
+    # [[cos(pi/2), sin(pi/2)],
+    #  [-sin(pi/2), cos(pi/2)]] = [[0, 1], [-1, 0]]
+    # Block 2 (Cos_1, Sin_1):
+    # [[-1, 0], [0, -1]]
+    expected_T_block1 = np.array([[0.0, 1.0], [-1.0, 0.0]])
+    expected_T_block2 = np.array([[-1.0, 0.0], [0.0, -1.0]])
+    expected_T = np.zeros((8, 8))
+    # data_1
+    expected_T[0:2, 0:2] = expected_T_block1
+    expected_T[2:4, 2:4] = expected_T_block2
+    # data_2
+    expected_T[4:6, 4:6] = expected_T_block1
+    expected_T[6:8, 6:8] = expected_T_block2
+    np.testing.assert_allclose(T_v, expected_T, atol=ATOL, rtol=RTOL)
+
+    # Only the first two states (one sin and one cos component) of each observed series are observed
+    expected_Z = np.zeros((2, 8))
+    expected_Z[0, 0] = 1.0
+    expected_Z[0, 2] = 1.0
+    expected_Z[1, 4] = 1.0
+    expected_Z[1, 6] = 1.0
+    np.testing.assert_allclose(Z_v, expected_Z, atol=ATOL, rtol=RTOL)
+
+    np.testing.assert_allclose(R_v, np.eye(8), atol=ATOL, rtol=RTOL)
+
+    Q_diag = np.diag(Q_v)
+    expected_Q_diag = np.r_[np.full(4, 0.1**2), np.full(4, 0.8**2)]
+    np.testing.assert_allclose(Q_diag, expected_Q_diag, atol=ATOL, rtol=RTOL)
+
+
+def test_add_two_frequency_seasonality_different_observed(rng):
+    mod1 = st.FrequencySeasonality(
+        season_length=4,
+        n=2,  # saturated
+        name="freq1",
+        innovations=True,
+        observed_state_names=["data_1"],
+    )
+    mod2 = st.FrequencySeasonality(
+        season_length=6,
+        n=1,  # unsaturated
+        name="freq2",
+        innovations=True,
+        observed_state_names=["data_2"],
+    )
+
+    mod = (mod1 + mod2).build(verbose=False)
+
+    params = {
+        "freq1": np.array([1.0, 0.0, 0.0], dtype=config.floatX),
+        "freq2": np.array([3.0, 0.0], dtype=config.floatX),
+        "sigma_freq1": np.array(0.0, dtype=config.floatX),
+        "sigma_freq2": 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=4 * 6 * 3)
+
+    assert_pattern_repeats(y[:, 0], 4, atol=ATOL, rtol=RTOL)
+    assert_pattern_repeats(y[:, 1], 6, atol=ATOL, rtol=RTOL)
+
+    assert mod.state_names == [
+        "freq1_Cos_0[data_1]",
+        "freq1_Sin_0[data_1]",
+        "freq1_Cos_1[data_1]",
+        "freq1_Sin_1[data_1]",
+        "freq2_Cos_0[data_2]",
+        "freq2_Sin_0[data_2]",
+    ]
+
+    assert mod.shock_names == [
+        "freq1_Cos_0[data_1]",
+        "freq1_Sin_0[data_1]",
+        "freq1_Cos_1[data_1]",
+        "freq1_Sin_1[data_1]",
+        "freq2_Cos_0[data_2]",
+        "freq2_Sin_0[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_v, T_v = fn(
+        freq1=np.array([1.0, 0.0, 1.2], dtype=config.floatX),
+        freq2=np.array([3.0, 0.0], dtype=config.floatX),
+    )
+
+    # Make sure the extra 0 in from the first component (the saturated state) is there!
+    np.testing.assert_allclose(np.array([1.0, 0.0, 1.2, 0.0, 3.0, 0.0]), x0_v, atol=ATOL, rtol=RTOL)
+
+    # Transition matrix is block diagonal: 4x4 for freq1, 2x2 for freq2
+    # freq1: n=4, lambdas = 2*pi*1/6, 2*pi*2/6
+    lam1 = 2 * np.pi * 1 / 4
+    lam2 = 2 * np.pi * 2 / 4
+    freq1_T1 = np.array([[np.cos(lam1), np.sin(lam1)], [-np.sin(lam1), np.cos(lam1)]])
+    freq1_T2 = np.array([[np.cos(lam2), np.sin(lam2)], [-np.sin(lam2), np.cos(lam2)]])
+    freq1_T = np.zeros((4, 4))
+
+    # freq2: n=4, lambdas = 2*pi*1/6
+    lam3 = 2 * np.pi * 1 / 6
+    freq2_T = np.array([[np.cos(lam3), np.sin(lam3)], [-np.sin(lam3), np.cos(lam3)]])
+
+    freq1_T[0:2, 0:2] = freq1_T1
+    freq1_T[2:4, 2:4] = freq1_T2
+
+    expected_T = np.zeros((6, 6))
+    expected_T[0:4, 0:4] = freq1_T
+    expected_T[4:6, 4:6] = freq2_T
+
+    np.testing.assert_allclose(expected_T, T_v, atol=ATOL, rtol=RTOL)