bugfixes and touchup

Author: jessegrabowski

pymc_extras/statespace/core/statespace.py

@@ -2347,7 +2347,7 @@ def _build_forecast_model(
             cov_dims = ["data_time", ALL_STATE_DIM, ALL_STATE_AUX_DIM]
 
         with pm.Model(coords=temp_coords) as forecast_model:
-            (_, _, *matrices), grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(
+            _, grouped_outputs = self._kalman_filter_outputs_from_dummy_graph(
                 data_dims=["data_time", OBS_STATE_DIM],
             )
 
@@ -2374,9 +2374,20 @@ def _build_forecast_model(
                 "P0_slice", cov_frozen[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
             )
 
-            # Get matrices with n_timesteps set to forecast length for time-varying models
-            # Note: matrices already has x0, P0 skipped from _kalman_filter_outputs_from_dummy_graph
-            forecast_matrices = self._insert_constant_timestep(matrices, len(forecast_index))
+            # Get fresh matrices with n_timesteps placeholder still intact.
+            # Build for the full timeline (training + forecast) so that time-varying matrices
+            # continue at the correct phase, then slice to keep only the forecast portion.
+            n_train = len(time_index)
+            n_total = n_train + len(forecast_index)
+
+            full_matrices = self._insert_constant_timestep(self.unpack_statespace(), n_total)
+            _, _, *forecast_matrices = full_matrices
+
+            forecast_names = MATRIX_NAMES[2:]  # c, d, T, Z, R, H, Q
+            forecast_matrices = [
+                m[n_train:] if m.ndim == (2 if name in VECTOR_VALUED else 3) else m
+                for m, name in zip(forecast_matrices, forecast_names)
+            ]
 
             _ = LinearGaussianStateSpace(
                 "forecast",
@@ -2640,6 +2651,12 @@ def impulse_response_function(
         -------
         pm.InferenceData
             An Arviz InferenceData object containing impulse response function in a variable named "irf".
+
+        Notes
+        -----
+        For models with time-varying transition matrices, the IRF is computed starting at phase 0 of the
+        time-varying cycle. This means the response represents the effect of a shock occurring at the first
+        modeled state, T(0).
         """
         options = [shock_size, shock_cov, shock_trajectory]
         n_options = sum(x is not None for x in options)

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

@@ -452,14 +452,15 @@ def _build_time_varying_transition(self) -> TensorVariable:
 
         T_hold = pt.eye(n)
 
-        time_idx = pt.arange(self.n_timesteps)
-        is_rotation_step = pt.eq(time_idx % self.duration, self.duration - 1)
+        # Build one complete cycle: [I, I, ..., I, T_rotate] of length `duration`
+        # Then tile to cover n_timesteps
+        cycle_matrices = [T_hold for _ in range(self.duration - 1)] + [T_rotate]
+        T_cycle = pt.stack(cycle_matrices)  # (duration, n, n)
 
-        return pt.where(
-            is_rotation_step[:, None, None],
-            pt.broadcast_to(T_rotate, (self.n_timesteps, n, n)),
-            pt.broadcast_to(T_hold, (self.n_timesteps, n, n)),
-        )
+        n_cycles = (self.n_timesteps + self.duration - 1) // self.duration  # ceiling division
+        T_tiled = pt.tile(T_cycle, (n_cycles, 1, 1))
+
+        return T_tiled[: self.n_timesteps]
 
     def _build_transition_matrix(self) -> TensorVariable:
         """Build the full transition matrix, handling multivariate via block_diag."""
@@ -522,15 +523,18 @@ def _build_initial_state_simple(self, initial_params: TensorVariable) -> TensorV
             if use_tv:
                 return initial_params
             else:
-                return pt.extra_ops.repeat(initial_params, self.duration, axis=0)
+                # Static mode: state is d blocks of s elements each
+                # Tile the full season vector d times
+                return pt.tile(initial_params, self.duration)
         else:
             if use_tv:
                 return initial_params.ravel()
             else:
-                return pt.extra_ops.repeat(initial_params, self.duration, axis=1).ravel()
+                # Static mode: tile each row (endog) d times, then ravel
+                return pt.tile(initial_params, (1, self.duration)).ravel()
 
     def _apply_start_state_shift(
-        self, initial_state: TensorVariable, T: TensorVariable
+        self, initial_state: TensorVariable, T: TensorVariable | None
     ) -> TensorVariable:
         """Shift initial state to account for start_state offset."""
         if self.start_idx == 0:
@@ -569,7 +573,7 @@ def _build_selection_and_state_cov(self) -> None:
         k_posdef = self._k_posdef_per_endog()
 
         R = pt.zeros((k_states, k_posdef))[0, 0].set(1.0)
-        self.ssm["selection", :, :] = pt.join(0, *[R for _ in range(k_endog_effective)])
+        self.ssm["selection", :, :] = pt.linalg.block_diag(*[R for _ in range(k_endog_effective)])
 
         sigma = self.make_and_register_variable(
             f"sigma_{self.name}",
@@ -606,8 +610,10 @@ def make_symbolic_graph(self) -> None:
         else:
             initial_state = self._build_initial_state_simple(initial_params)
 
-        # Apply start_state shift (handles time-varying vs static internally)
-        T_for_shift = self._build_static_transition()
+        # Apply start_state shift
+        T_for_shift = (
+            None if self._uses_time_varying_transition else self._build_static_transition()
+        )
         initial_state = self._apply_start_state_shift(initial_state, T_for_shift)
 
         self.ssm["initial_state", :] = initial_state

tests/statespace/core/test_statespace.py

@@ -1358,12 +1358,32 @@ def test_sample_unconditional_posterior(self, ss_mod_time_varying, idata_time_va
 
     @pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
     @pytest.mark.filterwarnings("ignore:No start date provided")
-    def test_forecast(self, ss_mod_time_varying, idata_time_varying):
-        result = ss_mod_time_varying.forecast(idata_time_varying, periods=10)
+    @pytest.mark.parametrize(
+        "periods", [10, 50], ids=["shorter_than_training", "longer_than_training"]
+    )
+    def test_forecast(self, ss_mod_time_varying, idata_time_varying, periods):
+        n_obs = 40  # must match pymc_mod_time_varying fixture
+        result = ss_mod_time_varying.forecast(idata_time_varying, periods=periods)
+
         assert "forecast_latent" in result
         assert "forecast_observed" in result
-        assert result["forecast_latent"].shape[2] == 10
+        assert result["forecast_latent"].dims == ("chain", "draw", "time", "state")
+        assert result["forecast_observed"].dims == ("chain", "draw", "time", "observed_state")
+        assert result["forecast_latent"].shape[2] == periods
         assert not np.any(np.isnan(result["forecast_latent"].values))
+        assert not np.any(np.isnan(result["forecast_observed"].values))
+
+        # Value check: the model has y_t = d_t + Z @ x_t with Z=[[1]] and H=0 (no obs noise),
+        # so forecast_observed - forecast_latent = d_t = slope * t.
+        # Forecast matrices are phase-aligned: they continue from n_obs, not from 0.
+        latent = result["forecast_latent"].values  # (chain, draw, time, state)
+        observed = result["forecast_observed"].values  # (chain, draw, time, obs)
+        slope = idata_time_varying.posterior["slope"].values  # (chain, draw)
+
+        intercepts = observed[..., 0] - latent[..., 0]  # (chain, draw, time)
+        expected = slope[..., None] * np.arange(n_obs, n_obs + periods)[None, None, :]
+
+        assert_allclose(intercepts, expected, atol=1e-5, rtol=1e-5)
 
     @pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
     def test_impulse_response_function(self, ss_mod_time_varying, idata_time_varying):

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

@@ -55,7 +55,9 @@ def test_time_seasonality(s, d, innovations, remove_first_state, rng):
 
 @pytest.mark.parametrize("d", [1, 3])
 @pytest.mark.parametrize("start_state", [0, 2, "state_2"])
-def test_time_seasonality_start_state(d, start_state, rng):
+@pytest.mark.parametrize("use_time_varying", [True, False], ids=["time_varying", "static"])
+@pytest.mark.parametrize("remove_first_state", [True, False])
+def test_time_seasonality_start_state(d, start_state, use_time_varying, remove_first_state, rng):
     s = 4
     state_names = [f"state_{i}" for i in range(s)]
 
@@ -65,13 +67,18 @@ def test_time_seasonality_start_state(d, start_state, rng):
         innovations=False,
         name="season",
         state_names=state_names,
-        remove_first_state=True,
+        remove_first_state=remove_first_state,
         start_state=start_state,
+        use_time_varying=use_time_varying,
     )
 
-    params = np.array([1.0, 2.0, 3.0], dtype=config.floatX)
-    implied_gamma0 = -params.sum()
-    default_seasons = [implied_gamma0, params[0], params[1], params[2]]
+    if remove_first_state:
+        params = np.array([1.0, 2.0, 3.0], dtype=config.floatX)
+        implied_gamma0 = -params.sum()
+        default_seasons = [implied_gamma0, params[0], params[1], params[2]]
+    else:
+        params = np.array([1.0, 2.0, 3.0, 4.0], dtype=config.floatX)
+        default_seasons = [params[0], params[1], params[2], params[3]]
 
     start_idx = state_names.index(start_state) if isinstance(start_state, str) else start_state
     expected_seasons = default_seasons[start_idx:] + default_seasons[:start_idx]
@@ -175,6 +182,12 @@ def test_time_seasonality_multiple_observed(rng, d, remove_first_state):
     np.testing.assert_allclose(T_v, expected_T, atol=ATOL, rtol=RTOL)
     np.testing.assert_allclose(Q_v, np.array([[0.1**2, 0.0], [0.0, 0.8**2]]), atol=ATOL, rtol=RTOL)
 
+    k_states_per_endog = d * (s - 1) if remove_first_state else d * s
+    Z0 = np.zeros((1, k_states_per_endog))
+    Z0[0, 0] = 1.0
+    expected_Z = np.block([[Z0, np.zeros_like(Z0)], [np.zeros_like(Z0), Z0]])
+    np.testing.assert_allclose(Z_v, expected_Z, atol=ATOL, rtol=RTOL)
+
 
 @pytest.mark.parametrize("d", [2, 3])
 @pytest.mark.parametrize("remove_first_state", [True, False])
@@ -204,6 +217,13 @@ def test_time_seasonality_multiple_observed_time_varying(rng, d, remove_first_st
     assert_pattern_repeats(y[:, 0], s * d, atol=ATOL, rtol=RTOL)
     assert_pattern_repeats(y[:, 1], s * d, atol=ATOL, rtol=RTOL)
 
+    k_states_per_endog = (s - 1) if remove_first_state else s
+    Z0 = np.zeros((1, k_states_per_endog))
+    Z0[0, 0] = 1.0
+    expected_Z = np.block([[Z0, np.zeros_like(Z0)], [np.zeros_like(Z0), Z0]])
+    Z_v = mod.ssm["design"].eval()
+    np.testing.assert_allclose(Z_v, expected_Z, atol=ATOL, rtol=RTOL)
+
 
 def test_time_seasonality_shared_states():
     mod = st.TimeSeasonality(
@@ -231,6 +251,34 @@ def test_time_seasonality_shared_states():
     np.testing.assert_allclose(R, np.array([[1.0], [0.0], [0.0]]))
 
 
+@pytest.mark.parametrize("use_time_varying", [True, False], ids=["time_varying", "static"])
+def test_time_seasonality_shared_states_with_duration(rng, use_time_varying):
+    s, d = 4, 3
+    mod = st.TimeSeasonality(
+        season_length=s,
+        duration=d,
+        innovations=False,
+        name="season",
+        observed_state_names=["data_1", "data_2"],
+        remove_first_state=True,
+        share_states=True,
+        use_time_varying=use_time_varying,
+    )
+
+    assert mod.k_endog == 2
+    expected_k_states = (s - 1) if use_time_varying else d * (s - 1)
+    assert mod.k_states == expected_k_states
+    assert mod.k_posdef == 0
+
+    x0 = np.array([1.0, 2.0, 3.0], dtype=config.floatX)
+    params = {"params_season": x0}
+    _, y = simulate_from_numpy_model(mod, rng, params, steps=s * d * 4)
+
+    assert_pattern_repeats(y[:, 0], s * d, atol=ATOL, rtol=RTOL)
+    assert_pattern_repeats(y[:, 1], s * d, atol=ATOL, rtol=RTOL)
+    np.testing.assert_allclose(y[:, 0], y[:, 1], atol=ATOL, rtol=RTOL)
+
+
 def test_add_mixed_shared_not_shared_time_seasonality():
     shared = st.TimeSeasonality(
         season_length=3,
@@ -394,8 +442,10 @@ def test_add_time_varying_and_static_seasonality(rng):
     x, y = simulate_from_numpy_model(mod, rng, params, steps=steps)
 
     # Combined output should have period = LCM of individual periods
-    # For testing, verify the model runs without error and output shape is correct
+    # LCM(s1*d1, s2*d2) = LCM(12, 10) = 60
     assert y.shape == (steps,)
+    assert not np.any(np.isnan(y))
+    assert_pattern_repeats(y, 60, atol=ATOL, rtol=RTOL)
 
 
 @pytest.mark.parametrize("n", [1, 2, 3, None])
@@ -409,8 +459,9 @@ def test_frequency_seasonality(n, s, rng):
     x0 = rng.normal(size=mod.n_coefs).astype(config.floatX)
     params = {"params_season": x0, "sigma_season": 0.0}
 
-    decimal = s_str.split(".") if "." in (s_str := str(s)) else "0"
-    T = int(s * 10 ** len(decimal))
+    s_str = str(s)
+    n_decimal = len(s_str.split(".")[1]) if "." in s_str else 0
+    T = int(s * 10**n_decimal)
 
     x, y = simulate_from_numpy_model(mod, rng, params, steps=2 * T)
     assert_pattern_repeats(y, T, atol=ATOL, rtol=RTOL)