add .skb.iter_cv_splits() (#1653)

Co-authored-by: Riccardo Cappuzzo <7548232+rcap107@users.noreply.github.com>

Author: jeromedockes

CHANGES.rst

@@ -13,6 +13,10 @@ New features
 ------------
 - TableReport now displays the mean statistic for boolean columns.
   :pr:`1647` by :user:`Abdelhakim Benechehab <abenechehab>`.
+- :meth:`DataOp.skb.iter_cv_splits` iterates over the training and testing
+  environments produced by a CV splitter -- similar to
+  :meth:`DataOp.skb.train_test_split` but for multiple cross-validation splits.
+  :pr:`1653` by :user:`Jérôme Dockès <jeromedockes>`.
 
 Changes
 -------

doc/api_reference.py

@@ -237,6 +237,7 @@
                     "DataOp.skb.make_grid_search",
                     "DataOp.skb.make_randomized_search",
                     "DataOp.skb.if_else",
+                    "DataOp.skb.iter_cv_splits",
                     "DataOp.skb.iter_learners_grid",
                     "DataOp.skb.iter_learners_randomized",
                     "DataOp.skb.mark_as_X",

skrub/_data_ops/_estimator.py

@@ -4,8 +4,10 @@
 from sklearn import model_selection
 from sklearn.base import BaseEstimator, TransformerMixin, clone
 from sklearn.exceptions import NotFittedError
+from sklearn.model_selection import KFold, check_cv
 from sklearn.utils.validation import check_is_fitted
 
+from .. import _dataframe as sbd
 from .. import _join_utils
 from ._choosing import BaseNumericChoice, get_default
 from ._data_ops import Apply, check_subsampled_X_y_shape
@@ -25,7 +27,7 @@
 from ._inspection import describe_params
 from ._parallel_coord import DEFAULT_COLORSCALE, plot_parallel_coord
 from ._subsampling import env_with_subsampling
-from ._utils import X_NAME, Y_NAME, _CloudPickle, attribute_error
+from ._utils import KFOLD_5, X_NAME, Y_NAME, _CloudPickle, attribute_error
 
 _FITTING_METHODS = ["fit", "fit_transform"]
 _SKLEARN_SEARCH_FITTED_ATTRIBUTES_TO_COPY = [
@@ -675,6 +677,39 @@ def train_test_split(
     return result
 
 
+def iter_cv_splits(data_op, environment, *, keep_subsampling=False, cv=KFOLD_5):
+    """Yield splits of an environment into training an testing environments.
+
+    This functionality is exposed to users through the
+    ``DataOp.skb.iter_cv_splits()`` method. See the corresponding docstring for
+    details and examples.
+    """
+    if cv is KFOLD_5:
+        cv = KFold(5)
+    cv = check_cv(cv)
+    environment = env_with_subsampling(data_op, environment, keep_subsampling)
+    X, y = _compute_Xy(data_op, environment)
+    for train_idx, test_idx in cv.split(X, y):
+        X_train, X_test = sbd.select_rows(X, train_idx), sbd.select_rows(X, test_idx)
+        train_env = {**environment, X_NAME: X_train}
+        test_env = {**environment, X_NAME: X_test}
+        split_info = {
+            "train": train_env,
+            "test": test_env,
+            "X_train": X_train,
+            "X_test": X_test,
+        }
+        if y is not None:
+            y_train, y_test = sbd.select_rows(y, train_idx), sbd.select_rows(
+                y, test_idx
+            )
+            train_env[Y_NAME] = y_train
+            test_env[Y_NAME] = y_test
+            split_info["y_train"] = y_train
+            split_info["y_test"] = y_test
+        yield split_info
+
+
 class ParamSearch(_CloudPickleDataOp, BaseEstimator):
     """Learner that evaluates a skrub DataOp with hyperparameter tuning.
 

skrub/_data_ops/_skrub_namespace.py

@@ -19,7 +19,13 @@
     check_name,
     deferred,
 )
-from ._estimator import ParamSearch, SkrubLearner, cross_validate, train_test_split
+from ._estimator import (
+    ParamSearch,
+    SkrubLearner,
+    cross_validate,
+    iter_cv_splits,
+    train_test_split,
+)
 from ._evaluation import (
     choices,
     clone,
@@ -33,7 +39,7 @@
     full_report,
 )
 from ._subsampling import SubsamplePreviews, env_with_subsampling
-from ._utils import NULL, attribute_error
+from ._utils import KFOLD_5, NULL, attribute_error
 
 
 def _var_values_provided(data_op, environment):
@@ -1528,9 +1534,9 @@ def train_test_split(
 
             - train: a dictionary containing the training environment
             - test: a dictionary containing the test environment
-            - X_train: the value of the variable marked with ``skb.mark_as_x()`` in
+            - X_train: the value of the variable marked with ``skb.mark_as_X()`` in
               the train environment
-            - X_test: the value of the variable marked with ``skb.mark_as_x()`` in
+            - X_test: the value of the variable marked with ``skb.mark_as_X()`` in
               the test environment
             - y_train: the value of the variable marked with ``skb.mark_as_y()`` in
               the train environment, if there is one (may not be the case for
@@ -1584,6 +1590,71 @@ def train_test_split(
             **split_func_kwargs,
         )
 
+    def iter_cv_splits(self, environment=None, *, keep_subsampling=False, cv=KFOLD_5):
+        """Yield splits of an environment into training and testing environments.
+
+        Parameters
+        ----------
+        environment : dict, optional
+            The environment (dict mapping variable names to values) containing the
+            full data. If ``None`` (the default), the data is retrieved from the
+            DataOp.
+
+        keep_subsampling : bool, default=False
+            If True, and if subsampling has been configured (see
+            :meth:`DataOp.skb.subsample`), use a subsample of the data. By
+            default subsampling is not applied and all the data is used.
+
+        cv : int, cross-validation generator or iterable, default=KFold(5)
+            The default is 5-fold without shuffling. Can be a cross-validation
+            splitter, an iterable yielding pairs of (train, test) indices, or an
+            int to specify the number of folds for KFold splitting.
+
+        Yields
+        ------
+        dict
+            For each split, a dict is produced, containing the following keys:
+
+            - train: a dictionary containing the training environment
+            - test: a dictionary containing the test environment
+            - X_train: the value of the variable marked with ``skb.mark_as_X()`` in
+              the train environment
+            - X_test: the value of the variable marked with ``skb.mark_as_X()`` in
+              the test environment
+            - y_train: the value of the variable marked with ``skb.mark_as_y()`` in
+              the train environment, if there is one (may not be the case for
+              unsupervised learning).
+            - y_test: the value of the variable marked with ``skb.mark_as_y()`` in
+              the test environment, if there is one (may not be the case for
+              unsupervised learning).
+
+        Examples
+        --------
+        >>> import skrub
+        >>> from sklearn.dummy import DummyClassifier
+        >>> from sklearn.metrics import accuracy_score
+
+        >>> orders = skrub.var("orders")
+        >>> X = orders.skb.drop("delayed").skb.mark_as_X()
+        >>> y = orders["delayed"].skb.mark_as_y()
+        >>> delayed = X.skb.apply(skrub.TableVectorizer()).skb.apply(
+        ...     DummyClassifier(), y=y
+        ... )
+        >>> df = skrub.datasets.toy_orders().orders
+        >>> accuracies = []
+        >>> for split in delayed.skb.iter_cv_splits({"orders": df}, cv=3):
+        ...     learner = delayed.skb.make_learner().fit(split["train"])
+        ...     prediction = learner.predict(split["test"])
+        ...     accuracies.append(accuracy_score(split["y_test"], prediction))
+        >>> accuracies
+        [1.0, 0.0, 1.0]
+        """
+        if environment is None:
+            environment = self.get_data()
+        yield from iter_cv_splits(
+            self._data_op, environment, keep_subsampling=keep_subsampling, cv=cv
+        )
+
     def make_grid_search(self, *, fitted=False, keep_subsampling=False, **kwargs):
         """Find the best parameters with grid search.
 

skrub/_data_ops/_utils.py

@@ -12,6 +12,7 @@
 class Sentinels(enum.Enum):
     NULL = "NULL"
     OPTIONAL_VALUE = "value"
+    KFOLD_5 = "KFold(5)"
 
     def __repr__(self):
         return self.value
@@ -22,6 +23,7 @@ def __str__(self):
 
 NULL = Sentinels.NULL
 OPTIONAL_VALUE = Sentinels.OPTIONAL_VALUE
+KFOLD_5 = Sentinels.KFOLD_5
 
 
 def simple_repr(data_op):

skrub/_data_ops/tests/test_estimators.py

@@ -422,6 +422,31 @@ def test_train_test_split(with_y):
         assert e.skb.eval() == [7, 6, 5, 4, 3, 2, 1, 0]
 
 
+def test_iter_cv_splits():
+    X = skrub.X(np.arange(5) * 10)
+    splits = X.skb.iter_cv_splits()
+    s = next(splits)
+    assert list(s["X_train"]) == list(s["train"]["_skrub_X"]) == [10, 20, 30, 40]
+    assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [0]
+    s = next(splits)
+    assert list(s["X_train"]) == list(s["train"]["_skrub_X"]) == [0, 20, 30, 40]
+    assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [10]
+
+    X = skrub.X(np.arange(4) * 10)
+    y = skrub.y(np.arange(4) * -10)
+    splits = skrub.as_data_op((X, y)).skb.iter_cv_splits(cv=4)
+    s = next(splits)
+    assert list(s["X_train"]) == list(s["train"]["_skrub_X"]) == [10, 20, 30]
+    assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [0]
+    assert list(s["y_train"]) == list(s["train"]["_skrub_y"]) == [-10, -20, -30]
+    assert list(s["y_test"]) == list(s["test"]["_skrub_y"]) == [0]
+    s = next(splits)
+    assert list(s["X_train"]) == list(s["train"]["_skrub_X"]) == [0, 20, 30]
+    assert list(s["X_test"]) == list(s["test"]["_skrub_X"]) == [10]
+    assert list(s["y_train"]) == list(s["train"]["_skrub_y"]) == [0, -20, -30]
+    assert list(s["y_test"]) == list(s["test"]["_skrub_y"]) == [-10]
+
+
 def test_train_test_split_splitter_renaming():
     # TODO remove when `splitter` is removed in 0.7.0
     X = skrub.X(list(range(10)))

skrub/_dataframe/_common.py

@@ -103,6 +103,7 @@
     "sample",
     "head",
     "slice",
+    "select_rows",
     "replace",
     "with_columns",
     "abs",
@@ -1355,6 +1356,26 @@ def _slice_polars(obj, *start_stop):
     return obj.slice(start, stop - start)
 
 
+@dispatch
+def select_rows(obj, idx):
+    return np.asarray(obj)[list(idx)]
+
+
+@select_rows.specialize("pandas")
+def _select_rows_pandas(obj, idx):
+    return obj.iloc[list(idx)]
+
+
+@select_rows.specialize("polars")
+def _select_rows_polars(obj, idx):
+    idx = list(idx)
+    if not idx:
+        # polars changed from interpreting indexing with an empty list as list
+        # of columns to list of row indices at some point.
+        return obj.head(0)
+    return obj[idx]
+
+
 @dispatch
 def replace(col, old, new):
     raise_dispatch_unregistered_type(col, kind="Series")

skrub/_dataframe/tests/test_common.py

@@ -35,6 +35,7 @@ def test_not_implemented():
         "copy_index",
         "index",
         "with_columns",
+        "select_rows",
     }
     for func_name in sorted(set(ns.__all__) - has_default_impl):
         func = getattr(ns, func_name)
@@ -786,6 +787,23 @@ def test_slice(df_module, obj, s):
     assert_array_equal(out, expected)
 
 
+@pytest.mark.parametrize("obj", ["column", "dataframe"])
+@pytest.mark.parametrize("idx", [[], [1], [2, 1]])
+def test_select_rows(df_module, obj, idx):
+    out = ns.select_rows(getattr(df_module, f"example_{obj}"), idx)
+    if obj == "dataframe":
+        out = ns.col(out, "float-col")
+    out = ns.to_numpy(out)
+    expected = ns.to_numpy(df_module.example_column)[idx]
+    assert_array_equal(out, expected)
+
+
+def test_select_rows_array():
+    a = np.arange(6).reshape((3, 2))
+    assert_array_equal(ns.select_rows(a, (2, 1)), a[[2, 1], :])
+    assert_array_equal(ns.select_rows(a[0], (1, 0)), a[0, [1, 0]])
+
+
 def test_is_null(df_module):
     s = ns.pandas_convert_dtypes(df_module.make_column("", [0, None, 2, None, 4]))
     expected = df_module.make_column("", [False, True, False, True, False])