Update optimize_acqf_homotopy and refactor mixed optimizer dispatch

Summary: see title. This refactors dispatch to mixed optimizers into a shared utility for botorch and ax. This also cleans up optimize_acqf_homotopy by removing fixed_features_list and leveraging the new dispatch util. This adds support for using optimize_acqf_mixed_alternating.

Reviewed By: bletham

Differential Revision: D91913317

Author: sdaulton

botorch/optim/optimize_homotopy.py

@@ -5,7 +5,8 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable
+from collections.abc import Callable, Mapping, Sequence
+from itertools import product
 from typing import Any
 
 import torch
@@ -14,6 +15,10 @@
 from botorch.optim.homotopy import Homotopy
 from botorch.optim.initializers import TGenInitialConditions
 from botorch.optim.optimize import optimize_acqf, optimize_acqf_mixed
+from botorch.optim.optimize_mixed import (
+    optimize_acqf_mixed_alternating,
+    should_use_mixed_alternating_optimizer,
+)
 from torch import Tensor
 
 
@@ -64,7 +69,9 @@ def optimize_acqf_homotopy(
     inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
     equality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
     nonlinear_inequality_constraints: list[tuple[Callable, bool]] | None = None,
-    fixed_features_list: list[dict[int, float]] | None = None,
+    fixed_features: dict[int, float] | None = None,
+    discrete_dims: Mapping[int, Sequence[float]] | None = None,
+    cat_dims: Mapping[int, Sequence[float]] | None = None,
     post_processing_func: Callable[[Tensor], Tensor] | None = None,
     batch_initial_conditions: Tensor | None = None,
     gen_candidates: TGenCandidates | None = None,
@@ -124,10 +131,19 @@ def optimize_acqf_homotopy(
             Using non-linear inequality constraints also requires that ``batch_limit``
             is set to 1, which will be done automatically if not specified in
             ``options``.
-        fixed_features_list: A list of maps ``{feature_index: value}``. The i-th
-            item represents the fixed_feature for the i-th optimization. If
-            ``fixed_features_list`` is provided, ``optimize_acqf_mixed`` is invoked.
-            All indices (``feature_index``) should be non-negative.
+        fixed_features: A map ``{feature_index: value}`` for features that should
+            be fixed to a particular value during generation. Used with
+            ``optimize_acqf`` or ``optimize_acqf_mixed_alternating``.
+        discrete_dims: A dictionary mapping indices of discrete and binary
+            dimensions to a list of allowed values for that dimension. If provided
+            along with ``cat_dims``, the optimizer is chosen based on the number
+            of discrete combinations: ``optimize_acqf_mixed_alternating`` is used
+            if there are more than 10 combinations, otherwise ``optimize_acqf_mixed``
+            is used.
+        cat_dims: A dictionary mapping indices of categorical dimensions
+            to a list of allowed values for that dimension. If provided
+            along with ``discrete_dims``, the optimizer is chosen based on the
+            number of discrete combinations (see ``discrete_dims``).
         post_processing_func: A function that post-processes an optimization
             result appropriately (i.e., according to ``round-trip``
             transformations).
@@ -151,27 +167,69 @@ def optimize_acqf_homotopy(
         ic_gen_kwargs: Additional keyword arguments passed to function specified by
             ``ic_generator``
     """
-    shared_optimize_acqf_kwargs = {
-        "num_restarts": num_restarts,
-        "inequality_constraints": inequality_constraints,
-        "equality_constraints": equality_constraints,
-        "nonlinear_inequality_constraints": nonlinear_inequality_constraints,
-        "return_best_only": False,  # False to make n_restarts persist through homotopy.
-        "gen_candidates": gen_candidates,
-        "ic_generator": ic_generator,
-        "timeout_sec": timeout_sec,
-        "retry_on_optimization_warning": retry_on_optimization_warning,
-        **ic_gen_kwargs,
-    }
-
-    if fixed_features_list and len(fixed_features_list) > 1:
+    # Determine which optimization function to use based on the arguments.
+    # This logic uses the shared `should_use_mixed_alternating_optimizer` utility
+    # which is also used by `determine_optimizer` in Ax.
+    use_mixed_alternating = should_use_mixed_alternating_optimizer(
+        discrete_dims=discrete_dims,
+        cat_dims=cat_dims,
+    )
+
+    if use_mixed_alternating:
+        # Use optimize_acqf_mixed_alternating for mixed discrete/continuous problems
+        # with many discrete combinations.
+        optimization_fn = optimize_acqf_mixed_alternating
+        shared_optimize_acqf_kwargs = {
+            "num_restarts": num_restarts,
+            "inequality_constraints": inequality_constraints,
+            "equality_constraints": equality_constraints,
+            "discrete_dims": discrete_dims,
+            "cat_dims": cat_dims,
+            "fixed_features": fixed_features,
+        }
+        fixed_features_kwargs = {}
+    elif discrete_dims is not None or cat_dims is not None:
+        # Use optimize_acqf_mixed for mixed problems with few discrete combinations.
+        # Build fixed_features_list from discrete_dims and cat_dims.
+
+        all_discrete_dims = {**(discrete_dims or {}), **(cat_dims or {})}
+        dim_indices = sorted(all_discrete_dims.keys())
+        value_lists = [all_discrete_dims[idx] for idx in dim_indices]
+        fixed_features_list = [
+            dict(zip(dim_indices, combo)) for combo in product(*value_lists)
+        ]
         optimization_fn = optimize_acqf_mixed
+        shared_optimize_acqf_kwargs = {
+            "num_restarts": num_restarts,
+            "inequality_constraints": inequality_constraints,
+            "equality_constraints": equality_constraints,
+            "nonlinear_inequality_constraints": nonlinear_inequality_constraints,
+            "return_best_only": False,  # to make n_restarts persist through homotopy.
+            "gen_candidates": gen_candidates,
+            "ic_generator": ic_generator,
+            "timeout_sec": timeout_sec,
+            "retry_on_optimization_warning": retry_on_optimization_warning,
+            **ic_gen_kwargs,
+        }
         fixed_features_kwargs = {"fixed_features_list": fixed_features_list}
     else:
         optimization_fn = optimize_acqf
-        fixed_features_kwargs = {
-            "fixed_features": fixed_features_list[0] if fixed_features_list else None
+        shared_optimize_acqf_kwargs = {
+            "num_restarts": num_restarts,
+            "inequality_constraints": inequality_constraints,
+            "equality_constraints": equality_constraints,
+            "nonlinear_inequality_constraints": nonlinear_inequality_constraints,
+            "return_best_only": False,  # to make n_restarts persist through homotopy.
+            "gen_candidates": gen_candidates,
+            "ic_generator": ic_generator,
+            "timeout_sec": timeout_sec,
+            "retry_on_optimization_warning": retry_on_optimization_warning,
+            **ic_gen_kwargs,
         }
+        if fixed_features is not None:
+            fixed_features_kwargs = {"fixed_features": fixed_features}
+        else:
+            fixed_features_kwargs = {}
 
     candidate_list, acq_value_list = [], []
     if q > 1:
@@ -183,16 +241,34 @@ def optimize_acqf_homotopy(
         homotopy.restart()
 
         while not homotopy.should_stop:
-            candidates, acq_values = optimization_fn(
-                acq_function=acq_function,
-                bounds=bounds,
-                q=1,
-                options=options,
-                batch_initial_conditions=candidates,
-                raw_samples=q_raw_samples,
-                **fixed_features_kwargs,
-                **shared_optimize_acqf_kwargs,
-            )
+            if use_mixed_alternating:
+                # optimize_acqf_mixed_alternating handles its own initialization
+                # and always returns best only, so we get shape (1, d).
+                candidates, acq_values = optimization_fn(
+                    acq_function=acq_function,
+                    bounds=bounds,
+                    q=1,
+                    options=options,
+                    raw_samples=q_raw_samples or raw_samples,
+                    post_processing_func=post_processing_func,
+                    **fixed_features_kwargs,
+                    **shared_optimize_acqf_kwargs,
+                )
+                # Reshape to (1, 1, d) to match optimize_acqf output shape and
+                # ensure acq_values is 1-d for prune_candidates.
+                candidates = candidates.unsqueeze(0)
+                acq_values = acq_values.view(-1)
+            else:
+                candidates, acq_values = optimization_fn(
+                    acq_function=acq_function,
+                    bounds=bounds,
+                    q=1,
+                    options=options,
+                    batch_initial_conditions=candidates,
+                    raw_samples=q_raw_samples,
+                    **fixed_features_kwargs,
+                    **shared_optimize_acqf_kwargs,
+                )
 
             homotopy.step()
 
@@ -208,19 +284,36 @@ def optimize_acqf_homotopy(
             ).unsqueeze(1)
 
         # Optimize one more time with the final options
-        candidates, acq_values = optimization_fn(
-            acq_function=acq_function,
-            bounds=bounds,
-            q=1,
-            options=final_options,
-            raw_samples=q_raw_samples,
-            batch_initial_conditions=candidates,
-            **fixed_features_kwargs,
-            **shared_optimize_acqf_kwargs,
-        )
+        if use_mixed_alternating:
+            candidates, acq_values = optimization_fn(
+                acq_function=acq_function,
+                bounds=bounds,
+                q=1,
+                options=final_options,
+                raw_samples=raw_samples,
+                post_processing_func=post_processing_func,
+                **fixed_features_kwargs,
+                **shared_optimize_acqf_kwargs,
+            )
+            # Reshape to (1, 1, d) to match optimize_acqf output shape.
+            candidates = candidates.unsqueeze(0)
+            acq_values = acq_values.view(-1)
+        else:
+            candidates, acq_values = optimization_fn(
+                acq_function=acq_function,
+                bounds=bounds,
+                q=1,
+                options=final_options,
+                raw_samples=q_raw_samples,
+                batch_initial_conditions=candidates,
+                **fixed_features_kwargs,
+                **shared_optimize_acqf_kwargs,
+            )
 
         # Post-process the candidates and grab the best candidate
-        if post_processing_func is not None:
+        # Note: post_processing_func is already applied within
+        # optimize_acqf_mixed_alternating, so we skip it here in that case.
+        if post_processing_func is not None and not use_mixed_alternating:
             candidates = post_processing_func(candidates)
             acq_values = acq_function(candidates)
 

botorch/optim/optimize_mixed.py

@@ -58,6 +58,14 @@
 # Convergence is defined as the improvements of one discrete, followed by a scalar
 # optimization yield less than ``CONVERGENCE_TOL`` improvements.
 
+# Threshold for choosing between optimize_acqf_mixed and
+# optimize_acqf_mixed_alternating in mixed (not fully discrete) search spaces.
+ALTERNATING_OPTIMIZER_THRESHOLD = 10
+
+# For fully discrete search spaces.
+MAX_CHOICES_ENUMERATE = 10_000
+MAX_CARDINALITY_FOR_LOCAL_SEARCH = 100
+
 SUPPORTED_OPTIONS = {
     "initialization_strategy",
     "tol",
@@ -74,6 +82,41 @@
 SUPPORTED_INITIALIZATION = {"continuous_relaxation", "equally_spaced", "random"}
 
 
+def should_use_mixed_alternating_optimizer(
+    discrete_dims: Mapping[int, Sequence[float]] | None = None,
+    cat_dims: Mapping[int, Sequence[float]] | None = None,
+) -> bool:
+    r"""Determine whether to use ``optimize_acqf_mixed_alternating`` for a mixed
+    (not fully discrete) search space based on the number of discrete combinations.
+
+    For mixed search spaces, if there are more than ``ALTERNATING_OPTIMIZER_THRESHOLD``
+    combinations of discrete choices, we use ``optimize_acqf_mixed_alternating``,
+    which alternates between continuous and discrete optimization steps. Otherwise,
+    we use ``optimize_acqf_mixed``, which enumerates all discrete combinations and
+    optimizes the continuous features with discrete features being fixed.
+
+    Args:
+        discrete_dims: A dictionary mapping indices of discrete (ordinal)
+            dimensions to their respective sets of values provided as a sequence.
+        cat_dims: A dictionary mapping indices of categorical dimensions
+            to their respective sets of values provided as a sequence.
+
+    Returns:
+        ``True`` if ``optimize_acqf_mixed_alternating`` should be used, ``False``
+        if ``optimize_acqf_mixed`` should be used instead.
+    """
+    if discrete_dims is None and cat_dims is None:
+        return False
+
+    n_combos = 1
+    for values in (discrete_dims or {}).values():
+        n_combos *= len(values)
+    for values in (cat_dims or {}).values():
+        n_combos *= len(values)
+
+    return n_combos > ALTERNATING_OPTIMIZER_THRESHOLD
+
+
 def _setup_continuous_relaxation(
     discrete_dims: dict[int, list[float]],
     max_discrete_values: int,