Extend TimeSeriesImputer to handle multiple series

docs/source/release_notes.rst

@@ -17,6 +17,7 @@ Release Notes
     * Enhancements
         * Added support for prediction intervals for VARMAX regressor :pr:`4267`
         * Integrated multiseries time series into AutoMLSearch :pr:`4270`
+        * Extended TimeSeriesImputer to handle multiple series :pr:`4291`
     * Fixes
         * Fixed error when stacking data with no exogenous variables :pr:`4275`
     * Changes

evalml/pipelines/components/component_base.py

@@ -256,7 +256,9 @@ def _handle_nullable_types(self, X=None, y=None):
         Args:
             X (pd.DataFrame, optional): Input data to a component of shape [n_samples, n_features].
                 May contain nullable types.
-            y (pd.Series, optional): The target of length [n_samples]. May contain nullable types.
+            y (pd.Series or pd.DataFrame, optional): The target of length [n_samples] or the
+                unstacked target for a multiseries problem of length [n_samples, n_features*n_series].
+                May contain nullable types.
 
         Returns:
             X, y with any incompatible nullable types downcasted to compatible equivalents.
@@ -278,5 +280,4 @@ def _handle_nullable_types(self, X=None, y=None):
                 handle_boolean_nullable=y_bool_incompatible,
                 handle_integer_nullable=y_int_incompatible,
             )
-
         return X, y

evalml/pipelines/components/transformers/imputers/time_series_imputer.py

@@ -93,6 +93,7 @@
         self._backwards_cols = None
         self._interpolate_cols = None
         self._impute_target = None
