Make prediction intervals name more robust for sktime (#4283)

* Make pred intervals name more robust
alteryx · Aug 18, 2023 · 24ba211 · 24ba211
1 parent 2ad46d3
commit 24ba211
Show file tree

Hide file tree

Showing 4 changed files with 8 additions and 4 deletions.
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
@@ -4,6 +4,7 @@ Release Notes
     * Enhancements
     * Fixes
     * Changes
+        * Updated ``ARIMARegressor`` to be compatible with sktime v0.22.0 and beyond :pr:`4283`
     * Documentation Changes
         * Removed erroneous warnings from Data Checks User Guide page and removed ``tqdm`` warning in all notebooks :pr:`4274`
     * Testing Changes

diff --git a/evalml/pipelines/components/estimators/regressors/arima_regressor.py b/evalml/pipelines/components/estimators/regressors/arima_regressor.py
@@ -229,8 +229,9 @@ def _parse_prediction_intervals(
         y_pred_intervals: pd.DataFrame,
         conf_int: float,
     ) -> Tuple[pd.Series, pd.Series]:
-        preds_lower = y_pred_intervals.loc(axis=1)[("Coverage", conf_int, "lower")]
-        preds_upper = y_pred_intervals.loc(axis=1)[("Coverage", conf_int, "upper")]
+        coverage_name = y_pred_intervals.columns[0][0]
+        preds_lower = y_pred_intervals.loc(axis=1)[(coverage_name, conf_int, "lower")]
+        preds_upper = y_pred_intervals.loc(axis=1)[(coverage_name, conf_int, "upper")]
         preds_lower.name = None
         preds_upper.name = None
         return preds_lower, preds_upper

diff --git a/evalml/tests/component_tests/test_arima_regressor.py b/evalml/tests/component_tests/test_arima_regressor.py
@@ -391,8 +391,10 @@ def test_arima_regressor_respects_use_covariates(mock_predict, mock_fit, ts_data
 
 
 @pytest.mark.parametrize("no_features", [True, False])
-def test_arima_regressor_prediction_intervals(no_features, ts_data):
+@pytest.mark.parametrize("y_name", [None, "y"])
+def test_arima_regressor_prediction_intervals(no_features, ts_data, y_name):
     X_train, X_test, y_train = ts_data(no_features=no_features)
+    y_train.name = y_name
 
     clf = ARIMARegressor(use_covariates=not no_features)
 

diff --git a/evalml/tests/dependency_update_check/latest_dependency_versions.txt b/evalml/tests/dependency_update_check/latest_dependency_versions.txt
@@ -29,7 +29,7 @@ scikit-optimize==0.9.0
 scipy==1.10.1
 seaborn==0.12.2
 shap==0.42.1
-sktime==0.21.1
+sktime==0.22.0
 statsmodels==0.14.0
 texttable==1.6.7
 tomli==2.0.1