Expand test coverage of forecast option combinations

Author: jessegrabowski

pymc_experimental/statespace/core/statespace.py

@@ -1502,16 +1502,21 @@ def sample_statespace_matrices(
 
     @staticmethod
     def _validate_forecast_args(
-        time_index: pd.Index,
+        time_index: pd.RangeIndex | pd.DatetimeIndex,
         start: int | pd.Timestamp,
         periods: int | None = None,
         end: int | pd.Timestamp = None,
         scenario: pd.DataFrame | np.ndarray | None = None,
         use_scenario_index: bool = False,
         verbose: bool = True,
     ):
-        if start not in time_index:
-            raise ValueError("start must be in the data index used to fit the model.")
+        if isinstance(start, pd.Timestamp) and start not in time_index:
+            raise ValueError("Datetime start must be in the data index used to fit the model.")
+        elif isinstance(start, int):
+            if abs(start) > len(time_index):
+                raise ValueError(
+                    "Integer start must be within the range of the data index used to fit the model."
+                )
         if periods is None and end is None:
             raise ValueError("Must specify one of either periods or end")
         if periods is not None and end is not None:
@@ -1559,6 +1564,22 @@ def _validate_scenario_data(
         name: str | None = None,
         verbose=True,
     ):
+        """
+        Validate the scenario data provided to the forecast method by checking that it has the correct shape and
+        dimensions.
+
+        Parameters
+        ----------
+        scenario
+        name
+        verbose
+
+        Returns
+        -------
+        scenario: pd.DataFrame | np.ndarray | dict[str, pd.DataFrame | np.ndarray]
+            Scenario data, validated and potentially modified.
+
+        """
         if not self._needs_exog_data:
             return scenario
 
@@ -1758,18 +1779,27 @@ def get_or_create_index(x, time_index, start=None):
             forecast_index = None
 
             if is_datetime:
-                freq = time_index.inferred_freq
+                freq = time_index.freq
                 if isinstance(start, int):
                     start = time_index[start]
+                if isinstance(end, int):
+                    raise ValueError(
+                        "end must be a timestamp if using a datetime index. To specify a number of "
+                        "timesteps from the start date, use the periods argument instead."
+                    )
                 if end is not None:
                     forecast_index = pd.date_range(start, end=end, freq=freq)
                 if periods is not None:
-                    # date_range include both start and end, but we're going to pop off the start later (it will be
-                    # interpreted as x0). So we need to add 1 to the periods so the user gets "periods" number of
-                    # forecasts back
+                    # date_range includes both the start and end date, but we're going to pop off the start later
+                    # (it will be interpreted as x0). So we need to add 1 to the periods so the user gets "periods"
+                    # number of forecasts back
                     forecast_index = pd.date_range(start, periods=periods + 1, freq=freq)
 
             else:
+                # If the user provided a positive integer as start, directly interpret it as the start time. If its
+                # negative, interpret it as a positional index.
+                if start < 0:
+                    start = time_index[start]
                 if end is not None:
                     forecast_index = pd.RangeIndex(start, end, step=1, dtype="int")
                 if periods is not None:

tests/statespace/test_statespace.py

@@ -114,6 +114,18 @@ def pymc_mod(ss_mod):
     return pymc_mod
 
 
+@pytest.fixture(scope="session")
+def ss_mod_no_exog(rng):
+    ll = st.LevelTrendComponent(order=2, innovations_order=1)
+    return ll.build()
+
+
+@pytest.fixture(scope="session")
+def ss_mod_no_exog_dt(rng):
+    ll = st.LevelTrendComponent(order=2, innovations_order=1)
+    return ll.build()
+
+
 @pytest.fixture(scope="session")
 def exog_ss_mod(rng):
     ll = st.LevelTrendComponent()
@@ -143,6 +155,42 @@ def exog_pymc_mod(exog_ss_mod, rng):
     return m
 
 
+@pytest.fixture(scope="session")
+def pymc_mod_no_exog(ss_mod_no_exog, rng):
+    y = pd.DataFrame(rng.normal(size=(100, 1)).astype(floatX), columns=["y"])
+
+    with pm.Model(coords=ss_mod_no_exog.coords) as m:
+        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
+        P0_sigma = pm.Exponential("P0_sigma", 1)
+        P0 = pm.Deterministic(
+            "P0", pt.eye(ss_mod_no_exog.k_states) * P0_sigma, dims=["state", "state_aux"]
+        )
+        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
+        ss_mod_no_exog.build_statespace_graph(y)
+
+    return m
+
+
+@pytest.fixture(scope="session")
+def pymc_mod_no_exog_dt(ss_mod_no_exog_dt, rng):
+    y = pd.DataFrame(
+        rng.normal(size=(100, 1)).astype(floatX),
+        columns=["y"],
+        index=pd.date_range("2020-01-01", periods=100, freq="D"),
+    )
+
+    with pm.Model(coords=ss_mod_no_exog_dt.coords) as m:
+        initial_trend = pm.Normal("initial_trend", dims=["trend_state"])
+        P0_sigma = pm.Exponential("P0_sigma", 1)
+        P0 = pm.Deterministic(
+            "P0", pt.eye(ss_mod_no_exog_dt.k_states) * P0_sigma, dims=["state", "state_aux"]
+        )
+        sigma_trend = pm.Exponential("sigma_trend", 1, dims=["trend_shock"])
+        ss_mod_no_exog_dt.build_statespace_graph(y)
+
+    return m
+
+
 @pytest.fixture(scope="session")
 def idata(pymc_mod, rng):
     with pymc_mod:
@@ -162,6 +210,24 @@ def idata_exog(exog_pymc_mod, rng):
     return idata
 
 
+@pytest.fixture(scope="session")
+def idata_no_exog(pymc_mod_no_exog, rng):
+    with pymc_mod_no_exog:
+        idata = pm.sample(draws=10, tune=0, chains=1, random_seed=rng)
+        idata_prior = pm.sample_prior_predictive(draws=10, random_seed=rng)
+    idata.extend(idata_prior)
+    return idata
+
+
+@pytest.fixture(scope="session")
+def idata_no_exog_dt(pymc_mod_no_exog_dt, rng):
+    with pymc_mod_no_exog_dt:
+        idata = pm.sample(draws=10, tune=0, chains=1, random_seed=rng)
+        idata_prior = pm.sample_prior_predictive(draws=10, random_seed=rng)
+    idata.extend(idata_prior)
+    return idata
+
+
 def test_invalid_filter_name_raises():
     msg = "The following are valid filter types: " + ", ".join(list(FILTER_FACTORY.keys()))
     with pytest.raises(NotImplementedError, match=msg):
@@ -664,28 +730,75 @@ def test_invalid_scenarios():
         ss_mod._validate_scenario_data(scenario)
 
 
+@pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
 @pytest.mark.parametrize("filter_output", ["predicted", "filtered", "smoothed"])
-def test_forecast(filter_output, ss_mod, idata, rng):
-    time_idx = idata.posterior.coords["time"].values
-    forecast_idata = ss_mod.forecast(
-        idata, start=time_idx[-1], periods=10, filter_output=filter_output, random_seed=rng
+@pytest.mark.parametrize(
+    "mod_name, idata_name, start, end, periods",
+    [
+        ("ss_mod_no_exog", "idata_no_exog", None, None, 10),
+        ("ss_mod_no_exog", "idata_no_exog", -1, None, 10),
+        ("ss_mod_no_exog", "idata_no_exog", 10, None, 10),
+        ("ss_mod_no_exog", "idata_no_exog", 10, 21, None),
+        ("ss_mod_no_exog_dt", "idata_no_exog_dt", None, None, 10),
+        ("ss_mod_no_exog_dt", "idata_no_exog_dt", -1, None, 10),
+        ("ss_mod_no_exog_dt", "idata_no_exog_dt", 10, None, 10),
+        ("ss_mod_no_exog_dt", "idata_no_exog_dt", 10, "2020-01-21", None),
+        ("ss_mod_no_exog_dt", "idata_no_exog_dt", "2020-03-01", "2020-03-11", None),
+        ("ss_mod_no_exog_dt", "idata_no_exog_dt", "2020-03-01", None, 10),
+    ],
+    ids=[
+        "range_default",
+        "range_negative",
+        "range_int",
+        "range_end",
+        "datetime_default",
+        "datetime_negative",
+        "datetime_int",
+        "datetime_int_end",
+        "datetime_datetime_end",
+        "datetime_datetime",
+    ],
+)
+def test_forecast(filter_output, mod_name, idata_name, start, end, periods, rng, request):
+    mod = request.getfixturevalue(mod_name)
+    idata = request.getfixturevalue(idata_name)
+    time_idx = mod._get_fit_time_index()
+    is_datetime = isinstance(time_idx, pd.DatetimeIndex)
+
+    if isinstance(start, str):
+        t0 = pd.Timestamp(start)
+    elif isinstance(start, int):
+        t0 = time_idx[start]
+    else:
+        t0 = time_idx[-1]
+
+    delta = time_idx.freq if is_datetime else 1
+
+    forecast_idata = mod.forecast(
+        idata, start=start, end=end, periods=periods, filter_output=filter_output, random_seed=rng
     )
 
-    assert forecast_idata.coords["time"].values.shape == (10,)
+    forecast_idx = forecast_idata.coords["time"].values
+    forecast_idx = pd.DatetimeIndex(forecast_idx) if is_datetime else pd.Index(forecast_idx)
+
+    assert forecast_idx.shape == (10,)
     assert forecast_idata.forecast_latent.dims == ("chain", "draw", "time", "state")
     assert forecast_idata.forecast_observed.dims == ("chain", "draw", "time", "observed_state")
 
     assert not np.any(np.isnan(forecast_idata.forecast_latent.values))
     assert not np.any(np.isnan(forecast_idata.forecast_observed.values))
 
+    assert forecast_idx[0] == (t0 + delta)
+
 
 @pytest.mark.filterwarnings("ignore:No time index found on the supplied data.")
-def test_forecast_with_exog_data(rng, exog_ss_mod, idata_exog):
+@pytest.mark.parametrize("start", [None, -1, 10])
+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.iloc[5, 0] = 1e9
 
     forecast_idata = exog_ss_mod.forecast(
-        idata_exog, periods=10, random_seed=rng, scenario=scenario
+        idata_exog, start=start, periods=10, random_seed=rng, scenario=scenario
     )
 
     components = exog_ss_mod.extract_components_from_idata(forecast_idata)