added and modified statespace tests

Author: Dekermanjian

tests/statespace/test_statespace.py

@@ -128,32 +128,51 @@ def ss_mod_no_exog_dt(rng):
 
 
 @pytest.fixture(scope="session")
-def exog_ss_mod(rng):
-    ll = st.LevelTrendComponent()
-    reg = st.RegressionComponent(name="exog", state_names=["a", "b", "c"])
-    mod = (ll + reg).build(verbose=False)
+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),
+            "y": rng.normal(size=(10,)),
+        }
+    )
 
-    return mod
+    df.loc[[1, 3, 9], ["y"]] = np.nan
+    return df.set_index("date")
 
 
 @pytest.fixture(scope="session")
-def exog_pymc_mod(exog_ss_mod, rng):
-    y = rng.normal(size=(100, 1)).astype(floatX)
-    X = rng.normal(size=(100, 3)).astype(floatX)
+def exog_ss_mod(exog_data):
+    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
+        innovations=False,  # Typically fixed effect (no stochastic evolution)
+        state_names=exog_data[["x1"]].columns.tolist(),
+    )
 
-    with pm.Model(coords=exog_ss_mod.coords) as m:
-        exog_data = pm.Data("data_exog", X)
-        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
-        P0_sigma = pm.Exponential("P0_sigma", 1)
-        P0 = pm.Deterministic(
-            "P0", pt.eye(exog_ss_mod.k_states) * P0_sigma, dims=["state", "state_aux"]
+    combined_model = level_trend + exog
+    return combined_model.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=["trend_state"])
+
+        data_exog = pm.Data(
+            "data_exog", exog_data["x1"].values[:, None], dims=["time", "exog_state"]
         )
-        beta_exog = pm.Normal("beta_exog", dims=["exog_state"])
+        beta_exog = pm.Normal("beta_exog", mu=0, sigma=1, dims=["exog_state"])
 
-        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
-        exog_ss_mod.build_statespace_graph(y, save_kalman_filter_outputs_in_idata=True)
+        exog_ss_mod.build_statespace_graph(exog_data["y"])
 
-    return m
+    return struct_model
 
 
 @pytest.fixture(scope="session")
@@ -844,10 +863,14 @@ 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:No time index found on the supplied data.")
-@pytest.mark.parametrize("start", [None, -1, 10])
+@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, 3)), columns=["a", "b", "c"])
+    scenario = pd.DataFrame(np.zeros((10, 1)), columns=["x1"])
     scenario.iloc[5, 0] = 1e9
 
     forecast_idata = exog_ss_mod.forecast(
@@ -856,14 +879,14 @@ 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="LevelTrend[level]")
-    betas = components.forecast_latent.sel(state=["exog[a]", "exog[b]", "exog[c]"])
+    betas = components.forecast_latent.sel(state=["exog[x1]"])
 
     scenario.index.name = "time"
     scenario_xr = (
         scenario.unstack()
         .to_xarray()
         .rename({"level_0": "state"})
-        .assign_coords(state=["exog[a]", "exog[b]", "exog[c]"])
+        .assign_coords(state=["exog[x1]"])
     )
 
     regression_effect = forecast_idata.forecast_observed.isel(observed_state=0) - level
@@ -872,91 +895,25 @@ 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:Provided data contains missing values")
 @pytest.mark.filterwarnings("ignore:The RandomType SharedVariables")
-def test_foreacast_valid_index(rng):
+@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_foreacast_valid_index(exog_pymc_mod, exog_ss_mod, exog_data):
     # Regression test for issue reported at  https://github.com/pymc-devs/pymc-extras/issues/424
