FIX: clean code to lint and type-check errors

Author: thibaultcordier

mapie/classification.py

@@ -13,15 +13,15 @@
 from sklearn.utils.validation import (_check_y, _num_samples, check_is_fitted,
                                       indexable)
 
-from ._machine_precision import EPSILON
-from ._typing import ArrayLike, NDArray
-from .estimator.classification.estimator import EnsembleClassifier
-from .metrics import classification_mean_width_score
-from .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_verbose,
-                    compute_quantiles)
-from .conformity_scores.utils_classification_conformity_scores import (
+from mapie._machine_precision import EPSILON
+from mapie._typing import ArrayLike, NDArray
+from mapie.estimator import EnsembleClassifier
+from mapie.metrics import classification_mean_width_score
+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_verbose,
+                         compute_quantiles)
+from mapie.conformity_scores.utils import (
     get_true_label_position
 )
 
@@ -988,9 +988,9 @@ def _split_data(
         sample_weight,
         groups,
         size_raps
-    ) -> Tuple[ArrayLike, ArrayLike, ArrayLike, ArrayLike, NDArray, ArrayLike]:
-
+    ):
         """Split data for raps method
+
         Parameters
         ----------
         X: ArrayLike
@@ -1013,15 +1013,15 @@ def _split_data(
 
         Returns
         -------
-        Tuple[ArrayLike, ArrayLike, ArrayLike, NDArray, Optional[NDArray],
-        Optional[ArrayLike]]
+        Tuple[NDArray, NDArray, NDArray, NDArray, Optional[NDArray],
+        Optional[NDArray]]
 
-            - ArrayLike of shape (n_samples, n_features)
-            - ArrayLike of shape (n_samples,)
-            - ArrayLike of shape (n_samples,)
-            - ArrayLike of shape (n_samples,)
+            - NDArray of shape (n_samples, n_features)
+            - NDArray of shape (n_samples,)
+            - NDArray of shape (n_samples,)
+            - NDArray of shape (n_samples,)
+            - NDArray of shape (n_samples,)
             - NDArray of shape (n_samples,)
-            - ArrayLike of shape (n_samples,)
         """
         raps_split = ShuffleSplit(
             1, test_size=size_raps, random_state=self.random_state
@@ -1096,15 +1096,25 @@ def fit(
             The model itself.
         """
         # Checks
-        (estimator, cv, X, y, y_enc, sample_weight, groups, n_samples) = (
-            self._check_fit_parameter(X, y, sample_weight, groups)
-        )
+        (estimator,
+         cv,
+         X,
+         y,
+         y_enc,
+         sample_weight,
+         groups,
+         n_samples) = self._check_fit_parameter(X, y, sample_weight, groups)
 
         if self.method == "raps":
             (X, y_enc, y, n_samples, sample_weight, groups) = self._split_data(
                 X, y_enc, sample_weight, groups, size_raps
             )
 
+        # Cast
+        X, y_enc, y = cast(NDArray, X), cast(NDArray, y_enc), cast(NDArray, y)
+        sample_weight = cast(NDArray, sample_weight)
+        groups = cast(NDArray, groups)
+
         # Work
         self.estimator_ = EnsembleClassifier(
             estimator,
@@ -1117,7 +1127,7 @@ def fit(
         )
 
         self.estimator_ = self.estimator_.fit(
-            X, y, y_enc, sample_weight, groups,
+            X, y, y_enc=y_enc, sample_weight=sample_weight, groups=groups,
             **fit_params
         )
 

mapie/conformity_scores/conformity_scores.py

@@ -5,7 +5,7 @@
 
 from mapie._compatibility import np_nanquantile
 from mapie._typing import ArrayLike, NDArray
-from mapie.estimator.regression.interface import EnsembleEstimator
+from mapie.estimator import EnsembleRegressor
 
 
 class ConformityScore(metaclass=ABCMeta):
@@ -317,7 +317,7 @@ def _beta_optimize(
     def get_bounds(
         self,
         X: ArrayLike,
-        estimator: EnsembleEstimator,
+        estimator: EnsembleRegressor,
         conformity_scores: NDArray,
         alpha_np: NDArray,
         ensemble: bool = False,

mapie/conformity_scores/utils_classification_conformity_scores.py renamed to mapie/conformity_scores/utils.py

@@ -0,0 +1,9 @@
+from .interface import EnsembleEstimator
+from .regressor import EnsembleRegressor
+from .classifier import EnsembleClassifier
+
+__all__ = [
+    "EnsembleEstimator",
+    "EnsembleRegressor",
+    "EnsembleClassifier",
+]

mapie/estimator/__init__.py

@@ -10,7 +10,7 @@
 from sklearn.utils.validation import _num_samples, check_is_fitted
 
 from mapie._typing import ArrayLike, NDArray
-from mapie.estimator.classification.interface import EnsembleEstimator
+from mapie.estimator.interface import EnsembleEstimator
 from mapie.utils import (
     check_no_agg_cv,
     fit_estimator,
@@ -294,7 +294,7 @@ def fit(
         self,
         X: ArrayLike,
         y: ArrayLike,
-        y_enc: ArrayLike,
+        y_enc: Optional[ArrayLike] = None,
         sample_weight: Optional[ArrayLike] = None,
         groups: Optional[ArrayLike] = None,
         **fit_params,
@@ -341,6 +341,9 @@ def fit(
         self.use_split_method_ = check_no_agg_cv(X, self.cv, self.no_agg_cv_)
         estimator = self.estimator
         n_samples = _num_samples(y)
+        if y_enc is None:
+            raise ValueError
+        y_enc = cast(NDArray, y_enc)
 
         # Computation
         if cv == "prefit":
@@ -378,25 +381,26 @@ def fit(
 
     def predict_proba_calib(
         self,
-        X: ArrayLike,
-        y: ArrayLike,
-        y_enc: ArrayLike,
-        groups: Optional[ArrayLike] = None,
-    ) -> Tuple[NDArray, ArrayLike, ArrayLike]:
+        X: NDArray,
+        y: NDArray,
+        y_enc: NDArray,
+        groups: Optional[NDArray] = None,
+        **predict_params
+    ) -> Tuple[NDArray, NDArray, NDArray]:
         """
         Perform predictions on X : the calibration set.
 
         Parameters
         ----------
-        X: ArrayLike of shape (n_samples_test, n_features)
+        X: NDArray of shape (n_samples_test, n_features)
             Input data
 
-        y: Optional[ArrayLike] of shape (n_samples_test,)
+        y: Optional[NDArray] of shape (n_samples_test,)
             Input labels.
 
             By default ``None``.
 
-        groups: Optional[ArrayLike] of shape (n_samples_test,)
+        groups: Optional[NDArray] of shape (n_samples_test,)
             Group labels for the samples used while splitting the dataset into
             train/test set.
 
@@ -439,7 +443,7 @@ def predict_proba_calib(
                 # are not used during calibration
                 self.k_ = self.k_[val_indices]
                 y_pred_proba = y_pred_proba[val_indices]
-                # y_enc = y_enc[val_indices]
+                y_enc = y_enc[val_indices]
                 y = cast(NDArray, y)[val_indices]
 
         return y_pred_proba, y, y_enc
@@ -448,7 +452,8 @@ def predict(
         self,
         X: ArrayLike,
         alpha_np: ArrayLike = [],
-        agg_scores: Any = None
+        agg_scores: Any = None,
+        **predict_params
     ) -> NDArray:
         """
         Predict target from X. It also computes the prediction per train sample

mapie/estimator/classification/__init__.py

@@ -0,0 +1,39 @@
+from __future__ import annotations
+
+from abc import ABCMeta, abstractmethod
+
+from mapie._typing import ArrayLike
+
+
+class EnsembleEstimator(metaclass=ABCMeta):
+    """
+    This class implements methods to handle the training and usage of the
+    estimator. This estimator can be unique or composed by cross validated
+    estimators.
+    """
+
+    @abstractmethod
+    def fit(
+        self,
+        X: ArrayLike,
+        y: ArrayLike,
+        **kwargs
+    ) -> EnsembleEstimator:
+        """
+        Fit the base estimator under the ``single_estimator_`` attribute.
+        Fit all cross-validated estimator clones
+        and rearrange them into a list, the ``estimators_`` attribute.
+        Out-of-fold conformity scores are stored under
+        the ``conformity_scores_`` attribute.
+        """
+
+    @abstractmethod
+    def predict(
+        self,
+        X: ArrayLike,
+        **kwargs
+    ):
+        """
+        Predict target from X. It also computes the prediction per train sample
+        for each test sample according to ``self.method``.
+        """