return feasible candidate if there is one and return_best_only=True (#2778)

Summary:
Pull Request resolved: #2778

as title.

If `return_best_only=True`, we only need one candidate that satisfies the parameter constraints. In that case, we shouldn't error out if we did find a feasible point.

Reviewed By: Balandat

Differential Revision: D71509592

fbshipit-source-id: 4a1dae0db371452a0e05752c6b0128d0b3cd6df6

Author: sdaulton

botorch/generation/gen.py

@@ -20,16 +20,12 @@
 import numpy.typing as npt
 import torch
 from botorch.acquisition import AcquisitionFunction
-from botorch.exceptions.errors import (
-    CandidateGenerationError,
-    OptimizationGradientError,
-)
+from botorch.exceptions.errors import OptimizationGradientError
 from botorch.exceptions.warnings import OptimizationWarning
 from botorch.generation.utils import _remove_fixed_features_from_optimization
 from botorch.logging import logger
 from botorch.optim.parameter_constraints import (
     _arrayify,
-    evaluate_feasibility,
     make_scipy_bounds,
     make_scipy_linear_constraints,
     make_scipy_nonlinear_inequality_constraints,
@@ -264,23 +260,6 @@ def f(x):
         fixed_features=fixed_features,
     )
 
-    # SLSQP can sometimes fail to produce a feasible candidate. Check for
-    # feasibility and error out if necessary.
-    if not (
-        is_feasible := evaluate_feasibility(
-            X=candidates,
-            inequality_constraints=inequality_constraints,
-            equality_constraints=equality_constraints,
-            nonlinear_inequality_constraints=nonlinear_inequality_constraints,
-        )
-    ).all():
-        raise CandidateGenerationError(
-            f"The {method} optimizer produced infeasible candidates. "
-            f"{(~is_feasible).sum().item()} out of {is_feasible.numel()} batches "
-            "of candidates were infeasible. Please make sure the constraints are "
-            "satisfiable and relax them if needed. "
-        )
-
     clamped_candidates = columnwise_clamp(
         X=candidates, lower=lower_bounds, upper=upper_bounds, raise_on_violation=True
     )

botorch/optim/optimize.py

@@ -35,6 +35,7 @@
     gen_one_shot_kg_initial_conditions,
     TGenInitialConditions,
 )
+from botorch.optim.parameter_constraints import evaluate_feasibility
 from botorch.optim.stopping import ExpMAStoppingCriterion
 from torch import Tensor
 
@@ -354,6 +355,15 @@ def _optimize_acqf_batch(opt_inputs: OptimizeAcqfInputs) -> tuple[Tensor, Tensor
         ),
     )
 
+    gen_kwargs = {}
+    for constraint_name in [
+        "inequality_constraints",
+        "equality_constraints",
+        "nonlinear_inequality_constraints",
+    ]:
+        if (constraint := getattr(opt_inputs, constraint_name)) is not None:
+            gen_kwargs[constraint_name] = constraint
+
     def _optimize_batch_candidates() -> tuple[Tensor, Tensor, list[Warning]]:
         batch_candidates_list: list[Tensor] = []
         batch_acq_values_list: list[Tensor] = []
@@ -370,15 +380,6 @@ def _optimize_batch_candidates() -> tuple[Tensor, Tensor, list[Warning]]:
         upper_bounds = None if bounds[1].isinf().all() else bounds[1]
         gen_options = {k: v for k, v in options.items() if k not in INIT_OPTION_KEYS}
 
-        gen_kwargs = {}
-        for constraint_name in [
-            "inequality_constraints",
-            "equality_constraints",
-            "nonlinear_inequality_constraints",
-        ]:
-            if (constraint := getattr(opt_inputs, constraint_name)) is not None:
-                gen_kwargs[constraint_name] = constraint
-
         for i, batched_ics_ in enumerate(batched_ics):
             # optimize using random restart optimization
             with warnings.catch_warnings(record=True) as ws:
@@ -471,7 +472,27 @@ def _optimize_batch_candidates() -> tuple[Tensor, Tensor, list[Warning]]:
             ]
             batch_acq_values = torch.cat(acq_values_list, dim=0)
 
