Allow MapieRegressor to use K-fold iterator variants with stratification and groups.

AUTHORS.rst

@@ -36,5 +36,6 @@ Contributors
 * Arthur Phan <[email protected]>
 * Rafael Saraiva <[email protected]>
 * Mehdi Elion <[email protected]>
+* Pierre de Fréminville <[email protected]>
 
 To be continued ...

HISTORY.rst

@@ -4,6 +4,9 @@ History
 
 ##### (##########)
 ------------------
+* Allow the use of `y` and `groups` arguments in cross validator methods `get_n_splits`
+  and `split` to enable more cv-split variants for :class:`MapieRegressor`
+  (e.g. :class:`GroupKFold`, stratified continuous split).
 
 0.8.0 (2024-01-03)
 ------------------

Makefile

@@ -1,7 +1,7 @@
 .PHONY: tests doc build
 
-lint:	
-	flake8 . --exclude=doc
+lint:
+	flake8 . --exclude=doc,.venv
 
 type-check:
 	mypy mapie

mapie/classification.py

@@ -1047,6 +1047,7 @@ def fit(
         y: ArrayLike,
         sample_weight: Optional[ArrayLike] = None,
         size_raps: Optional[float] = .2,
+        groups: Optional[ArrayLike] = None,
     ) -> MapieClassifier:
         """
         Fit the base estimator or use the fitted base estimator.
@@ -1074,6 +1075,11 @@ def fit(
 
             By default ``.2``.
 
+        groups: Optional[ArrayLike] of shape (n_samples,)
+            Group labels for the samples used while splitting the dataset into
+            train/test set.
+
+            By default ``None``.
 
         Returns
         -------
@@ -1163,7 +1169,9 @@ def fit(
                     k,
                     sample_weight,
                 )
-                for k, (train_index, val_index) in enumerate(cv.split(X))
+                for k, (train_index, val_index) in enumerate(
+                    cv.split(X, y_enc, groups)
+                )
             )
             (
                 self.estimators_,

mapie/estimator/estimator.py

@@ -322,7 +322,11 @@ def _pred_multi(self, X: ArrayLike) -> NDArray:
         y_pred_multi = self._aggregate_with_mask(y_pred_multi, self.k_)
         return y_pred_multi
 
-    def predict_calib(self, X: ArrayLike) -> NDArray:
+    def predict_calib(
+            self,
+            X: ArrayLike,
+            y: Optional[ArrayLike] = None,
+            groups: Optional[ArrayLike] = None) -> NDArray:
         """
         Perform predictions on X : the calibration set.
 
@@ -331,6 +335,17 @@ def predict_calib(self, X: ArrayLike) -> NDArray:
         X: ArrayLike of shape (n_samples_test, n_features)
             Input data
 
+        y: ArrayLike of shape (n_samples,)
+            Input labels.
+
+            By default ``None``.
+
+        groups: Optional[ArrayLike] of shape (n_samples,)
+            Group labels for the samples used while splitting the dataset into
+            train/test set.
+
+            By default ``None``.
+
         Returns
         -------
         NDArray of shape (n_samples_test, 1)
@@ -349,15 +364,17 @@ def predict_calib(self, X: ArrayLike) -> NDArray:
                     delayed(self._predict_oof_estimator)(
                         estimator, X, calib_index,
                     )
-                    for (_, calib_index), estimator in zip(cv.split(X),
-                                                           self.estimators_)
+                    for (_, calib_index), estimator in zip(
+                        cv.split(X, y, groups),
+                        self.estimators_
+                    )
                 )
                 predictions, indices = map(
                     list, zip(*outputs)
                 )
                 n_samples = _num_samples(X)
                 pred_matrix = np.full(
-                    shape=(n_samples, cv.get_n_splits(X)),
+                    shape=(n_samples, cv.get_n_splits(X, y, groups)),
                     fill_value=np.nan,
                     dtype=float,
                 )
@@ -377,6 +394,7 @@ def fit(
         X: ArrayLike,
         y: ArrayLike,
         sample_weight: Optional[ArrayLike] = None,
+        groups: Optional[ArrayLike] = None
     ) -> EnsembleRegressor:
         """
         Fit the base estimator under the ``single_estimator_`` attribute.
@@ -397,6 +415,11 @@ def fit(
             Sample weights. If None, then samples are equally weighted.
             By default ``None``.
 
+        groups: Optional[ArrayLike] of shape (n_samples,)
+            Group labels for the samples used while splitting the dataset into
+            train/test set.
+            By default ``None``.
+
         Returns
         -------
         EnsembleRegressor
@@ -423,7 +446,7 @@ def fit(
             )
             cv = cast(BaseCrossValidator, cv)
             self.k_ = np.full(
-                shape=(n_samples, cv.get_n_splits(X, y)),
+                shape=(n_samples, cv.get_n_splits(X, y, groups)),
                 fill_value=np.nan,
                 dtype=float,
             )
@@ -434,7 +457,7 @@ def fit(
                     delayed(self._fit_oof_estimator)(
                         clone(estimator), X, y, train_index, sample_weight
                     )
-                    for train_index, _ in cv.split(X)
+                    for train_index, _ in cv.split(X, y, groups)
                 )
                 # In split-CP, we keep only the model fitted on train dataset
                 if self.use_split_method_:

mapie/estimator/interface.py

@@ -21,6 +21,7 @@ def fit(
         X: ArrayLike,
         y: ArrayLike,
         sample_weight: Optional[ArrayLike] = None,
+        groups: Optional[ArrayLike] = None,
     ) -> EnsembleEstimator:
         """
         Fit the base estimator under the ``single_estimator_`` attribute.
