Pr 451 - modified and added tests to statespace (pymc-devs#466)

* Use `set_data` in forecast

* Ignore new numpy matmul warnings in tests

* Tracking down data bug

* resolved merge conflicts environment-test.yml

* added and modified statespace tests

* added logic to update static shape of target when forecasting with exogenous variables

* make sure updated dummy target has correct dimensions

* Revert test env name change

* Add some checks to `test_foreacast_valid_index`

---------

Co-authored-by: jessegrabowski <[email protected]>

Author: 2 people

pymc_extras/statespace/core/statespace.py

@@ -1699,6 +1699,10 @@ def _validate_forecast_args(
             raise ValueError(
                 "Must specify one of either periods or end unless use_scenario_index=True"
             )
+        if periods is None and end is None and not use_scenario_index:
+            raise ValueError(
+                "Must specify one of either periods or end unless use_scenario_index=True"
+            )
         if periods is not None and end is not None:
             raise ValueError("Must specify exactly one of either periods or end")
         if scenario is None and use_scenario_index:
@@ -2246,23 +2250,58 @@ def forecast(
             use_scenario_index=use_scenario_index,
         )
         scenario = self._finalize_scenario_initialization(scenario, forecast_index)
+        temp_coords = self._fit_coords.copy()
 
-        forecast_model = self._build_forecast_model(
-            time_index=time_index,
-            t0=t0,
-            forecast_index=forecast_index,
-            scenario=scenario,
-            filter_output=filter_output,
-            mvn_method=mvn_method,
-        )
+        dims = None
+        if all([dim in temp_coords for dim in [filter_time_dim, ALL_STATE_DIM, OBS_STATE_DIM]]):
+            dims = [TIME_DIM, ALL_STATE_DIM, OBS_STATE_DIM]
 
-        with forecast_model:
-            if scenario is not None:
-                dummy_obs_data = np.zeros((len(forecast_index), self.k_endog))
-                pm.set_data(
-                    scenario | {"data": dummy_obs_data},
-                    coords={"data_time": np.arange(len(forecast_index))},
-                )
+        t0_idx = np.flatnonzero(time_index == t0)[0]
+
+        temp_coords["data_time"] = time_index
+        temp_coords[TIME_DIM] = forecast_index
+
+        mu_dims, cov_dims = None, None
+        if all([dim in self._fit_coords for dim in [TIME_DIM, ALL_STATE_DIM, ALL_STATE_AUX_DIM]]):
+            mu_dims = ["data_time", ALL_STATE_DIM]
+            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(
+                scenario=scenario,
+                data_dims=["data_time", OBS_STATE_DIM],
+            )
+
+            for name in self.data_names:
+                if name in scenario.keys():
+                    pm.set_data(
+                        {"data": np.zeros((len(forecast_index), self.k_endog))},
+                        coords={"data_time": np.arange(len(forecast_index))},
+                    )
+                    break
+
+            group_idx = FILTER_OUTPUT_TYPES.index(filter_output)
+            mu, cov = grouped_outputs[group_idx]
+
+            x0 = pm.Deterministic(
+                "x0_slice", mu[t0_idx], dims=mu_dims[1:] if mu_dims is not None else None
+            )
+            P0 = pm.Deterministic(
+                "P0_slice", cov[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
+            )
+
+            _ = LinearGaussianStateSpace(
+                "forecast",
+                x0,
+                P0,
+                *matrices,
+                steps=len(forecast_index),
+                dims=dims,
+                mode=self._fit_mode,
+                sequence_names=self.kalman_filter.seq_names,
+                k_endog=self.k_endog,
+                append_x0=False,
+            )
 
         forecast_model.rvs_to_initial_values = {
             k: None for k in forecast_model.rvs_to_initial_values.keys()

tests/statespace/core/test_statespace.py

@@ -152,28 +152,23 @@ def exog_data(rng):
 
     df.loc[[1, 3, 9], ["y"]] = np.nan
     return df.set_index("date")
-
-
-@pytest.fixture(scope="session")
-def exog_data_mv(rng):
+def exog_data(rng):
     # simulate data
     df = pd.DataFrame(
         {
             "date": pd.date_range(start="2023-05-01", end="2023-05-10", freq="D"),
             "x1": rng.choice(2, size=10, replace=True).astype(float),
-            "y1": rng.normal(size=(10,)),
-            "y2": rng.normal(size=(10,)),
+            "y": rng.normal(size=(10,)),
         }
     )
 
-    df.loc[[1, 3, 9], ["y1"]] = np.nan
-    df.loc[[3, 5, 7], ["y2"]] = np.nan
+    df.loc[[1, 3, 9], ["y"]] = np.nan
     return df.set_index("date")
 
 
 @pytest.fixture(scope="session")
 def exog_ss_mod(exog_data):
-    level_trend = st.LevelTrendComponent(name="trend", order=1, innovations_order=[0])
+    level_trend = st.LevelTrendComponent(order=1, innovations_order=[0])
     exog = st.RegressionComponent(
         name="exog",  # Name of this exogenous variable component
         k_exog=1,  # Only one exogenous variable now
@@ -185,83 +180,23 @@ def exog_ss_mod(exog_data):
     return combined_model.build()
 
 
-@pytest.fixture(scope="session")
-def exog_ss_mod_mv(exog_data_mv):
-    level_trend = st.LevelTrendComponent(
-        name="trend", order=1, innovations_order=[0], observed_state_names=["y1", "y2"]
-    )
-    exog = st.RegressionComponent(
-        name="exog",  # Name of this exogenous variable component
-        k_exog=1,  # Only one exogenous variable now
-        innovations=False,  # Typically fixed effect (no stochastic evolution)
-        state_names=exog_data_mv[["x1"]].columns.tolist(),
-        observed_state_names=["y1", "y2"],
-    )
-
-    combined_model = level_trend + exog
-    return combined_model.build()
-
-
-@pytest.fixture(scope="session")
-def ss_mod_multi_component(rng):
-    ll = st.LevelTrendComponent(
-        name="trend", order=2, innovations_order=1, observed_state_names=["y1", "y2"]
-    )
-    exog = st.RegressionComponent(
-        name="exog",
-        innovations=True,
-        state_names=["x1"],
-    )
-    ar = st.AutoregressiveComponent(observed_state_names=["y1"])
-    cycle = st.CycleComponent(cycle_length=2, observed_state_names=["y1", "y2"], innovations=True)
-    season = st.TimeSeasonality(season_length=2, observed_state_names=["y1"], innovations=True)
-
-    fseason = st.FrequencySeasonality(
-        season_length=2, observed_state_names=["y1"], innovations=True
-    )
-    measure_error = st.MeasurementError(observed_state_names=["y1", "y2"])
-    return (ll + exog + ar + cycle + season + fseason + measure_error).build()
-
-
 @pytest.fixture(scope="session")
 def exog_pymc_mod(exog_ss_mod, exog_data):
     # define pymc model
     with pm.Model(coords=exog_ss_mod.coords) as struct_model:
         P0_diag = pm.Gamma("P0_diag", alpha=2, beta=4, dims=["state"])
         P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=["state", "state_aux"])
 
-        initial_trend = pm.Normal("initial_trend", mu=[0], sigma=[0.005], dims=["state_trend"])
+        initial_trend = pm.Normal("initial_trend", mu=[0], sigma=[0.005], dims=["trend_state"])
 
         data_exog = pm.Data(
-            "data_exog", exog_data["x1"].values[:, None], dims=["time", "state_exog"]
+            "data_exog", exog_data["x1"].values[:, None], dims=["time", "exog_state"]
         )
-        beta_exog = pm.Normal("beta_exog", mu=0, sigma=1, dims=["state_exog"])
+        beta_exog = pm.Normal("beta_exog", mu=0, sigma=1, dims=["exog_state"])
 
-        exog_ss_mod.build_statespace_graph(exog_data["y"], save_kalman_filter_outputs_in_idata=True)
+        exog_ss_mod.build_statespace_graph(exog_data["y"])
 
     return struct_model
-
-
-@pytest.fixture(scope="session")
-def exog_pymc_mod_mv(exog_ss_mod_mv, exog_data_mv):
-    # define pymc model
-    with pm.Model(coords=exog_ss_mod_mv.coords) as struct_model:
-        P0_diag = pm.Gamma("P0_diag", alpha=2, beta=4, dims=["state"])
-        P0 = pm.Deterministic("P0", pt.diag(P0_diag), dims=["state", "state_aux"])
-
-        initial_trend = pm.Normal(
-            "initial_trend", mu=[0], sigma=[0.005], dims=["endog_trend", "state_trend"]
-        )
-
-        data_exog = pm.Data(
-            "data_exog", exog_data_mv["x1"].values[:, None], dims=["time", "state_exog"]
-        )
-        beta_exog = pm.Normal("beta_exog", mu=0, sigma=1, dims=["endog_exog", "state_exog"])
-
-        exog_ss_mod_mv.build_statespace_graph(
-            exog_data_mv[["y1", "y2"]], save_kalman_filter_outputs_in_idata=True
-        )
-
     return struct_model
 
 
@@ -1036,13 +971,19 @@ def test_forecast(filter_output, mod_name, idata_name, start, end, periods, rng,
     assert forecast_idx[0] == (t0 + delta)
 
 
+@pytest.mark.filterwarnings("ignore:Provided data contains missing values")
+@pytest.mark.filterwarnings("ignore:The RandomType SharedVariables")
 @pytest.mark.filterwarnings("ignore:Provided data contains missing values")
 @pytest.mark.filterwarnings("ignore:The RandomType SharedVariables")
 @pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
 @pytest.mark.filterwarnings("ignore:Skipping `CheckAndRaise` Op")
 @pytest.mark.filterwarnings("ignore:No frequency was specific on the data's DateTimeIndex.")
 @pytest.mark.parametrize("start", [None, -1, 5])
+@pytest.mark.filterwarnings("ignore:Skipping `CheckAndRaise` Op")
+@pytest.mark.filterwarnings("ignore:No frequency was specific on the data's DateTimeIndex.")
+@pytest.mark.parametrize("start", [None, -1, 5])
 def test_forecast_with_exog_data(rng, exog_ss_mod, idata_exog, start):
+    scenario = pd.DataFrame(np.zeros((10, 1)), columns=["x1"])
     scenario = pd.DataFrame(np.zeros((10, 1)), columns=["x1"])
     scenario.iloc[5, 0] = 1e9
 
@@ -1051,7 +992,7 @@ def test_forecast_with_exog_data(rng, exog_ss_mod, idata_exog, start):
     )
 
     components = exog_ss_mod.extract_components_from_idata(forecast_idata)
-    level = components.forecast_latent.sel(state="trend[level]")
+    level = components.forecast_latent.sel(state="LevelTrend[level]")
     betas = components.forecast_latent.sel(state=["exog[x1]"])
 
     scenario.index.name = "time"
@@ -1060,6 +1001,7 @@ def test_forecast_with_exog_data(rng, exog_ss_mod, idata_exog, start):
         .to_xarray()
         .rename({"level_0": "state"})
         .assign_coords(state=["exog[x1]"])
+        .assign_coords(state=["exog[x1]"])
     )
 
     regression_effect = forecast_idata.forecast_observed.isel(observed_state=0) - level
@@ -1068,138 +1010,6 @@ def test_forecast_with_exog_data(rng, exog_ss_mod, idata_exog, start):
     assert_allclose(regression_effect, regression_effect_expected)
 
 
-@pytest.mark.filterwarnings("ignore:Provided data contains missing values")
-@pytest.mark.filterwarnings("ignore:The RandomType SharedVariables")
-@pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
-@pytest.mark.filterwarnings("ignore:Skipping `CheckAndRaise` Op")
-@pytest.mark.filterwarnings("ignore:No frequency was specific on the data's DateTimeIndex.")
-@pytest.mark.parametrize("start", [None, -1, 5])
-def test_forecast_with_exog_data_mv(rng, exog_ss_mod_mv, idata_exog_mv, start):
-    scenario = pd.DataFrame(np.zeros((10, 1)), columns=["x1"])
-    scenario.iloc[5, 0] = 1e9
-
-    forecast_idata = exog_ss_mod_mv.forecast(
-        idata_exog_mv, start=start, periods=10, random_seed=rng, scenario=scenario
-    )
-
-    components = exog_ss_mod_mv.extract_components_from_idata(forecast_idata)
-    level_y1 = components.forecast_latent.sel(state="trend[level[y1]]")
-    level_y2 = components.forecast_latent.sel(state="trend[level[y2]]")
-    betas_y1 = components.forecast_latent.sel(state=["exog[x1[y1]]"])
-    betas_y2 = components.forecast_latent.sel(state=["exog[x1[y2]]"])
-
-    scenario.index.name = "time"
-    scenario_xr_y1 = (
-        scenario.unstack()
-        .to_xarray()
-        .rename({"level_0": "state"})
-        .assign_coords(state=["exog[x1[y1]]"])
-    )
-
-    scenario_xr_y2 = (
-        scenario.unstack()
-        .to_xarray()
-        .rename({"level_0": "state"})
-        .assign_coords(state=["exog[x1[y2]]"])
-    )
-
-    regression_effect_y1 = forecast_idata.forecast_observed.isel(observed_state=0) - level_y1
-    regression_effect_expected_y1 = (betas_y1 * scenario_xr_y1).sum(dim=["state"])
-
-    regression_effect_y2 = forecast_idata.forecast_observed.isel(observed_state=1) - level_y2
-    regression_effect_expected_y2 = (betas_y2 * scenario_xr_y2).sum(dim=["state"])
-
-    np.testing.assert_allclose(regression_effect_y1, regression_effect_expected_y1)
-    np.testing.assert_allclose(regression_effect_y2, regression_effect_expected_y2)
-
-
-@pytest.mark.filterwarnings("ignore:Provided data contains missing values")
-@pytest.mark.filterwarnings("ignore:The RandomType SharedVariables")
-@pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
-@pytest.mark.filterwarnings("ignore:Skipping `CheckAndRaise` Op")
-@pytest.mark.filterwarnings("ignore:No frequency was specific on the data's DateTimeIndex.")
-def test_build_forecast_model(rng, exog_ss_mod, exog_pymc_mod, exog_data, idata_exog):
-    data_before_build_forecast_model = {d.name: d.get_value() for d in exog_pymc_mod.data_vars}
-
-    scenario = pd.DataFrame(
-        {
-            "date": pd.date_range(start="2023-05-11", end="2023-05-20", freq="D"),
-            "x1": rng.choice(2, size=10, replace=True).astype(float),
-        }
-    )
-    scenario.set_index("date", inplace=True)
-
-    time_index = exog_ss_mod._get_fit_time_index()
-    t0, forecast_index = exog_ss_mod._build_forecast_index(
-        time_index=time_index,
-        start=exog_data.index[-1],
-        end=scenario.index[-1],
-        scenario=scenario,
-    )
-
-    test_forecast_model = exog_ss_mod._build_forecast_model(
-        time_index=time_index,
-        t0=t0,
-        forecast_index=forecast_index,
-        scenario=scenario,
-        filter_output="predicted",
-        mvn_method="svd",
-    )
-
-    frozen_shared_inputs = [
-        inpt
-        for inpt in graph_inputs([test_forecast_model.x0_slice, test_forecast_model.P0_slice])
-        if isinstance(inpt, SharedVariable)
-        and not isinstance(inpt.get_value(), np.random.Generator)
-    ]
-
-    assert (
-        len(frozen_shared_inputs) == 0
-    )  # check there are no non-random generator SharedVariables in the frozen inputs
-
-    unfrozen_shared_inputs = [
-        inpt
-        for inpt in graph_inputs([test_forecast_model.forecast_combined])
-        if isinstance(inpt, SharedVariable)
-        and not isinstance(inpt.get_value(), np.random.Generator)
-    ]
-
-    # Check that there is one (in this case) unfrozen shared input and it corresponds to the exogenous data
-    assert len(unfrozen_shared_inputs) == 1
-    assert unfrozen_shared_inputs[0].name == "data_exog"
-
-    data_after_build_forecast_model = {d.name: d.get_value() for d in test_forecast_model.data_vars}
-
-    with test_forecast_model:
-        dummy_obs_data = np.zeros((len(forecast_index), exog_ss_mod.k_endog))
-        pm.set_data(
-            {"data_exog": scenario} | {"data": dummy_obs_data},
-            coords={"data_time": np.arange(len(forecast_index))},
-        )
-        idata_forecast = pm.sample_posterior_predictive(
-            idata_exog, var_names=["x0_slice", "P0_slice"]
-        )
-
-    np.testing.assert_allclose(
-        unfrozen_shared_inputs[0].get_value(), scenario["x1"].values.reshape((-1, 1))
-    )  # ensure the replaced data matches the exogenous data
-
-    for k in data_before_build_forecast_model.keys():
-        assert (  # check that the data needed to init the forecasts doesn't change
-            data_before_build_forecast_model[k].mean() == data_after_build_forecast_model[k].mean()
-        )
-
-    # Check that the frozen states and covariances correctly match the sliced index
-    np.testing.assert_allclose(
-        idata_exog.posterior["predicted_covariance"].sel(time=t0).mean(("chain", "draw")).values,
-        idata_forecast.posterior_predictive["P0_slice"].mean(("chain", "draw")).values,
-    )
-    np.testing.assert_allclose(
-        idata_exog.posterior["predicted_state"].sel(time=t0).mean(("chain", "draw")).values,
-        idata_forecast.posterior_predictive["x0_slice"].mean(("chain", "draw")).values,
-    )
-
-
 @pytest.mark.filterwarnings("ignore:Provided data contains missing values")
 @pytest.mark.filterwarnings("ignore:The RandomType SharedVariables")
 @pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
@@ -1231,16 +1041,3 @@ def test_foreacast_valid_index(exog_pymc_mod, exog_ss_mod, exog_data):
 
     assert forecasts.forecast_latent.shape[2] == n_periods
     assert forecasts.forecast_observed.shape[2] == n_periods
-
-
-def test_param_dims_coords(ss_mod_multi_component):
-    for param in ss_mod_multi_component.param_names:
-        shape = ss_mod_multi_component.param_info[param]["shape"]
-        dims = ss_mod_multi_component.param_dims.get(param, None)
-        if len(shape) == 0:
-            assert dims is None
-            continue
-        for i, s in zip(shape, dims):
-            assert i == len(
-                ss_mod_multi_component.coords[s]
-            ), f"Mismatch between shape {i} and dimension {s}"