Deprecating objective in favor of posterior_transform for `MultiO…

…bjectiveAnalyticAcquisitionFunction`

Summary: This commit deprecates `AnalyticMultiOutputObjective` and replaces it with a posterior transform.

Differential Revision: D43047711

fbshipit-source-id: 99b51e226842d61551b95baf5152cf6efb5dd29b

botorch/acquisition/multi_objective/analytic.py

@@ -25,10 +25,7 @@
 
 import torch
 from botorch.acquisition.acquisition import AcquisitionFunction
-from botorch.acquisition.multi_objective.objective import (
-    AnalyticMultiOutputObjective,
-    IdentityAnalyticMultiOutputObjective,
-)
+from botorch.acquisition.objective import IdentityPosteriorTransform, PosteriorTransform
 from botorch.exceptions.errors import UnsupportedError
 from botorch.models.model import Model
 from botorch.utils.multi_objective.box_decompositions.non_dominated import (
@@ -43,23 +40,30 @@ class MultiObjectiveAnalyticAcquisitionFunction(AcquisitionFunction):
     r"""Abstract base class for Multi-Objective batch acquisition functions."""
 
     def __init__(
-        self, model: Model, objective: Optional[AnalyticMultiOutputObjective] = None
+        self,
+        model: Model,
+        posterior_transform: Optional[PosteriorTransform] = None,
+        **kwargs,
     ) -> None:
         r"""Constructor for the MultiObjectiveAnalyticAcquisitionFunction base class.
 
         Args:
             model: A fitted model.
-            objective: An AnalyticMultiOutputObjective (optional).
+            posterior_transform: A PosteriorTransform (optional).
         """
         super().__init__(model=model)
-        if objective is None:
-            objective = IdentityAnalyticMultiOutputObjective()
-        elif not isinstance(objective, AnalyticMultiOutputObjective):
+        posterior_transform = self._deprecate_acqf_objective(
+            posterior_transform=posterior_transform,
+            objective=kwargs.get("objective"),
+        )
+        if posterior_transform is None:
+            posterior_transform = IdentityPosteriorTransform()
+        elif not isinstance(posterior_transform, PosteriorTransform):
             raise UnsupportedError(
-                "Only objectives of type AnalyticMultiOutputObjective are supported "
-                "for Multi-Objective analytic acquisition functions."
+                "Only a posterior_transform of type PosteriorTransform is "
+                "supported for Multi-Objective analytic acquisition functions."
             )
-        self.objective = objective
+        self.posterior_transform = posterior_transform
 
     @abstractmethod
     def forward(self, X: Tensor) -> Tensor:
@@ -81,7 +85,8 @@ def __init__(
         model: Model,
         ref_point: List[float],
         partitioning: NondominatedPartitioning,
-        objective: Optional[AnalyticMultiOutputObjective] = None,
+        posterior_transform: Optional[PosteriorTransform] = None,
+        **kwargs,
     ) -> None:
         r"""Expected Hypervolume Improvement supporting m>=2 outcomes.
 
@@ -118,7 +123,7 @@ def __init__(
             partitioning: A `NondominatedPartitioning` module that provides the non-
                 dominated front and a partitioning of the non-dominated space in hyper-
                 rectangles.
-            objective: An `AnalyticMultiOutputObjective`.
+            posterior_transform: A `PosteriorTransform`.
         """
         # TODO: we could refactor this __init__ logic into a
         # HypervolumeAcquisitionFunction Mixin
@@ -138,7 +143,7 @@ def __init__(
             raise ValueError(
                 "At least one pareto point must be better than the reference point."
             )
-        super().__init__(model=model, objective=objective)
+        super().__init__(model=model, posterior_transform=posterior_transform, **kwargs)
         self.register_buffer("ref_point", ref_point)
         self.partitioning = partitioning
         cell_bounds = self.partitioning.get_hypercell_bounds()
@@ -203,7 +208,7 @@ def nu(self, lower: Tensor, upper: Tensor, mu: Tensor, sigma: Tensor) -> None:
 
     @t_batch_mode_transform()
     def forward(self, X: Tensor) -> Tensor:
-        posterior = self.objective(self.model.posterior(X))
+        posterior = self.posterior_transform(self.model.posterior(X))
         mu = posterior.mean
         sigma = posterior.variance.clamp_min(1e-9).sqrt()
         # clamp here, since upper_bounds will contain `inf`s, which

botorch/acquisition/multi_objective/objective.py

@@ -6,6 +6,8 @@
 
 from __future__ import annotations
 
+import warnings
+
 from abc import abstractmethod
 from typing import List, Optional
 
@@ -14,10 +16,10 @@
     AcquisitionObjective,
     GenericMCObjective,
     MCAcquisitionObjective,
+    UnstandardizePosteriorTransform,
 )
 from botorch.exceptions.errors import BotorchError, BotorchTensorDimensionError
 from botorch.models.model import Model
-from botorch.models.transforms.outcome import Standardize
 from botorch.posteriors import GPyTorchPosterior
 from botorch.utils import apply_constraints
 from botorch.utils.transforms import normalize_indices
@@ -261,35 +263,35 @@ def forward(self, samples: Tensor, X: Optional[Tensor] = None) -> Tensor:
 
 
 class AnalyticMultiOutputObjective(AcquisitionObjective):
-    r"""Abstract base class for multi-output analyic objectives."""
-    # TODO: Refactor these as PosteriorTransform as well.
+    r"""Abstract base class for multi-output analyic objectives.
 
-    @abstractmethod
-    def forward(self, posterior: GPyTorchPosterior) -> GPyTorchPosterior:
-        r"""Transform the posterior
-
-        Args:
-            posterior: A posterior.
+    DEPRECATED - This will be removed in the next version.
 
-        Returns:
-            A transformed posterior.
-        """
-        pass  # pragma: no cover
+    """
+    pass  # pragma: no cover
 
 
 class IdentityAnalyticMultiOutputObjective(AnalyticMultiOutputObjective):
+    """DEPRECATED - This will be removed in the next version."""
+
+    def __init__(self):
+        warnings.warn(
+            "IdentityAnalyticMultiOutputObjective is deprecated. "
+            "Use IdentityPosteriorTransform instead.",
+            DeprecationWarning,
+        )
+        super().__init__()
+
     def forward(self, posterior: GPyTorchPosterior) -> GPyTorchPosterior:
         return posterior
 
 
-class UnstandardizeAnalyticMultiOutputObjective(AnalyticMultiOutputObjective):
+class UnstandardizeAnalyticMultiOutputObjective(
+    UnstandardizePosteriorTransform, AnalyticMultiOutputObjective
+):
     r"""Objective that unstandardizes the posterior.
 
-    TODO: remove this when MultiTask models support outcome transforms.
-
-    Example:
-        >>> unstd_objective = UnstandardizeAnalyticMultiOutputObjective(Y_mean, Y_std)
-        >>> unstd_posterior = unstd_objective(posterior)
+    DEPRECATED - This will be removed in the next version.
     """
 
     def __init__(self, Y_mean: Tensor, Y_std: Tensor) -> None:
@@ -300,18 +302,12 @@ def __init__(self, Y_mean: Tensor, Y_std: Tensor) -> None:
             Y_std: `m`-dim tensor of outcome standard deviations
 
         """
-        if Y_mean.ndim > 1 or Y_std.ndim > 1:
-            raise BotorchTensorDimensionError(
-                "Y_mean and Y_std must both be 1-dimensional, but got "
-                f"{Y_mean.ndim} and {Y_std.ndim}"
-            )
-        super().__init__()
-        self.outcome_transform = Standardize(m=Y_mean.shape[0]).to(Y_mean)
-        Y_std_unsqueezed = Y_std.unsqueeze(0)
-        self.outcome_transform.means = Y_mean.unsqueeze(0)
-        self.outcome_transform.stdvs = Y_std_unsqueezed
-        self.outcome_transform._stdvs_sq = Y_std_unsqueezed.pow(2)
-        self.outcome_transform.eval()
+        warnings.warn(
+            "UnstandardizeAnalyticMultiOutputObjective is deprecated. "
+            "Use UnstandardizePosteriorTransform instead.",
+            DeprecationWarning,
+        )
+        super().__init__(Y_mean=Y_mean, Y_std=Y_std)
 
     def forward(self, posterior: GPyTorchPosterior) -> Tensor:
         return self.outcome_transform.untransform_posterior(posterior)

botorch/acquisition/objective.py

@@ -14,8 +14,9 @@
 from typing import Callable, List, Optional, TYPE_CHECKING, Union
 
 import torch
-from botorch.exceptions.errors import UnsupportedError
+from botorch.exceptions.errors import BotorchTensorDimensionError, UnsupportedError
 from botorch.models.model import Model
+from botorch.models.transforms.outcome import Standardize
 from botorch.posteriors.gpytorch import GPyTorchPosterior, scalarize_posterior
 from botorch.sampling import IIDNormalSampler, MCSampler
 from botorch.utils import apply_constraints
@@ -75,6 +76,14 @@ def forward(self, posterior) -> Posterior:
 from botorch.models.deterministic import DeterministicModel  # noqa
 
 
+class IdentityPosteriorTransform(PosteriorTransform):
+    def evaluate(self, Y: Tensor) -> Tensor:
+        return Y
+
+    def forward(self, posterior: GPyTorchPosterior) -> GPyTorchPosterior:
+        return posterior
+
+
 class ScalarizedPosteriorTransform(PosteriorTransform):
     r"""An affine posterior transform for scalarizing multi-output posteriors.
 
@@ -258,6 +267,44 @@ def forward(self, posterior: GPyTorchPosterior) -> GPyTorchPosterior:
         return GPyTorchPosterior(distribution=new_mvn)
 
 
+class UnstandardizePosteriorTransform(PosteriorTransform):
+    r"""Posterior transform that unstandardizes the posterior.
+
+    TODO: remove this when MultiTask models support outcome transforms.
+
+    Example:
+        >>> unstd_transform = UnstandardizePosteriorTransform(Y_mean, Y_std)
+        >>> unstd_posterior = unstd_transform(posterior)
+    """
+
+    def __init__(self, Y_mean: Tensor, Y_std: Tensor) -> None:
+        r"""Initialize objective.
+
+        Args:
+            Y_mean: `m`-dim tensor of outcome means
+            Y_std: `m`-dim tensor of outcome standard deviations
+
+        """
+        if Y_mean.ndim > 1 or Y_std.ndim > 1:
+            raise BotorchTensorDimensionError(
+                "Y_mean and Y_std must both be 1-dimensional, but got "
+                f"{Y_mean.ndim} and {Y_std.ndim}"
+            )
+        super().__init__()
+        self.outcome_transform = Standardize(m=Y_mean.shape[0]).to(Y_mean)
+        Y_std_unsqueezed = Y_std.unsqueeze(0)
+        self.outcome_transform.means = Y_mean.unsqueeze(0)
+        self.outcome_transform.stdvs = Y_std_unsqueezed
+        self.outcome_transform._stdvs_sq = Y_std_unsqueezed.pow(2)
+        self.outcome_transform.eval()
+
+    def evaluate(self, Y: Tensor) -> Tensor:
+        return self.outcome_transform(Y)
+
+    def forward(self, posterior: GPyTorchPosterior) -> Tensor:
+        return self.outcome_transform.untransform_posterior(posterior)
+
+
 class MCAcquisitionObjective(Module, ABC):
     r"""Abstract base class for MC-based objectives.
 

test/acquisition/multi_objective/test_analytic.py

@@ -10,16 +10,15 @@
     ExpectedHypervolumeImprovement,
     MultiObjectiveAnalyticAcquisitionFunction,
 )
-from botorch.acquisition.multi_objective.objective import (
-    AnalyticMultiOutputObjective,
-    IdentityAnalyticMultiOutputObjective,
-    IdentityMCMultiOutputObjective,
-)
+from botorch.acquisition.multi_objective.objective import IdentityMCMultiOutputObjective
+from botorch.acquisition.objective import IdentityPosteriorTransform, PosteriorTransform
 from botorch.exceptions.errors import BotorchError, UnsupportedError
+from botorch.posteriors import GPyTorchPosterior
 from botorch.utils.multi_objective.box_decompositions.non_dominated import (
     NondominatedPartitioning,
 )
 from botorch.utils.testing import BotorchTestCase, MockModel, MockPosterior
+from torch import Tensor
 
 
 class DummyMultiObjectiveAnalyticAcquisitionFunction(
@@ -29,8 +28,11 @@ def forward(self, X):
         pass
 
 
-class DummyAnalyticMultiOutputObjective(AnalyticMultiOutputObjective):
-    def forward(self, samples):
+class DummyPosteriorTransform(PosteriorTransform):
+    def evaluate(self, Y: Tensor) -> Tensor:
+        pass
+
+    def forward(self, posterior: GPyTorchPosterior) -> GPyTorchPosterior:
         pass
 
 
@@ -43,17 +45,17 @@ def test_init(self):
         mm = MockModel(MockPosterior(mean=torch.rand(2, 1)))
         # test default init
         acqf = DummyMultiObjectiveAnalyticAcquisitionFunction(model=mm)
-        self.assertIsInstance(acqf.objective, IdentityAnalyticMultiOutputObjective)
+        self.assertIsInstance(acqf.posterior_transform, IdentityPosteriorTransform)
         # test custom init
-        objective = DummyAnalyticMultiOutputObjective()
+        posterior_transform = DummyPosteriorTransform()
         acqf = DummyMultiObjectiveAnalyticAcquisitionFunction(
-            model=mm, objective=objective
+            model=mm, posterior_transform=posterior_transform
         )
-        self.assertEqual(acqf.objective, objective)
+        self.assertEqual(acqf.posterior_transform, posterior_transform)
         # test unsupported objective
         with self.assertRaises(UnsupportedError):
             DummyMultiObjectiveAnalyticAcquisitionFunction(
-                model=mm, objective=IdentityMCMultiOutputObjective()
+                model=mm, posterior_transform=IdentityMCMultiOutputObjective()
             )
         acqf = DummyMultiObjectiveAnalyticAcquisitionFunction(model=mm)
         # test set_X_pending