@@ -41,6 +42,11 @@ def fit(
             Sample weights. If None, then samples are equally weighted.
             By default ``None``.
 
+        groups: Optional[ArrayLike] of shape (n_samples,)
+            Group labels for the samples used while splitting the dataset into
+            train/test set.
+            By default ``None``.
+
         Returns
         -------
         EnsembleRegressor

mapie/regression/regression.py

@@ -392,6 +392,7 @@ def _check_fit_parameters(
         X: ArrayLike,
         y: ArrayLike,
         sample_weight: Optional[ArrayLike] = None,
+        groups: Optional[ArrayLike] = None
     ):
         """
         Perform several checks on class parameters.
@@ -407,6 +408,11 @@ def _check_fit_parameters(
         sample_weight: Optional[NDArray] of shape (n_samples,)
             Non-null sample weights.
 
+        groups: Optional[ArrayLike] of shape (n_samples,)
+            Group labels for the samples used while splitting the dataset into
+            train/test set.
+            By default ``None``.
+
         Raises
         ------
         ValueError
@@ -449,14 +455,21 @@ def _check_fit_parameters(
         X = cast(NDArray, X)
         y = cast(NDArray, y)
         sample_weight = cast(Optional[NDArray], sample_weight)
+        groups = cast(Optional[NDArray], groups)
 
-        return estimator, cs_estimator, agg_function, cv, X, y, sample_weight
+        return (
+            estimator, cs_estimator,
+            agg_function, cv,
+            X, y,
+            sample_weight, groups
+        )
 
     def fit(
         self,
         X: ArrayLike,
         y: ArrayLike,
         sample_weight: Optional[ArrayLike] = None,
+        groups: Optional[ArrayLike] = None
     ) -> MapieRegressor:
         """
         Fit estimator and compute conformity scores used for
@@ -484,6 +497,11 @@ def fit(
 
             By default ``None``.
 
+        groups: Optional[ArrayLike] of shape (n_samples,)
+            Group labels for the samples used while splitting the dataset into
+            train/test set.
+            By default ``None``.
+
         Returns
         -------
         MapieRegressor
@@ -496,7 +514,8 @@ def fit(
          cv,
          X,
          y,
-         sample_weight) = self._check_fit_parameters(X, y, sample_weight)
+         sample_weight,
+         groups) = self._check_fit_parameters(X, y, sample_weight, groups)
 
         self.estimator_ = EnsembleRegressor(
             estimator,
@@ -509,10 +528,12 @@ def fit(
             self.verbose
         )
         # Fit the prediction function
-        self.estimator_ = self.estimator_.fit(X, y, sample_weight)
+        self.estimator_ = self.estimator_.fit(
+            X, y, sample_weight=sample_weight, groups=groups
+        )
 
         # Predict on calibration data
-        y_pred = self.estimator_.predict_calib(X)
+        y_pred = self.estimator_.predict_calib(X, y=y, groups=groups)
 
         # Compute the conformity scores (manage jk-ab case)
         self.conformity_scores_ = \

mapie/subsample.py

@@ -56,7 +56,7 @@ def __init__(
         self.random_state = random_state
 
     def split(
-        self, X: NDArray
+        self, X: NDArray, *args: Any, **kargs: Any
     ) -> Generator[Tuple[NDArray, NDArray], None, None]:
         """
         Generate indices to split data into training and test sets.
@@ -89,7 +89,8 @@ def split(
             test_index = np.setdiff1d(indices, train_index)
             yield train_index, test_index
 
-    def get_n_splits(self, *args: Any, **kargs: Any) -> int:
+    def get_n_splits(
+            self, *args: Any, **kargs: Any) -> int:
         """
         Returns the number of splitting iterations in the cross-validator.
 
@@ -154,7 +155,7 @@ def __init__(
         self.random_state = random_state
 
     def split(
-        self, X: NDArray
+        self, X: NDArray, *args: Any, **kargs: Any
     ) -> Generator[Tuple[NDArray, NDArray], None, None]:
         """
         Generate indices to split data into training and test sets.

mapie/tests/test_regression.py

@@ -361,6 +361,28 @@ def test_results_with_constant_sample_weights(strategy: str) -> None:
     np.testing.assert_allclose(y_pis1, y_pis2)
 
 
+@pytest.mark.parametrize("strategy", [*STRATEGIES])
+def test_results_with_constant_groups(strategy: str) -> None:
+    """
+    Test predictions when groups are None
+    or constant with different values.
+    """
+    n_samples = len(X)
+    mapie0 = MapieRegressor(**STRATEGIES[strategy])
+    mapie1 = MapieRegressor(**STRATEGIES[strategy])
+    mapie2 = MapieRegressor(**STRATEGIES[strategy])
+    mapie0.fit(X, y, groups=None)
+    mapie1.fit(X, y, groups=np.ones(shape=n_samples))
+    mapie2.fit(X, y, groups=np.ones(shape=n_samples) * 5)
+    y_pred0, y_pis0 = mapie0.predict(X, alpha=0.05)
+    y_pred1, y_pis1 = mapie1.predict(X, alpha=0.05)
+    y_pred2, y_pis2 = mapie2.predict(X, alpha=0.05)
+    np.testing.assert_allclose(y_pred0, y_pred1)
+    np.testing.assert_allclose(y_pred1, y_pred2)
+    np.testing.assert_allclose(y_pis0, y_pis1)
+    np.testing.assert_allclose(y_pis1, y_pis2)
+
+
 @pytest.mark.parametrize("strategy", [*STRATEGIES])
 def test_prediction_between_low_up(strategy: str) -> None:
     """Test that prediction lies between low and up prediction intervals."""

mapie/utils.py

@@ -214,6 +214,8 @@ def check_no_agg_cv(
     X: ArrayLike,
     cv: Union[int, str, BaseCrossValidator, BaseShuffleSplit],
     no_agg_cv_array: list,
+    y: Optional[ArrayLike] = None,
+    groups: Optional[ArrayLike] = None
 ) -> bool:
     """
     Check if cross-validator is ``"prefit"``, ``"split"`` or any split
@@ -230,6 +232,17 @@ def check_no_agg_cv(
     no_agg_cv_array: list
         List of all non-aggregated cv methods.
 
+    y: Opional[ArrayLike] of shape (n_samples,)
+        Input labels.
+
+        By default ``None``.
+
+    groups: Optional[ArrayLike] of shape (n_samples,)
+        Group labels for the samples used while splitting the dataset into
+        train/test set.
+
+        By default ``None``.
+
     Returns
     -------
     bool
@@ -240,7 +253,7 @@ def check_no_agg_cv(
     elif isinstance(cv, int):
         return cv == 1
     elif hasattr(cv, "get_n_splits"):
-        return cv.get_n_splits(X) == 1
+        return cv.get_n_splits(X, y, groups) == 1
     else:
         raise ValueError(
             "Invalid cv argument. "
@@ -598,7 +611,7 @@ def check_lower_upper_bounds(
     if any_final_inversion:
         warnings.warn(
             "WARNING: The predictions have issues.\n"
-            + "The upper predictions are lower than"
+            + "The upper predictions are lower than "
             + "the lower predictions at some points."
         )