UPD: modify the user warning with a better function call

scikit-learn-contrib · Dec 21, 2023 · fc0522f · fc0522f
1 parent e83389d
commit fc0522f
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 5 deletions.
diff --git a/mapie/regression/regression.py b/mapie/regression/regression.py
@@ -2,6 +2,8 @@
 
 from typing import Iterable, Optional, Tuple, Union, cast
 
+import warnings
+
 import numpy as np
 from sklearn.base import BaseEstimator, RegressorMixin
 from sklearn.linear_model import LinearRegression
@@ -592,9 +594,10 @@ def predict(
 
         else:
             if optimize_beta and self.method != 'enbpi':
-                raise UserWarning(
-                    "Beta optimisation should only be used for "
-                    "method='enbpi'."
+                warnings.warn(
+                    "WARNING: Beta optimisation should only be used for "
+                    "method='enbpi'.",
+                    UserWarning
                 )
 
             n = len(self.conformity_scores_)

diff --git a/mapie/tests/test_regression.py b/mapie/tests/test_regression.py
@@ -641,8 +641,10 @@ def test_beta_optimize_user_warning() -> None:
     """
     Test that a UserWarning is displayed when optimize_beta is used.
     """
-    mapie_reg = MapieRegressor().fit(X, y)
-    with pytest.raises(
+    mapie_reg = MapieRegressor(
+        conformity_score=AbsoluteConformityScore(sym=False)
+    ).fit(X, y)
+    with pytest.warns(
         UserWarning, match=r"Beta optimisation should only be used for*",
     ):
         mapie_reg.predict(X, alpha=0.05, optimize_beta=True)