create utility for padded batched pareto frontiers

botorch/acquisition/multi_objective/analytic.py

@@ -31,7 +31,9 @@
 )
 from botorch.exceptions.errors import UnsupportedError
 from botorch.models.model import Model
-from botorch.utils.multi_objective.box_decomposition import NondominatedPartitioning
+from botorch.utils.multi_objective.box_decompositions.non_dominated import (
+    NondominatedPartitioning,
+)
 from botorch.utils.transforms import t_batch_mode_transform
 from torch import Tensor
 from torch.distributions import Normal
@@ -139,7 +141,7 @@ def __init__(
         super().__init__(model=model, objective=objective)
         self.register_buffer("ref_point", ref_point)
         self.partitioning = partitioning
-        cell_bounds = self.partitioning.get_hypercell_bounds(ref_point=self.ref_point)
+        cell_bounds = self.partitioning.get_hypercell_bounds()
         self.register_buffer("cell_lower_bounds", cell_bounds[0])
         self.register_buffer("cell_upper_bounds", cell_bounds[1])
         # create indexing tensor of shape `2^m x m`

botorch/acquisition/multi_objective/monte_carlo.py

@@ -31,7 +31,9 @@
 from botorch.exceptions.errors import UnsupportedError
 from botorch.models.model import Model
 from botorch.sampling.samplers import MCSampler, SobolQMCNormalSampler
-from botorch.utils.multi_objective.box_decomposition import NondominatedPartitioning
+from botorch.utils.multi_objective.box_decompositions.non_dominated import (
+    NondominatedPartitioning,
+)
 from botorch.utils.objective import apply_constraints_nonnegative_soft
 from botorch.utils.torch import BufferDict
 from botorch.utils.transforms import concatenate_pending_points, t_batch_mode_transform
@@ -150,7 +152,7 @@ def __init__(
         self.constraints = constraints
         self.eta = eta
         self.register_buffer("ref_point", ref_point)
-        cell_bounds = partitioning.get_hypercell_bounds(ref_point=self.ref_point)
+        cell_bounds = partitioning.get_hypercell_bounds()
         self.register_buffer("cell_lower_bounds", cell_bounds[0])
         self.register_buffer("cell_upper_bounds", cell_bounds[1])
         self.q = -1

botorch/acquisition/utils.py

@@ -25,7 +25,9 @@
 from botorch.exceptions.warnings import SamplingWarning
 from botorch.models.model import Model
 from botorch.sampling.samplers import IIDNormalSampler, MCSampler, SobolQMCNormalSampler
-from botorch.utils.multi_objective.box_decomposition import NondominatedPartitioning
+from botorch.utils.multi_objective.box_decompositions.non_dominated import (
+    NondominatedPartitioning,
+)
 from botorch.utils.transforms import squeeze_last_dim
 from torch import Tensor
 from torch.quasirandom import SobolEngine
@@ -122,19 +124,23 @@ def get_acquisition_function(
         )
     elif acquisition_function_name == "qEHVI":
         # pyre-fixme [16]: `Model` has no attribute `train_targets`
-        if "ref_point" not in kwargs:
+        try:
+            ref_point = kwargs["ref_point"]
+        except KeyError:
             raise ValueError("`ref_point` must be specified in kwargs for qEHVI")
-        if "Y" not in kwargs:
+        try:
+            Y = kwargs["Y"]
+        except KeyError:
             raise ValueError("`Y` must be specified in kwargs for qEHVI")
-        ref_point = kwargs["ref_point"]
-        Y = kwargs.get("Y")
         # get feasible points
         if constraints is not None:
             feas = torch.stack([c(Y) <= 0 for c in constraints], dim=-1).all(dim=-1)
             Y = Y[feas]
         obj = objective(Y)
         partitioning = NondominatedPartitioning(
-            num_outcomes=obj.shape[-1], Y=obj, alpha=kwargs.get("alpha", 0.0)
+            ref_point=torch.as_tensor(ref_point, dtype=Y.dtype, device=Y.device),
+            Y=obj,
+            alpha=kwargs.get("alpha", 0.0),
         )
         return moo_monte_carlo.qExpectedHypervolumeImprovement(
             model=model,

botorch/utils/multi_objective/__init__.py

@@ -4,7 +4,6 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from botorch.utils.multi_objective.box_decomposition import NondominatedPartitioning
 from botorch.utils.multi_objective.hypervolume import Hypervolume
 from botorch.utils.multi_objective.pareto import is_non_dominated
 from botorch.utils.multi_objective.scalarization import get_chebyshev_scalarization
@@ -14,5 +13,4 @@
     "get_chebyshev_scalarization",
     "is_non_dominated",
     "Hypervolume",
-    "NondominatedPartitioning",
 ]