Merge pull request #53 from simai-ml/small_changes

small changes

Author: vtaquet

doc/quick_start.rst

@@ -66,18 +66,30 @@ The estimated prediction intervals can then be plotted as follows.
     
     from matplotlib import pyplot as plt
     from mapie.metrics import coverage_score
+
+    coverage_scores = [
+        coverage_score(y, y_preds[:, 1, i], y_preds[:, 2, i])
+        for i, _ in enumerate(alpha)
+    ]
+
     plt.xlabel("x")
     plt.ylabel("y")
     plt.scatter(X, y, alpha=0.3)
     plt.plot(X, y_preds[:, 0, 0], color="C1")
     order = np.argsort(X[:, 0])
     plt.plot(X[order], y_preds[order][:, 1, 1], color="C1", ls="--")
     plt.plot(X[order], y_preds[order][:, 2, 1], color="C1", ls="--")
-    plt.fill_between(X[order].ravel(), y_preds[:, 1, 0][order].ravel(), y_preds[:, 2, 0][order].ravel(), alpha=0.2)
-    coverage_scores = [coverage_score(y, y_preds[:, 1, i], y_preds[:, 2, i]) for i, _ in enumerate(alpha)]
+    plt.fill_between(
+        X[order].ravel(),
+        y_preds[:, 1, 0][order].ravel(),
+        y_preds[:, 2, 0][order].ravel(),
+        alpha=0.2
+    )
     plt.title(
-        f"Target and effective coverages for alpha={alpha[0]:.2f}: ({1-alpha[0]:.3f}, {coverage_scores[0]:.3f})\n" +
-        f"Target and effective coverages for alpha={alpha[1]:.2f}: ({1-alpha[1]:.3f}, {coverage_scores[1]:.3f})"
+        f"Target and effective coverages for "
+        f"alpha={alpha[0]:.2f}: ({1-alpha[0]:.3f}, {coverage_scores[0]:.3f})\n"
+        f"Target and effective coverages for "
+        f"alpha={alpha[1]:.2f}: ({1-alpha[1]:.3f}, {coverage_scores[1]:.3f})"
     )
     plt.show()
 

examples/plot_toy_model.py

@@ -22,6 +22,11 @@
 mapie.fit(X, y)
 y_preds = mapie.predict(X)
 
+coverage_scores = [
+    coverage_score(y, y_preds[:, 1, i], y_preds[:, 2, i])
+    for i, _ in enumerate(alpha)
+]
+
 plt.xlabel("x")
 plt.ylabel("y")
 plt.scatter(X, y, alpha=0.3)
@@ -35,10 +40,6 @@
     y_preds[:, 2, 0][order].ravel(),
     alpha=0.2
 )
-coverage_scores = [
-    coverage_score(y, y_preds[:, 1, i], y_preds[:, 2, i])
-    for i, _ in enumerate(alpha)
-]
 plt.title(
     f"Target and effective coverages for "
     f"alpha={alpha[0]:.2f}: ({1-alpha[0]:.3f}, {coverage_scores[0]:.3f})\n"

mapie/estimators.py

@@ -293,9 +293,9 @@ def _check_cv(
             if cv >= 2:
                 return KFold(n_splits=cv)
         if (
-            isinstance(cv, KFold) or
-            isinstance(cv, LeaveOneOut) or
-            cv == "prefit"
+            isinstance(cv, KFold)
+            or isinstance(cv, LeaveOneOut)
+            or cv == "prefit"
         ):
             return cv
         raise ValueError(
@@ -441,12 +441,11 @@ def _fit_and_predict_oof_model(
         """
         X_train, y_train, X_val = X[train_index], y[train_index], X[val_index]
         if sample_weight is None:
-            sample_weight_train = None
+            estimator = fit_estimator(estimator, X_train, y_train)
         else:
-            sample_weight_train = sample_weight[train_index]
-        estimator = fit_estimator(
-            estimator, X_train, y_train, sample_weight_train
-        )
+            estimator = fit_estimator(
+                estimator, X_train, y_train, sample_weight[train_index]
+            )
         y_pred = estimator.predict(X_val)
         val_id = np.full_like(y_pred, k)
         return estimator, y_pred, val_id, val_index
@@ -555,7 +554,7 @@ def predict(self, X: ArrayLike) -> np.ndarray:
         """
         # Checks
         check_is_fitted(
-            self, 
+            self,
             [
                 "n_features_in_",
                 "single_estimator_",