diff --git a/pyproject.toml b/pyproject.toml index 4e275b4..44da694 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/mother/optimization/core.py b/src/mother/optimization/core.py index befef5e..129b771 100644 --- a/src/mother/optimization/core.py +++ b/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 diff --git a/test/unit/test_model_tuner.py b/test/unit/test_model_tuner.py index bb71e29..d0dac43 100644 --- a/test/unit/test_model_tuner.py +++ b/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: diff --git a/uv.lock b/uv.lock index b6e564c..6f03a5b 100644 --- a/uv.lock +++ b/uv.lock @@ -2871,7 +2871,7 @@ requires-dist = [ { name = "leidenalg", marker = "extra == 'report'", specifier = ">=0.10.2,<0.11" }, { name = "matplotlib", marker = "extra == 'report'", specifier = ">3.7.0" }, { name = "numpy", specifier = ">=2.2,<2.3" }, - { name = "optuna", specifier = ">=4.2.0,<5" }, + { name = "optuna", specifier = ">=4.9.0,<5" }, { name = "pandas", specifier = ">=2.2,<3.0" }, { name = "plotly", marker = "extra == 'report'", specifier = ">=6.0.1,<7" }, { name = "pydantic-settings", specifier = ">=2.4.0,<3" },