Update entropy search sampler

Author: benmltu

botorch/acquisition/multi_objective/joint_entropy_search.py

@@ -23,20 +23,20 @@
 
 import torch
 from botorch import settings
-from botorch.acquisition.acquisition import AcquisitionFunction
+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
 from botorch.posteriors.gpytorch import GPyTorchPosterior
-from botorch.sampling.samplers import SobolQMCNormalSampler
+from botorch.sampling.normal import SobolQMCNormalSampler
 from botorch.utils.transforms import concatenate_pending_points, t_batch_mode_transform
 from torch import Tensor
 from torch.distributions import Normal
 
 
-class LowerBoundMultiObjectiveEntropySearch(AcquisitionFunction):
+class LowerBoundMultiObjectiveEntropySearch(AcquisitionFunction, MCSamplerMixin):
     r"""Abstract base class for the lower bound multi-objective entropy search
     acquisition functions.
     """
@@ -74,6 +74,8 @@ def __init__(
                 estimate.
         """
         super().__init__(model=model)
+        sampler = SobolQMCNormalSampler(sample_shape=torch.Size([num_samples]))
+        MCSamplerMixin.__init__(self, sampler=sampler)
         # Batch GP models (e.g. fantasized models) are not currently supported
         if isinstance(model, ModelListGP):
             train_X = model.models[0].train_inputs[0]
@@ -119,9 +121,6 @@ def __init__(
                 + "."
             )
 
-        self.sampler = SobolQMCNormalSampler(
-            num_samples=num_samples, collapse_batch_dims=True
-        )
         self.set_X_pending(X_pending)
 
     @abstractmethod
@@ -420,7 +419,7 @@ def _compute_monte_carlo_variables(
                 samples.
         """
         # `num_mc_samples x batch_shape x q x num_pareto_samples x 1 x M`
-        samples = self.sampler(posterior)
+        samples = self.get_posterior_samples(posterior)
 
         # `num_mc_samples x batch_shape x q x num_pareto_samples`
         if self.model.num_outputs == 1:

botorch/acquisition/multi_objective/max_value_entropy_search.py

@@ -338,7 +338,7 @@ def _compute_monte_carlo_variables(
         """
 
         # `num_mc_samples x batch_shape x q x 1 x M`
-        samples = self.sampler(posterior)
+        samples = self.get_posterior_samples(posterior)
 
         # `num_mc_samples x batch_shape x q`
         if self.model.num_outputs == 1:

test/acquisition/multi_objective/test_joint_entropy_search.py

@@ -17,7 +17,7 @@
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.transforms.outcome import Standardize
-from botorch.sampling.samplers import SobolQMCNormalSampler
+from botorch.sampling.normal import SobolQMCNormalSampler
 from botorch.utils.testing import BotorchTestCase
 
 

test/acquisition/test_joint_entropy_search.py

@@ -12,7 +12,7 @@
 from botorch.models.gp_regression import SingleTaskGP
 from botorch.models.model_list_gp_regression import ModelListGP
 from botorch.models.transforms.outcome import Standardize
-from botorch.sampling.samplers import SobolQMCNormalSampler
+from botorch.sampling.normal import SobolQMCNormalSampler
 from botorch.utils.testing import BotorchTestCase