made slight change with _build_forecast_model and created a test case

Author: Dekermanjian

pymc_extras/statespace/core/statespace.py

@@ -2080,11 +2080,11 @@ def _build_forecast_model(
                 for data_var in forecast_model.data_vars
             }
 
-            replacements_diff = np.setdiff1d(
+            missing_data_vars = np.setdiff1d(
                 ar1=[*self.data_names, "data"], ar2=[k.name for k, _ in sub_dict.items()]
             )
-            if replacements_diff.size > 0:
-                raise ValueError(f"{replacements_diff} data used for fitting not found!")
+            if missing_data_vars.size > 0:
+                raise ValueError(f"{missing_data_vars} data used for fitting not found!")
 
             mu_frozen, cov_frozen = graph_replace([mu, cov], replace=sub_dict, strict=True)
 
@@ -2095,13 +2095,6 @@ def _build_forecast_model(
                 "P0_slice", cov_frozen[t0_idx], dims=cov_dims[1:] if cov_dims is not None else None
             )
 
-            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))},
-                )
-
             _ = LinearGaussianStateSpace(
                 "forecast",
                 x0,
@@ -2263,6 +2256,14 @@ def forecast(
             mvn_method=mvn_method,
         )
 
+        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))},
+                )
+
         forecast_model.rvs_to_initial_values = {
             k: None for k in forecast_model.rvs_to_initial_values.keys()
         }

tests/statespace/core/test_statespace.py

@@ -895,6 +895,58 @@ 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.")
+def test_build_forecast_model(rng, exog_ss_mod, exog_pymc_mod, exog_data):
+    data_before_build_forecast_model = {d.name: d.get_value() for d in exog_pymc_mod.data_vars}
+
+    scenario1 = 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),
+        }
+    )
+    scenario1.set_index("date", inplace=True)
+
+    scenario2 = pd.DataFrame(
+        {
+            "date": pd.date_range(start="2023-05-11", end="2023-05-20", freq="D"),
+            "x1": np.zeros(shape=(10,)),
+        }
+    )
+    scenario2.set_index("date", inplace=True)
+
+    for scenario in [scenario1, scenario2]:
+        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="smoothed",
+            mvn_method="svd",
+        )
+
+        data_after_build_forecast_model = {
+            d.name: d.get_value() for d in test_forecast_model.data_vars
+        }
+        for k in data_before_build_forecast_model.keys():
+            assert (
+                data_before_build_forecast_model[k].mean()
+                == data_after_build_forecast_model[k].mean()
+            )
+
+
 @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.")