UPD: add missing attributes + keep quantiles_ attribute

scikit-learn-contrib · Jul 4, 2024 · 7d053b7 · 7d053b7
1 parent 59586a9
commit 7d053b7
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 0 deletions.
diff --git a/mapie/classification.py b/mapie/classification.py
@@ -140,12 +140,18 @@ class MapieClassifier(BaseEstimator, ClassifierMixin):
     estimator_: EnsembleClassifier
         Sklearn estimator that handle all that is related to the estimator.
 
+    conformity_score_function_: BaseClassificationScore
+        Score function that handle all that is related to conformity scores.
+
     n_features_in_: int
         Number of features passed to the fit method.
 
     conformity_scores_: ArrayLike of shape (n_samples_train)
         The conformity scores used to calibrate the prediction sets.
 
+    quantiles_: ArrayLike of shape (n_alpha)
+        The quantiles estimated from ``conformity_scores_`` and alpha values.
+
     References
     ----------
     [1] Mauricio Sadinle, Jing Lei, and Larry Wasserman.
@@ -741,4 +747,6 @@ def predict(
             **kwargs
         )
 
+        self.quantiles_ = self.conformity_score_function_.quantiles_
+
         return y_pred, prediction_sets
diff --git a/mapie/regression/regression.py b/mapie/regression/regression.py
@@ -165,6 +165,9 @@ class MapieRegressor(BaseEstimator, RegressorMixin):
     estimator_: EnsembleRegressor
         Sklearn estimator that handle all that is related to the estimator.
 
+    conformity_score_function_: BaseRegressionScore
+        Score function that handle all that is related to conformity scores.
+
     conformity_scores_: ArrayLike of shape (n_samples_train,)
         Conformity scores between ``y_train`` and ``y_pred``.