Add : **predict_params in fit and predict method for Mapie Regression

mapie/estimator/classifier.py

@@ -448,7 +448,7 @@ def predict(
         self,
         X: ArrayLike,
         agg_scores: Optional[str] = None,
-        **predict_params
+        **predict_params,
     ) -> NDArray:
         """
         Predict target from X. It also computes the prediction per train sample

mapie/estimator/interface.py

@@ -37,4 +37,18 @@ def predict(
         """
         Predict target from X. It also computes the prediction per train sample
         for each test sample according to ``self.method``.
+
+        Parameters
+        ----------
+        X: ArrayLike of shape (n_samples, n_features)
+            Test data.
+
+        **kwargs : dict
+            Additional fit and predict parameters.
+
+        Returns
+        -------
+        Tuple[NDArray, NDArray]
+            - Predictions
+            - Predictions sets
         """

mapie/estimator/regressor.py

@@ -233,6 +233,7 @@ def _predict_oof_estimator(
         estimator: RegressorMixin,
         X: ArrayLike,
         val_index: ArrayLike,
+        **predict_params
     ) -> Tuple[NDArray, ArrayLike]:
         """
         Perform predictions on a single out-of-fold model on a validation set.
@@ -248,14 +249,17 @@ def _predict_oof_estimator(
         val_index: ArrayLike of shape (n_samples_val)
             Validation data indices.
 
+        **predict_params : dict
+            Additional predict parameters.
+
         Returns
         -------
         Tuple[NDArray, ArrayLike]
             Predictions of estimator from val_index of X.
         """
         X_val = _safe_indexing(X, val_index)
         if _num_samples(X_val) > 0:
-            y_pred = estimator.predict(X_val)
+            y_pred = estimator.predict(X_val, **predict_params)
         else:
             y_pred = np.array([])
         return y_pred, val_index
@@ -306,7 +310,7 @@ def _aggregate_with_mask(
         else:
             raise ValueError("The value of self.agg_function is not correct")
 
-    def _pred_multi(self, X: ArrayLike) -> NDArray:
+    def _pred_multi(self, X: ArrayLike, **predict_params) -> NDArray:
         """
         Return a prediction per train sample for each test sample, by
         aggregation with matrix ``k_``.
@@ -316,12 +320,15 @@ def _pred_multi(self, X: ArrayLike) -> NDArray:
         X: ArrayLike of shape (n_samples_test, n_features)
             Input data
 
+        **predict_params : dict
+            Additional predict parameters.
+
         Returns
         -------
         NDArray of shape (n_samples_test, n_samples_train)
         """
         y_pred_multi = np.column_stack(
-            [e.predict(X) for e in self.estimators_]
+            [e.predict(X, **predict_params) for e in self.estimators_]
         )
         # At this point, y_pred_multi is of shape
         # (n_samples_test, n_estimators_). The method
@@ -334,7 +341,8 @@ def predict_calib(
         self,
         X: ArrayLike,
         y: Optional[ArrayLike] = None,
-        groups: Optional[ArrayLike] = None
+        groups: Optional[ArrayLike] = None,
+        **predict_params
     ) -> NDArray:
         """
         Perform predictions on X : the calibration set.
@@ -355,6 +363,9 @@ def predict_calib(
 
             By default ``None``.
 
+        **predict_params : dict
+            Additional predict parameters.
+
         Returns
         -------
         NDArray of shape (n_samples_test, 1)
@@ -371,7 +382,7 @@ def predict_calib(
                 cv = cast(BaseCrossValidator, self.cv)
                 outputs = Parallel(n_jobs=self.n_jobs, verbose=self.verbose)(
                     delayed(self._predict_oof_estimator)(
-                        estimator, X, calib_index,
+                        estimator, X, calib_index, **predict_params
                     )
                     for (_, calib_index), estimator in zip(
                         cv.split(X, y, groups),
@@ -404,7 +415,7 @@ def fit(
         y: ArrayLike,
         sample_weight: Optional[ArrayLike] = None,
         groups: Optional[ArrayLike] = None,
-        **fit_params,
+        **fit_params
     ) -> EnsembleRegressor:
         """
         Fit the base estimator under the ``single_estimator_`` attribute.
@@ -526,6 +537,9 @@ def predict(
             predictions (3 arrays). If ``False`` the method return the
             simple predictions only.
 
+        **predict_params : dict
+            Additional predict parameters.
+
         Returns
         -------
         Tuple[NDArray, NDArray, NDArray]
@@ -535,15 +549,15 @@ def predict(
         """
         check_is_fitted(self, self.fit_attributes)
 
-        y_pred = self.single_estimator_.predict(X)
+        y_pred = self.single_estimator_.predict(X, **predict_params)
         if not return_multi_pred and not ensemble:
             return y_pred
 
         if self.method in self.no_agg_methods_ or self.use_split_method_:
             y_pred_multi_low = y_pred[:, np.newaxis]
             y_pred_multi_up = y_pred[:, np.newaxis]
         else:
-            y_pred_multi = self._pred_multi(X)
+            y_pred_multi = self._pred_multi(X, **predict_params)
 
             if self.method == "minmax":
                 y_pred_multi_low = np.min(y_pred_multi, axis=1, keepdims=True)

mapie/regression/quantile_regression.py

@@ -547,7 +547,6 @@ def fit(
              The model itself.
         """
         self.cv = self._check_cv(cast(str, self.cv))
-
         # Initialization
         self.estimators_: List[RegressorMixin] = []
         if self.cv == "prefit":
@@ -649,6 +648,7 @@ def predict(
         optimize_beta: bool = False,
         allow_infinite_bounds: bool = False,
         symmetry: Optional[bool] = True,
+        **predict_params,
     ) -> Union[NDArray, Tuple[NDArray, NDArray]]:
         """
         Predict target on new samples with confidence intervals.
@@ -676,6 +676,9 @@ def predict(
             each residuals separatly or to use the maximum of the two
             combined.
 
+        predict_params : dict
+            Additional predict parameters.
+
         Returns
         -------
         Union[NDArray, Tuple[NDArray, NDArray]]
@@ -699,7 +702,7 @@ def predict(
             dtype=float,
         )
         for i, est in enumerate(self.estimators_):
-            y_preds[i] = est.predict(X)
+            y_preds[i] = est.predict(X, **predict_params)
         check_lower_upper_bounds(y_preds[0], y_preds[1], y_preds[2])
         if symmetry:
             quantile = np.full(

mapie/regression/regression.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import warnings
-from typing import Iterable, Optional, Tuple, Union, cast
+from typing import Any, Iterable, Optional, Tuple, Union, cast
 
 import numpy as np
 from sklearn.base import BaseEstimator, RegressorMixin
@@ -228,6 +228,7 @@ def __init__(
         verbose: int = 0,
         conformity_score: Optional[ConformityScore] = None,
         random_state: Optional[Union[int, np.random.RandomState]] = None,
+        predict_params: Optional[bool] = False
     ) -> None:
         self.estimator = estimator
         self.method = method
@@ -238,6 +239,7 @@ def __init__(
         self.verbose = verbose
         self.conformity_score = conformity_score
         self.random_state = random_state
+        self.predict_params = predict_params
 
     def _check_parameters(self) -> None:
         """
@@ -467,7 +469,7 @@ def fit(
         y: ArrayLike,
         sample_weight: Optional[ArrayLike] = None,
         groups: Optional[ArrayLike] = None,
-        **fit_params,
+        **kwargs: Any
     ) -> MapieRegressor:
         """
         Fit estimator and compute conformity scores used for
@@ -500,14 +502,19 @@ def fit(
             train/test set.
             By default ``None``.
 
-        **fit_params : dict
-            Additional fit parameters.
+        kwargs : dict
+            Additional fit and predict parameters.
 
         Returns
         -------
         MapieRegressor
             The model itself.
         """
+        fit_params = kwargs.pop('fit_params', {})
+        predict_params = kwargs.pop('predict_params', {})
+
+        if len(predict_params) > 0:
+            self.predict_params = True
         # Checks
         (estimator,
          self.conformity_score_function_,
@@ -534,7 +541,8 @@ def fit(
         )
 
         # Predict on calibration data
-        y_pred = self.estimator_.predict_calib(X, y=y, groups=groups)
+        y_pred = self.estimator_.predict_calib(X, y=y, groups=groups,
+                                               **predict_params)
 
         # Compute the conformity scores (manage jk-ab case)
         self.conformity_scores_ = \
@@ -551,6 +559,7 @@ def predict(
         alpha: Optional[Union[float, Iterable[float]]] = None,
         optimize_beta: bool = False,
         allow_infinite_bounds: bool = False,
+        **predict_params
     ) -> Union[NDArray, Tuple[NDArray, NDArray]]:
         """
         Predict target on new samples with confidence intervals.
@@ -600,6 +609,9 @@ def predict(
 
             By default ``False``.
 
+        predict_params : dict
+            Additional predict parameters.
+
         Returns
         -------
         Union[NDArray, Tuple[NDArray, NDArray]]
@@ -609,6 +621,24 @@ def predict(
                 - [:, 0, :]: Lower bound of the prediction interval.
                 - [:, 1, :]: Upper bound of the prediction interval.
         """
+
+        if self.predict_params is True:
+            warnings.warn(
+                f"Be careful that predict_params: '{predict_params}' "
+                "is used in fit method",
+                UserWarning
+            )
+
+        elif (len(predict_params) > 0 and
+              self.predict_params is False and
+              self.cv != "prefit"):
+            raise ValueError(
+                f"Using 'predict_param' '{predict_params}' "
+                f"without having used it in the fit method. "
+                f"Please ensure '{predict_params}' "
+                f"is used in the fit method before calling predict."
+            )
+
         # Checks
         check_is_fitted(self, self.fit_attributes)
         self._check_ensemble(ensemble)
@@ -617,7 +647,7 @@ def predict(
         # If alpha is None, predict the target without confidence intervals
         if alpha is None:
             y_pred = self.estimator_.predict(
-                X, ensemble, return_multi_pred=False
+                X, ensemble, return_multi_pred=False, **predict_params
             )
             return np.array(y_pred)
 

mapie/regression/time_series_regression.py

@@ -405,6 +405,7 @@ def predict(
         alpha: Optional[Union[float, Iterable[float]]] = None,
         optimize_beta: bool = False,
         allow_infinite_bounds: bool = False,
+        **predict_params
     ) -> Union[NDArray, Tuple[NDArray, NDArray]]:
         """
         Predict target on new samples with confidence intervals.
@@ -439,6 +440,9 @@ def predict(
         allow_infinite_bounds: bool
             Allow infinite prediction intervals to be produced.
 
+        predict_params : dict
+            Additional predict parameters.
+
         Returns
         -------
         Union[NDArray, Tuple[NDArray, NDArray]]
@@ -450,15 +454,16 @@ def predict(
         """
         if alpha is None:
             super().predict(
-                X, ensemble=ensemble, alpha=alpha, optimize_beta=optimize_beta
+                X, ensemble=ensemble, alpha=alpha, optimize_beta=optimize_beta,
+                **predict_params
             )
 
         if self.method == "aci":
             alpha = self._get_alpha(alpha)
 
         return super().predict(
             X, ensemble=ensemble, alpha=alpha, optimize_beta=optimize_beta,
-            allow_infinite_bounds=allow_infinite_bounds
+            allow_infinite_bounds=allow_infinite_bounds, **predict_params
         )
 
     def _more_tags(self):