Source code for ax.benchmark.methods.modular_botorch

# 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.

# pyre-strict

from typing import Any, Dict, Optional, Type, Union

from ax.benchmark.benchmark_method import (
    BenchmarkMethod,
    get_benchmark_scheduler_options,
)
from ax.modelbridge.generation_node import GenerationStep
from ax.modelbridge.generation_strategy import GenerationStrategy
from ax.modelbridge.registry import Models
from ax.models.torch.botorch_modular.model import SurrogateSpec
from ax.service.scheduler import SchedulerOptions
from botorch.acquisition.acquisition import AcquisitionFunction
from botorch.acquisition.analytic import LogExpectedImprovement
from botorch.acquisition.logei import qLogNoisyExpectedImprovement
from botorch.acquisition.multi_objective.logei import (
    qLogNoisyExpectedHypervolumeImprovement,
)
from botorch.acquisition.multi_objective.monte_carlo import (
    qNoisyExpectedHypervolumeImprovement,
)
from botorch.models.fully_bayesian import SaasFullyBayesianSingleTaskGP
from botorch.models.model import Model


model_names_abbrevations: Dict[str, str] = {
    SaasFullyBayesianSingleTaskGP.__name__: "SAAS",
}
acqf_name_abbreviations: Dict[str, str] = {
    qLogNoisyExpectedImprovement.__name__: "qLogNEI",
    qNoisyExpectedHypervolumeImprovement.__name__: "qNEHVI",
    qLogNoisyExpectedHypervolumeImprovement.__name__: "qLogNEHVI",
    LogExpectedImprovement.__name__: "LogEI",
}


[docs]def get_sobol_botorch_modular_acquisition( model_cls: Type[Model], acquisition_cls: Type[AcquisitionFunction], distribute_replications: bool, scheduler_options: Optional[SchedulerOptions] = None, name: Optional[str] = None, num_sobol_trials: int = 5, model_gen_kwargs: Optional[Dict[str, Any]] = None, ) -> BenchmarkMethod: """Get a `BenchmarkMethod` that uses Sobol followed by MBM. Args: model_cls: BoTorch model class, e.g. SingleTaskGP acquisition_cls: Acquisition function class, e.g. `qLogNoisyExpectedImprovement`. distribute_replications: Whether to use multiple machines scheduler_options: Passed as-is to scheduler. Default: `get_benchmark_scheduler_options()`. name: Name that will be attached to the `GenerationStrategy`. num_sobol_trials: Number of Sobol trials; if the scheduler_options specify to use `BatchTrial`s, then this refers to the number of `BatchTrial`s. model_gen_kwargs: Passed to the BoTorch `GenerationStep` and ultimately to the BoTorch `Model`. Example: >>> # A simple example >>> from ax.benchmark.methods.sobol_botorch_modular import ( ... get_sobol_botorch_modular_acquisition ... ) >>> from ax.benchmark.benchmark_method import get_benchmark_scheduler_options >>> >>> method = get_sobol_botorch_modular_acquisition( ... model_cls=SingleTaskGP, ... acquisition_cls=qLogNoisyExpectedImprovement, ... distribute_replications=False, ... ) >>> # Pass sequential=False to BoTorch's optimize_acqf >>> batch_method = get_sobol_botorch_modular_acquisition( ... model_cls=SingleTaskGP, ... acquisition_cls=qLogNoisyExpectedImprovement, ... distribute_replications=False, ... scheduler_options=get_benchmark_scheduler_options( ... batch_size=5, ... ), ... model_gen_kwargs={ ... "model_gen_options": { ... "optimizer_kwargs": {"sequential": False} ... } ... }, ... num_sobol_trials=1, ... ) """ model_kwargs: Dict[ str, Union[Type[AcquisitionFunction], Dict[str, SurrogateSpec], bool] ] = { "botorch_acqf_class": acquisition_cls, "surrogate_specs": {"BoTorch": SurrogateSpec(botorch_model_class=model_cls)}, } model_name = model_names_abbrevations.get(model_cls.__name__, model_cls.__name__) acqf_name = acqf_name_abbreviations.get( acquisition_cls.__name__, acquisition_cls.__name__ ) # Historically all benchmarks were sequential, so sequential benchmarks # don't get anything added to their name, for continuity batch_suffix = "" if ( scheduler_options is not None and (batch_size := scheduler_options.batch_size) is not None ): if batch_size > 1: batch_suffix = f"_q{batch_size}" name = name or f"MBM::{model_name}_{acqf_name}{batch_suffix}" generation_strategy = GenerationStrategy( name=name, steps=[ GenerationStep( model=Models.SOBOL, num_trials=num_sobol_trials, min_trials_observed=num_sobol_trials, ), GenerationStep( model=Models.BOTORCH_MODULAR, num_trials=-1, model_kwargs=model_kwargs, model_gen_kwargs=model_gen_kwargs or {}, ), ], ) return BenchmarkMethod( name=generation_strategy.name, generation_strategy=generation_strategy, scheduler_options=scheduler_options or get_benchmark_scheduler_options(), distribute_replications=distribute_replications, )