doc & tests

Author: jeromedockes

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,9 +4,10 @@
 from sklearn import model_selection
 from sklearn.base import BaseEstimator, TransformerMixin, clone
 from sklearn.exceptions import NotFittedError
-from sklearn.model_selection import check_cv
+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
@@ -26,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 = [
@@ -676,12 +677,20 @@ def train_test_split(
     return result
 
 
-def iter_cv_splits(data_op, environment, *, keep_subsampling=False, cv=5):
+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 = X[train_idx], X[test_idx]
+        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 = {
@@ -691,7 +700,9 @@ def iter_cv_splits(data_op, environment, *, keep_subsampling=False, cv=5):
             "X_test": X_test,
         }
         if y is not None:
-            y_train, y_test = y[train_idx], y[test_idx]
+            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

skrub/_data_ops/_skrub_namespace.py

@@ -39,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):
@@ -1534,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
@@ -1590,7 +1590,65 @@ def train_test_split(
             **split_func_kwargs,
         )
 
-    def iter_cv_splits(self, environment=None, *, keep_subsampling=False, cv=5):
+    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(

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",
@@ -1374,6 +1375,21 @@ 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):
+    return obj[list(idx)]
+
+
 @dispatch
 def replace(col, old, new):
     raise _raise(col, kind="Series")

skrub/_dataframe/tests/test_common.py

@@ -34,6 +34,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)
@@ -783,6 +784,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])