Non-linear inequality constraints

Summary:
See title.

This diff isn't ready for review by any means, but is an attempt to temporarily unblock facebook#769.

Differential Revision: D33853834

fbshipit-source-id: 5c00203ff438226010693cb04a2694c0dbb0cea7

Author: dme65

Diff for: ax/models/torch/botorch.py

@@ -84,6 +84,7 @@
         Tensor,
         Optional[Tuple[Tensor, Tensor]],
         Optional[Tuple[Tensor, Tensor]],
+        Optional[List[Callable]],
         Optional[Dict[int, float]],
         Optional[TConfig],
         Optional[Dict[int, float]],

Diff for: ax/models/torch/botorch_defaults.py

@@ -304,7 +304,9 @@ def scipy_optimizer(
     n: int,
     inequality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None,
     equality_constraints: Optional[List[Tuple[Tensor, Tensor, float]]] = None,
+    nonlinear_inequality_constraints: Optional[List[Callable]] = None,
     fixed_features: Optional[Dict[int, float]] = None,
+    batch_initial_conditions: Optional[Tensor] = None,
     rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
     **kwargs: Any,
 ) -> Tuple[Tensor, Tensor]:
@@ -321,6 +323,14 @@ def scipy_optimizer(
         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`
+        nonlinear_inequality_constraints: A list of callables with that represent
+            non-linear inequality constraints of the form `callable(x) >= 0`. Each
+            callable is expected to take a `(num_restarts) x q x d`-dim tensor as an
+            input and return a `(num_restarts) x q`-dim tensor with the constraint
+            values. The constraints will later be passed to SLSQP. You need to pass in
+            `batch_initial_conditions` in this case.
+        batch_initial_conditions: A tensor to specify the initial conditions. Set
+            this if you do not want to use default initialization strategy.
         fixed_features: A map {feature_index: value} for features that should
             be fixed to a particular value during generation.
         rounding_func: A function that rounds an optimization result
@@ -352,7 +362,9 @@ def scipy_optimizer(
         options=kwargs,
         inequality_constraints=inequality_constraints,
         equality_constraints=equality_constraints,
+        nonlinear_inequality_constraints=nonlinear_inequality_constraints,
         fixed_features=fixed_features,
+        batch_initial_conditions=batch_initial_conditions,
         sequential=sequential,
         post_processing_func=rounding_func,
     )
@@ -365,6 +377,7 @@ def recommend_best_observed_point(
     objective_weights: Tensor,
     outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
     linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
+    nonlinear_inequality_constraints: Optional[List[Callable]] = None,
     fixed_features: Optional[Dict[int, float]] = None,
     model_gen_options: Optional[TConfig] = None,
     target_fidelities: Optional[Dict[int, float]] = None,
@@ -385,6 +398,8 @@ def recommend_best_observed_point(
         linear_constraints: A tuple of (A, b). For k linear constraints on
             d-dimensional x, A is (k x d) and b is (k x 1) such that
             A x <= b.
+        nonlinear_inequality_constraints: A list of callables with that represent
+            non-linear inequality constraints of the form `callable(x) >= 0`.
         fixed_features: A map {feature_index: value} for features that
             should be fixed to a particular value in the best point.
         model_gen_options: A config dictionary that can contain
@@ -396,6 +411,11 @@ def recommend_best_observed_point(
     Returns:
         A d-array of the best point, or None if no feasible point was observed.
     """
+    if nonlinear_inequality_constraints:
+        raise NotImplementedError(
+            "`nonlinear_inequality_constraints` aren't supported by "
+            "`recommend_best_observed_point`."
+        )
     if target_fidelities:
         raise NotImplementedError(
             "target_fidelities not implemented for base BotorchModel"
@@ -421,6 +441,7 @@ def recommend_best_out_of_sample_point(
     objective_weights: Tensor,
     outcome_constraints: Optional[Tuple[Tensor, Tensor]] = None,
     linear_constraints: Optional[Tuple[Tensor, Tensor]] = None,
+    nonlinear_inequality_constraints: Optional[List[Callable]] = None,
     fixed_features: Optional[Dict[int, float]] = None,
     model_gen_options: Optional[TConfig] = None,
     target_fidelities: Optional[Dict[int, float]] = None,
@@ -442,6 +463,8 @@ def recommend_best_out_of_sample_point(
         linear_constraints: A tuple of (A, b). For k linear constraints on
             d-dimensional x, A is (k x d) and b is (k x 1) such that
             A x <= b.
+        nonlinear_inequality_constraints: A list of callables with that represent
+            non-linear inequality constraints of the form `callable(x) >= 0`.
         fixed_features: A map {feature_index: value} for features that
             should be fixed to a particular value in the best point.
         model_gen_options: A config dictionary that can contain
@@ -483,6 +506,8 @@ def recommend_best_out_of_sample_point(
     # (including transforming constraints b/c of fixed features)
     if inequality_constraints is not None:
         raise UnsupportedError("Inequality constraints are not supported!")
+    if nonlinear_inequality_constraints is not None:
+        raise UnsupportedError("Non-linear inequality constraints are not supported!")
 
     return_best_only = optimizer_options.get("return_best_only", True)
     bounds_ = torch.tensor(bounds, dtype=model.dtype, device=model.device)