diff --git a/ax/benchmark/benchmark.py b/ax/benchmark/benchmark.py index 6daee1e55da..c48709092cd 100644 --- a/ax/benchmark/benchmark.py +++ b/ax/benchmark/benchmark.py @@ -65,6 +65,19 @@ def compute_score_trace( return score_trace.clip(min=0, max=100) +def _create_benchmark_experiment( + problem: BenchmarkProblemBase, method_name: str +) -> Experiment: + """Creates an empty experiment for the given problem and method.""" + return Experiment( + name=f"{problem.name}|{method_name}_{int(time())}", + search_space=problem.search_space, + optimization_config=problem.optimization_config, + tracking_metrics=problem.tracking_metrics, + runner=problem.runner, + ) + + def benchmark_replication( problem: BenchmarkProblemBase, method: BenchmarkMethod, @@ -79,13 +92,7 @@ def benchmark_replication( from `botorch.utils.sampling`. """ - experiment = Experiment( - name=f"{problem.name}|{method.name}_{int(time())}", - search_space=problem.search_space, - optimization_config=problem.optimization_config, - tracking_metrics=problem.tracking_metrics, - runner=problem.runner, - ) + experiment = _create_benchmark_experiment(problem=problem, method_name=method.name) scheduler = Scheduler( experiment=experiment, diff --git a/ax/benchmark/methods/choose_generation_strategy.py b/ax/benchmark/methods/choose_generation_strategy.py index f4f57ec2b18..d4792d329a3 100644 --- a/ax/benchmark/methods/choose_generation_strategy.py +++ b/ax/benchmark/methods/choose_generation_strategy.py @@ -3,32 +3,4 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import Optional - -from ax.benchmark.benchmark_method import ( - BenchmarkMethod, - get_sequential_optimization_scheduler_options, -) -from ax.benchmark.benchmark_problem import BenchmarkProblemBase -from ax.modelbridge.dispatch_utils import choose_generation_strategy -from ax.service.scheduler import SchedulerOptions - - -def get_choose_generation_strategy_method( - problem: BenchmarkProblemBase, - scheduler_options: Optional[SchedulerOptions] = None, - distribute_replications: bool = False, -) -> BenchmarkMethod: - generation_strategy = choose_generation_strategy( - search_space=problem.search_space, - optimization_config=problem.optimization_config, - num_trials=problem.num_trials, - ) - - return BenchmarkMethod( - name=f"ChooseGenerationStrategy::{problem.name}", - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - distribute_replications=distribute_replications, - ) +# TODO: remove this file in the next diff diff --git a/ax/benchmark/methods/gpei_and_moo.py b/ax/benchmark/methods/gpei_and_moo.py index 4b803ddaec1..d66ee238e16 100644 --- a/ax/benchmark/methods/gpei_and_moo.py +++ b/ax/benchmark/methods/gpei_and_moo.py @@ -3,58 +3,4 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import Optional - -from ax.benchmark.benchmark_method import ( - BenchmarkMethod, - get_sequential_optimization_scheduler_options, -) -from ax.modelbridge.generation_strategy import GenerationStep, GenerationStrategy -from ax.modelbridge.registry import Models -from ax.service.scheduler import SchedulerOptions - - -def get_gpei_default( - scheduler_options: Optional[SchedulerOptions] = None, -) -> BenchmarkMethod: - generation_strategy = GenerationStrategy( - name="SOBOL+GPEI::default", - steps=[ - GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5), - GenerationStep( - model=Models.GPEI, - num_trials=-1, - max_parallelism=1, - ), - ], - ) - - return BenchmarkMethod( - name=generation_strategy.name, - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - ) - - -def get_moo_default( - scheduler_options: Optional[SchedulerOptions] = None, -) -> BenchmarkMethod: - generation_strategy = GenerationStrategy( - name="SOBOL+MOO::default", - steps=[ - GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5), - GenerationStep( - model=Models.MOO, - num_trials=-1, - max_parallelism=1, - ), - ], - ) - - return BenchmarkMethod( - name=generation_strategy.name, - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - ) +# TODO: delete this file in the next diff diff --git a/ax/benchmark/methods/modular_botorch.py b/ax/benchmark/methods/modular_botorch.py index be380c3e679..4c889516da9 100644 --- a/ax/benchmark/methods/modular_botorch.py +++ b/ax/benchmark/methods/modular_botorch.py @@ -3,7 +3,7 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import Any, Dict, Optional, Type +from typing import Dict, Optional, Type, Union from ax.benchmark.benchmark_method import ( BenchmarkMethod, @@ -13,228 +13,33 @@ from ax.modelbridge.registry import Models from ax.models.torch.botorch_modular.surrogate import Surrogate from ax.service.scheduler import SchedulerOptions -from ax.utils.common.constants import Keys from botorch.acquisition.acquisition import AcquisitionFunction -from botorch.acquisition.monte_carlo import qNoisyExpectedImprovement -from botorch.acquisition.multi_objective.monte_carlo import ( - qNoisyExpectedHypervolumeImprovement, -) -from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP -from botorch.models.gp_regression import FixedNoiseGP - - -def get_sobol_botorch_modular_fixed_noise_gp_qnei( - scheduler_options: Optional[SchedulerOptions] = None, -) -> BenchmarkMethod: - model_gen_kwargs = { - "model_gen_options": { - Keys.OPTIMIZER_KWARGS: { - "num_restarts": 50, - "raw_samples": 1024, - }, - Keys.ACQF_KWARGS: { - "prune_baseline": True, - }, - } - } - - generation_strategy = GenerationStrategy( - name="SOBOL+BOTORCH_MODULAR::FixedNoiseGP_qNoisyExpectedImprovement", - steps=[ - GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5), - GenerationStep( - model=Models.BOTORCH_MODULAR, - num_trials=-1, - max_parallelism=1, - model_kwargs={ - "surrogate": Surrogate(FixedNoiseGP), - "botorch_acqf_class": qNoisyExpectedImprovement, - }, - model_gen_kwargs=model_gen_kwargs, - ), - ], - ) - - return BenchmarkMethod( - name=generation_strategy.name, - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - ) - - -def get_sobol_botorch_modular_fixed_noise_gp_qnehvi( - scheduler_options: Optional[SchedulerOptions] = None, -) -> BenchmarkMethod: - model_gen_kwargs = { - "model_gen_options": { - Keys.OPTIMIZER_KWARGS: { - "num_restarts": 50, - "raw_samples": 1024, - }, - Keys.ACQF_KWARGS: { - "prune_baseline": True, - "qmc": True, - "mc_samples": 512, - }, - } - } - - generation_strategy = GenerationStrategy( - name="SOBOL+BOTORCH_MODULAR::FixedNoiseGP_qNoisyExpectedHypervolumeImprovement", - steps=[ - GenerationStep( - model=Models.SOBOL, - num_trials=5, - min_trials_observed=5, - ), - GenerationStep( - model=Models.BOTORCH_MODULAR, - num_trials=-1, - max_parallelism=1, - model_kwargs={ - "surrogate": Surrogate(FixedNoiseGP), - "botorch_acqf_class": qNoisyExpectedHypervolumeImprovement, - }, - model_gen_kwargs=model_gen_kwargs, - ), - ], - ) - - return BenchmarkMethod( - name=generation_strategy.name, - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - ) - - -def get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnei( - scheduler_options: Optional[SchedulerOptions] = None, - distribute_replications: bool = True, -) -> BenchmarkMethod: # noqa - return get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp( - qNoisyExpectedImprovement, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - distribute_replications=distribute_replications, - ) - - -def get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp( - botorch_acqf_class: Type[AcquisitionFunction], - scheduler_options: Optional[SchedulerOptions] = None, - distribute_replications: bool = True, -) -> BenchmarkMethod: # noqa - generation_strategy = GenerationStrategy( - name="SOBOL+BOTORCH_MODULAR::SaasFullyBayesianSingleTaskGP_" - + botorch_acqf_class.__name__, # noqa - steps=[ - GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5), - GenerationStep( - model=Models.BOTORCH_MODULAR, - num_trials=-1, - max_parallelism=1, - model_kwargs={ - "surrogate": Surrogate( - botorch_model_class=SaasFullyBayesianSingleTaskGP - ), - "botorch_acqf_class": botorch_acqf_class, - }, - ), - ], - ) - - return BenchmarkMethod( - name=generation_strategy.name, - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - distribute_replications=distribute_replications, - ) - - -def get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnehvi( - scheduler_options: Optional[SchedulerOptions] = None, - distribute_replications: bool = True, -) -> BenchmarkMethod: # noqa - generation_strategy = GenerationStrategy( - name="SOBOL+BOTORCH_MODULAR::SaasFullyBayesianSingleTaskGP_qNoisyExpectedHypervolumeImprovement", # noqa - steps=[ - GenerationStep( - model=Models.SOBOL, - num_trials=5, - min_trials_observed=5, - ), - GenerationStep( - model=Models.BOTORCH_MODULAR, - num_trials=-1, - max_parallelism=1, - model_kwargs={ - "surrogate": Surrogate( - botorch_model_class=SaasFullyBayesianSingleTaskGP - ), - "botorch_acqf_class": qNoisyExpectedHypervolumeImprovement, - }, - ), - ], - ) - return BenchmarkMethod( - name=generation_strategy.name, - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - distribute_replications=distribute_replications, - ) - - -def get_sobol_botorch_modular_default( - scheduler_options: Optional[SchedulerOptions] = None, - distribute_replications: bool = False, -) -> BenchmarkMethod: - generation_strategy = GenerationStrategy( - name="SOBOL+BOTORCH_MODULAR::default", - steps=[ - GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5), - GenerationStep( - model=Models.BOTORCH_MODULAR, - num_trials=-1, - max_parallelism=1, - ), - ], - ) - - return BenchmarkMethod( - name=generation_strategy.name, - generation_strategy=generation_strategy, - scheduler_options=scheduler_options - or get_sequential_optimization_scheduler_options(), - distribute_replications=distribute_replications, - ) +from botorch.models.model import Model def get_sobol_botorch_modular_acquisition( + model_cls: Type[Model], acquisition_cls: Type[AcquisitionFunction], - acquisition_options: Optional[Dict[str, Any]] = None, scheduler_options: Optional[SchedulerOptions] = None, distribute_replications: bool = False, + name: Optional[str] = None, ) -> BenchmarkMethod: + model_kwargs: Dict[str, Union[Type[AcquisitionFunction], Surrogate]] = { + "botorch_acqf_class": acquisition_cls + } + model_kwargs["surrogate"] = Surrogate(model_cls) + generation_strategy = GenerationStrategy( - name=f"SOBOL+BOTORCH_MODULAR::{acquisition_cls.__name__}", + name=f"MBM::{model_cls.__name__}_{acquisition_cls.__name__}" + if name is None + else name, steps=[ - GenerationStep( - model=Models.SOBOL, - num_trials=5, - min_trials_observed=5, - ), + GenerationStep(model=Models.SOBOL, num_trials=5, min_trials_observed=5), GenerationStep( model=Models.BOTORCH_MODULAR, num_trials=-1, max_parallelism=1, - model_kwargs={ - "botorch_acqf_class": acquisition_cls, - "acquisition_options": acquisition_options, - }, + model_kwargs=model_kwargs, ), ], ) diff --git a/ax/benchmark/methods/registry.py b/ax/benchmark/methods/registry.py new file mode 100644 index 00000000000..b19679f90a5 --- /dev/null +++ b/ax/benchmark/methods/registry.py @@ -0,0 +1,98 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass +from typing import Any, Callable, Dict, Optional, Type + +from ax.benchmark.benchmark_method import BenchmarkMethod +from ax.benchmark.methods.modular_botorch import get_sobol_botorch_modular_acquisition +from ax.service.scheduler import SchedulerOptions + +from botorch.acquisition import AcquisitionFunction +from botorch.acquisition.analytic import LogExpectedImprovement +from botorch.acquisition.logei import qLogNoisyExpectedImprovement +from botorch.acquisition.multi_objective.monte_carlo import ( + qNoisyExpectedHypervolumeImprovement, +) +from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP +from botorch.models.gp_regression import FixedNoiseGP, SingleTaskGP +from botorch.models.model import Model + + +@dataclass +class BenchmarkMethodRegistryEntry: + factory_fn: Callable[..., BenchmarkMethod] + factory_kwargs: Dict[str, Any] + + +def mbm_method_factory( + model_cls: Type[Model], + acquisition_cls: Type[AcquisitionFunction], +) -> BenchmarkMethodRegistryEntry: + model_names = { + SaasFullyBayesianSingleTaskGP.__name__: "SAAS", + } + acqf_names = { + qLogNoisyExpectedImprovement.__name__: "qLogNEI", + qNoisyExpectedHypervolumeImprovement.__name__: "qNEHVI", + LogExpectedImprovement.__name__: "LogEI", + } + + model_name = model_names.get(model_cls.__name__, model_cls.__name__) + acqf_name = acqf_names.get(acquisition_cls.__name__, acquisition_cls.__name__) + + method = BenchmarkMethodRegistryEntry( + factory_fn=get_sobol_botorch_modular_acquisition, + factory_kwargs={ + "acquisition_cls": acquisition_cls, + "model_cls": model_cls, + "name": f"MBM::{model_name}_{acqf_name}", + }, + ) + return method + + +# TODO: unit test that registry key matches method name +BENCHMARK_METHOD_REGISTRY: Dict[str, BenchmarkMethodRegistryEntry] = { + "MBM::FixedNoiseGP_qLogNEI": mbm_method_factory( + model_cls=FixedNoiseGP, acquisition_cls=qLogNoisyExpectedImprovement + ), + "MBM::FixedNoiseGP_qNEHVI": mbm_method_factory( + model_cls=FixedNoiseGP, acquisition_cls=qNoisyExpectedHypervolumeImprovement + ), + "MBM::SAAS_qLogNEI": mbm_method_factory( + model_cls=SaasFullyBayesianSingleTaskGP, + acquisition_cls=qLogNoisyExpectedImprovement, + ), + # MBM, choose_generation_strategy default for single-objective optimization + "MBM::SingleTaskGP_qLogNEI": mbm_method_factory( + model_cls=SingleTaskGP, acquisition_cls=qLogNoisyExpectedImprovement + ), + # Not a default + "MBM::SingleTaskGP_LogEI": mbm_method_factory( + model_cls=SingleTaskGP, acquisition_cls=LogExpectedImprovement + ), + # MBM, choose_generation_strategy default for multi-objective optimization + "MBM::SingleTaskGP_qNEHVI": mbm_method_factory( + model_cls=SingleTaskGP, acquisition_cls=qNoisyExpectedHypervolumeImprovement + ), + "MBM::SAAS_qNEHVI": mbm_method_factory( + model_cls=SaasFullyBayesianSingleTaskGP, + acquisition_cls=qNoisyExpectedHypervolumeImprovement, + ), +} + + +def get_method( + method_name: str, + distribute_replications: bool = False, + scheduler_options: Optional[SchedulerOptions] = None, +) -> BenchmarkMethod: + entry = BENCHMARK_METHOD_REGISTRY[method_name] + return entry.factory_fn( + **entry.factory_kwargs, + distribute_replications=distribute_replications, + scheduler_options=scheduler_options, + ) diff --git a/ax/benchmark/tests/test_benchmark.py b/ax/benchmark/tests/test_benchmark.py index 0e856b0ddf4..25ff875fcb6 100644 --- a/ax/benchmark/tests/test_benchmark.py +++ b/ax/benchmark/tests/test_benchmark.py @@ -17,15 +17,8 @@ ) from ax.benchmark.benchmark_problem import SingleObjectiveBenchmarkProblem from ax.benchmark.benchmark_result import BenchmarkResult -from ax.benchmark.methods.gpei_and_moo import get_gpei_default -from ax.benchmark.methods.modular_botorch import ( - get_sobol_botorch_modular_acquisition, - get_sobol_botorch_modular_default, - get_sobol_botorch_modular_fixed_noise_gp_qnehvi, - get_sobol_botorch_modular_fixed_noise_gp_qnei, - get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnehvi, - get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnei, -) +from ax.benchmark.methods.modular_botorch import get_sobol_botorch_modular_acquisition +from ax.benchmark.methods.registry import get_method from ax.modelbridge.modelbridge_utils import extract_search_space_digest from ax.service.utils.scheduler_options import SchedulerOptions from ax.storage.json_store.load import load_experiment @@ -41,7 +34,8 @@ ) from ax.utils.testing.core_stubs import get_dataset, get_experiment from ax.utils.testing.mock import fast_botorch_optimize -from botorch.acquisition.max_value_entropy_search import qMaxValueEntropy +from botorch.acquisition.logei import qLogNoisyExpectedImprovement +from botorch.models.gp_regression import SingleTaskGP from botorch.test_functions.synthetic import Branin @@ -147,25 +141,18 @@ def test_replication_mbm(self) -> None: for method, problem in [ ( get_sobol_botorch_modular_acquisition( - acquisition_cls=qMaxValueEntropy, + model_cls=SingleTaskGP, + acquisition_cls=qLogNoisyExpectedImprovement, scheduler_options=get_sequential_optimization_scheduler_options(), ), get_single_objective_benchmark_problem(infer_noise=False, num_trials=6), ), ( - get_sobol_botorch_modular_fixed_noise_gp_qnei(), - get_single_objective_benchmark_problem(infer_noise=False, num_trials=6), - ), - ( - get_sobol_botorch_modular_fixed_noise_gp_qnehvi(), + get_method("MBM::FixedNoiseGP_qNEHVI", distribute_replications=False), get_multi_objective_benchmark_problem(infer_noise=False, num_trials=6), ), ( - get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnei(), - get_single_objective_benchmark_problem(num_trials=6), - ), - ( - get_sobol_botorch_modular_saas_fully_bayesian_single_task_gp_qnehvi(), + get_method("MBM::SAAS_qNEHVI", distribute_replications=False), get_multi_objective_benchmark_problem(num_trials=6), ), ]: @@ -177,22 +164,6 @@ def test_replication_mbm(self) -> None: ) self.assertTrue(np.all(res.score_trace <= 100)) - @fast_botorch_optimize - def test_replication_mbm_default(self) -> None: - method = get_sobol_botorch_modular_default() - for problem in [ - get_single_objective_benchmark_problem(infer_noise=False, num_trials=6), - get_multi_objective_benchmark_problem(infer_noise=False, num_trials=6), - get_single_objective_benchmark_problem(num_trials=6), - get_multi_objective_benchmark_problem(num_trials=6), - ]: - with self.subTest(problem=problem): - res = benchmark_replication(problem=problem, method=method, seed=0) - self.assertEqual( - problem.num_trials, - len(not_none(res.experiment).trials), - ) - def test_replication_moo_sobol(self) -> None: problem = get_multi_objective_benchmark_problem() @@ -235,7 +206,12 @@ def test_benchmark_one_method_problem(self) -> None: def test_benchmark_multiple_problems_methods(self) -> None: aggs = benchmark_multiple_problems_methods( problems=[get_single_objective_benchmark_problem(num_trials=6)], - methods=[get_sobol_benchmark_method(), get_gpei_default()], + methods=[ + get_sobol_benchmark_method(), + get_sobol_botorch_modular_acquisition( + model_cls=SingleTaskGP, acquisition_cls=qLogNoisyExpectedImprovement + ), + ], seeds=(0, 1), ) @@ -251,7 +227,9 @@ def test_timeout(self) -> None: num_trials=1000, # Unachievable num_trials ) - generation_strategy = get_gpei_default().generation_strategy + generation_strategy = get_sobol_botorch_modular_acquisition( + model_cls=SingleTaskGP, acquisition_cls=qLogNoisyExpectedImprovement + ).generation_strategy method = BenchmarkMethod( name=generation_strategy.name, diff --git a/ax/benchmark/tests/test_methods.py b/ax/benchmark/tests/test_methods.py index 6326ed7b376..0224eae116a 100644 --- a/ax/benchmark/tests/test_methods.py +++ b/ax/benchmark/tests/test_methods.py @@ -11,32 +11,37 @@ from ax.benchmark.problems.registry import get_problem from ax.modelbridge.registry import Models from ax.utils.common.testutils import TestCase +from ax.utils.common.typeutils import not_none from ax.utils.testing.mock import fast_botorch_optimize +from botorch.acquisition.analytic import LogExpectedImprovement from botorch.acquisition.knowledge_gradient import qKnowledgeGradient +from botorch.models.gp_regression import SingleTaskGP class TestMethods(TestCase): def test_mbm_acquisition(self) -> None: method = get_sobol_botorch_modular_acquisition( - scheduler_options=get_sequential_optimization_scheduler_options(), + model_cls=SingleTaskGP, acquisition_cls=qKnowledgeGradient, - acquisition_options={"num_fantasies": 16}, + scheduler_options=get_sequential_optimization_scheduler_options(), ) - self.assertEqual(method.name, "SOBOL+BOTORCH_MODULAR::qKnowledgeGradient") + self.assertEqual(method.name, "MBM::SingleTaskGP_qKnowledgeGradient") gs = method.generation_strategy sobol, kg = gs._steps self.assertEqual(kg.model, Models.BOTORCH_MODULAR) - model_kwargs = kg.model_kwargs - # pyre-fixme[16]: Optional type has no attribute `__getitem__`. + model_kwargs = not_none(kg.model_kwargs) self.assertEqual(model_kwargs["botorch_acqf_class"], qKnowledgeGradient) - self.assertEqual(model_kwargs["acquisition_options"], {"num_fantasies": 16}) + self.assertEqual( + model_kwargs["surrogate"].botorch_model_class.__name__, "SingleTaskGP" + ) @fast_botorch_optimize def test_benchmark_replication_runs(self) -> None: problem = get_problem(problem_name="ackley4") method = get_sobol_botorch_modular_acquisition( + model_cls=SingleTaskGP, scheduler_options=get_sequential_optimization_scheduler_options(), - acquisition_cls=qKnowledgeGradient, + acquisition_cls=LogExpectedImprovement, ) n_sobol_trials = method.generation_strategy._steps[0].num_trials # Only run one non-Sobol trial diff --git a/ax/benchmark/tests/test_methods_registry.py b/ax/benchmark/tests/test_methods_registry.py new file mode 100644 index 00000000000..1bbdaa76579 --- /dev/null +++ b/ax/benchmark/tests/test_methods_registry.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# +# This source code is licensed under the MIT license found in the +# LICENSE file in the root directory of this source tree. + + +from ax.benchmark.methods.registry import BENCHMARK_METHOD_REGISTRY, get_method +from ax.modelbridge.registry import Models +from ax.utils.common.testutils import TestCase +from ax.utils.common.typeutils import not_none +from botorch.acquisition.logei import qLogNoisyExpectedImprovement +from botorch.acquisition.multi_objective.monte_carlo import ( + qNoisyExpectedHypervolumeImprovement, +) +from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP +from botorch.models.gp_regression import FixedNoiseGP + + +class TestRegistry(TestCase): + def test_names_match(self) -> None: + """Check that all names in the registry match the corresponding method.""" + for name in BENCHMARK_METHOD_REGISTRY.keys(): + method = get_method(name) + self.assertEqual(name, method.name) + + def test_mbm_methods(self) -> None: + test_cases = { + "MBM::FixedNoiseGP_qLogNEI": (FixedNoiseGP, qLogNoisyExpectedImprovement), + "MBM::SAAS_qNEHVI": ( + SaasFullyBayesianSingleTaskGP, + qNoisyExpectedHypervolumeImprovement, + ), + } + for name, (model_cls, acqf_cls) in test_cases.items(): + with self.subTest(name=name): + method = get_method(name) + gs = method.generation_strategy + sobol_step, botorch_step = gs._steps + self.assertIs(botorch_step.model, Models.BOTORCH_MODULAR) + model_kwargs = not_none(botorch_step.model_kwargs) + self.assertIs(model_kwargs["surrogate"].botorch_model_class, model_cls) + self.assertIs(model_kwargs["botorch_acqf_class"], acqf_cls)