Added support for prediction intervals for VARMAX regressor (#4267)

* Initial commit

* Updated release notes

* Fixed random state handling

* Updated tests

* Moved number of simulations as constant.

docs/source/release_notes.rst

@@ -19,6 +19,7 @@ Release Notes
         * Added stacking and unstacking utility functions to work with multiseries data :pr:`4250`
         * Added multiseries regression pipeline class :pr:`4256`
         * Added multiseries VARMAX regressor :pr:`4238`
+        * Added support for prediction intervals for VARMAX regressor :pr:`4267`
     * Fixes
         * Added support for pandas 2 :pr:`4216`
         * Fixed bug where time series pipelines would fail due to MASE needing `y_train` when scoring :pr:`4258`

evalml/pipelines/components/estimators/regressors/exponential_smoothing_regressor.py

@@ -27,6 +27,8 @@ class ExponentialSmoothingRegressor(Estimator):
         random_seed (int): Seed for the random number generator. Defaults to 0.
     """
 
+    _N_REPETITIONS = 400
+
     name = "Exponential Smoothing Regressor"
     hyperparameter_ranges = {
         "trend": [None, "additive"],
@@ -167,7 +169,7 @@ def get_prediction_intervals(
         # anchor represents where the simulations should start from (forecasting is done from the "end")
         y_pred = self._component_obj._fitted_forecaster.simulate(
             nsimulations=X.shape[0],
-            repetitions=400,
+            repetitions=self._N_REPETITIONS,
             anchor="end",
             random_state=self.parameters["random_state"],
         )

evalml/pipelines/components/estimators/regressors/varmax_regressor.py

@@ -34,6 +34,8 @@ class VARMAXRegressor(Estimator):
             solely be based off of the datetimes and target values. Defaults to True.
     """
 
+    _N_REPETITIONS = 400
+
     name = "VARMAX Regressor"
     hyperparameter_ranges = {
         "p": Integer(0, 10),
@@ -215,11 +217,51 @@ def get_prediction_intervals(
             predictions (pd.Series): Not used for VARMAX regressor.
 
         Returns:
-            dict: Prediction intervals, keys are in the format {coverage}_lower or {coverage}_upper.
+            dict[dict]: A dict of prediction intervals, where the dict is in the format {series_id: {coverage}_lower or {coverage}_upper}.
         """
-        raise NotImplementedError(
-            "VARMAX does not have prediction intervals implemented yet.",
+        if coverage is None:
+            coverage = [0.95]
+
+        X, y = self._manage_woodwork(X, y)
+        use_exog = (
+            # If exogenous variables were used during training
+            self._component_obj._fitted_forecaster.model.exog is not None
+            and self.use_covariates
+        )
+        if use_exog:
+            X = X.ww.select(exclude=["Datetime"])
+            X = convert_bool_to_double(X)
+        # Accesses the fitted statsmodels model within sktime
+        # nsimulations represents how many steps should be simulated
+        # repetitions represents the number of simulations that should be run (confusing, I know)
+        # anchor represents where the simulations should start from (forecasting is done from the "end")
+        y_pred = self._component_obj._fitted_forecaster.simulate(
+            nsimulations=X.shape[0],
+            repetitions=self._N_REPETITIONS,
+            anchor="end",
+            random_state=self.random_seed,
+            exog=X if use_exog else None,
         )
+        prediction_interval_result = {}
+        # Access the target column names (i.e. the series_id values) that the VARMAX component obj was fitted on
+        for series in self._component_obj._fitted_forecaster.model.endog_names:
+            series_result = {}
+            series_preds = y_pred[[col for col in y_pred.columns if series in col]]
+            for conf_int in coverage:
+                prediction_interval_lower = series_preds.quantile(
+                    q=round((1 - conf_int) / 2, 3),
+                    axis="columns",
+                )
+                prediction_interval_upper = series_preds.quantile(
+                    q=round((1 + conf_int) / 2, 3),
+                    axis="columns",
+                )
+                prediction_interval_lower.index = X.index
+                prediction_interval_upper.index = X.index
+                series_result[f"{conf_int}_lower"] = prediction_interval_lower
+                series_result[f"{conf_int}_upper"] = prediction_interval_upper
+            prediction_interval_result[series] = series_result
+        return prediction_interval_result
 
     @property
     def feature_importance(self) -> np.ndarray:

evalml/tests/component_tests/test_varmax_regressor.py

@@ -260,26 +260,10 @@ def test_varmax_regressor_respects_use_covariates(
     assert "X" not in mock_predict.call_args.kwargs
 
 
-def test_varmax_regressor_prediction_intervals_not_implemented_yet_error(
-    ts_multiseries_data,
-):
-    X_train, X_test, y_train = ts_multiseries_data(n_series=2)
-
-    clf = VARMAXRegressor()
-
-    clf.fit(X_train, y_train)
-    with pytest.raises(
-        NotImplementedError,
-        match="VARMAX does not have prediction intervals implemented yet.",
-    ):
-        clf.get_prediction_intervals(X_test)
-
-
 def test_varmax_regressor_can_forecast_arbitrary_dates_no_covariates(
     ts_multiseries_data,
 ):
     X, _, y = ts_multiseries_data(n_series=2)
-
     X_train, X_test, y_train, y_test = split_data(
         X,
         y,
@@ -329,3 +313,46 @@ def test_varmax_regressor_can_forecast_arbitrary_dates_past_holdout(
     varmax.fit(X_train, y_train)
 
     varmax.predict(X_test)
+
+
+@pytest.mark.parametrize("use_X_train", [True, False])
+@pytest.mark.parametrize("coverages", [None, [0.85], [0.95, 0.90, 0.85]])
+@pytest.mark.parametrize("use_covariates", [True, False])
+def test_varmax_regressor_prediction_intervals(
+    use_covariates,
+    coverages,
+    use_X_train,
+    ts_multiseries_data,
+):
+    X_train, X_test, y_train = ts_multiseries_data(no_features=not use_covariates)
+
+    clf = VARMAXRegressor(use_covariates=use_covariates)
+    clf.fit(X=X_train if use_X_train else None, y=y_train)
+
+    # Check we are not using exogenous variables if use_covariates=False, even if X_test is passed.
+    if not use_covariates or not use_X_train:
+        with patch.object(
+            clf._component_obj._fitted_forecaster,
+            "simulate",
+        ) as mock_simulate:
+            clf.get_prediction_intervals(X_test, None, coverages)
+            assert mock_simulate.call_args[1]["exog"] is None
+
+    results_coverage = clf.get_prediction_intervals(X_test, None, coverages)
+    predictions = clf.predict(X_test)
+
+    series_id_targets = list(results_coverage.keys())
+    for series in series_id_targets:
+        series_results_coverage = results_coverage[series]
+        conf_ints = list(series_results_coverage.keys())
+        data = list(series_results_coverage.values())
+
+        assert len(conf_ints) == (len(coverages) if coverages is not None else 1) * 2
+        assert len(data) == (len(coverages) if coverages is not None else 1) * 2
+
+        for interval in coverages if coverages is not None else [0.95]:
+            conf_int_lower = f"{interval}_lower"
+            conf_int_upper = f"{interval}_upper"
+
+            assert (series_results_coverage[conf_int_upper] > predictions[series]).all()
+            assert (predictions[series] > series_results_coverage[conf_int_lower]).all()