ENH: V1 cqr implmentation (#576)

ENH: V1 cqr implmentation

Author: Valentin-Laurent

mapie_v1/_utils.py

@@ -54,9 +54,11 @@ def check_if_X_y_different_from_fit(
 
 def make_intervals_single_if_single_alpha(
     intervals: NDArray,
-    alphas: List[float]
+    alphas: Union[float, List[float]]
 ) -> NDArray:
-    if len(alphas) == 1:
+    if isinstance(alphas, float):
+        return intervals[:, :, 0]
+    if isinstance(alphas, list) and len(alphas) == 1:
         return intervals[:, :, 0]
     return intervals
 

mapie_v1/integration_tests/tests/test_regression.py

@@ -10,6 +10,7 @@
 from sklearn.ensemble import RandomForestRegressor
 from sklearn.linear_model import QuantileRegressor
 from sklearn.ensemble import GradientBoostingRegressor
+from sklearn.model_selection import train_test_split
 
 from mapie.subsample import Subsample
 from mapie._typing import ArrayLike
@@ -306,22 +307,29 @@ def test_intervals_and_predictions_exact_equality_jackknife(params_jackknife):
     )
     gbr_models.append(estimator_)
 
+sample_weight_train = train_test_split(
+    X,
+    y,
+    sample_weight,
+    test_size=0.4,
+    random_state=RANDOM_STATE
+)[-2]
 
 params_test_cases_quantile = [
     {
         "v0": {
             "alpha": 0.2,
             "cv": "split",
             "method": "quantile",
-            "calib_size": 0.3,
+            "calib_size": 0.4,
             "sample_weight": sample_weight,
             "random_state": RANDOM_STATE,
         },
         "v1": {
             "confidence_level": 0.8,
             "prefit": False,
-            "test_size": 0.3,
-            "fit_params": {"sample_weight": sample_weight},
+            "test_size": 0.4,
+            "fit_params": {"sample_weight": sample_weight_train},
             "random_state": RANDOM_STATE,
         },
     },
@@ -330,15 +338,15 @@ def test_intervals_and_predictions_exact_equality_jackknife(params_jackknife):
             "estimator": gbr_models,
             "cv": "prefit",
             "method": "quantile",
-            "calib_size": 0.3,
+            "calib_size": 0.2,
             "sample_weight": sample_weight,
             "optimize_beta": True,
             "random_state": RANDOM_STATE,
         },
         "v1": {
             "estimator": gbr_models,
             "prefit": True,
-            "test_size": 0.3,
+            "test_size": 0.2,
             "fit_params": {"sample_weight": sample_weight},
             "minimize_interval_width": True,
             "random_state": RANDOM_STATE,
@@ -418,12 +426,16 @@ def compare_model_predictions_and_intervals(
     v1_params: Dict = {},
     prefit: bool = False,
     test_size: Optional[float] = None,
+    sample_weight: Optional[ArrayLike] = None,
     random_state: int = 42,
 ) -> None:
 
     if test_size is not None:
         X_train, X_conf, y_train, y_conf = train_test_split_shuffle(
-            X, y, test_size=test_size, random_state=random_state
+            X,
+            y,
+            test_size=test_size,
+            random_state=random_state,
         )
     else:
         X_train, X_conf, y_train, y_conf = X, X, y, y

mapie_v1/regression.py

@@ -3,14 +3,15 @@
 from typing_extensions import Self
 
 import numpy as np
-from sklearn.linear_model import LinearRegression, QuantileRegressor
+from sklearn.linear_model import LinearRegression
 from sklearn.base import RegressorMixin, clone
 from sklearn.model_selection import BaseCrossValidator
+from sklearn.pipeline import Pipeline
 
 from mapie.subsample import Subsample
 from mapie._typing import ArrayLike, NDArray
 from mapie.conformity_scores import BaseRegressionScore
-from mapie.regression import MapieRegressor
+from mapie.regression import MapieRegressor, MapieQuantileRegressor
 from mapie.utils import check_estimator_fit_predict
 from mapie_v1.conformity_scores._utils import (
     check_and_select_regression_conformity_score,
@@ -904,12 +905,29 @@ class ConformalizedQuantileRegressor:
 
     def __init__(
         self,
-        estimator: RegressorMixin = QuantileRegressor(),
-        confidence_level: Union[float, List[float]] = 0.9,
-        conformity_score: Union[str, BaseRegressionScore] = "absolute",
-        random_state: Optional[Union[int, np.random.RandomState]] = None,
+        estimator: Optional[
+            Union[
+                RegressorMixin,
+                Pipeline,
+                List[Union[RegressorMixin, Pipeline]]
+            ]
+        ] = None,
+        confidence_level: float = 0.9,
+        prefit: bool = False,
     ) -> None:
-        pass
+
+        self._alpha = 1 - confidence_level
+        self.prefit = prefit
+
+        cv: str = "prefit" if prefit else "split"
+        self._mapie_quantile_regressor = MapieQuantileRegressor(
+            estimator=estimator,
+            method="quantile",
+            cv=cv,
+            alpha=self._alpha,
+        )
+
+        self._sample_weight: Optional[NDArray] = None
 
     def fit(
         self,
@@ -937,6 +955,27 @@ def fit(
         Self
             The fitted ConformalizedQuantileRegressor instance.
         """
+
+        if self.prefit:
+            raise ValueError(
+                "The estimators are already fitted, the .fit() method should"
+                " not be called with prefit=True."
+            )
+
+        if fit_params:
+            fit_params_ = copy.deepcopy(fit_params)
+            self._sample_weight = fit_params_.pop("sample_weight", None)
+        else:
+            fit_params_ = {}
+
+        self._mapie_quantile_regressor._initialize_fit_conformalize()
+        self._mapie_quantile_regressor._fit_estimators(
+                                            X=X_train,
+                                            y=y_train,
+                                            sample_weight=self._sample_weight,
+                                            **fit_params_,
+        )
+
         return self
 
     def conformalize(
@@ -969,6 +1008,14 @@ def conformalize(
             The ConformalizedQuantileRegressor instance with calibrated
             prediction intervals.
         """
+        self.predict_params = predict_params if predict_params else {}
+
+        self._mapie_quantile_regressor.conformalize(
+            X_conf,
+            y_conf,
+            **self.predict_params
+        )
+
         return self
 
     def predict_set(
@@ -1007,7 +1054,18 @@ def predict_set(
             Prediction intervals with shape `(n_samples, 2)`, with lower
             and upper bounds for each sample.
         """
-        return np.ndarray(0)
+        _, intervals = self._mapie_quantile_regressor.predict(
+            X,
+            optimize_beta=minimize_interval_width,
+            allow_infinite_bounds=allow_infinite_bounds,
+            symmetry=symmetric_intervals,
+            **self.predict_params
+        )
+
+        return make_intervals_single_if_single_alpha(
+            intervals,
+            self._alpha
+        )
 
     def predict(
         self,
@@ -1026,7 +1084,9 @@ def predict(
         NDArray
             Array of point predictions with shape `(n_samples,)`.
         """
-        return np.ndarray(0)
+        estimator = self._mapie_quantile_regressor
+        predictions, _ = estimator.predict(X, **self.predict_params)
+        return predictions
 
 
 class GibbsConformalRegressor: