From de201de716f9cc817666512e114368c1cef75175 Mon Sep 17 00:00:00 2001
From: David Eriksson <deriksson@meta.com>
Date: Tue, 5 Nov 2024 20:49:58 -0800
Subject: [PATCH] Skip fixed feature enumerations in optimize_acqf_mixed that
 can't satisfy the parameter constraints

Summary: When using `optimize_acqf_mixed`, some combinations of fixed features may result in the parameter constraints being impossible to satisfy. This causes `optimize_acqf` to error out. This diff skips the combinations of fixed features where the parameter constraints are impossible to satisfy.

Differential Revision: D65514819
---
 botorch/optim/optimize.py   | 52 +++++++++++++++++++++++++------------
 test/optim/test_optimize.py | 35 +++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 17 deletions(-)

diff --git a/botorch/optim/optimize.py b/botorch/optim/optimize.py
index 61e44d32b6..eb98bc6372 100644
--- a/botorch/optim/optimize.py
+++ b/botorch/optim/optimize.py
@@ -25,6 +25,7 @@
     qHypervolumeKnowledgeGradient,
 )
 from botorch.exceptions import InputDataError, UnsupportedError
+from botorch.exceptions.errors import CandidateGenerationError
 from botorch.exceptions.warnings import OptimizationWarning
 from botorch.generation.gen import gen_candidates_scipy, TGenCandidates
 from botorch.logging import logger
@@ -938,27 +939,44 @@ def optimize_acqf_mixed(
 
     if q == 1:
         ff_candidate_list, ff_acq_value_list = [], []
+        num_candidate_generation_failures = 0
         for fixed_features in fixed_features_list:
-            candidate, acq_value = optimize_acqf(
-                acq_function=acq_function,
-                bounds=bounds,
-                q=q,
-                num_restarts=num_restarts,
-                raw_samples=raw_samples,
-                options=options or {},
-                inequality_constraints=inequality_constraints,
-                equality_constraints=equality_constraints,
-                nonlinear_inequality_constraints=nonlinear_inequality_constraints,
-                fixed_features=fixed_features,
-                post_processing_func=post_processing_func,
-                batch_initial_conditions=batch_initial_conditions,
-                ic_generator=ic_generator,
-                return_best_only=True,
-                **ic_gen_kwargs,
-            )
+            try:
+                candidate, acq_value = optimize_acqf(
+                    acq_function=acq_function,
+                    bounds=bounds,
+                    q=q,
+                    num_restarts=num_restarts,
+                    raw_samples=raw_samples,
+                    options=options or {},
+                    inequality_constraints=inequality_constraints,
+                    equality_constraints=equality_constraints,
+                    nonlinear_inequality_constraints=nonlinear_inequality_constraints,
+                    fixed_features=fixed_features,
+                    post_processing_func=post_processing_func,
+                    batch_initial_conditions=batch_initial_conditions,
+                    ic_generator=ic_generator,
+                    return_best_only=True,
+                    **ic_gen_kwargs,
+                )
+            except CandidateGenerationError:
+                # if candidate generation fails, we skip this candidate
+                num_candidate_generation_failures += 1
+                continue
             ff_candidate_list.append(candidate)
             ff_acq_value_list.append(acq_value)
 
+        if len(ff_candidate_list) == 0:
+            raise CandidateGenerationError(
+                "Candidate generation failed for all `fixed_features`."
+            )
+        elif num_candidate_generation_failures > 0:
+            warnings.warn(
+                "Candidate generation failed for "
+                f"{num_candidate_generation_failures} `fixed_features`.",
+                OptimizationWarning,
+                stacklevel=3,
+            )
         ff_acq_values = torch.stack(ff_acq_value_list)
         best = torch.argmax(ff_acq_values)
         return ff_candidate_list[best], ff_acq_values[best]
diff --git a/test/optim/test_optimize.py b/test/optim/test_optimize.py
index c634ce4ef9..58f41f4a4e 100644
--- a/test/optim/test_optimize.py
+++ b/test/optim/test_optimize.py
@@ -23,6 +23,7 @@
     qHypervolumeKnowledgeGradient,
 )
 from botorch.exceptions import InputDataError, UnsupportedError
+from botorch.exceptions.errors import CandidateGenerationError
 from botorch.exceptions.warnings import OptimizationWarning
 from botorch.generation.gen import gen_candidates_scipy, gen_candidates_torch
 from botorch.models import SingleTaskGP
@@ -1588,6 +1589,40 @@ def test_optimize_acqf_one_shot_large_q(self):
                 raw_samples=10,
             )
 
+    def test_optimize_acqf_mixed_ff_with_constraint(self):
+        mock_acq_function = MockAcquisitionFunction()
+        with self.assertWarnsRegex(
+            OptimizationWarning,
+            "Candidate generation failed for 1 `fixed_features`.",
+        ):
+            optimize_acqf_mixed(
+                acq_function=mock_acq_function,
+                q=1,
+                fixed_features_list=[{0: 0}, {0: 1}],
+                bounds=torch.stack([torch.zeros(3), 4 * torch.ones(3)]),
+                num_restarts=2,
+                raw_samples=10,
+                inequality_constraints=[
+                    (torch.zeros(1), torch.ones(1), 1)
+                ],  # x[0] >= 1
+            )
+        # No fixed features satisfy the constraint
+        with self.assertRaisesRegex(
+            CandidateGenerationError,
+            "Candidate generation failed for all `fixed_features`.",
+        ):
+            optimize_acqf_mixed(
+                acq_function=mock_acq_function,
+                q=1,
+                fixed_features_list=[{0: 0}],
+                bounds=torch.stack([torch.zeros(3), 4 * torch.ones(3)]),
+                num_restarts=2,
+                raw_samples=10,
+                inequality_constraints=[
+                    (torch.zeros(1), torch.ones(1), 1)
+                ],  # x[0] >= 1
+            )
+
 
 class TestOptimizeAcqfDiscrete(BotorchTestCase):
     def test_optimize_acqf_discrete(self):