+    # SLSQP can sometimes fail to produce a feasible candidate. Check for
+    # feasibility and error out if necessary.
+    is_feasible = evaluate_feasibility(
+        X=batch_candidates,
+        inequality_constraints=gen_kwargs.get("inequality_constraints"),
+        equality_constraints=gen_kwargs.get("equality_constraints"),
+        nonlinear_inequality_constraints=gen_kwargs.get(
+            "nonlinear_inequality_constraints"
+        ),
+    )
+    infeasible = ~is_feasible
+    if (opt_inputs.return_best_only and (not is_feasible.any())) or infeasible.all():
+        raise CandidateGenerationError(
+            f"The optimizer produced infeasible candidates. "
+            f"{(~is_feasible).sum().item()} out of {is_feasible.numel()} batches "
+            "of candidates were infeasible. Please make sure the constraints are "
+            "satisfiable and relax them if needed. "
+        )
     if opt_inputs.return_best_only:
+        # filter for feasible candidates
+        batch_acq_values[infeasible] = -float("inf")
         best = torch.argmax(batch_acq_values.view(-1), dim=0)
         batch_candidates = batch_candidates[best]
         batch_acq_values = batch_acq_values[best]

test/generation/test_gen.py

@@ -11,10 +11,7 @@
 
 import torch
 from botorch.acquisition import qExpectedImprovement, qKnowledgeGradient
-from botorch.exceptions.errors import (
-    CandidateGenerationError,
-    OptimizationGradientError,
-)
+from botorch.exceptions.errors import OptimizationGradientError
 from botorch.exceptions.warnings import OptimizationWarning
 from botorch.fit import fit_gpytorch_mll
 from botorch.generation.gen import (
@@ -389,27 +386,6 @@ def test_gen_candidates_scipy_invalid_method(self) -> None:
                 upper_bounds=1,
             )
 
-    def test_gen_candidates_scipy_infeasible_candidates(self) -> None:
-        # Check for error when infeasible candidates are generated.
-        ics = torch.rand(2, 3, 1, device=self.device)
-        with mock.patch(
-            "botorch.generation.gen.minimize_with_timeout",
-            return_value=OptimizeResult(x=ics.view(-1).cpu().numpy()),
-        ), self.assertRaisesRegex(
-            CandidateGenerationError, "infeasible candidates. 2 out of 2"
-        ):
-            gen_candidates_scipy(
-                initial_conditions=ics,
-                acquisition_function=MockAcquisitionFunction(),
-                inequality_constraints=[
-                    (  # X[..., 0] >= 2.0, which is infeasible.
-                        torch.tensor([0], device=self.device),
-                        torch.tensor([1.0], device=self.device),
-                        2.0,
-                    )
-                ],
-            )
-
 
 class TestRandomRestartOptimization(TestBaseCandidateGeneration):
     def test_random_restart_optimization(self):

test/optim/test_optimize.py

@@ -341,7 +341,7 @@ def test_optimize_acqf_sequential(
                 gcs_return_vals = [
                     (
                         torch.tensor(
-                            [[[1.1, 2.1, 3.1]]], device=self.device, dtype=dtype
+                            [[[1.1, 2.1, 4.0]]], device=self.device, dtype=dtype
                         ),
                         torch.tensor([i], device=self.device, dtype=dtype),
                     )
@@ -357,11 +357,20 @@ def test_optimize_acqf_sequential(
                 if mock_gen_candidates is mock_gen_candidates_scipy:
                     # x[2] * 4 >= 5
                     inequality_constraints = [
-                        (torch.tensor([2]), torch.tensor([4]), torch.tensor(5))
+                        (
+                            torch.tensor([2], dtype=torch.long, device=self.device),
+                            torch.tensor([4], device=self.device),
+                            torch.tensor(5, device=self.device),
+                        )
                     ]
                     equality_constraints = [
-                        (torch.tensor([0, 1]), torch.ones(2), torch.tensor(4.0))
+                        (
+                            torch.tensor([0, 1], dtype=torch.long, device=self.device),
+                            torch.ones(2, device=self.device),
+                            torch.tensor(16.0, device=self.device),
+                        )
                     ]
+                    equality_constraints = None
                 # gen_candidates_torch does not support constraints
                 else:
                     inequality_constraints = None
@@ -1183,6 +1192,77 @@ def __call__(self, x, f):
         self.assertEqual(f_obj(x2), 2.0)
         self.assertEqual(f_np_wrapper.call_count, 2)
 
+    def _test_optimize_acqf_infeasible_candidates(
+        self, mock_gen_batch_initial_conditions, q, num_restarts, ics
+    ):
+        mock_acq_function = MockAcquisitionFunction()
+        mock_gen_batch_initial_conditions.side_effect = [ics for _ in range(2)]
+        with mock.patch(
+            "botorch.generation.gen.minimize_with_timeout",
+            return_value=OptimizeResult(x=ics.view(-1).cpu().numpy()),
+        ):
+            candidates, _ = optimize_acqf(
+                acq_function=mock_acq_function,
+                bounds=torch.tensor(
+                    [[0.0], [3.0]], dtype=ics.dtype, device=self.device
+                ),
+                q=q,
+                num_restarts=num_restarts,
+                raw_samples=1,
+                inequality_constraints=[
+                    (  # X[..., 0] >= 2.0, which is infeasible.
+                        torch.tensor([0], device=self.device),
+                        torch.tensor([1.0], dtype=ics.dtype, device=self.device),
+                        2.0,
+                    )
+                ],
+                sequential=False,
+                gen_candidates=gen_candidates_scipy,
+                initial_conditions=ics,
+                acquisition_function=MockAcquisitionFunction(),
+            )
+            return candidates
+
+    @mock.patch("botorch.optim.optimize.gen_batch_initial_conditions")
+    def test_optimize_acqf_all_infeasible_candidates(
+        self, mock_gen_batch_initial_conditions
+    ) -> None:
+        # Check for error when all batches of candidates are infeasible w.r.t
+        # parameter constraints.
+        q = 3
+        num_restarts = 2
+        for dtype in (torch.float, torch.double):
+            with self.assertRaisesRegex(
+                CandidateGenerationError, "infeasible candidates. 2 out of 2"
+            ):
+                self._test_optimize_acqf_infeasible_candidates(
+                    mock_gen_batch_initial_conditions=mock_gen_batch_initial_conditions,
+                    q=q,
+                    num_restarts=num_restarts,
+                    ics=torch.rand(num_restarts, q, 1, dtype=dtype, device=self.device),
+                )
+
+    @mock.patch("botorch.optim.optimize.gen_batch_initial_conditions")
+    def test_optimize_acqf_some_infeasible_candidates(
+        self, mock_gen_batch_initial_conditions
+    ) -> None:
+        # Check that no error is raised when the first batch of candidates
+        # contains points that are infeasible w.r.t parameter constraints, but
+        # second batch is feasible
+        q = 3
+        num_restarts = 2
+
+        for dtype in (torch.float, torch.double):
+            ics = torch.rand(num_restarts, q, 1, dtype=dtype, device=self.device)
+            ics[1, ..., 0] = 3.0  # make second batch feasible
+            candidates = self._test_optimize_acqf_infeasible_candidates(
+                mock_gen_batch_initial_conditions=mock_gen_batch_initial_conditions,
+                q=q,
+                num_restarts=num_restarts,
+                ics=ics,
+            )
+            self.assertTrue(torch.equal(candidates, ics[1]))
+
 
 class TestAllOptimizers(BotorchTestCase):
     def test_raises_with_negative_fixed_features(self) -> None:
@@ -1708,7 +1788,9 @@ def test_optimize_acqf_one_shot_large_q(self):
     def test_optimize_acqf_mixed_ff_with_constraint(self):
         mock_acq_function = MockAcquisitionFunction()
         bounds = torch.stack([torch.zeros(3), 4 * torch.ones(3)])
-        ineq_constraints = [(torch.zeros(1), torch.ones(1), 1)]  # x[0] >= 1
+        ineq_constraints = [
+            (torch.zeros(1, dtype=torch.long), torch.ones(1), 1)
+        ]  # x[0] >= 1
         with self.assertWarnsRegex(
             OptimizationWarning,
             "Candidate generation failed for 1 combinations of `fixed_features`. "