-
-    index = pd.date_range(start="2023-05-01", end="2025-01-29", freq="D")
-    T, k = len(index), 2
-    data = np.zeros((T, k))
-    idx = rng.choice(T, size=10, replace=False)
-    cols = rng.choice(k, size=10, replace=True)
-
-    data[idx, cols] = 1
-
-    df_holidays = pd.DataFrame(data, index=index, columns=["Holiday 1", "Holiday 2"])
-
-    data = rng.normal(size=(T, 1))
-    nan_locs = rng.choice(T, size=10, replace=False)
-    data[nan_locs] = np.nan
-    y = pd.DataFrame(data, index=index, columns=["sales"])
-
-    level_trend = st.LevelTrendComponent(order=1, innovations_order=[0])
-    weekly_seasonality = st.TimeSeasonality(
-        season_length=7,
-        state_names=["Sun", "Mon", "Tues", "Wed", "Thu", "Fri", "Sat"],
-        innovations=True,
-        remove_first_state=False,
-    )
-    quarterly_seasonality = st.FrequencySeasonality(season_length=365, n=2, innovations=True)
-    ar1 = st.AutoregressiveComponent(order=1)
-    me = st.MeasurementError()
-
-    exog = st.RegressionComponent(
-        name="exog",  # Name of this exogenous variable component
-        k_exog=2,  # Only one exogenous variable now
-        innovations=False,  # Typically fixed effect (no stochastic evolution)
-        state_names=df_holidays.columns.tolist(),
-    )
-
-    combined_model = level_trend + weekly_seasonality + quarterly_seasonality + me + ar1 + exog
-    ss_mod = combined_model.build()
-
-    with pm.Model(coords=ss_mod.coords) as struct_model:
-        P0_diag = pm.Gamma("P0_diag", alpha=2, beta=10, 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=["trend_state"])
-        #     sigma_trend = pm.Gamma("sigma_trend", alpha=2, beta=1, dims=["trend_shock"])  # Applied to the level only
-
-        Seasonal_coefs = pm.ZeroSumNormal(
-            "Seasonal[s=7]_coefs", sigma=0.5, dims=["Seasonal[s=7]_state"]
-        )  # DOW dev. from weekly mean
-        sigma_Seasonal = pm.Gamma(
-            "sigma_Seasonal[s=7]", alpha=2, beta=1
-        )  # How much this dev. can dev.
-
-        Frequency_coefs = pm.Normal(
-            "Frequency[s=365, n=2]", mu=0, sigma=0.5, dims=["Frequency[s=365, n=2]_state"]
-        )  # amplitudes in short-term (weekly noise culprit)
-        sigma_Frequency = pm.Gamma(
-            "sigma_Frequency[s=365, n=2]", alpha=2, beta=1
-        )  # smoothness & adaptability over time
-
-        ar_params = pm.Laplace("ar_params", mu=0, b=0.2, dims=["ar_lag"])
-        sigma_ar = pm.Gamma("sigma_ar", alpha=2, beta=1)
-
-        sigma_measurement_error = pm.HalfStudentT("sigma_MeasurementError", nu=3, sigma=1)
-
-        data_exog = pm.Data("data_exog", df_holidays.values, dims=["time", "exog_state"])
-        beta_exog = pm.Normal("beta_exog", mu=0, sigma=1, dims=["exog_state"])
-
-        ss_mod.build_statespace_graph(y, mode="JAX")
-
+    with exog_pymc_mod:
         idata = pm.sample_prior_predictive()
 
-    post = ss_mod.sample_conditional_prior(idata)
-
     # Define start date and forecast period
-    start_date, n_periods = pd.to_datetime("2024-4-15"), 8
+    start_date, n_periods = pd.to_datetime("2023-05-05"), 5
 
     # Extract exogenous data for the forecast period
     scenario = {
         "data_exog": pd.DataFrame(
-            df_holidays.loc[start_date:].iloc[:n_periods], columns=df_holidays.columns
+            exog_data[["x1"]].loc[start_date:].iloc[:n_periods], columns=exog_data[["x1"]].columns
         )
     }
 
     # Generate the forecast
-    forecasts = ss_mod.forecast(idata.prior, scenario=scenario, use_scenario_index=True)
+    forecasts = exog_ss_mod.forecast(idata.prior, scenario=scenario, use_scenario_index=True)