project point to feasible space via quadratic programming (#3010)

Summary:

see title. This is particularly useful for resolving numerical issues with Ax when it checks parameter constraints.

Differential Revision: D82328877

Author: sdaulton

botorch/optim/optimize.py

@@ -35,7 +35,10 @@
     gen_one_shot_kg_initial_conditions,
     TGenInitialConditions,
 )
-from botorch.optim.parameter_constraints import evaluate_feasibility
+from botorch.optim.parameter_constraints import (
+    evaluate_feasibility,
+    project_to_feasible_space_via_slsqp,
+)
 from botorch.optim.stopping import ExpMAStoppingCriterion
 from torch import Tensor
 
@@ -513,15 +516,47 @@ def _optimize_batch_candidates() -> tuple[Tensor, Tensor, list[Warning]]:
 
     # SLSQP can sometimes fail to produce a feasible candidate. Check for
     # feasibility and error out if necessary.
+    # if there are equality constraints, project the candidate to the feasible set
+    equality_constraints = gen_kwargs.get("equality_constraints")
+    inequality_constraints = gen_kwargs.get("inequality_constraints")
+    nonlinear_inequality_constraints = gen_kwargs.get(
+        "nonlinear_inequality_constraints"
+    )
     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"
-        ),
+        inequality_constraints=inequality_constraints,
+        equality_constraints=equality_constraints,
+        nonlinear_inequality_constraints=nonlinear_inequality_constraints,
     )
     infeasible = ~is_feasible
+    if nonlinear_inequality_constraints is None and infeasible.any():
+        projected_candidates = project_to_feasible_space_via_slsqp(
+            X=batch_candidates[infeasible],
+            bounds=opt_inputs.bounds,
+            equality_constraints=equality_constraints,
+            inequality_constraints=inequality_constraints,
+        )
+        if opt_inputs.post_processing_func is not None:
+            projected_candidates = opt_inputs.post_processing_func(projected_candidates)
+        batch_candidates[infeasible] = projected_candidates
+        # recompute AF values for projected points
+        with torch.no_grad():
+            batch_acq_values[infeasible] = torch.cat(
+                [
+                    opt_inputs.acq_function(cand)
+                    for cand in projected_candidates.split(batch_limit, dim=0)
+                ],
+                dim=0,
+            )
+        # re-evaluate feasibility
+        is_feasible = evaluate_feasibility(
+            X=batch_candidates,
+            inequality_constraints=inequality_constraints,
+            equality_constraints=equality_constraints,
+            nonlinear_inequality_constraints=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. "

botorch/optim/parameter_constraints.py

@@ -18,14 +18,15 @@
 import numpy.typing as npt
 import torch
 from botorch.exceptions.errors import CandidateGenerationError, UnsupportedError
-from scipy.optimize import Bounds
+from botorch.optim.utils import columnwise_clamp
+from scipy.optimize import Bounds, minimize
 from torch import Tensor
 
 
 ScipyConstraintDict = dict[
     str, Union[str, Callable[[np.ndarray], float], Callable[[np.ndarray], np.ndarray]]
 ]
-CONST_TOL = 1e-6
+CONST_TOL = 1e-8
 
 
 def make_scipy_bounds(
@@ -691,3 +692,69 @@ def evaluate_feasibility(
                 tolerance=tolerance,
             )
     return is_feasible
+
+
+def project_to_feasible_space_via_slsqp(
+    X: Tensor,
+    bounds: Tensor,
+    inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
+    equality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
+) -> Tensor:
+    """Project X onto the feasible space by solving a quadratic program.
+
+    This uses SLSQP with gradients to solve the quadratic program.
+
+    Args:
+        X: A `(batch_shape x) n x d`-dim tensor of inptus.
+        bounds: A `2 x d`-dim tensor of lower and upper bounds.
+        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`. `indices` and
+            `coefficients` should be torch tensors. See the docstring of
+            `make_scipy_linear_constraints` for an example. Only intra-point
+            constraints are supported and `indices` should be a 1-d tensor.
+        equality_constraints: A list of tuples (indices, coefficients, rhs).
+
+    Returns:
+        A `(batch_shape x) n x d`-dim tensor of  projected values.
+    """
+    if inequality_constraints is None and equality_constraints is None:
+        return X
+    bounds_scipy = make_scipy_bounds(
+        X=X, lower_bounds=bounds[0], upper_bounds=bounds[1]
+    )
+    constraints = make_scipy_linear_constraints(
+        shapeX=X.shape,
+        inequality_constraints=inequality_constraints,
+        equality_constraints=equality_constraints,
+    )
+    # Define squared distance objective
+    X_np = X.flatten().detach().cpu().numpy()
+
+    def objective(x: np.ndarray):
+        return 0.5 * np.sum((x - X_np) ** 2)
+
+    def grad_objective(x: np.ndarray):
+        return x - X_np
+
+    x0 = (
+        columnwise_clamp(X=X, lower=bounds[0], upper=bounds[1], raise_on_violation=True)
+        .detach()
+        .cpu()
+        .numpy()
+        .flatten()
+    )
+    result = minimize(
+        fun=objective,
+        x0=x0,
+        method="SLSQP",
+        jac=grad_objective,
+        bounds=bounds_scipy,
+        constraints=constraints,
+        tol=CONST_TOL,
+    )
+
+    if not result.success:
+        raise RuntimeError(f"Optimization failed: {result.message}")
+
+    return torch.from_numpy(result.x).to(X).view(X.shape)