GPSampler as Default Optuna Sampler

pyproject.toml

@@ -25,7 +25,7 @@ dependencies = [
   "catboost>=1.2.9,<=1.2.10",
   "feature-engine>=1.8.0,<2",
   "numpy>=2.2,<2.3",
-  "optuna>=4.2.0,<5",
+  "optuna>=4.9.0,<5",
   "pandas>=2.2,<3.0",
   "pydantic-settings>=2.4.0,<3",
   "rdkit>=2024.9.1",

src/mother/optimization/core.py

@@ -125,6 +125,36 @@ class MotherTuner:
              -> Pipeline:
     """
 
+    @staticmethod
+    def _build_tpe_sampler(
+        seed: int,
+        n_startup_trials: int,
+    ) -> optuna.samplers.TPESampler:
+        return optuna.samplers.TPESampler(
+            multivariate=True,
+            group=True,
+            constant_liar=True,
+            seed=seed,
+            n_startup_trials=n_startup_trials,
+        )
+
+    @staticmethod
+    def _build_gp_sampler(
+        seed: int,
+        n_startup_trials: int,
+    ) -> optuna.samplers.GPSampler:
+        return optuna.samplers.GPSampler(
+            seed=seed,
+            independent_sampler=optuna.samplers.TPESampler(
+                seed=seed,
+                n_startup_trials=n_startup_trials,
+            ),
+            n_startup_trials=n_startup_trials,
+            deterministic_objective=False,
+            # constant liar is always True for GPSampler,
+            # so we don't need to set it explicitly
+        )
+
     def __init__(
         self,
         scorer: typing.Union[typing.Callable, str],
@@ -143,15 +173,30 @@ def __init__(
         self.early_stopping_optuna: bool = early_stopping_optuna
         self.tuning_direction: typing.Union[StudyDirection, str] = tuning_direction
         self.scorer: typing.Callable = skl_metrics.get_scorer(scorer)
+
         if sampler is None:
-            module_logger.debug("Setting up default sampler TPE")
-            self.sampler = optuna.samplers.TPESampler(
-                multivariate=kwargs.get("multivariate", True),
-                group=True,
-                constant_liar=True,
-                seed=seed,
-                n_startup_trials=n_startup_trials,
-            )
+            if torch_available:
+                try:
+                    module_logger.debug("torch available — using GPSampler as default")
+                    self.sampler = self._build_gp_sampler(
+                        seed=seed,
+                        n_startup_trials=n_startup_trials,
+                    )
+                except Exception as gpsampler_error:
+                    module_logger.warning(
+                        "GPSampler initialization failed (%s). Falling back to TPESampler.",
+                        gpsampler_error,
+                    )
+                    self.sampler = self._build_tpe_sampler(
+                        seed=seed,
+                        n_startup_trials=n_startup_trials,
+                    )
+            else:
+                module_logger.debug("torch not available — falling back to TPESampler")
+                self.sampler = self._build_tpe_sampler(
+                    seed=seed,
+                    n_startup_trials=n_startup_trials,
+                )
         else:
             self.sampler = sampler
 

test/unit/test_model_tuner.py

@@ -1,4 +1,5 @@
 import numpy as np
+import optuna
 import pandas as pd
 import pytest
 from optuna.samplers import RandomSampler
@@ -208,6 +209,114 @@ def test_tune_mother_model_classification_with_catboost(
         assert probabilities.shape[1] == n_classes, f"Multiclass probabilities should have {n_classes} columns."
 
 
+class TestGetCallbacks:
+    def test_early_stopping_enabled_returns_callback_when_torch_available(self, monkeypatch):
+        """get_callbacks should return terminator callback when enabled and torch is available."""
+        monkeypatch.setattr("mother.optimization.core.torch_available", True)
+        tuner = MotherTuner(
+            scorer=make_scorer(mean_squared_error, greater_is_better=False),
+            early_stopping_optuna=True,
+        )
+
+        callbacks = tuner.get_callbacks()
+
+        assert callbacks is not None
+        assert len(callbacks) == 1
+
+    def test_early_stopping_disabled_returns_none(self):
+        tuner = MotherTuner(
+            scorer=make_scorer(mean_squared_error, greater_is_better=False),
+            early_stopping_optuna=False,
+        )
+
+        callbacks = tuner.get_callbacks()
+
+        assert callbacks is None
+
+    def test_early_stopping_enabled_torch_missing_returns_none_with_warning(self, monkeypatch, caplog):
+        monkeypatch.setattr("mother.optimization.core.torch_available", False)
+        tuner = MotherTuner(
+            scorer=make_scorer(mean_squared_error, greater_is_better=False),
+            early_stopping_optuna=True,
+        )
+
+        with caplog.at_level("WARNING"):
+            callbacks = tuner.get_callbacks()
+
+        assert callbacks is None
+        assert "Torch not installed" in caplog.text
+
+
+class TestDefaultSamplerSelection:
+    def test_gp_sampler_when_torch_available(self, monkeypatch):
+        """GPSampler should be selected when torch is available."""
+        pytest.importorskip("torch")
+        pytest.importorskip("scipy")
+        monkeypatch.setattr("mother.optimization.core.torch_available", True)
+        tuner = MotherTuner(
+            scorer=make_scorer(mean_squared_error, greater_is_better=False),
+        )
+        assert isinstance(tuner.sampler, optuna.samplers.GPSampler)
+        assert isinstance(tuner.sampler._independent_sampler, optuna.samplers.TPESampler)
+
+    def test_tpe_sampler_when_torch_unavailable(self, monkeypatch):
+        """TPESampler should be selected when torch is not available."""
+        monkeypatch.setattr("mother.optimization.core.torch_available", False)
+        tuner = MotherTuner(
+            scorer=make_scorer(mean_squared_error, greater_is_better=False),
+        )
+        assert isinstance(tuner.sampler, optuna.samplers.TPESampler)
+
+    def test_tpe_sampler_fallback_when_gp_sampler_init_fails(self, monkeypatch):
+        """If GPSampler init fails, MotherTuner should fall back to TPESampler."""
+
+        def _raise_gp_init_error(*args, **kwargs):
+            raise RuntimeError("gp init failed")
+
+        monkeypatch.setattr("mother.optimization.core.torch_available", True)
+        monkeypatch.setattr("optuna.samplers.GPSampler", _raise_gp_init_error)
+
+        tuner = MotherTuner(
+            scorer=make_scorer(mean_squared_error, greater_is_better=False),
+        )
+        assert isinstance(tuner.sampler, optuna.samplers.TPESampler)
+
+    def test_gp_sampler_warns_for_dynamic_search_space(self, caplog):
+        """GPSampler should warn that dynamic search spaces fall back to independent sampling."""
+        pytest.importorskip("torch")
+        pytest.importorskip("scipy")
+
+        sampler = optuna.samplers.GPSampler(
+            seed=0,
+            n_startup_trials=1,
+            warn_independent_sampling=True,
+        )
+        study = optuna.create_study(direction="minimize", sampler=sampler)
+
+        def objective(trial):
+            # Alternate parameter names across trials to force a dynamic search space.
+            x = trial.suggest_float("x", 0.0, 1.0)
+            if trial.number % 2 == 0:
+                trial.suggest_float("even_only", 0.0, 1.0)
+            else:
+                trial.suggest_float("odd_only", 0.0, 1.0)
+            return x
+
+        with caplog.at_level("WARNING", logger="optuna.samplers._gp.sampler"):
+            study.optimize(objective, n_trials=6)
+
+        assert "dynamic search space is not supported by GPSampler" in caplog.text
+
+    def test_custom_sampler_used_directly(self):
+        """A user-provided sampler should be used as-is."""
+        custom = RandomSampler(seed=0)
+        tuner = MotherTuner(
+            scorer=make_scorer(mean_squared_error, greater_is_better=False),
+            sampler=custom,
+        )
+        assert tuner.sampler is custom
+
+
 @pytest.mark.slow
 @pytest.mark.serial
 class TestTuneMotherModels: