compute_best_feasible_f utility (#1931)

Summary:
Pull Request resolved: #1931

This commit separates out the `compute_best_feasible_f` function from `qNEI` as a utility function in order to use it in the `input_constructors` and `get_acquisition_function` (follow-up).

Reviewed By: Balandat

Differential Revision: D47365085

fbshipit-source-id: 2d8203544bd646638d9ea6c20789d0700f600c0b

Author: SebastianAment

botorch/acquisition/monte_carlo.py

@@ -27,7 +27,6 @@
 from typing import Any, Callable, List, Optional, Protocol, Tuple, Union
 
 import torch
-from botorch import acquisition
 from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
 from botorch.acquisition.cached_cholesky import CachedCholeskyMCAcquisitionFunction
 from botorch.acquisition.objective import (
@@ -36,7 +35,10 @@
     MCAcquisitionObjective,
     PosteriorTransform,
 )
-from botorch.acquisition.utils import prune_inferior_points
+from botorch.acquisition.utils import (
+    compute_best_feasible_objective,
+    prune_inferior_points,
+)
 from botorch.exceptions.errors import UnsupportedError
 from botorch.models.model import Model
 from botorch.sampling.base import MCSampler
@@ -591,46 +593,24 @@ def _get_samples_and_objectives(self, X: Tensor) -> Tuple[Tensor, Tensor]:
         return samples, obj
 
     def _compute_best_feasible_objective(self, samples: Tensor, obj: Tensor) -> Tensor:
-        """
+        r"""Computes best feasible objective value from samples.
+
         Args:
             samples: `sample_shape x batch_shape x q x m`-dim posterior samples.
             obj: A `sample_shape x batch_shape x q`-dim Tensor of MC objective values.
 
         Returns:
             A `sample_shape x batch_shape x 1`-dim Tensor of best feasible objectives.
         """
-        if self._constraints is not None:
-            # is_feasible is sample_shape x batch_shape x q
-            is_feasible = compute_smoothed_constraint_indicator(
-                constraints=self._constraints, samples=samples, eta=self._eta
-            )
-            is_feasible = is_feasible > 0.5  # due to smooth approximation
-            if is_feasible.any():
-                obj = torch.where(is_feasible, obj, -torch.inf)
-            else:  # if there are no feasible observations, estimate a lower
-                # bound on the objective by sampling convex combinations of X_baseline.
-                convex_weights = torch.rand(
-                    32,
-                    self.X_baseline.shape[-2],
-                    dtype=self.X_baseline.dtype,
-                    device=self.X_baseline.device,
-                )
-                weights_sum = convex_weights.sum(dim=0, keepdim=True)
-                convex_weights = convex_weights / weights_sum
-                # infeasible cost M is such that -M < min_x f(x), thus
-                # 0 < min_x f(x) - (-M), so we should take -M as a lower
-                # bound on the best feasible objective
-                return -acquisition.utils.get_infeasible_cost(
-                    X=convex_weights @ self.X_baseline,
-                    model=self.model,
-                    objective=self.objective,
-                    posterior_transform=self.posterior_transform,
-                ).expand(*obj.shape[:-1], 1)
-
-        # we don't need to differentiate through X_baseline for now, so taking
-        # the regular max over the n points to get best_f is fine
-        with torch.no_grad():
-            return obj.amax(dim=-1, keepdim=True)
+        return compute_best_feasible_objective(
+            samples=samples,
+            obj=obj,
+            constraints=self._constraints,
+            model=self.model,
+            objective=self.objective,
+            posterior_transform=self.posterior_transform,
+            X_baseline=self.X_baseline,
+        )
 
 
 class qProbabilityOfImprovement(SampleReducingMCAcquisitionFunction):

botorch/acquisition/utils.py

@@ -32,6 +32,7 @@
     FastNondominatedPartitioning,
     NondominatedPartitioning,
 )
+from botorch.utils.objective import compute_feasibility_indicator
 from botorch.utils.sampling import optimize_posterior_samples
 from botorch.utils.transforms import is_fully_bayesian
 from torch import Tensor
@@ -213,6 +214,113 @@ def get_acquisition_function(
     )
 
 
+def compute_best_feasible_objective(
+    samples: Tensor,
+    obj: Tensor,
+    constraints: Optional[List[Callable[[Tensor], Tensor]]],
+    model: Optional[Model] = None,
+    objective: Optional[MCAcquisitionObjective] = None,
+    posterior_transform: Optional[PosteriorTransform] = None,
+    X_baseline: Optional[Tensor] = None,
+    infeasible_obj: Optional[Tensor] = None,
+) -> Tensor:
+    """Computes the largest `obj` value that is feasible under the `constraints`. If
+    `constraints` is None, returns the best unconstrained objective value.
+
+    When no feasible observations exist and `infeasible_obj` is not `None`, returns
+    `infeasible_obj` (potentially reshaped). When no feasible observations exist and
+    `infeasible_obj` is `None`, uses `model`, `objective`, `posterior_transform`, and
+    `X_baseline` to infer and return an `infeasible_obj` `M` s.t. `M < min_x f(x)`.
+
+    Args:
+        samples: `(sample_shape) x batch_shape x q x m`-dim posterior samples.
+        obj: A `(sample_shape) x batch_shape x q`-dim Tensor of MC objective values.
+        constraints: A list of constraint callables which map posterior samples to
+            a scalar. The associated constraint is considered satisfied if this
+            scalar is less than zero.
+        model: A Model, only required when there are no feasible observations.
+        objective: An MCAcquisitionObjective, only optionally used when there are no
+            feasible observations.
+        posterior_transform: A PosteriorTransform, only optionally used when there are
+            no feasible observations.
+        X_baseline: A `batch_shape x d`-dim Tensor of baseline points, only required
+            when there are no feasible observations.
+        infeasible_obj: A Tensor to be returned when no feasible points exist.
+
+    Returns:
+        A `(sample_shape) x batch_shape x 1`-dim Tensor of best feasible objectives.
+    """
+    if constraints is None:  # unconstrained case
+        # we don't need to differentiate through X_baseline for now, so taking
+        # the regular max over the n points to get best_f is fine
+        with torch.no_grad():
+            return obj.amax(dim=-1, keepdim=True)
+
+    is_feasible = compute_feasibility_indicator(
+        constraints=constraints, samples=samples
+    )  # sample_shape x batch_shape x q
+    if is_feasible.any():
+        obj = torch.where(is_feasible, obj, -torch.inf)
+        with torch.no_grad():
+            return obj.amax(dim=-1, keepdim=True)
+
+    elif infeasible_obj is not None:
+        return infeasible_obj.expand(*obj.shape[:-1], 1)
+
+    else:
+        if model is None:
+            raise ValueError(
+                "Must specify `model` when no feasible observation exists."
+            )
+        if X_baseline is None:
+            raise ValueError(
+                "Must specify `X_baseline` when no feasible observation exists."
+            )
+        return _estimate_objective_lower_bound(
+            model=model,
+            objective=objective,
+            posterior_transform=posterior_transform,
+            X=X_baseline,
+        ).expand(*obj.shape[:-1], 1)
+
+
+def _estimate_objective_lower_bound(
+    model: Model,
+    objective: Optional[MCAcquisitionObjective],
+    posterior_transform: Optional[PosteriorTransform],
+    X: Tensor,
+) -> Tensor:
+    """Estimates a lower bound on the objective values by evaluating the model at convex
+    combinations of `X`, returning the 6-sigma lower bound of the computed statistics.
+
+    Args:
+        model: A fitted model.
+        objective: An MCAcquisitionObjective with `m` outputs.
+        posterior_transform: A PosteriorTransform.
+        X: A `n x d`-dim Tensor of design points from which to draw convex combinations.
+
+    Returns:
+        A `m`-dimensional Tensor of lower bounds of the objectives.
+    """
+    convex_weights = torch.rand(
+        32,
+        X.shape[-2],
+        dtype=X.dtype,
+        device=X.device,
+    )
+    weights_sum = convex_weights.sum(dim=0, keepdim=True)
+    convex_weights = convex_weights / weights_sum
+    # infeasible cost M is such that -M < min_x f(x), thus
+    # 0 < min_x f(x) - (-M), so we should take -M as a lower
+    # bound on the best feasible objective
+    return -get_infeasible_cost(
+        X=convex_weights @ X,
+        model=model,
+        objective=objective,
+        posterior_transform=posterior_transform,
+    )
+
+
 def get_infeasible_cost(
     X: Tensor,
     model: Model,

botorch/utils/objective.py

@@ -95,14 +95,37 @@ def apply_constraints_nonnegative_soft(
     return obj.clamp_min(0).mul(w)  # Enforce non-negativity of obj, apply constraints.
 
 
+def compute_feasibility_indicator(
+    constraints: Optional[List[Callable[[Tensor], Tensor]]],
+    samples: Tensor,
+) -> Tensor:
+    r"""Computes the feasibility of a list of constraints given posterior samples.
+
+    Args:
+        constraints: A list of callables, each mapping a batch_shape x q x m`-dim Tensor
+            to a `batch_shape x q`-dim Tensor, where negative values imply feasibility.
+        samples: A batch_shape x q x m`-dim Tensor of posterior samples.
+
+    Returns:
+        A `batch_shape x q`-dim tensor of Boolean feasibility values.
+    """
+    ind = torch.ones(samples.shape[:-1], dtype=torch.bool, device=samples.device)
+    if constraints is not None:
+        for constraint in constraints:
+            ind = ind.logical_and(constraint(samples) < 0)
+    return ind
+
+
 def compute_smoothed_constraint_indicator(
     constraints: List[Callable[[Tensor], Tensor]],
     samples: Tensor,
     eta: Union[Tensor, float],
 ) -> Tensor:
-    r"""Computes the feasibility indicator of a list of constraints given posterior
-    samples, using a sigmoid to smoothly approximate the feasibility indicator
-    of each individual constraint to ensure differentiability and high gradient signal.
+    r"""Computes the smoothed feasibility indicator of a list of constraints.
+
+    Given posterior samples, using a sigmoid to smoothly approximate the feasibility
+    indicator of each individual constraint to ensure differentiability and high
+    gradient signal.
 
     Args:
         constraints: A list of callables, each mapping a Tensor of size `b x q x m`

test/acquisition/test_utils.py

@@ -19,6 +19,7 @@
     ScalarizedPosteriorTransform,
 )
 from botorch.acquisition.utils import (
+    compute_best_feasible_objective,
     expand_trace_observations,
     get_acquisition_function,
     get_infeasible_cost,
@@ -575,7 +576,83 @@ def test_GetUnknownAcquisitionFunction(self):
             )
 
 
-class TestGetInfeasibleCost(BotorchTestCase):
+class TestConstraintUtils(BotorchTestCase):
+    def test_compute_best_feasible_objective(self):
+        for dtype in (torch.float, torch.double):
+            with self.subTest(dtype=dtype):
+                tkwargs = {"dtype": dtype, "device": self.device}
+                n = 5
+                X = torch.arange(n, **tkwargs).view(-1, 1)
+                means = torch.arange(n, **tkwargs).view(-1, 1)
+                samples = means
+                variances = torch.tensor(
+                    [0.09, 0.25, 0.36, 0.25, 0.09], **tkwargs
+                ).view(-1, 1)
+                mm = MockModel(
+                    MockPosterior(mean=means, variance=variances, samples=samples)
+                )
+
+                # testing all feasible points
+                obj = means.squeeze(-1)
+                constraints = [lambda samples: -torch.ones_like(samples[..., 0])]
+                best_f = compute_best_feasible_objective(
+                    samples=means, obj=obj, constraints=constraints
+                )
+                self.assertAllClose(best_f, obj.amax(dim=-1, keepdim=True))
+
+                # testing with some infeasible points
+                con_cutoff = 3.0
+                best_f = compute_best_feasible_objective(
+                    samples=means,
+                    obj=obj,
+                    constraints=[
+                        lambda samples: samples[..., 0] - (con_cutoff + 1 / 2)
+                    ],
+                )
+                # only first three points are feasible
+                self.assertAllClose(best_f, torch.tensor([con_cutoff], **tkwargs))
+
+                # testing with no feasible points and infeasible obj
+                infeasible_obj = torch.tensor(torch.pi, **tkwargs)
+                best_f = compute_best_feasible_objective(
+                    samples=means,
+                    obj=obj,
+                    constraints=[lambda X: torch.ones_like(X[..., 0])],
+                    infeasible_obj=infeasible_obj,
+                )
+                self.assertAllClose(best_f, infeasible_obj.unsqueeze(0))
+
+                # testing with no feasible points and not infeasible obj
+                def objective(Y, X):
+                    return Y.squeeze(-1) - 5.0
+
+                best_f = compute_best_feasible_objective(
+                    samples=means,
+                    obj=obj,
+                    constraints=[lambda X: torch.ones_like(X[..., 0])],
+                    model=mm,
+                    X_baseline=X,
+                    objective=objective,
+                )
+                self.assertAllClose(
+                    best_f, -get_infeasible_cost(X=X, model=mm, objective=objective)
+                )
+
+                with self.assertRaisesRegex(ValueError, "Must specify `model`"):
+                    best_f = compute_best_feasible_objective(
+                        samples=means,
+                        obj=obj,
+                        constraints=[lambda X: torch.ones_like(X[..., 0])],
+                        X_baseline=X,
+                    )
+                with self.assertRaisesRegex(ValueError, "Must specify `X_baseline`"):
+                    best_f = compute_best_feasible_objective(
+                        samples=means,
+                        obj=obj,
+                        constraints=[lambda X: torch.ones_like(X[..., 0])],
+                        model=mm,
+                    )
+
     def test_get_infeasible_cost(self):
         for dtype in (torch.float, torch.double):
             tkwargs = {"dtype": dtype, "device": self.device}