+        self._y_all_null_cols = None
         super().__init__(
             parameters=parameters,
             component_obj=None,
@@ -137,11 +138,17 @@
         self._backwards_cols = _filter_cols("backwards_fill", X)
         self._interpolate_cols = _filter_cols("interpolate", X)
 
-        if y is not None:
+        if isinstance(y, pd.Series):
             y = infer_feature_types(y)
             if y.isnull().any():
                 self._impute_target = self.parameters["target_impute_strategy"]
 
+        elif isinstance(y, pd.DataFrame):
+            y = infer_feature_types(y)
+            y_nan_ratio = y.isna().sum() / y.shape[0]
+            self._y_all_null_cols = y_nan_ratio[y_nan_ratio == 1].index.tolist()
+            if y.isnull().values.any():
+                self._impute_target = self.parameters["target_impute_strategy"]
         return self
 
     def transform(self, X, y=None):
@@ -212,19 +219,33 @@
             new_ltypes.update(new_int_ltypes)
         X_not_all_null.ww.init(schema=original_schema, logical_types=new_ltypes)
 
-        y_imputed = pd.Series(y)
+        y_imputed = (
+            y.ww.drop(self._y_all_null_cols)
+            if isinstance(y, pd.DataFrame)
+            else pd.Series(y)
+        )
         if y is not None and len(y) > 0:
             if self._impute_target == "forwards_fill":
-                y_imputed = y.pad()
+                y_imputed = y_imputed.pad()
                 y_imputed.bfill(inplace=True)
             elif self._impute_target == "backwards_fill":
-                y_imputed = y.bfill()
+                y_imputed = y_imputed.bfill()
                 y_imputed.pad(inplace=True)
             elif self._impute_target == "interpolate":
-                y_imputed = y.interpolate()
+                y_imputed = y_imputed.interpolate()
                 y_imputed.bfill(inplace=True)
             # Re-initialize woodwork with the downcast logical type
-            y_imputed = ww.init_series(y_imputed, logical_type=y.ww.logical_type)
+            if isinstance(y, pd.Series):
+                y_imputed = ww.init_series(y_imputed, logical_type=y.ww.logical_type)
+            else:
+                y_original_schema = y.ww.schema.get_subset_schema(
+                    list(y_imputed.columns),
+                )
+                y_new_ltypes = {
+                    col: _determine_non_nullable_equivalent(ltype)
+                    for col, ltype in y_original_schema.logical_types.items()
+                }
+                y_imputed.ww.init(schema=y_original_schema, logical_types=y_new_ltypes)
 
         return X_not_all_null, y_imputed
 
@@ -234,15 +255,22 @@
         Args:
             X (pd.DataFrame, optional): Input data to a component of shape [n_samples, n_features].
                 May contain nullable types.
-            y (pd.Series, optional): The target of length [n_samples]. May contain nullable types.
+            y (pd.Series or pd.DataFrame, optional): The target of length [n_samples] or the
+                unstacked target for a multiseries problem of length [n_samples, n_features*n_series].
+                May contain nullable types.
 
         Returns:
             X, y with any incompatible nullable types downcasted to compatible equivalents when interpolate is used. Is NoOp otherwise.
         """
         if self._impute_target == "interpolate":
             # For BooleanNullable, we have to avoid Categorical columns
             # since the category dtype also has incompatibilities with linear interpolate, which is expected
-            if isinstance(y.ww.logical_type, BooleanNullable):
+            # TODO: Avoid categorical columns for BooleanNullable in multiseries when
+            #       multiseries timeseries supports categorical
+            if isinstance(y, pd.Series) and isinstance(
+                y.ww.logical_type,
+                BooleanNullable,
+            ):
                 y = ww.init_series(y, Double)
             else:
                 _, y = super()._handle_nullable_types(None, y)

evalml/tests/component_tests/test_time_series_imputer.py

@@ -722,3 +722,78 @@
         _, nullable_series = imputer._handle_nullable_types(None, nullable_series)
 
     nullable_series.interpolate()
+
+
+@pytest.mark.parametrize(
+    "nans_present",
+    [True, False],
+)
+def test_time_series_imputer_multiseries(
+    multiseries_ts_data_unstacked,
+    nans_present,
+):
+    X, y = multiseries_ts_data_unstacked
+    imputer = TimeSeriesImputer(target_impute_strategy="interpolate")
+
+    if nans_present:
+        for count, col in enumerate(y, start=1):
+            y[col][count] = np.nan
+
+    imputer.fit(X, y)
+    assert imputer._y_all_null_cols == []
+
+    _, y_imputed = imputer.transform(X, y)
+    assert isinstance(y_imputed, pd.DataFrame)
+
+    y_expected = pd.DataFrame({f"target_{i}": range(i, 100, 5) for i in range(5)})
+    assert_frame_equal(y_imputed, y_expected, check_dtype=False)
+
+
+@pytest.mark.parametrize(
+    "num_nan_cols",
+    [1, 2, 3],
+)
+@pytest.mark.parametrize(
+    "nan_in_other_cols",
+    [True, False],
+)
+def test_time_series_imputer_multiseries_some_columns_all_nan(
+    multiseries_ts_data_unstacked,
+    num_nan_cols,
+    nan_in_other_cols,
+):
+    X, y = multiseries_ts_data_unstacked
+    imputer = TimeSeriesImputer(target_impute_strategy="interpolate")
+
+    for count, col in enumerate(y, start=1):
+        if count <= num_nan_cols:
+            y[col] = np.nan
+        if count == num_nan_cols and not nan_in_other_cols:
+            break
+        else:
+            y[col][count] = np.nan
+
+    imputer.fit(X, y)
+    _, y_imputed = imputer.transform(X, y)
+
+    y_expected = pd.DataFrame(
+        {f"target_{i}": range(i, 100, 5) for i in range(num_nan_cols, 5)},
+    )
+    assert_frame_equal(y_imputed, y_expected, check_dtype=False)
+
+
+def test_imputer_multiseries_drops_columns_with_all_nan(multiseries_ts_data_unstacked):
+    X, y = multiseries_ts_data_unstacked
+    for col in y:
+        y[col] = np.nan
+    imputer = TimeSeriesImputer(target_impute_strategy="interpolate")
+    imputer.fit(X, y)
+    assert imputer._y_all_null_cols == y.columns.tolist()
+    _, y_imputed = imputer.transform(X, y)
+    expected = y.drop(y.columns.tolist(), axis=1)
+    assert_frame_equal(
+        y_imputed,
+        expected,
+        check_column_type=False,
+        check_index_type=False,
+    )

evalml/tests/utils_tests/test_nullable_type_utils.py

@@ -53,11 +53,24 @@
 
 
 @pytest.mark.parametrize(
-    "downcast_util, data_type",
-    [(_downcast_nullable_X, "X"), (_downcast_nullable_y, "y")],
+    "downcast_util, data_type, y_type",
+    [
+        (_downcast_nullable_X, "X", "series"),
+        (_downcast_nullable_y, "y", "series"),
+        (_downcast_nullable_y, "y", "dataframe"),
+    ],
 )
-def test_downcast_utils_handle_woodwork_not_init(X_y_binary, downcast_util, data_type):
-    X, y = X_y_binary
+def test_downcast_utils_handle_woodwork_not_init(
+    X_y_binary,
+    multiseries_ts_data_unstacked,
+    downcast_util,
+    data_type,
+    y_type,
+):
+    if y_type == "series":
+        X, y = X_y_binary
+    else:
+        X, y = multiseries_ts_data_unstacked
     # Remove woodwork types
     if data_type == "X":
         data = X.copy()

evalml/utils/nullable_type_utils.py

@@ -1,4 +1,4 @@
-import woodwork as ww
+import pandas as pd
 from woodwork.logical_types import AgeNullable, BooleanNullable, IntegerNullable
 
 DOWNCAST_TYPE_DICT = {
@@ -48,7 +48,7 @@
         to other dtypes via Woodwork logical type transformations.
 
     Args:
-        y (pd.Series): Target data of shape [n_samples] whose nullable types will be changed.
+        y (pd.Series or pd.DataFrame): Target data of shape [n_samples] or [n_samples, n_features*n_series] whose nullable types will be changed.
         handle_boolean_nullable (bool, optional): Whether or not to downcast data with BooleanNullable logical types.
         handle_integer_nullable (bool, optional): Whether or not to downcast data with IntegerNullable or AgeNullable logical types.
 
@@ -57,16 +57,20 @@
         y with any incompatible nullable types downcasted to compatible equivalents.
     """
     if y.ww.schema is None:
-        y = ww.init_series(y)
+        y.ww.init()
 
     incompatible_logical_types = _get_incompatible_nullable_types(
         handle_boolean_nullable,
         handle_integer_nullable,
     )
 
-    if isinstance(y.ww.logical_type, tuple(incompatible_logical_types)):
-        new_ltype = _determine_downcast_type(y)
-        return y.ww.set_logical_type(new_ltype)
+    if isinstance(y, pd.DataFrame):
+        y = _downcast_nullable_X(y)
+
+    else:
+        if isinstance(y.ww.logical_type, tuple(incompatible_logical_types)):
+            new_ltype = _determine_downcast_type(y)
+            return y.ww.set_logical_type(new_ltype)
 
     return y