Add input constructor for LogCEI (#2973)

Summary:
Pull Request resolved: #2973

Similar to LogPOF, the input constructor converts a constraints_tuple (A,b) to a dictionary of input constraints. Also infers best_f from a training set, similar to LogEI.

The input constructor does not throw an error if objective_index lies in the constraints, as this is already handled by the LogCEI constructor (this is included as a test case).

Reviewed By: SebastianAment

Differential Revision: D80183196

fbshipit-source-id: dc6490fdc23da340765efdd13c40f792d4ef90af

Author: Ethan Che

botorch/acquisition/input_constructors.py

@@ -20,6 +20,7 @@
 from botorch.acquisition.active_learning import qNegIntegratedPosteriorVariance
 from botorch.acquisition.analytic import (
     ExpectedImprovement,
+    LogConstrainedExpectedImprovement,
     LogExpectedImprovement,
     LogNoisyExpectedImprovement,
     LogProbabilityOfFeasibility,
@@ -361,28 +362,63 @@ def construct_inputs_pof(
     Returns:
         A dict mapping kwarg names of the constructor to values.
     """
-    # Construct a dictionary of the form `{i: [lower, upper]}`,
-    # where `i` is the output index, and `lower` and `upper` are
-    # lower and upper bounds on that output (resp. interpreted
-    # as -Inf / Inf if None).
-    weights, bounds = constraints_tuple
-    constraints_dict = {}
-    for w, b in zip(weights, bounds):
-        nonzero_w = w.nonzero()
-        if nonzero_w.numel() != 1:
-            raise BotorchError(
-                "LogProbabilityOfFeasibility only support constraints on single"
-                " outcomes."
-            )
-        i = nonzero_w.item()
-        w_i = w[i]
-        is_ub = torch.sign(w_i) == 1.0
-        b = b.item()
-        bounds = (None, b / w_i) if is_ub else (b / w_i, None)
-        constraints_dict[i] = bounds
+    # Construct a constraint dictionary from constraint_tuple
+    constraints_dict = _construct_constraint_dict_from_tuple(
+        constraints_tuple, LogProbabilityOfFeasibility
+    )
+
     return {"model": model, "constraints": constraints_dict}
 
 
+@acqf_input_constructor(LogConstrainedExpectedImprovement)
+def construct_inputs_logcei(
+    model: Model,
+    training_data: MaybeDict[SupervisedDataset],
+    objective_index: int,
+    constraints_tuple: tuple[Tensor, Tensor],
+    best_f: float | Tensor | None = None,
+    maximize: bool = True,
+) -> dict[str, Any]:
+    r"""Construct kwargs for the log constrained expected improvement
+    acquisition function.
+
+    Args:
+        model: The model to be used in the acquisition function.
+        training_data: Dataset(s) used to train the model.
+            Used to determine default value for `best_f`.
+        objective_index: The index of the objective.
+        constraints_tuple: A tuple of `(A, b)`. For `k` outcome constraints
+            and `m` outputs at `f(x)``, `A` is `k x m` and `b` is `k x 1` such
+            that `A f(x) <= b`.
+        best_f: Either a scalar or a `b`-dim Tensor (batch mode) representing
+                the best feasible function value observed so far (assumed noiseless).
+        maximize: If True, consider the problem a maximization problem.
+
+    Returns:
+        A dict mapping kwarg names of the constructor to values.
+    """
+
+    # If no best_f provided, compute it from the training data
+    # For LogCEI, posterior_transform is not used.
+    if best_f is None:
+        best_f = get_best_f_analytic(
+            training_data=training_data,
+        )
+
+    # Construct a constraint dictionary from constraint_tuple
+    constraints_dict = _construct_constraint_dict_from_tuple(
+        constraints_tuple, LogConstrainedExpectedImprovement
+    )
+
+    return {
+        "model": model,
+        "best_f": best_f,
+        "objective_index": objective_index,
+        "constraints": constraints_dict,
+        "maximize": maximize,
+    }
+
+
 @acqf_input_constructor(UpperConfidenceBound)
 def construct_inputs_ucb(
     model: Model,
@@ -1984,3 +2020,30 @@ def _get_ref_point(
         ref_point = objective(objective_thresholds)
 
     return ref_point
+
+
+def _construct_constraint_dict_from_tuple(
+    constraints_tuple: tuple, acqf_class: type[AcquisitionFunction]
+) -> dict[str, Any]:
+    """
+    Construct a dictionary of the form `{i: [lower, upper]}`,
+    where `i` is the output index, and `lower` and `upper` are
+    lower and upper bounds on that output (resp. interpreted
+    as -Inf / Inf if None).
+    """
+    weights, bounds = constraints_tuple
+    constraints_dict = {}
+    for w, b in zip(weights, bounds):
+        nonzero_w = w.nonzero()
+        if nonzero_w.numel() != 1:
+            raise BotorchError(
+                f"{acqf_class.__name__} only support constraints on single outcomes."
+            )
+        i = nonzero_w.item()
+        w_i = w[i]
+        is_ub = torch.sign(w_i) == 1.0
+        b = b.item()
+        bounds = (None, b / w_i) if is_ub else (b / w_i, None)
+        constraints_dict[i] = bounds
+
+    return constraints_dict

test/acquisition/test_input_constructors.py

@@ -24,6 +24,7 @@
 from botorch.acquisition.active_learning import qNegIntegratedPosteriorVariance
 from botorch.acquisition.analytic import (
     ExpectedImprovement,
+    LogConstrainedExpectedImprovement,
     LogExpectedImprovement,
     LogNoisyExpectedImprovement,
     LogProbabilityOfFeasibility,
@@ -475,6 +476,72 @@ def test_construct_inputs_noisy_ei(self) -> None:
                 with self.assertRaisesRegex(ValueError, "Field `X` must be shared"):
                     c(model=mock_model, training_data=self.multiX_multiY)
 
+    def test_construct_inputs_LogCEI(self) -> None:
+        c = get_acqf_input_constructor(LogConstrainedExpectedImprovement)
+        mock_model = self.mock_model
+        constraints_tuple = [torch.tensor([[0.0, 1.0]]), torch.tensor([[2.0]])]
+        constraints = {1: (None, 2.0)}
+        best_f_expected = self.blockX_blockY[0].Y.squeeze().max()
+        objective_index = 0
+        # test that best_f is inferred from training data
+        # test constraint tuple
+        kwargs = c(
+            model=mock_model,
+            objective_index=objective_index,
+            training_data=self.blockX_blockY,
+            constraints_tuple=constraints_tuple,
+            maximize=False,
+        )
+        self.assertEqual(
+            set(kwargs.keys()),
+            {"model", "best_f", "objective_index", "constraints", "maximize"},
+        )
+        self.assertIs(kwargs["model"], mock_model)
+        self.assertEqual(kwargs["objective_index"], objective_index)
+        self.assertEqual(kwargs["constraints"], constraints)
+        self.assertEqual(kwargs["best_f"], best_f_expected)
+        self.assertFalse(kwargs["maximize"])
+        # test that best_f overrides default from training data
+        # test that negative constraints work
+        constraints_tuple = [torch.tensor([[0.0, -1.0]]), torch.tensor([[-2.0]])]
+        constraints = {1: (2.0, None)}
+        kwargs = c(
+            model=mock_model,
+            objective_index=objective_index,
+            training_data=self.blockX_blockY,
+            best_f=0.1,
+            constraints_tuple=constraints_tuple,
+        )
+        self.assertIs(kwargs["model"], mock_model)
+        self.assertEqual(kwargs["objective_index"], objective_index)
+        self.assertEqual(kwargs["constraints"], constraints)
+        self.assertEqual(kwargs["best_f"], 0.1)
+        self.assertTrue(kwargs["maximize"])
+        # test that constraints on multiple outcomes raises an exception
+        with self.assertRaisesRegex(
+            BotorchError,
+            "LogConstrainedExpectedImprovement only support constraints on single"
+            " outcomes.",
+        ):
+            c(
+                model=mock_model,
+                objective_index=objective_index,
+                training_data=self.blockX_blockY,
+                constraints_tuple=[torch.tensor([[1.0, 1.0]]), torch.tensor([[2.0]])],
+            )
+        # test that if objective_index coincides with constraints raises a value error
+        with self.assertRaisesRegex(
+            ValueError,
+            "Output corresponding to objective should not be a constraint.",
+        ):
+            kwargs = c(
+                model=mock_model,
+                objective_index=1,
+                training_data=self.blockX_blockY,
+                constraints_tuple=[torch.tensor([[0.0, -1.0]]), torch.tensor([[-2.0]])],
+            )
+            LogConstrainedExpectedImprovement(**kwargs)
+
     def test_construct_inputs_eubo(self) -> None:
         """test input constructor for analytical EUBO and MC qEUBO"""
 
@@ -1845,6 +1912,17 @@ def setUp(self, suppress_input_warnings: bool = True) -> None:
             },
         )
 
+        self.cases["LogCEI"] = (
+            [LogConstrainedExpectedImprovement],
+            {
+                "model": st_soo_model,
+                "objective_index": 0,
+                "training_data": self.blockX_blockY,
+                "constraints": {0: [-5, 5]},
+                **constraints_tuple_dict,
+            },
+        )
+
         def constraint(X: Tensor) -> Tensor:
             return X[..., 0].abs() - 5