Support inequality constraints & X_avoid support in optimize_acqf_dis…

…crete

Summary: These were previously handled in `Acquisition.optimize` in Ax. Pushing it down to BoTorch makes the functionality more broadly useful and helps simplify `Acquisition.optimize` (next diff).

Differential Revision: D64841997

botorch/optim/optimize.py

@@ -1000,6 +1000,9 @@ def optimize_acqf_discrete(
     choices: Tensor,
     max_batch_size: int = 2048,
     unique: bool = True,
+    X_avoid: Tensor | None = None,
+    duplicate_tol: float = 1e-6,
+    inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
 ) -> tuple[Tensor, Tensor]:
     r"""Optimize over a discrete set of points using batch evaluation.
 
@@ -1017,6 +1020,14 @@ def optimize_acqf_discrete(
             a large training set.
         unique: If True return unique choices, o/w choices may be repeated
             (only relevant if `q > 1`).
+        X_avoid: An `n x d` tensor of candidates that we aren't allowed to pick.
+            Any points within `duplicate_tol` of `X_avoid` will be removed from
+            the set of choices.
+        duplicate_tol: The tolerance for deduplication of choices using `X_avoid`.
+        inequality constraints: A list of tuples (indices, coefficients, rhs),
+            with each tuple encoding an inequality constraint of the form
+            `\sum_i (X[indices[i]] * coefficients[i]) >= rhs`.
+            Infeasible points will be removed from the set of choices.
 
     Returns:
         A two-element tuple containing
@@ -1029,8 +1040,33 @@ def optimize_acqf_discrete(
             "Discrete optimization is not supported for"
             "one-shot acquisition functions."
         )
-    if choices.numel() == 0:
-        raise InputDataError("`choices` must be non-emtpy.")
+    # This can be vectorized, but using a for-loop to avoid memory issues.
+    if X_avoid is not None and unique:
+        for x in X_avoid:
+            choices = choices[(choices - x).abs().max(dim=-1).values > duplicate_tol]
+    if inequality_constraints is not None:
+        choices = _filter_infeasible(
+            X=choices, inequality_constraints=inequality_constraints
+        )
+    len_choices = len(choices)
+    if len_choices == 0:
+        message = "`choices` must be non-empty."
+        if X_avoid is not None or inequality_constraints is not None:
+            message += (
+                " No feasible points remain after removing `X_avoid` and "
+                "filtering out infeasible points."
+            )
+        raise InputDataError(message)
+    elif len_choices < q:
+        warnings.warn(
+            (
+                f"Requested {q=} candidates from fully discrete search "
+                f"space, but only {len_choices} possible choices remain. "
+            ),
+            OptimizationWarning,
+            stacklevel=2,
+        )
+        q = len_choices
     choices_batched = choices.unsqueeze(-2)
     if q > 1:
         candidate_list, acq_value_list = [], []

test/optim/test_optimize.py

@@ -22,6 +22,7 @@
     qHypervolumeKnowledgeGradient,
 )
 from botorch.exceptions import InputDataError, UnsupportedError
+from botorch.exceptions.warnings import OptimizationWarning
 from botorch.generation.gen import gen_candidates_scipy, gen_candidates_torch
 from botorch.models import SingleTaskGP
 from botorch.models.model_list_gp_regression import ModelListGP
@@ -1556,7 +1557,7 @@ def test_optimize_acqf_discrete(self):
             mock_acq_function = SquaredAcquisitionFunction()
             mock_acq_function.set_X_pending(None)
             # ensure proper raising of errors if no choices
-            with self.assertRaisesRegex(InputDataError, "`choices` must be non-emtpy."):
+            with self.assertRaisesRegex(InputDataError, "`choices` must be non-empty."):
                 optimize_acqf_discrete(
                     acq_function=mock_acq_function,
                     q=q,
@@ -1613,14 +1614,51 @@ def test_optimize_acqf_discrete(self):
             self.assertAllClose(acq_value, expected_acq_value)
             self.assertAllClose(candidates, expected_candidates)
 
-        with self.assertRaises(UnsupportedError):
-            acqf = MockOneShotAcquisitionFunction()
+        acqf = MockOneShotAcquisitionFunction()
+        with self.assertRaisesRegex(UnsupportedError, "one-shot acquisition"):
             optimize_acqf_discrete(
                 acq_function=acqf,
                 q=1,
                 choices=torch.tensor([[0.5], [0.2]]),
             )
 
+    def test_optimize_acqf_discrete_X_avoid_and_constraints(self):
+        # Check that choices are filtered correctly using X_avoid and constraints.
+        tkwargs: dict[str, Any] = {"device": self.device, "dtype": torch.double}
+        mock_acq_function = SquaredAcquisitionFunction()
+        choices = torch.rand(2, 2, **tkwargs)
+        with self.assertRaisesRegex(InputDataError, "No feasible points"):
+            optimize_acqf_discrete(
+                acq_function=mock_acq_function,
+                q=1,
+                choices=choices,
+                X_avoid=choices,
+            )
+        with self.assertWarnsRegex(OptimizationWarning, "Requested q=2 candidates"):
+            candidates, _ = optimize_acqf_discrete(
+                acq_function=mock_acq_function,
+                q=2,
+                choices=choices,
+                X_avoid=choices[:1],
+            )
+        self.assertAllClose(candidates, choices[1:])
+        constraints = [
+            (  # X[..., 0] >= 1.0
+                torch.tensor([0], dtype=torch.long, device=self.device),
+                torch.tensor([1.0], **tkwargs),
+                1.0,
+            )
+        ]
+        choices[0, 0] = 1.0
+        with self.assertWarnsRegex(OptimizationWarning, "Requested q=2 candidates"):
+            candidates, _ = optimize_acqf_discrete(
+                acq_function=mock_acq_function,
+                q=2,
+                choices=choices,
+                inequality_constraints=constraints,
+            )
+        self.assertAllClose(candidates, choices[:1])
+
     def test_optimize_acqf_discrete_local_search(self):
         for q, dtype in itertools.product((1, 2), (torch.float, torch.double)):
             tkwargs = {"device": self.device, "dtype": dtype}