Source code for ax.benchmark.problems.registry

# 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

from ax.benchmark.benchmark_problem import (
    BenchmarkProblem,
    MultiObjectiveBenchmarkProblem,
    SingleObjectiveBenchmarkProblem,
)
from ax.benchmark.problems.hd_embedding import embed_higher_dimension
from ax.benchmark.problems.hpo.torchvision import PyTorchCNNTorchvisionBenchmarkProblem
from ax.benchmark.problems.synthetic.hss.jenatton import get_jenatton_benchmark_problem
from botorch.test_functions.multi_objective import BraninCurrin
from botorch.test_functions.synthetic import Ackley, Branin, Hartmann, Powell


[docs]@dataclass class BenchmarkProblemRegistryEntry: factory_fn: Callable[..., BenchmarkProblem] factory_kwargs: Dict[str, Any]
BENCHMARK_PROBLEM_REGISTRY = { "ackley": BenchmarkProblemRegistryEntry( factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic, factory_kwargs={ "test_problem_class": Ackley, "test_problem_kwargs": {}, "num_trials": 50, "infer_noise": True, }, ), "branin": BenchmarkProblemRegistryEntry( factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic, factory_kwargs={ "test_problem_class": Branin, "test_problem_kwargs": {}, "num_trials": 30, "infer_noise": True, }, ), "branin_currin": BenchmarkProblemRegistryEntry( factory_fn=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective, factory_kwargs={ "test_problem_class": BraninCurrin, "test_problem_kwargs": {}, "num_trials": 30, "infer_noise": True, }, ), "branin_currin30": BenchmarkProblemRegistryEntry( factory_fn=lambda n: embed_higher_dimension( problem=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective( test_problem_class=BraninCurrin, test_problem_kwargs={}, num_trials=100, infer_noise=True, ), total_dimensionality=n, ), factory_kwargs={"n": 30}, ), "hartmann6": BenchmarkProblemRegistryEntry( factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic, factory_kwargs={ "test_problem_class": Hartmann, "test_problem_kwargs": {"dim": 6}, "num_trials": 50, "infer_noise": True, }, ), "hartmann30": BenchmarkProblemRegistryEntry( factory_fn=lambda n: embed_higher_dimension( problem=SingleObjectiveBenchmarkProblem.from_botorch_synthetic( test_problem_class=Hartmann, test_problem_kwargs={"dim": 6}, num_trials=100, infer_noise=True, ), total_dimensionality=n, ), factory_kwargs={"n": 30}, ), "hpo_pytorch_cnn_MNIST": BenchmarkProblemRegistryEntry( factory_fn=PyTorchCNNTorchvisionBenchmarkProblem.from_dataset_name, factory_kwargs={"name": "MNIST", "num_trials": 50, "infer_noise": True}, ), "hpo_pytorch_cnn_FashionMNIST": BenchmarkProblemRegistryEntry( factory_fn=PyTorchCNNTorchvisionBenchmarkProblem.from_dataset_name, factory_kwargs={"name": "FashionMNIST", "num_trials": 50, "infer_noise": True}, ), "jenatton": BenchmarkProblemRegistryEntry( factory_fn=get_jenatton_benchmark_problem, factory_kwargs={"num_trials": 50, "infer_noise": True}, ), "powell": BenchmarkProblemRegistryEntry( factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic, factory_kwargs={ "test_problem_class": Powell, "test_problem_kwargs": {}, "num_trials": 50, "infer_noise": True, }, ), # Problems without inferred noise "branin_fixed_noise": BenchmarkProblemRegistryEntry( factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic, factory_kwargs={ "test_problem_class": Branin, "test_problem_kwargs": {}, "num_trials": 30, "infer_noise": False, }, ), "branin_currin_fixed_noise": BenchmarkProblemRegistryEntry( factory_fn=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective, factory_kwargs={ "test_problem_class": BraninCurrin, "test_problem_kwargs": {}, "num_trials": 30, "infer_noise": False, }, ), "branin_currin30_fixed_noise": BenchmarkProblemRegistryEntry( factory_fn=lambda n: embed_higher_dimension( problem=MultiObjectiveBenchmarkProblem.from_botorch_multi_objective( test_problem_class=BraninCurrin, test_problem_kwargs={}, num_trials=100, infer_noise=False, ), total_dimensionality=n, ), factory_kwargs={"n": 30}, ), "hartmann6_fixed_noise": BenchmarkProblemRegistryEntry( factory_fn=SingleObjectiveBenchmarkProblem.from_botorch_synthetic, factory_kwargs={ "test_problem_class": Hartmann, "test_problem_kwargs": {"dim": 6}, "num_trials": 50, "infer_noise": False, }, ), "hartmann30_fixed_noise": BenchmarkProblemRegistryEntry( factory_fn=lambda n: embed_higher_dimension( problem=SingleObjectiveBenchmarkProblem.from_botorch_synthetic( test_problem_class=Hartmann, test_problem_kwargs={"dim": 6}, num_trials=100, infer_noise=False, ), total_dimensionality=n, ), factory_kwargs={"n": 30}, ), "jenatton_fixed_noise": BenchmarkProblemRegistryEntry( factory_fn=get_jenatton_benchmark_problem, factory_kwargs={"num_trials": 50, "infer_noise": False}, ), }
[docs]def get_problem( problem_name: str, ) -> BenchmarkProblem: entry = BENCHMARK_PROBLEM_REGISTRY[problem_name] return entry.factory_fn(**entry.factory_kwargs)