add discrete acquisition function optimization for discrete spaces with gpbo (#79)

* add discrete acquisition function optimization for discrete spaces with gpbo

---------

Signed-off-by: Grossberger Lukas (CR/AIR2.2) <Lukas.Grossberger@de.bosch.com>

Author: LGro

blackboxopt/__init__.py

@@ -1,4 +1,4 @@
-__version__ = "4.8.5"
+__version__ = "4.9.0"
 
 from parameterspace import ParameterSpace
 

blackboxopt/optimizers/botorch_base.py

@@ -32,7 +32,7 @@
     from botorch.acquisition import AcquisitionFunction
     from botorch.exceptions import BotorchTensorDimensionWarning
     from botorch.models.model import Model
-    from botorch.optim import optimize_acqf
+    from botorch.optim import optimize_acqf, optimize_acqf_discrete
     from botorch.sampling.samplers import IIDNormalSampler
     from sklearn.impute import SimpleImputer
 
@@ -166,27 +166,50 @@ def to_numerical(
     return X, Y
 
 
-def init_af_opt_kwargs(af_opt_kwargs: Optional[dict]) -> dict:
-    """Provide default initialization for the acquisition function optimizer
-    configuration. Ensure that all mandatory fields are set to be used by BoTorch.
+def _acquisition_function_optimizer_factory(
+    search_space: ps.ParameterSpace,
+    af_opt_kwargs: Optional[dict],
+    torch_dtype: torch.dtype,
+) -> Callable[[AcquisitionFunction], Tuple[torch.Tensor, torch.Tensor]]:
+    """Prepare either BoTorch's `optimize_acqf_discrete` or `optimize_acqf` depending
+    on whether the search space is fully discrete or not and set required defaults if
+    not overridden by `af_opt_kwargs`.
 
     Args:
-        af_opt_kwargs: Acquisition function configuration.
+        search_space: Search space used for optimization.
+        af_opt_kwargs: Acquisition function optimizer configuration, e.g. containing
+            values for `n_samples` for discrete optimization, and `num_restarts`,
+            `raw_samples` for the continuous optimization case.
+        torch_dtype: Torch tensor type.
 
     Returns:
-        Acquisition function optimizer configuration, applicable for BoTorch's
-            optimizer.
+        Acquisition function optimizer that takes an acquisition function and returns a
+        candidate with its associate acquisition function value.
     """
+    kwargs = {} if af_opt_kwargs is None else af_opt_kwargs.copy()
 
-    af_opt_config = {} if af_opt_kwargs is None else af_opt_kwargs
-
-    # number of initial samples during AF optimization
-    af_opt_config.setdefault("raw_samples", 1024)
-
-    # number of restarts during AF optimization
-    af_opt_config.setdefault("num_restarts", 4)
-
-    return af_opt_config
+    is_fully_discrete_space = not any(
+        search_space[n]["parameter"].is_continuous
+        for n in search_space.get_parameter_names()
+    )
+    if is_fully_discrete_space:
+        choices = torch.Tensor(
+            [
+                search_space.to_numerical(search_space.sample())
+                for _ in range(kwargs.pop("n_samples", 5_000))
+            ]
+        ).to(dtype=torch_dtype)
+        return functools.partial(optimize_acqf_discrete, q=1, choices=choices, **kwargs)
+
+    return functools.partial(
+        optimize_acqf,
+        q=1,
+        # The numerical representation always lives on the unit hypercube
+        bounds=torch.tensor([[0, 1]] * len(search_space), dtype=torch_dtype).T,
+        num_restarts=kwargs.pop("num_restarts", 4),
+        raw_samples=kwargs.pop("raw_samples", 1024),
+        **kwargs,
+    )
 
 
 def filter_y_nans(
@@ -256,7 +279,8 @@ def __init__(
                 Providing a partially initialized class is possible with, e.g.
                 `functools.partial(UpperConfidenceBound, beta=6.0, maximize=False)`.
             af_optimizer_kwargs: Settings for acquisition function optimizer,
-                see `botorch.optim.optimize_acqf`.
+                see `botorch.optim.optimize_acqf` and in case the whole search space
+                is discrete: `botorch.optim.optimize_acqf_discrete`.
             num_initial_random_samples: Size of the initial space-filling design that
                 is used before starting BO. The points are sampled randomly in the
                 search space. If no random sampling is required, set it to 0.
@@ -286,7 +310,7 @@ def __init__(
 
         self.model = model
         self.acquisition_function_factory = acquisition_function_factory
-        self.af_opt_kwargs = init_af_opt_kwargs(af_optimizer_kwargs)
+        self.af_optimizer_kwargs = af_optimizer_kwargs
 
     def _create_fantasy_model(self, model: Model) -> Model:
         """Create model with the pending specifications and model based
@@ -327,7 +351,6 @@ def _generate_evaluation_specification(self):
         fantasy_model = self._create_fantasy_model(self.model)
         fantasy_model.eval()
 
-        # find next configuration by optimizing the acquisition function
         af = self.acquisition_function_factory(fantasy_model)
         if getattr(af, "maximize", False):
             raise ValueError(
@@ -338,17 +361,12 @@ def _generate_evaluation_specification(self):
                 "acquisition_function_factory init argument."
             )
 
-        # numerical representation always lives on hypercube
-        bounds = torch.tensor(
-            [[0, 1]] * len(self.search_space), dtype=self.torch_dtype
-        ).T
-
-        configuration, _ = optimize_acqf(
-            af,
-            bounds=bounds,
-            q=1,
-            **self.af_opt_kwargs,
+        acquisition_function_optimizer = _acquisition_function_optimizer_factory(
+            search_space=self.search_space,
+            af_opt_kwargs=self.af_optimizer_kwargs,
+            torch_dtype=self.torch_dtype,
         )
+        configuration, _ = acquisition_function_optimizer(af)
 
         return EvaluationSpecification(
             configuration=self.search_space.from_numerical(configuration[0]),

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "blackboxopt"
-version = "4.8.5"
+version = "4.9.0"
 description = "A common interface for blackbox optimization algorithms along with useful helpers like parallel optimization loops, analysis and visualization scripts."
 readme = "README.md"
 repository = "https://github.com/boschresearch/blackboxopt"

tests/optimizers/botorch_base_test.py

@@ -6,14 +6,17 @@
 from functools import partial
 
 import numpy as np
+import parameterspace as ps
 import pytest
 import torch
 from botorch.acquisition import UpperConfidenceBound
 from botorch.models import SingleTaskGP
+from botorch.optim import optimize_acqf, optimize_acqf_discrete
 
 from blackboxopt import ConstraintsError, Evaluation, Objective
 from blackboxopt.optimizers.botorch_base import (
     SingleObjectiveBOTorchOptimizer,
+    _acquisition_function_optimizer_factory,
     filter_y_nans,
     impute_nans_with_constant,
     to_numerical,
@@ -57,6 +60,73 @@ def test_all_reference_tests(reference_test, seed):
     )
 
 
+def test_acquisition_function_optimizer_factory_with_continuous():
+    continuous_space = ps.ParameterSpace()
+    continuous_space.add(ps.ContinuousParameter("conti1", (0.0, 1.0)))
+    continuous_space.add(ps.ContinuousParameter("conti2", (-1.0, 1.0)))
+
+    af_opt = _acquisition_function_optimizer_factory(
+        continuous_space, af_opt_kwargs={}, torch_dtype=torch.float64
+    )
+
+    assert af_opt.func == optimize_acqf  # pylint: disable=no-member
+
+
+def test_acquisition_function_optimizer_factory_with_discrete_space():
+    discrete_space = ps.ParameterSpace()
+    discrete_space.add(ps.IntegerParameter("integ", (-5, 10)))
+    discrete_space.add(ps.OrdinalParameter("ordin", ("small", "medium", "large")))
+    discrete_space.add(ps.CategoricalParameter("categ", ("woof", "miaow", "moo")))
+
+    af_opt = _acquisition_function_optimizer_factory(
+        discrete_space, af_opt_kwargs={}, torch_dtype=torch.float64
+    )
+
+    assert af_opt.func == optimize_acqf_discrete  # pylint: disable=no-member
+
+
+def test_acquisition_function_optimizer_factory_with_mixed_space():
+    mixed_space = ps.ParameterSpace()
+    mixed_space.add(ps.OrdinalParameter("ordin", ("small", "medium", "large")))
+    mixed_space.add(ps.ContinuousParameter("conti", (0.0, 1.0)))
+
+    af_opt = _acquisition_function_optimizer_factory(
+        mixed_space, af_opt_kwargs={}, torch_dtype=torch.float64
+    )
+
+    assert af_opt.func == optimize_acqf  # pylint: disable=no-member
+
+
+def test_find_optimum_in_1d_discrete_space(seed):
+    space = ps.ParameterSpace()
+    space.add(ps.IntegerParameter("integ", (0, 2)))
+    batch_shape = torch.Size()
+    opt = SingleObjectiveBOTorchOptimizer(
+        search_space=space,
+        objective=Objective("loss", greater_is_better=False),
+        model=SingleTaskGP(
+            torch.empty((*batch_shape, 0, len(space)), dtype=torch.float64),
+            torch.empty((*batch_shape, 0, 1), dtype=torch.float64),
+        ),
+        acquisition_function_factory=partial(
+            UpperConfidenceBound, beta=1.0, maximize=False
+        ),
+        max_pending_evaluations=5,
+        seed=seed,
+    )
+
+    losses = []
+    for _ in range(10):
+        es = opt.generate_evaluation_specification()
+        loss = es.configuration["integ"] ** 2
+        losses.append(loss)
+        opt.report(es.create_evaluation(objectives={"loss": loss}))
+
+    assert (
+        sum(l == 0 for l in losses) > 5
+    ), "After figuring out the best of the three points, it should only propose that."
+
+
 def test_impute_nans_with_constant():
     x1_no_nans = torch.tensor([[0.1, 0.1], [0.7, 0.2], [1.0, 0.3]])
     x1_some_nans = torch.tensor([[0.1, 0.1], [0.7, float("nan")], [1.0, 0.3]])