Merge pull request #163 from muhlbach/main

Update sklearn.py by addind catch to OptunaSearchCV

Author: not522

optuna_integration/sklearn/sklearn.py

@@ -501,6 +501,11 @@ class OptunaSearchCV(BaseEstimator):
                 See the tutorial of `Callback for Study.optimize <https://optuna.readthedocs.io/en/stable/tutorial/20_recipes/007_optuna_callback.html#optuna-callback>`_
                 for how to use and implement callback functions.
 
+        catch:
+            A study continues to run even when a trial raises one of the exceptions specified
+            in this argument. Default is an empty tuple, i.e. the study will stop for any
+            exception except for :class:`~optuna.exceptions.TrialPruned`.
+
     Attributes:
         best_estimator_:
             Estimator that was chosen by the search. This is present only if
@@ -733,6 +738,7 @@ def __init__(
         timeout: float | None = None,
         verbose: int = 0,
         callbacks: list[Callable[[study_module.Study, FrozenTrial], None]] | None = None,
+        catch: Iterable[type[Exception]] | type[Exception] = (),
     ) -> None:
         _imports.check()
 
@@ -767,6 +773,7 @@ def __init__(
         self.timeout = timeout
         self.verbose = verbose
         self.callbacks = callbacks
+        self.catch = catch
 
     def _check_is_fitted(self) -> None:
         attributes = ["n_splits_", "sample_indices_", "scorer_", "study_"]
@@ -925,6 +932,7 @@ def fit(
             n_trials=self.n_trials,
             timeout=self.timeout,
             callbacks=self.callbacks,
+            catch=self.catch,
         )
 
         _logger.info("Finished hyperparameter search!")

tests/sklearn/test_sklearn.py

@@ -24,6 +24,8 @@
 from sklearn.exceptions import NotFittedError
 from sklearn.linear_model import LogisticRegression
 from sklearn.linear_model import SGDClassifier
+from sklearn.metrics import make_scorer
+from sklearn.metrics import r2_score
 from sklearn.model_selection import PredefinedSplit
 from sklearn.neighbors import KernelDensity
 from sklearn.tree import DecisionTreeRegressor
@@ -458,6 +460,47 @@ def test_callbacks() -> None:
         assert callback.call_count == n_trials
 
 
+@pytest.mark.parametrize("catch", [(ValueError,), ()])
+def test_catch(catch: tuple) -> None:
+
+    class MockDististribution(distributions.BaseDistribution):
+
+        def _contains(self) -> None:  # type: ignore
+            raise ValueError
+
+        def single(self) -> None:  # type: ignore
+            raise ValueError
+
+        def to_internal_repr(self) -> None:  # type: ignore
+            raise ValueError
+
+    est = SGDClassifier(max_iter=5, tol=1e-03)
+    X, y = make_blobs(n_samples=10)
+    param_dist = {"param": MockDististribution()}
+    n_trials = 3
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore", ExperimentalWarning)
+        optuna_search = OptunaSearchCV(
+            est,
+            param_dist,
+            cv=3,
+            max_iter=5,
+            n_trials=n_trials,
+            error_score=0,
+            refit=False,
+            scoring=make_scorer(r2_score),
+            catch=catch,
+        )
+
+    if catch:
+        optuna_search.fit(X, y)
+        assert optuna_search.n_trials_ == 3
+    else:
+        with pytest.raises(ValueError):
+            optuna_search.fit(X, y)
+
+
 @pytest.mark.filterwarnings("ignore::UserWarning")
 @patch("optuna_integration.sklearn.sklearn.cross_validate")
 def test_terminator_cv_score_reporting(mock: MagicMock) -> None: