Add Inequality Constraints to SEBO

ax/benchmark/tests/problems/test_mixed_integer_problems.py

@@ -61,7 +61,7 @@ def test_problems(self) -> None:
         cases: list[tuple[BenchmarkProblem, dict[str, float], torch.Tensor]] = [
             (
                 get_discrete_hartmann(),
-                {f"x{i+1}": 0.0 for i in range(6)},
+                {f"x{i + 1}": 0.0 for i in range(6)},
                 torch.zeros(6, dtype=torch.double),
             ),
             (
@@ -71,28 +71,28 @@ def test_problems(self) -> None:
             ),
             (
                 get_discrete_ackley(),
-                {f"x{i+1}": 0.0 for i in range(13)},
+                {f"x{i + 1}": 0.0 for i in range(13)},
                 torch.zeros(13, dtype=torch.double),
             ),
             (
                 get_discrete_ackley(),
                 {
-                    **{f"x{i+1}": 2 for i in range(0, 5)},
-                    **{f"x{i+1}": 4 for i in range(5, 10)},
-                    **{f"x{i+1}": 1.0 for i in range(10, 13)},
+                    **{f"x{i + 1}": 2 for i in range(0, 5)},
+                    **{f"x{i + 1}": 4 for i in range(5, 10)},
+                    **{f"x{i + 1}": 1.0 for i in range(10, 13)},
                 },
                 torch.ones(13, dtype=torch.double),
             ),
             (
                 get_discrete_rosenbrock(),
-                {f"x{i+1}": 0.0 for i in range(10)},
+                {f"x{i + 1}": 0.0 for i in range(10)},
                 torch.full((10,), -5.0, dtype=torch.double),
             ),
             (
                 get_discrete_rosenbrock(),
                 {
-                    **{f"x{i+1}": 3 for i in range(0, 6)},
-                    **{f"x{i+1}": 1.0 for i in range(6, 10)},
+                    **{f"x{i + 1}": 3 for i in range(0, 6)},
+                    **{f"x{i + 1}": 1.0 for i in range(6, 10)},
                 },
                 torch.full((10,), 10.0, dtype=torch.double),
             ),

ax/core/search_space.py

@@ -714,7 +714,7 @@ def _find_applicable_parameters(root: Parameter) -> set[str]:
         ):
             raise RuntimeError(
                 error_msg_prefix
-                + f"Parameters {applicable_paramers- set(parameters.keys())} are"
+                + f"Parameters {applicable_paramers - set(parameters.keys())} are"
                 " missing."
             )
 

ax/modelbridge/tests/test_prediction_utils.py

@@ -162,5 +162,5 @@ def _attach_completed_trials(ax_client: AxClient) -> None:
 
 # Test metric evaluation method
 def _evaluate_test_metrics(parameters: TParameterization) -> TEvaluationOutcome:
-    x = np.array([parameters.get(f"x{i+1}") for i in range(2)])
+    x = np.array([parameters.get(f"x{i + 1}") for i in range(2)])
     return {"test_metric1": (x[0] / x[1], 0.0), "test_metric2": (x[0] + x[1], 0.0)}

ax/models/torch/botorch_modular/sebo.py

@@ -233,15 +233,11 @@ def optimize(
             with the weight for each candidate.
         """
         if self.penalty_name == "L0_norm":
-            if inequality_constraints is not None:
-                raise NotImplementedError(
-                    "Homotopy does not support optimization with inequality "
-                    + "constraints. Use L1 penalty norm instead."
-                )
             candidates, expected_acquisition_value, weights = (
                 self._optimize_with_homotopy(
                     n=n,
                     search_space_digest=search_space_digest,
+                    inequality_constraints=inequality_constraints,
                     fixed_features=fixed_features,
                     rounding_func=rounding_func,
                     optimizer_options=optimizer_options,
@@ -269,6 +265,7 @@ def _optimize_with_homotopy(
         n: int,
         search_space_digest: SearchSpaceDigest,
         fixed_features: dict[int, float] | None = None,
+        inequality_constraints: list[tuple[Tensor, Tensor, float]] | None = None,
         rounding_func: Callable[[Tensor], Tensor] | None = None,
         optimizer_options: dict[str, Any] | None = None,
     ) -> tuple[Tensor, Tensor, Tensor]:
@@ -296,21 +293,31 @@ def _optimize_with_homotopy(
                 )
             ],
         )
-        batch_initial_conditions = get_batch_initial_conditions(
-            acq_function=self.acqf,
-            raw_samples=optimizer_options_with_defaults["raw_samples"],
-            X_pareto=self.acqf.X_baseline,
-            target_point=self.target_point,
-            bounds=bounds,
-            num_restarts=optimizer_options_with_defaults["num_restarts"],
-        )
+        if inequality_constraints is None:
+            batch_initial_conditions = get_batch_initial_conditions(
+                acq_function=self.acqf,
+                raw_samples=optimizer_options_with_defaults["raw_samples"],
+                X_pareto=self.acqf.X_baseline,
+                target_point=self.target_point,
+                bounds=bounds,
+                num_restarts=optimizer_options_with_defaults["num_restarts"],
+            )
+        else:
+            warnings.warn(
+                "Adopting Botorch's default for sampling initial conditions as "
+                "inequality constraints are not supported for custom SEBO "
+                "initialization based on pareto front perturbation."
+            )
+            batch_initial_conditions = None
+
         candidates, expected_acquisition_value = optimize_acqf_homotopy(
             q=n,
             acq_function=self.acqf,
             bounds=bounds,
             homotopy=homotopy,
             num_restarts=optimizer_options_with_defaults["num_restarts"],
             raw_samples=optimizer_options_with_defaults["raw_samples"],
+            inequality_constraints=inequality_constraints,
             post_processing_func=rounding_func,
             fixed_features=fixed_features,
             batch_initial_conditions=batch_initial_conditions,

ax/models/torch/tests/test_sebo.py

@@ -222,7 +222,7 @@ def test_optimize_l1(self, mock_optimize_acqf: Mock) -> None:
             optimizer_options=self.optimizer_options,
         )
 
-        args, kwargs = mock_optimize_acqf.call_args
+        kwargs = mock_optimize_acqf.call_args.kwargs
         self.assertEqual(kwargs["acq_function"], acquisition.acqf)
         self.assertEqual(kwargs["q"], 2)
         self.assertEqual(kwargs["inequality_constraints"], self.inequality_constraints)
@@ -248,7 +248,7 @@ def test_optimize_l0(self, mock_optimize_acqf_homotopy: Mock) -> None:
             optimizer_options=self.optimizer_options,
         )
 
-        args, kwargs = mock_optimize_acqf_homotopy.call_args
+        kwargs = mock_optimize_acqf_homotopy.call_args.kwargs
         self.assertEqual(kwargs["acq_function"], acquisition.acqf)
         self.assertEqual(kwargs["q"], 2)
         self.assertEqual(kwargs["post_processing_func"], self.rounding_func)
@@ -270,7 +270,7 @@ def test_optimize_l0(self, mock_optimize_acqf_homotopy: Mock) -> None:
             rounding_func=self.rounding_func,
             optimizer_options=self.optimizer_options,
         )
-        args, kwargs = mock_optimize_acqf_homotopy.call_args
+        kwargs = mock_optimize_acqf_homotopy.call_args.kwargs
         self.assertEqual(kwargs["acq_function"], acquisition2.acqf)
         self.assertEqual(kwargs["q"], 2)
         self.assertEqual(kwargs["post_processing_func"], self.rounding_func)
@@ -282,19 +282,22 @@ def test_optimize_l0(self, mock_optimize_acqf_homotopy: Mock) -> None:
             fixed_features=self.fixed_features,
             options={"penalty": "L0_norm", "target_point": self.target_point},
         )
-        with self.assertRaisesRegex(
-            NotImplementedError,
-            "Homotopy does not support optimization with inequality "
-            "constraints. Use L1 penalty norm instead.",
-        ):
-            acquisition.optimize(
-                n=2,
-                search_space_digest=self.search_space_digest,
-                inequality_constraints=self.inequality_constraints,
-                fixed_features=self.fixed_features,
-                rounding_func=self.rounding_func,
-                optimizer_options=self.optimizer_options,
-            )
+        acquisition.optimize(
+            n=2,
+            search_space_digest=self.search_space_digest,
+            inequality_constraints=self.inequality_constraints,
+            fixed_features=self.fixed_features,
+            rounding_func=self.rounding_func,
+            optimizer_options=self.optimizer_options,
+        )
+        kwargs = mock_optimize_acqf_homotopy.call_args.kwargs
+        self.assertEqual(kwargs["acq_function"], acquisition.acqf)
+        self.assertEqual(kwargs["q"], 2)
+        self.assertEqual(kwargs["inequality_constraints"], self.inequality_constraints)
+        self.assertEqual(kwargs["post_processing_func"], self.rounding_func)
+        self.assertEqual(kwargs["batch_initial_conditions"], None)
+        self.assertEqual(kwargs["num_restarts"], self.optimizer_options["num_restarts"])
+        self.assertEqual(kwargs["raw_samples"], self.optimizer_options["raw_samples"])
 
     def test_clamp_to_target(self) -> None:
         X = torch.tensor(

ax/service/tests/test_global_stopping.py

@@ -47,7 +47,7 @@ def get_ax_client_for_branin(
 
     def evaluate(self, parameters: TParameterization) -> dict[str, tuple[float, float]]:
         """Evaluates the parameters for branin experiment."""
-        x = np.array([parameters.get(f"x{i+1}") for i in range(2)])
+        x = np.array([parameters.get(f"x{i + 1}") for i in range(2)])
         # pyre-fixme[7]: Expected `Dict[str, Tuple[float, float]]` but got
         #  `Dict[str, Tuple[Union[float, ndarray], float]]`.
         return {"branin": (branin(x), 0.0)}

ax/service/tests/test_interactive_loop.py

@@ -30,7 +30,7 @@ def _elicit(
             parameterization_with_trial_index: tuple[TParameterization, int],
         ) -> tuple[int, TEvaluationOutcome] | None:
             parameterization, trial_index = parameterization_with_trial_index
-            x = np.array([parameterization.get(f"x{i+1}") for i in range(6)])
+            x = np.array([parameterization.get(f"x{i + 1}") for i in range(6)])
 
             return (
                 trial_index,
@@ -89,7 +89,7 @@ def _elicit(
             parameterization, trial_index = parameterization_with_trial_index
             time.sleep(0.15)  # Sleep to induce MaxParallelismException in loop
 
-            x = np.array([parameterization.get(f"x{i+1}") for i in range(6)])
+            x = np.array([parameterization.get(f"x{i + 1}") for i in range(6)])
 
             return (
                 trial_index,

ax/utils/sensitivity/sobol_measures.py

@@ -941,12 +941,12 @@ def _get_torch_model(
     """
     if not isinstance(model_bridge, TorchModelBridge):
         raise NotImplementedError(
-            f"{type(model_bridge) = }, but only TorchModelBridge is supported."
+            f"{type(model_bridge)=}, but only TorchModelBridge is supported."
         )
     model = model_bridge.model  # should be of type TorchModel
     if not (isinstance(model, BotorchModel) or isinstance(model, ModularBoTorchModel)):
         raise NotImplementedError(
-            f"{type(model_bridge.model) = }, but only "
+            f"{type(model_bridge.model)=}, but only "
             "Union[BotorchModel, ModularBoTorchModel] is supported."
         )
     return model

scripts/insert_api_refs.py

@@ -44,7 +44,7 @@ def replace_backticks(source_path, docs_path):
             for i, l in enumerate(lines):
                 match = re.search(pattern, l)
                 if match:
-                    print(f"{f}:{i+1} s/{match.group(0)}/{link}")
+                    print(f"{f}:{i + 1} s/{match.group(0)}/{link}")
                     lines[i] = re.sub(pattern, link, l)
             open(f, "w").writelines(lines)