FIX: avoid double inference when predicting in split conformal classification (#721)

* FIX: avoid double inference when predicting in split conformal classification:
 - move probabilities prediction to the main class instead of conformity scores classes
 - use those probabilities to compute y_pred
 - change EnsembleClassifier .predict function to .predict_agg_proba
 - improve check_proba_normalized following copilot suggestion

Author: Valentin-Laurent

mapie/classification.py

@@ -24,10 +24,12 @@
     check_target, check_and_select_conformity_score,
 )
 from mapie.estimator.classifier import EnsembleClassifier
-from mapie.utils import (_check_alpha, _check_alpha_and_n_samples, _check_cv,
-                         _check_estimator_classification, _check_n_features_in,
-                         _check_n_jobs, _check_null_weight, _check_predict_params,
-                         _check_verbose)
+from mapie.utils import (
+    _check_alpha, _check_alpha_and_n_samples, _check_cv,
+    _check_estimator_classification, _check_n_features_in,
+    _check_n_jobs, _check_null_weight, _check_predict_params,
+    _check_verbose, check_proba_normalized,
+)
 from mapie.utils import (
     _transform_confidence_level_to_alpha_list,
     _raise_error_if_fit_called_in_prefit_mode,
@@ -1055,7 +1057,12 @@ def predict(
         alpha = cast(Optional[NDArray], _check_alpha(alpha))
 
         # Estimate predictions
-        y_pred = self.estimator_.single_estimator_.predict(X, **predict_params)
+        y_pred_proba = self.estimator_.single_estimator_.predict_proba(
+            X,
+            **predict_params
+        )
+        y_pred_proba = check_proba_normalized(y_pred_proba, axis=1)
+        y_pred = self.label_encoder_.inverse_transform(np.argmax(y_pred_proba, axis=1))
         if alpha is None:
             return y_pred
 
@@ -1067,9 +1074,17 @@ def predict(
         _check_alpha_and_n_samples(alpha_np, n)
 
         # Estimate prediction sets
+        if self.estimator_.cv != "prefit":
+            y_pred_proba = self.estimator_.predict_agg_proba(
+                X,
+                agg_scores,
+                **predict_params
+            )
+
         prediction_sets = self.conformity_score_function_.predict_set(
             X, alpha_np,
-            estimator=self.estimator_,
+            y_pred_proba=y_pred_proba,
+            cv=self.estimator_.cv,
             conformity_scores=self.conformity_scores_,
             include_last_label=include_last_label,
             agg_scores=agg_scores,

mapie/conformity_scores/classification.py

@@ -4,7 +4,7 @@
 import numpy as np
 
 from mapie.conformity_scores.interface import BaseConformityScore
-from mapie.estimator.classifier import EnsembleClassifier
+from sklearn.model_selection import BaseCrossValidator
 
 from numpy.typing import ArrayLike, NDArray
 
@@ -53,7 +53,8 @@ def get_predictions(
         self,
         X: NDArray,
         alpha_np: NDArray,
-        estimator: EnsembleClassifier,
+        y_pred_proba: NDArray,
+        cv: Optional[Union[int, str, BaseCrossValidator]],
         **kwargs
     ) -> NDArray:
         """
@@ -84,7 +85,7 @@ def get_conformity_score_quantiles(
         self,
         conformity_scores: NDArray,
         alpha_np: NDArray,
-        estimator: EnsembleClassifier,
+        cv: Optional[Union[int, str, BaseCrossValidator]],
         **kwargs
     ) -> NDArray:
         """
@@ -116,7 +117,7 @@ def get_prediction_sets(
         y_pred_proba: NDArray,
         conformity_scores: NDArray,
         alpha_np: NDArray,
-        estimator: EnsembleClassifier,
+        cv: Optional[Union[int, str, BaseCrossValidator]],
         **kwargs
     ) -> NDArray:
         """
@@ -150,13 +151,14 @@ def get_sets(
         self,
         X: NDArray,
         alpha_np: NDArray,
-        estimator: EnsembleClassifier,
+        y_pred_proba: NDArray,
+        cv: Optional[Union[int, str, BaseCrossValidator]],
         conformity_scores: NDArray,
         **kwargs
     ) -> NDArray:
         """
         Compute classes of the prediction sets from the observed values,
-        the estimator of type ``EnsembleClassifier`` and the conformity scores.
+        the predicted probabilities and the conformity scores.
 
         Parameters
         ----------
@@ -167,8 +169,11 @@ def get_sets(
             NDArray of floats between 0 and 1, representing the uncertainty
             of the confidence set.
 
-        estimator: EnsembleClassifier
-            Estimator that is fitted to predict y from X.
+        y_pred_proba: NDArray
+            Predicted probabilities from the estimator.
+
+        cv: Optional[Union[int, str, BaseCrossValidator]]
+            Cross-validation strategy used by the estimator.
 
         conformity_scores: NDArray of shape (n_samples,)
             Conformity scores.
@@ -178,19 +183,19 @@ def get_sets(
         NDArray of shape (n_samples, n_classes, n_alpha)
             Prediction sets (Booleans indicate whether classes are included).
         """
+        # Choice of the quantile
         # Predict probabilities
         y_pred_proba = self.get_predictions(
-            X, alpha_np, estimator, **kwargs
+            X, alpha_np, y_pred_proba, cv, **kwargs
         )
 
-        # Choice of the quantile
         self.quantiles_ = self.get_conformity_score_quantiles(
-            conformity_scores, alpha_np, estimator, **kwargs
+            conformity_scores, alpha_np, cv, **kwargs
         )
 
         # Build prediction sets
         prediction_sets = self.get_prediction_sets(
-            y_pred_proba, conformity_scores, alpha_np, estimator, **kwargs
+            y_pred_proba, conformity_scores, alpha_np, cv, **kwargs
         )
 
         return prediction_sets
@@ -213,6 +218,12 @@ def predict_set(
         alpha_np: NDArray of shape (n_alpha, )
             Represents the uncertainty of the confidence set to produce.
 
+        y_pred_proba: NDArray
+            Predicted probabilities from the estimator.
+
+        cv: Optional[Union[int, str, BaseCrossValidator]]
+            Cross-validation strategy used by the estimator.
+
         **kwargs: dict
             Additional keyword arguments.
 
@@ -221,4 +232,6 @@ def predict_set(
         The output structure depend on the ``get_sets`` method.
             The prediction sets for each sample and each alpha level.
         """
-        return self.get_sets(X=X, alpha_np=alpha_np, **kwargs)
+        return self.get_sets(
+            X=X, alpha_np=alpha_np, **kwargs
+        )

mapie/conformity_scores/sets/aps.py

@@ -3,12 +3,12 @@
 import numpy as np
 from sklearn.utils import check_random_state
 from sklearn.preprocessing import label_binarize
+from sklearn.model_selection import BaseCrossValidator
 
 from mapie.conformity_scores.sets.naive import NaiveConformityScore
 from mapie.conformity_scores.sets.utils import (
-    check_include_last_label, check_proba_normalized
+    check_include_last_label
 )
-from mapie.estimator.classifier import EnsembleClassifier
 
 from mapie._machine_precision import EPSILON
 from numpy.typing import ArrayLike, NDArray
@@ -46,24 +46,28 @@ def get_predictions(
         self,
         X: NDArray,
         alpha_np: NDArray,
-        estimator: EnsembleClassifier,
+        y_pred_proba: NDArray,
+        cv: Optional[Union[int, str, BaseCrossValidator]],
         agg_scores: Optional[str] = "mean",
         **kwargs
     ) -> NDArray:
         """
-        Get predictions from an EnsembleClassifier.
+        Just processes the passed y_pred_proba.
 
         Parameters
         -----------
         X: NDArray of shape (n_samples, n_features)
-            Observed feature values.
+            Observed feature values (not used since predictions are passed).
 
         alpha_np: NDArray of shape (n_alpha,)
             NDArray of floats between ``0`` and ``1``, represents the
             uncertainty of the confidence interval.
 
-        estimator: EnsembleClassifier
-            Estimator that is fitted to predict y from X.
+        y_pred_proba: NDArray
+            Predicted probabilities from the estimator.
+
+        cv: Optional[Union[int, str, BaseCrossValidator]]
+            Cross-validation strategy used by the estimator.
 
         agg_scores: Optional[str]
             Method to aggregate the scores from the base estimators.
@@ -77,8 +81,6 @@ def get_predictions(
         NDArray
             Array of predictions.
         """
-        y_pred_proba = estimator.predict(X, agg_scores)
-        y_pred_proba = check_proba_normalized(y_pred_proba, axis=1)
         if agg_scores != "crossval":
             y_pred_proba = np.repeat(
                 y_pred_proba[:, :, np.newaxis], len(alpha_np), axis=2
@@ -171,7 +173,7 @@ def get_conformity_score_quantiles(
         self,
         conformity_scores: NDArray,
         alpha_np: NDArray,
-        estimator: EnsembleClassifier,
+        cv: Optional[Union[int, str, BaseCrossValidator]],
         agg_scores: Optional[str] = "mean",
         **kwargs
     ) -> NDArray:
@@ -187,8 +189,8 @@ def get_conformity_score_quantiles(
             NDArray of floats between 0 and 1, representing the uncertainty
             of the confidence interval.
 
-        estimator: EnsembleClassifier
-            Estimator that is fitted to predict y from X.
+        cv: Optional[Union[int, str, BaseCrossValidator]]
+            Cross-validation strategy used by the estimator.
 
         agg_scores: Optional[str]
             Method to aggregate the scores from the base estimators.
@@ -204,7 +206,7 @@ def get_conformity_score_quantiles(
         """
         n = len(conformity_scores)
 
-        if estimator.cv == "prefit" or agg_scores in ["mean"]:
+        if cv == "prefit" or agg_scores in ["mean"]:
             quantiles_ = _compute_quantiles(conformity_scores, alpha_np)
         else:
             quantiles_ = (n + 1) * (1 - alpha_np)
@@ -328,7 +330,7 @@ def get_prediction_sets(
         y_pred_proba: NDArray,
         conformity_scores: NDArray,
         alpha_np: NDArray,
-        estimator: EnsembleClassifier,
+        cv: Optional[Union[int, str, BaseCrossValidator]],
         agg_scores: Optional[str] = "mean",
         include_last_label: Optional[Union[bool, str]] = True,
         **kwargs
@@ -349,8 +351,8 @@ def get_prediction_sets(
             NDArray of floats between 0 and 1, representing the uncertainty
             of the confidence interval.
 
-        estimator: EnsembleClassifier
-            Estimator that is fitted to predict y from X.
+        cv: Optional[Union[int, str, BaseCrossValidator]]
+            Cross-validation strategy used by the estimator.
 
         agg_scores: Optional[str]
             Method to aggregate the scores from the base estimators.
@@ -398,7 +400,7 @@ def get_prediction_sets(
         include_last_label = check_include_last_label(include_last_label)
 
         # specify which thresholds will be used
-        if estimator.cv == "prefit" or agg_scores in ["mean"]:
+        if cv == "prefit" or agg_scores in ["mean"]:
             thresholds = self.quantiles_
         else:
             thresholds = conformity_scores.ravel()
@@ -414,7 +416,7 @@ def get_prediction_sets(
             )
         )
         # get the prediction set by taking all probabilities above the last one
-        if estimator.cv == "prefit" or agg_scores in ["mean"]:
+        if cv == "prefit" or agg_scores in ["mean"]:
             y_pred_included = np.greater_equal(
                 y_pred_proba - y_pred_proba_last, -EPSILON
             )
@@ -432,7 +434,7 @@ def get_prediction_sets(
                 thresholds,
                 **kwargs
             )
-        if estimator.cv == "prefit" or agg_scores in ["mean"]:
+        if cv == "prefit" or agg_scores in ["mean"]:
             prediction_sets = y_pred_included
         else:
             # compute the number of times the inequality is verified