ENH: implement train_conformalize_test_split (#629)

* Implement train_conformalize_test_split

* Add unit tests for train_conformalize_test_split

* Rename mapie_v1._utils into mapie_v1.utils.

* Use train_conformalize_test_split in regression documentation examples

---------

Co-authored-by: Valentin Laurent <[email protected]>

Author: 2 people

examples/regression/1-quickstart/plot_prefit.py

@@ -24,12 +24,12 @@
 import scipy
 from lightgbm import LGBMRegressor
 from matplotlib import pyplot as plt
-from sklearn.model_selection import train_test_split
 from sklearn.neural_network import MLPRegressor
 
 from mapie._typing import NDArray
 from mapie.metrics import regression_coverage_score
 from mapie_v1.regression import SplitConformalRegressor, ConformalizedQuantileRegressor
+from mapie_v1.utils import train_conformalize_test_split
 
 warnings.filterwarnings("ignore")
 
@@ -56,13 +56,13 @@ def f(x: NDArray) -> NDArray:
 X = np.linspace(0, 1, n_samples)
 y = f(X) + rng.normal(0, sigma, n_samples)
 
-# Train/validation/test split
-X_train_conformalize, X_test, y_train_conformalize, y_test = train_test_split(
-    X, y, test_size=1 / 10, random_state=RANDOM_STATE
-)
-X_train, X_conformalize, y_train, y_conformalize = train_test_split(
-    X_train_conformalize, y_train_conformalize,
-    test_size=1 / 9, random_state=RANDOM_STATE
+# Train/conformalize/test split
+(
+    X_train, X_conformalize, X_test, y_train, y_conformalize, y_test
+) = train_conformalize_test_split(
+    X, y,
+    train_size=0.8, conformalize_size=0.1, test_size=0.1,
+    random_state=RANDOM_STATE
 )
 
 

examples/regression/1-quickstart/plot_toy_model.py

@@ -8,20 +8,21 @@
 import numpy as np
 from matplotlib import pyplot as plt
 from sklearn.datasets import make_regression
-from sklearn.model_selection import train_test_split
 
 from mapie.metrics import regression_coverage_score
 from mapie_v1.regression import SplitConformalRegressor
+from mapie_v1.utils import train_conformalize_test_split
 
 RANDOM_STATE = 42
 
 X, y = make_regression(n_samples=500, n_features=1, noise=20, random_state=RANDOM_STATE)
 
-X_train, X_test_conformalize, y_train, y_test_conformalize = train_test_split(
-    X, y, test_size=0.4,  random_state=RANDOM_STATE
-)
-X_test, X_conformalize, y_test, y_conformalize = train_test_split(
-    X_test_conformalize, y_test_conformalize, test_size=0.5, random_state=RANDOM_STATE
+(
+    X_train, X_conformalize, X_test, y_train, y_conformalize, y_test
+) = train_conformalize_test_split(
+    X, y,
+    train_size=0.6, conformalize_size=0.2, test_size=0.2,
+    random_state=RANDOM_STATE
 )
 
 confidence_level = [0.95, 0.68]

examples/regression/2-advanced-analysis/plot_ResidualNormalisedScore_tutorial.py

@@ -23,6 +23,7 @@
 from mapie.conformity_scores import ResidualNormalisedScore
 from mapie.metrics import regression_coverage_score_v2, regression_ssc_score
 from mapie_v1.regression import SplitConformalRegressor
+from mapie_v1.utils import train_conformalize_test_split
 
 warnings.filterwarnings("ignore")
 
@@ -64,17 +65,15 @@
 
 np.array(X)
 np.array(y)
-X_train_conformalize, X_test, y_train_conformalize, y_test = train_test_split(
-    X,
-    y,
-    random_state=RANDOM_STATE,
-    test_size=0.02
-)
-X_train, X_conformalize, y_train, y_conformalize = train_test_split(
-    X_train_conformalize,
-    y_train_conformalize,
+
+(
+    X_train, X_conformalize, X_test, y_train, y_conformalize, y_test
+) = train_conformalize_test_split(
+    X, y,
+    train_size=0.7, conformalize_size=0.28, test_size=0.02,
     random_state=RANDOM_STATE
 )
+
 X_conformalize_prefit, X_res, y_conformalize_prefit, y_res = train_test_split(
     X_conformalize,
     y_conformalize,

examples/regression/2-advanced-analysis/plot_conformal_predictive_distribution.py

@@ -23,11 +23,11 @@
 from matplotlib import pyplot as plt
 from sklearn.datasets import make_regression
 from sklearn.linear_model import LinearRegression
-from sklearn.model_selection import train_test_split
 
 from mapie.conformity_scores import (AbsoluteConformityScore,
                                      ResidualNormalisedScore)
 from mapie_v1.regression import SplitConformalRegressor
+from mapie_v1.utils import train_conformalize_test_split
 
 warnings.filterwarnings('ignore')
 
@@ -43,12 +43,13 @@
 X, y = make_regression(
     n_samples=1000, n_features=1, noise=20, random_state=RANDOM_STATE
 )
-X_train_conformalize, X_test, y_train_conformalize, y_test = train_test_split(
-    X, y, test_size=0.3, random_state=RANDOM_STATE
-)
 
-X_train, X_conformalize, y_train, y_conformalize = train_test_split(
-    X_train_conformalize, y_train_conformalize, test_size=0.3, random_state=RANDOM_STATE
+(
+    X_train, X_conformalize, X_test, y_train, y_conformalize, y_test
+) = train_conformalize_test_split(
+    X, y,
+    train_size=0.6, conformalize_size=0.2, test_size=0.2,
+    random_state=RANDOM_STATE
 )
 
 

examples/regression/2-advanced-analysis/plot_cqr_symmetry_difference.py

@@ -11,10 +11,10 @@
 from matplotlib import pyplot as plt
 from sklearn.datasets import make_regression
 from sklearn.ensemble import GradientBoostingRegressor
-from sklearn.model_selection import train_test_split
 
 from mapie.metrics import regression_coverage_score
 from mapie_v1.regression import ConformalizedQuantileRegressor
+from mapie_v1.utils import train_conformalize_test_split
 
 RANDOM_STATE = 1
 
@@ -25,13 +25,14 @@
     n_samples=1000, n_features=1, noise=20, random_state=RANDOM_STATE
 )
 
-X_train_conformalize, X_test, y_train_conformalize, y_test = train_test_split(
-    X, y, test_size=0.3, random_state=RANDOM_STATE
+(
+    X_train, X_conformalize, X_test, y_train, y_conformalize, y_test
+) = train_conformalize_test_split(
+    X, y,
+    train_size=0.6, conformalize_size=0.2, test_size=0.2,
+    random_state=RANDOM_STATE
 )
 
-X_train, X_conformalize, y_train, y_conformalize = train_test_split(
-    X_train_conformalize, y_train_conformalize, test_size=0.3, random_state=RANDOM_STATE
-)
 
 # Define confidence level
 confidence_level = 0.8

mapie_v1/_utils.py

@@ -11,7 +11,7 @@
 from mapie._typing import ArrayLike, NDArray
 from mapie.classification import MapieClassifier
 from mapie.conformity_scores import BaseClassificationScore
-from mapie_v1._utils import (
+from mapie_v1.utils import (
     transform_confidence_level_to_alpha_list,
     prepare_params,
     cast_predictions_to_ndarray_tuple,

mapie_v1/classification.py

@@ -13,7 +13,7 @@
 from mapie.regression import MapieRegressor, MapieQuantileRegressor
 from mapie.utils import check_estimator_fit_predict
 from mapie_v1.conformity_scores._utils import check_and_select_conformity_score
-from mapie_v1._utils import (
+from mapie_v1.utils import (
     transform_confidence_level_to_alpha_list,
     check_if_param_in_allowed_values,
     check_cv_not_string,