Repository: sktime/sktime

Issue:
[ENH] Refactor/simplify `sktime.forecasting.model_selection._split.BaseSplitter._split_vectorized`
`sktime.forecasting.model_selection._split.BaseSplitter._split_vectorized` is very hard to read and understand.

This method was introduced in https://github.com/alan-turing-institute/sktime/pull/2599. During the review I proposed the following refactoring idea: https://github.com/alan-turing-institute/sktime/pull/2599#issuecomment-1126736598 (see comments that followed). For convenience, I copy my main suggestion here:

> I looked a bit closer inside of the sktime.forecasting.model_selection._split.BaseSplitter._split_vectorized. Again, I completely understand the result, for each series in the hierarchy you want to get train/test windows with numeration in the context of one big stacked time series. But I am worried about the implementation of this method. There are several enclosed loops, several temporary objects.

> For each time series in hierarchy you can easily get train/test windows that start counting from zero in isolation for each series. What you want to do next is shift those windows by the iloc indices of each series beginning. Say you have two windows in isolation: [0, 1, 3] (out of [0, 1, 2, 3]), and [0, 2] (out of [0, 1, 2]). The index of the whole series is [0, 1, 2, 3, 4, 5, 6]. The "anchors" are then 0 and 4. So your final windows are [0, 1, 3] + 0 = [0, 1, 3], and [0, 2] + 4 = [4, 6]. I know, it is just a toy example, but isn't it easy to get those iloc "anchors"?


Discussion:
agreed. There must be an easier way to do this...
Hi, I would like to work on this. I will create a PR in the next few days.
nice, thanks, @mateuja! Let me know if have any questions (since I wrote the original code).

Resolving patch:
diff --git a/sktime/forecasting/model_selection/_split.py b/sktime/forecasting/model_selection/_split.py
--- a/sktime/forecasting/model_selection/_split.py
+++ b/sktime/forecasting/model_selection/_split.py
@@ -96,8 +96,7 @@ def _get_end(y_index: pd.Index, fh: ForecastingHorizon) -> int:
     fh_offset = null if fh.is_all_in_sample() else fh[-1]
     if array_is_int(fh):
         return n_timepoints - fh_offset - 1
-    else:
-        return y_index.get_loc(y_index[-1] - fh_offset)
+    return y_index.get_loc(y_index[-1] - fh_offset)
 
 
 def _check_window_lengths(
@@ -180,10 +179,7 @@ def _inputs_are_supported(args: list) -> bool:
     -------
     True if all inputs are compatible, False otherwise
     """
-    if all_inputs_are_iloc_like(args) or all_inputs_are_time_like(args):
-        return True
-    else:
-        return False
+    return all_inputs_are_iloc_like(args) or all_inputs_are_time_like(args)
 
 
 def _check_inputs_for_compatibility(args: list) -> None:
@@ -348,7 +344,7 @@ def __init__(
         super(BaseSplitter, self).__init__()
 
     def split(self, y: ACCEPTED_Y_TYPES) -> SPLIT_GENERATOR_TYPE:
-        """Get iloc references to train/test slits of `y`.
+        """Get iloc references to train/test splits of `y`.
 
         Parameters
         ----------
@@ -394,7 +390,7 @@ def _split(self, y: pd.Index) -> SPLIT_GENERATOR_TYPE:
         """
         raise NotImplementedError("abstract method")
 
-    def _split_vectorized(self, y: pd.Index) -> SPLIT_GENERATOR_TYPE:
+    def _split_vectorized(self, y: pd.MultiIndex) -> SPLIT_GENERATOR_TYPE:
         """Get iloc references to train/test splits of `y`, for pd.MultiIndex.
 
         This applies _split per time series instance in the multiindex.
@@ -412,51 +408,22 @@ def _split_vectorized(self, y: pd.Index) -> SPLIT_GENERATOR_TYPE:
         test : 1D np.ndarray of dtype int
             Test window indices, iloc references to test indices in y
         """
-        # challenge is obtaining iloc references for *the original data frame*
-        #   todo: try to shorten this, there must be a quicker way
-        train_test_res = dict()
-        train_iloc = dict()
-        test_iloc = dict()
-        y = pd.DataFrame(index=y)
-        y_index_df = y.reset_index(-1)
-        y_index_df["__index"] = range(len(y_index_df))
-        y_index_inst = y_index_df.index.unique()
-        for idx in y_index_inst:
-            train_iloc[idx] = dict()
-            test_iloc[idx] = dict()
-        for idx in y_index_inst:
-            y_inst = y_index_df.loc[idx]
-            y_inst = y_inst.reset_index(drop=True).set_index(y_inst.columns[0])
-            y_inst_index = y_inst.index
-            train_test_res[idx] = list(self._split(y_inst_index))
-            for i, tt in enumerate(train_test_res[idx]):
-                train_iloc[idx][i] = tt[0]
-                test_iloc[idx][i] = tt[1]
-            for i, train in train_iloc[idx].items():
-                train_iloc[idx][i] = y_inst["__index"].iloc[train].values
-            for i, test in test_iloc[idx].items():
-                test_iloc[idx][i] = y_inst["__index"].iloc[test].values
-
-        train_multi = dict()
-        test_multi = dict()
-        for idx in y_index_inst:
-            for i in train_iloc[idx].keys():
-                if i not in train_multi.keys():
-                    train_multi[i] = train_iloc[idx][i]
-                    test_multi[i] = test_iloc[idx][i]
-                else:
-                    train_multi[i] = np.concatenate(
-                        (train_iloc[idx][i], train_multi[i])
-                    )
-                    test_multi[i] = np.concatenate((test_iloc[idx][i], test_multi[i]))
-        for i in train_multi.keys():
-            train_multi[i] = np.sort(train_multi[i])
-            test_multi[i] = np.sort(test_multi[i])
-
-        for i in train_multi.keys():
-            train = train_multi[i]
-            test = test_multi[i]
-            yield train, test
+        srs = pd.DataFrame(index=y).reset_index(-1).iloc[:, 0]
+        index = srs.index
+        anchors = pd.Series(range(len(srs)), index).groupby(index).first().tolist()
+        splits = (self._split(pd.Index(inst.values)) for _, inst in srs.groupby(index))
+
+        train = []
+        test = []
+        for split_inst, anchor in zip(splits, anchors):
+            train_inst, test_inst = zip(*split_inst)
+            train.append(tuple(indices + anchor for indices in train_inst))
+            test.append(tuple(indices + anchor for indices in test_inst))
+
+        train = map(np.concatenate, zip(*train))
+        test = map(np.concatenate, zip(*test))
+
+        yield from zip(train, test)
 
     def split_loc(self, y: ACCEPTED_Y_TYPES) -> Iterator[Tuple[pd.Index, pd.Index]]:
         """Get loc references to train/test splits of `y`.
@@ -666,8 +633,7 @@ def _get_train_window(
             return np.argwhere(
                 (y >= y[max(train_start, 0)]) & (y <= y[min(split_point, len(y)) - 1])
             ).flatten()
-        else:
-            return np.array([], dtype=int)
+        return np.array([], dtype=int)
 
 
 class CutoffSplitter(BaseSplitter):
@@ -792,8 +758,7 @@ def get_cutoffs(self, y: Optional[ACCEPTED_Y_TYPES] = None) -> np.ndarray:
         """
         if array_is_int(self.cutoffs):
             return check_cutoffs(self.cutoffs)
-        else:
-            return np.argwhere(y.index.isin(check_cutoffs(self.cutoffs))).flatten()
+        return np.argwhere(y.index.isin(check_cutoffs(self.cutoffs))).flatten()
 
     @classmethod
     def get_test_params(cls, parameter_set="default"):
@@ -813,8 +778,7 @@ def get_test_params(cls, parameter_set="default"):
             `MyClass(**params)` or `MyClass(**params[i])` creates a valid test instance.
             `create_test_instance` uses the first (or only) dictionary in `params`
         """
-        params = {"cutoffs": np.array([3, 7, 10])}
-        return params
+        return {"cutoffs": np.array([3, 7, 10])}
 
 
 class BaseWindowSplitter(BaseSplitter):
@@ -840,8 +804,7 @@ def __init__(
     def _initial_window(self):
         if hasattr(self, "initial_window"):
             return self.initial_window
-        else:
-            return None
+        return None
 
     def _split(self, y: pd.Index) -> SPLIT_GENERATOR_TYPE:
         n_timepoints = y.shape[0]
@@ -1478,7 +1441,6 @@ def _split_by_fh(
     if X is None:
         return y_train, y_test
 
-    else:
-        X_train = X.loc[train]
-        X_test = X.loc[test]
-        return y_train, y_test, X_train, X_test
+    X_train = X.loc[train]
+    X_test = X.loc[test]
+    return y_train, y_test, X_train, X_test


Test patch: