Wire equality constraints through generator and acquisition to BoTorch (facebook#5177)

Summary:
Pull Request resolved: facebook#5177

Thread equality constraints from TorchOptConfig through to BoTorch's
optimize_acqf and related optimizers. This is the key diff that connects
Ax's equality constraint representation to BoTorch's SLSQP-based optimizer.

- Add `_to_equality_constraints` in `torch/utils.py` — converts (A, b) tensor
  format to BoTorch's `(indices, coefficients, rhs)` format. No sign negation
  needed (equality is symmetric).
- Update `BoTorchGenerator.gen()` to pass equality constraints to
  `acqf.optimize()`.
- Add `equality_constraints` parameter to `Acquisition.optimize()` and forward
  to `optimize_acqf`, `optimize_acqf_mixed`, `optimize_acqf_mixed_alternating`.
- Raise `ValueError` for discrete optimizers and NSGA-II (unsupported).
- Update `validate_candidates` to check equality constraints.
- Update `_prune_irrelevant_parameters` and `_remove_infeasible_candidates`
  to handle equality constraints.
- Update `Surrogate.best_point` to pass equality constraints.

Reviewed By: bletham

Differential Revision: D100256482

fbshipit-source-id: 7739880ce77a3791587044a1d5c046c59c15e43d

Author: sdaulton

ax/generators/torch/botorch_modular/acquisition.py

@@ -596,6 +596,7 @@ def optimize(
         n: int,
         search_space_digest: SearchSpaceDigest,
         inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
+        equality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
         fixed_features: dict[int, float] | None = None,
         rounding_func: Callable[[Tensor], Tensor] | None = None,
         optimizer_options: dict[str, Any] | None = None,
@@ -612,6 +613,9 @@ def optimize(
             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``.
+            equality_constraints: A list of tuples (indices, coefficients, rhs),
+                with each tuple encoding an equality constraint of the form
+                ``sum_i (X[indices[i]] * coefficients[i]) = rhs``.
             fixed_features: A map `{feature_index: value}` for features that
                 should be fixed to a particular value during generation.
             rounding_func: A function that post-processes an optimization
@@ -664,8 +668,8 @@ def optimize(
         # Ax expects `optimize_acqf` to return tensors of a certain shape.
         if optimizer_options is not None:
             forbidden_optimizer_options = [
-                "equality_constraints",
-                "inequality_constraints",  # These should be constructed by Ax
+                "equality_constraints",  # Constructed by Ax
+                "inequality_constraints",  # Constructed by Ax
                 "batch_initial_conditions",
                 "return_best_only",
                 "return_full_tree",
@@ -716,6 +720,7 @@ def optimize(
                 bounds=bounds,
                 q=n,
                 inequality_constraints=inequality_constraints,
+                equality_constraints=equality_constraints,
                 fixed_features=fixed_features,
                 post_processing_func=rounding_func,
                 acq_function_sequence=self.acq_function_sequence,
@@ -727,6 +732,11 @@ def optimize(
             "optimize_acqf_discrete",
             "optimize_acqf_discrete_local_search",
         ):
+            if equality_constraints:
+                raise ValueError(
+                    "Equality constraints are not supported with discrete "
+                    f"optimizer '{optimizer}'."
+                )
             X_observed = self.X_observed
             if self.X_pending is not None:
                 if X_observed is None:
@@ -805,6 +815,7 @@ def optimize(
                     discrete_choices=discrete_choices
                 ),
                 inequality_constraints=inequality_constraints,
+                equality_constraints=equality_constraints,
                 post_processing_func=rounding_func,
                 **optimizer_options_with_defaults,
             )
@@ -832,9 +843,15 @@ def optimize(
                 post_processing_func=rounding_func,
                 fixed_features=fixed_features,
                 inequality_constraints=inequality_constraints,
+                equality_constraints=equality_constraints,
                 **optimizer_options_with_defaults,
             )
         elif optimizer == "optimize_with_nsgaii":
+            if equality_constraints:
+                raise ValueError(
+                    "Equality constraints are not supported with "
+                    "optimizer 'optimize_with_nsgaii'."
+                )
             if optimize_with_nsgaii is not None:
                 acqf = assert_is_instance(
                     self.acqf, MultiOutputAcquisitionFunctionWrapper
@@ -873,6 +890,7 @@ def optimize(
                     candidates=candidates,
                     search_space_digest=search_space_digest,
                     inequality_constraints=inequality_constraints,
+                    equality_constraints=equality_constraints,
                     fixed_features=fixed_features,
                 )
         # Validate candidates before returning
@@ -883,6 +901,7 @@ def optimize(
             inequality_constraints=inequality_constraints,
             feature_names=search_space_digest.feature_names,
             task_features=search_space_digest.task_features,
+            equality_constraints=equality_constraints,
         )
 
         n_candidates = candidates.shape[0]
@@ -986,6 +1005,7 @@ def _prune_irrelevant_parameters(
         candidates: Tensor,
         search_space_digest: SearchSpaceDigest,
         inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
+        equality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
         fixed_features: dict[int, float] | None = None,
     ) -> tuple[Tensor, Tensor]:
         r"""Prune irrelevant parameters from the candidates using BONSAI.
@@ -1092,6 +1112,7 @@ def _prune_irrelevant_parameters(
                         candidates=pruned_candidates,
                         indices=indices,
                         inequality_constraints=inequality_constraints,
+                        equality_constraints=equality_constraints,
                     )
                     if pruned_candidates.shape[0] == 0:
                         # no feasible points, continue to
@@ -1205,6 +1226,7 @@ def _remove_infeasible_candidates(
     candidates: Tensor,
     indices: Tensor,
     inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
+    equality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
 ) -> tuple[Tensor, Tensor]:
     r"""Filter out infeasible candidates, based on the parameter constraints.
 
@@ -1214,24 +1236,19 @@ def _remove_infeasible_candidates(
             in [0, d-1).
         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. When q=1, or when
-            applying the same constraint to each candidate in the batch
-            (intra-point constraint), `indices` should be a 1-d tensor.
-            For inter-point constraints, in which the constraint is applied to the
-            whole batch of candidates, `indices` must be a 2-d tensor, where
-            in each row `indices[i] =(k_i, l_i)` the first index `k_i` corresponds
-            to the `k_i`-th element of the `q`-batch and the second index `l_i`
-            corresponds to the `l_i`-th feature of that element.
+            `\sum_i (X[indices[i]] * coefficients[i]) >= rhs`.
+        equality_constraints: A list of tuples (indices, coefficients, rhs),
+            with each tuple encoding an equality constraint of the form
+            `\sum_i (X[indices[i]] * coefficients[i]) = rhs`.
 
     Returns:
-        A two-element tuple containing the filter candidates and indices.
+        A two-element tuple containing the filtered candidates and indices.
     """
-    if inequality_constraints is not None:
+    if inequality_constraints is not None or equality_constraints is not None:
         is_feasible = evaluate_feasibility(
             X=candidates,
             inequality_constraints=inequality_constraints,
+            equality_constraints=equality_constraints,
         )
         candidates = candidates[is_feasible]
         indices = indices[is_feasible]

ax/generators/torch/botorch_modular/generator.py

@@ -26,6 +26,7 @@
     ModelConfig,
 )
 from ax.generators.torch.utils import (
+    _to_equality_constraints,
     _to_inequality_constraints,
     get_feature_importances_from_botorch_model,
     get_rounding_func,
@@ -401,6 +402,9 @@ def gen(
             inequality_constraints=_to_inequality_constraints(
                 linear_constraints=torch_opt_config.linear_constraints
             ),
+            equality_constraints=_to_equality_constraints(
+                equality_constraints=torch_opt_config.equality_constraints
+            ),
             fixed_features=torch_opt_config.fixed_features,
             rounding_func=botorch_rounding_func,
             optimizer_options=assert_is_instance(

ax/generators/torch/botorch_modular/surrogate.py

@@ -42,6 +42,7 @@
     use_model_list,
 )
 from ax.generators.torch.utils import (
+    _to_equality_constraints,
     _to_inequality_constraints,
     pick_best_out_of_sample_point_acqf_class,
     predict_from_model,
@@ -1059,6 +1060,9 @@ def best_out_of_sample_point(
             inequality_constraints=_to_inequality_constraints(
                 linear_constraints=torch_opt_config.linear_constraints
             ),
+            equality_constraints=_to_equality_constraints(
+                equality_constraints=torch_opt_config.equality_constraints
+            ),
             fixed_features=torch_opt_config.fixed_features,
         )
         return candidates[0], acqf_value

ax/generators/torch/botorch_modular/utils.py

@@ -892,8 +892,9 @@ def validate_candidates(
     inequality_constraints: list[tuple[Tensor, Tensor, float]] | None,
     feature_names: list[str] | None = None,
     task_features: list[int] | None = None,
+    equality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
 ) -> None:
-    """Validate candidates satisfy bounds, discrete, and inequality constraints.
+    """Validate candidates satisfy bounds, discrete, and linear constraints.
 
     Args:
         candidates: A `n x d`-dim Tensor of candidates to validate.
@@ -906,6 +907,9 @@ def validate_candidates(
         task_features: Optional list of task feature indices to skip discrete value
             validation for. Task features can be fixed to new task values via
             fixed_features that are not in the search space's discrete_choices.
+        equality_constraints: A list of tuples (indices, coefficients, rhs),
+            representing equality constraints of the form
+            `sum_i (X[indices[i]] * coefficients[i]) = rhs`.
 
     Raises:
         CandidateGenerationError: If any candidate violates constraints.
@@ -959,3 +963,17 @@ def validate_candidates(
                 f"Infeasible candidate indices: {infeasible_indices}. "
                 f"Number of constraints: {len(inequality_constraints)}."
             )
+
+    # 4. Equality constraint validation
+    if equality_constraints:
+        is_feasible = evaluate_feasibility(
+            X=candidates.unsqueeze(-2),  # Add q dimension
+            equality_constraints=equality_constraints,
+        )
+        if not is_feasible.all():
+            infeasible_indices = torch.where(~is_feasible)[0].tolist()
+            raise CandidateGenerationError(
+                f"Candidates violate equality constraints. "
+                f"Infeasible candidate indices: {infeasible_indices}. "
+                f"Number of constraints: {len(equality_constraints)}."
+            )