Rework botorch/acquisition/multi_objective directory structure

botorch/acquisition/multi_objective/__init__.py

@@ -4,19 +4,23 @@
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
 
-from botorch.acquisition.multi_objective.analytic import (
-    ExpectedHypervolumeImprovement,
+from botorch.acquisition.multi_objective.analytic import ExpectedHypervolumeImprovement
+from botorch.acquisition.multi_objective.base import (
     MultiObjectiveAnalyticAcquisitionFunction,
+    MultiObjectiveMCAcquisitionFunction,
 )
 from botorch.acquisition.multi_objective.hypervolume_knowledge_gradient import (
     qHypervolumeKnowledgeGradient,
     qMultiFidelityHypervolumeKnowledgeGradient,
 )
+from botorch.acquisition.multi_objective.logei import (
+    qLogExpectedHypervolumeImprovement,
+    qLogNoisyExpectedHypervolumeImprovement,
+)
 from botorch.acquisition.multi_objective.max_value_entropy_search import (
     qMultiObjectiveMaxValueEntropy,
 )
 from botorch.acquisition.multi_objective.monte_carlo import (
-    MultiObjectiveMCAcquisitionFunction,
     qExpectedHypervolumeImprovement,
     qNoisyExpectedHypervolumeImprovement,
 )
@@ -33,18 +37,20 @@
 
 
 __all__ = [
+    "ExpectedHypervolumeImprovement",
     "get_default_partitioning_alpha",
+    "IdentityMCMultiOutputObjective",
+    "MCMultiOutputObjective",
+    "MOMF",
+    "MultiObjectiveAnalyticAcquisitionFunction",
+    "MultiObjectiveMCAcquisitionFunction",
     "prune_inferior_points_multi_objective",
     "qExpectedHypervolumeImprovement",
     "qHypervolumeKnowledgeGradient",
+    "qLogExpectedHypervolumeImprovement",
+    "qLogNoisyExpectedHypervolumeImprovement",
     "qMultiFidelityHypervolumeKnowledgeGradient",
-    "qNoisyExpectedHypervolumeImprovement",
-    "MOMF",
     "qMultiObjectiveMaxValueEntropy",
-    "ExpectedHypervolumeImprovement",
-    "IdentityMCMultiOutputObjective",
-    "MCMultiOutputObjective",
-    "MultiObjectiveAnalyticAcquisitionFunction",
-    "MultiObjectiveMCAcquisitionFunction",
+    "qNoisyExpectedHypervolumeImprovement",
     "WeightedMCMultiOutputObjective",
 ]

botorch/acquisition/multi_objective/analytic.py

@@ -19,14 +19,14 @@
 
 from __future__ import annotations
 
-from abc import abstractmethod
 from itertools import product
 from typing import Optional
 
 import torch
-from botorch.acquisition.acquisition import AcquisitionFunction
+from botorch.acquisition.multi_objective.base import (
+    MultiObjectiveAnalyticAcquisitionFunction,
+)
 from botorch.acquisition.objective import PosteriorTransform
-from botorch.exceptions.errors import UnsupportedError
 from botorch.models.model import Model
 from botorch.utils.multi_objective.box_decompositions.non_dominated import (
     NondominatedPartitioning,
@@ -36,45 +36,6 @@
 from torch.distributions import Normal
 
 
-class MultiObjectiveAnalyticAcquisitionFunction(AcquisitionFunction):
-    r"""Abstract base class for Multi-Objective batch acquisition functions."""
-
-    def __init__(
-        self,
-        model: Model,
-        posterior_transform: Optional[PosteriorTransform] = None,
-    ) -> None:
-        r"""Constructor for the MultiObjectiveAnalyticAcquisitionFunction base class.
-
-        Args:
-            model: A fitted model.
-            posterior_transform: A PosteriorTransform (optional).
-        """
-        super().__init__(model=model)
-        if posterior_transform is None or isinstance(
-            posterior_transform, PosteriorTransform
-        ):
-            self.posterior_transform = posterior_transform
-        else:
-            raise UnsupportedError(
-                "Only a posterior_transform of type PosteriorTransform is "
-                "supported for Multi-Objective analytic acquisition functions."
-            )
-
-    @abstractmethod
-    def forward(self, X: Tensor) -> Tensor:
-        r"""Takes in a `batch_shape x 1 x d` X Tensor of t-batches with `1` `d`-dim
-        design point each, and returns a Tensor with shape `batch_shape'`, where
-        `batch_shape'` is the broadcasted batch shape of model and input `X`.
-        """
-        pass  # pragma: no cover
-
-    def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
-        raise UnsupportedError(
-            "Analytic acquisition functions do not account for X_pending yet."
-        )
-
-
 class ExpectedHypervolumeImprovement(MultiObjectiveAnalyticAcquisitionFunction):
     def __init__(
         self,

botorch/acquisition/multi_objective/base.py

@@ -0,0 +1,153 @@
+#!/usr/bin/env python3
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+r"""
+Base classes for multi-objective acquisition functions.
+"""
+
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import Callable, Optional, Union
+
+import torch
+from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
+from botorch.acquisition.multi_objective.objective import (
+    IdentityMCMultiOutputObjective,
+    MCMultiOutputObjective,
+)
+from botorch.acquisition.objective import PosteriorTransform
+from botorch.exceptions.errors import UnsupportedError
+from botorch.models.model import Model
+from botorch.models.transforms.input import InputPerturbation
+from botorch.sampling.base import MCSampler
+from torch import Tensor
+
+
+class MultiObjectiveAnalyticAcquisitionFunction(AcquisitionFunction):
+    r"""Abstract base class for Multi-Objective batch acquisition functions."""
+
+    def __init__(
+        self,
+        model: Model,
+        posterior_transform: Optional[PosteriorTransform] = None,
+    ) -> None:
+        r"""Constructor for the MultiObjectiveAnalyticAcquisitionFunction base class.
+
+        Args:
+            model: A fitted model.
+            posterior_transform: A PosteriorTransform (optional).
+        """
+        super().__init__(model=model)
+        if posterior_transform is None or isinstance(
+            posterior_transform, PosteriorTransform
+        ):
+            self.posterior_transform = posterior_transform
+        else:
+            raise UnsupportedError(
+                "Only a posterior_transform of type PosteriorTransform is "
+                "supported for Multi-Objective analytic acquisition functions."
+            )
+
+    @abstractmethod
+    def forward(self, X: Tensor) -> Tensor:
+        r"""Takes in a `batch_shape x 1 x d` X Tensor of t-batches with `1` `d`-dim
+        design point each, and returns a Tensor with shape `batch_shape'`, where
+        `batch_shape'` is the broadcasted batch shape of model and input `X`.
+        """
+        pass  # pragma: no cover
+
+    def set_X_pending(self, X_pending: Optional[Tensor] = None) -> None:
+        raise UnsupportedError(
+            "Analytic acquisition functions do not account for X_pending yet."
+        )
+
+
+class MultiObjectiveMCAcquisitionFunction(AcquisitionFunction, MCSamplerMixin, ABC):
+    r"""Abstract base class for Multi-Objective batch acquisition functions.
+
+    NOTE: This does not inherit from `MCAcquisitionFunction` to avoid circular imports.
+
+    Args:
+        _default_sample_shape: The `sample_shape` for the default sampler.
+    """
+
+    _default_sample_shape = torch.Size([128])
+
+    def __init__(
+        self,
+        model: Model,
+        sampler: Optional[MCSampler] = None,
+        objective: Optional[MCMultiOutputObjective] = None,
+        constraints: Optional[list[Callable[[Tensor], Tensor]]] = None,
+        eta: Union[Tensor, float] = 1e-3,
+        X_pending: Optional[Tensor] = None,
+    ) -> None:
+        r"""Constructor for the `MultiObjectiveMCAcquisitionFunction` base class.
+
+        Args:
+            model: A fitted model.
+            sampler: The sampler used to draw base samples. If not given,
+                a sampler is generated using `get_sampler`.
+                NOTE: For posteriors that do not support base samples,
+                a sampler compatible with intended use case must be provided.
+                See `ForkedRNGSampler` and `StochasticSampler` as examples.
+            objective: The MCMultiOutputObjective under which the samples are
+                evaluated. Defaults to `IdentityMCMultiOutputObjective()`.
+            constraints: A list of callables, each mapping a Tensor of dimension
+                `sample_shape x batch-shape x q x m` to a Tensor of dimension
+                `sample_shape x batch-shape x q`, where negative values imply
+                feasibility.
+            eta: The temperature parameter for the sigmoid function used for the
+                differentiable approximation of the constraints. In case of a float the
+                same eta is used for every constraint in constraints. In case of a
+                tensor the length of the tensor must match the number of provided
+                constraints. The i-th constraint is then estimated with the i-th
+                eta value.
+            X_pending:  A `m x d`-dim Tensor of `m` design points that have
+                points that have been submitted for function evaluation
+                but have not yet been evaluated.
+        """
+        super().__init__(model=model)
+        MCSamplerMixin.__init__(self, sampler=sampler)
+        if objective is None:
+            objective = IdentityMCMultiOutputObjective()
+        elif not isinstance(objective, MCMultiOutputObjective):
+            raise UnsupportedError(
+                "Only objectives of type MCMultiOutputObjective are supported for "
+                "Multi-Objective MC acquisition functions."
+            )
+        if (
+            hasattr(model, "input_transform")
+            and isinstance(model.input_transform, InputPerturbation)
+            and constraints is not None
+        ):
+            raise UnsupportedError(
+                "Constraints are not supported with input perturbations, due to"
+                "sample q-batch shape being different than that of the inputs."
+                "Use a composite objective that applies feasibility weighting to"
+                "samples before calculating the risk measure."
+            )
+        self.add_module("objective", objective)
+        self.constraints = constraints
+        if constraints:
+            if type(eta) is not Tensor:
+                eta = torch.full((len(constraints),), eta)
+            self.register_buffer("eta", eta)
+        self.X_pending = None
+        if X_pending is not None:
+            self.set_X_pending(X_pending)
+
+    @abstractmethod
+    def forward(self, X: Tensor) -> Tensor:
+        r"""Takes in a `batch_shape x q x d` X Tensor of t-batches with `q` `d`-dim
+        design points each, and returns a Tensor with shape `batch_shape'`, where
+        `batch_shape'` is the broadcasted batch shape of model and input `X`. Should
+        utilize the result of `set_X_pending` as needed to account for pending function
+        evaluations.
+        """
+        pass  # pragma: no cover

botorch/acquisition/multi_objective/hypervolume_knowledge_gradient.py

@@ -29,8 +29,8 @@
 from botorch.acquisition.cost_aware import CostAwareUtility
 from botorch.acquisition.decoupled import DecoupledAcquisitionFunction
 from botorch.acquisition.knowledge_gradient import ProjectedAcquisitionFunction
+from botorch.acquisition.multi_objective.base import MultiObjectiveMCAcquisitionFunction
 from botorch.acquisition.multi_objective.monte_carlo import (
-    MultiObjectiveMCAcquisitionFunction,
     qExpectedHypervolumeImprovement,
 )
 from botorch.acquisition.multi_objective.objective import MCMultiOutputObjective

botorch/acquisition/multi_objective/joint_entropy_search.py

@@ -25,7 +25,6 @@
 from botorch import settings
 from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
 from botorch.exceptions.errors import UnsupportedError
-
 from botorch.models.model import Model
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.utils import fantasize as fantasize_flag

botorch/acquisition/multi_objective/logei.py

@@ -14,7 +14,7 @@
 
 import torch
 from botorch.acquisition.logei import TAU_MAX, TAU_RELU
-from botorch.acquisition.multi_objective import MultiObjectiveMCAcquisitionFunction
+from botorch.acquisition.multi_objective.base import MultiObjectiveMCAcquisitionFunction
 from botorch.acquisition.multi_objective.objective import MCMultiOutputObjective
 from botorch.models.model import Model
 from botorch.sampling.base import MCSampler
@@ -60,7 +60,7 @@ def __init__(
         objective: Optional[MCMultiOutputObjective] = None,
         constraints: Optional[list[Callable[[Tensor], Tensor]]] = None,
         X_pending: Optional[Tensor] = None,
-        eta: Optional[Union[Tensor, float]] = 1e-2,
+        eta: Union[Tensor, float] = 1e-2,
         fat: bool = True,
         tau_relu: float = TAU_RELU,
         tau_max: float = TAU_MAX,
@@ -333,7 +333,7 @@ def __init__(
         objective: Optional[MCMultiOutputObjective] = None,
         constraints: Optional[list[Callable[[Tensor], Tensor]]] = None,
         X_pending: Optional[Tensor] = None,
-        eta: Optional[Union[Tensor, float]] = 1e-3,
+        eta: Union[Tensor, float] = 1e-3,
         prune_baseline: bool = False,
         alpha: float = 0.0,
         cache_pending: bool = True,

botorch/acquisition/multi_objective/max_value_entropy_search.py

@@ -20,17 +20,14 @@
 from __future__ import annotations
 
 from math import pi
-
 from typing import Callable, Optional, Union
 
 import torch
 from botorch.acquisition.max_value_entropy_search import qMaxValueEntropy
+from botorch.acquisition.multi_objective.base import MultiObjectiveMCAcquisitionFunction
 from botorch.acquisition.multi_objective.joint_entropy_search import (
     LowerBoundMultiObjectiveEntropySearch,
 )
-from botorch.acquisition.multi_objective.monte_carlo import (
-    MultiObjectiveMCAcquisitionFunction,
-)
 from botorch.models.converter import (
     batched_multi_output_to_single_output,
     model_list_to_batched,