class_docstr: MC-based batch Noisy Expected Improvement.
  
  This function does not assume a `best_f` is known (which would require
  noiseless observations). Instead, it uses samples from the joint posterior
  over the `q` test points and previously observed points. The improvement
  over previously observed points is computed for each sample and averaged.
  
  `qNEI(X) = E(max(max Y - max Y_baseline, 0))`, where
  `(Y, Y_baseline) ~ f((X, X_baseline)), X = (x_1,...,x_q)`
  
  Example:
      >>> model = SingleTaskGP(train_X, train_Y)
      >>> sampler = SobolQMCNormalSampler(1024)
      >>> qNEI = qNoisyExpectedImprovement(model, train_X, sampler)
      >>> qnei = qNEI(test_X)
class_docstr_tok: ["MC-based","batch","Noisy","Expected","Improvement",".","This","function","does","not","assume","a","`","best_f","`","is","known","(","which","would","require","noiseless","observations",")",".","Instead",",","it","uses","samples","from","the","joint","posterior","over","the","`","q","`","test","points","and","previously","observed","points",".","The","improvement","over","previously","observed","points","is","computed","for","each","sample","and","averaged",".","`","qNEI","(","X",")","=","E","(","max","(","max","Y","-","max","Y_baseline",",","0",")",")","`",",","where","`","(","Y",",","Y_baseline",")","~","f","(","(","X",",","X_baseline",")",")",",","X","=","(","x_1",",","...",",","x_q",")","`","Example",":",">",">",">","model","=","SingleTaskGP","(","train_X",",","train_Y",")",">",">",">","sampler","=","SobolQMCNormalSampler","(","1024",")",">",">",">","qNEI","=","qNoisyExpectedImprovement","(","model",",","train_X",",","sampler",")",">",">",">","qnei","=","qNEI","(","test_X",")"]
docstr: q-Noisy Expected Improvement.
  
  Args:
      model: A fitted model.
      X_baseline: A `batch_shape x r x d`-dim Tensor of `r` design points
          that have already been observed. These points are considered as
          the potential best design point.
      sampler: The sampler used to draw base samples. See `MCAcquisitionFunction`
          more details.
      objective: The MCAcquisitionObjective under which the samples are
          evaluated. Defaults to `IdentityMCObjective()`.
          NOTE: `ConstrainedMCObjective` for outcome constraints is deprecated in
          favor of passing the `constraints` directly to this constructor.
      posterior_transform: A PosteriorTransform (optional).
      X_pending: A `batch_shape x 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. Concatenated into `X` upon
          forward call. Copied and set to have no gradient.
      prune_baseline: If True, remove points in `X_baseline` that are
          highly unlikely to be the best point. This can significantly
          improve performance and is generally recommended. In order to
          customize pruning parameters, instead manually call
          `botorch.acquisition.utils.prune_inferior_points` on `X_baseline`
          before instantiating the acquisition function.
      cache_root: A boolean indicating whether to cache the root
          decomposition over `X_baseline` and use low-rank updates.
      constraints: A list of constraint callables which map a Tensor of posterior
          samples of dimension `sample_shape x batch-shape x q x m`-dim to a
          `sample_shape x batch-shape x q`-dim Tensor. The associated constraints
          are considered satisfied if the output is less than zero.
      eta: Temperature parameter(s) governing the smoothness of the sigmoid
          approximation to the constraint indicators. For more details, on this
          parameter, see the docs of `compute_smoothed_feasibility_indicator`.
      marginalize_dim: The dimension to marginalize over.
  
  TODO: similar to qNEHVI, when we are using sequential greedy candidate
  selection, we could incorporate pending points X_baseline and compute
  the incremental qNEI from the new point. This would greatly increase
  efficiency for large batches.
docstr_tok: ["q-Noisy","Expected","Improvement",".","Args",":","model",":","A","fitted","model",".","X_baseline",":","A","`","batch_shape","x","r","x","d","`","-dim","Tensor","of","`","r","`","design","points","that","have","already","been","observed",".","These","points","are","considered","as","the","potential","best","design","point",".","sampler",":","The","sampler","used","to","draw","base","samples",".","See","`","MCAcquisitionFunction","`","more","details",".","objective",":","The","MCAcquisitionObjective","under","which","the","samples","are","evaluated",".","Defaults","to","`","IdentityMCObjective","(",")","`",".","NOTE",":","`","ConstrainedMCObjective","`","for","outcome","constraints","is","deprecated","in","favor","of","passing","the","`","constraints","`","directly","to","this","constructor",".","posterior_transform",":","A","PosteriorTransform","(","optional",")",".","X_pending",":","A","`","batch_shape","x","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",".","Concatenated","into","`","X","`","upon","forward","call",".","Copied","and","set","to","have","no","gradient",".","prune_baseline",":","If","True",",","remove","points","in","`","X_baseline","`","that","are","highly","unlikely","to","be","the","best","point",".","This","can","significantly","improve","performance","and","is","generally","recommended",".","In","order","to","customize","pruning","parameters",",","instead","manually","call","`","botorch.acquisition.utils.prune_inferior_points","`","on","`","X_baseline","`","before","instantiating","the","acquisition","function",".","cache_root",":","A","boolean","indicating","whether","to","cache","the","root","decomposition","over","`","X_baseline","`","and","use","low-rank","updates",".","constraints",":","A","list","of","constraint","callables","which","map","a","Tensor","of","posterior","samples","of","dimension","`","sample_shape","x","batch-shape","x","q","x","m","`","-dim","to","a","`","sample_shape","x","batch-shape","x","q","`","-dim","Tensor",".","The","associated","constraints","are","considered","satisfied","if","the","output","is","less","than","zero",".","eta",":","Temperature","parameter","(","s",")","governing","the","smoothness","of","the","sigmoid","approximation","to","the","constraint","indicators",".","For","more","details",",","on","this","parameter",",","see","the","docs","of","`","compute_smoothed_feasibility_indicator","`",".","marginalize_dim",":","The","dimension","to","marginalize","over",".","TODO",":","similar","to","qNEHVI",",","when","we","are","using","sequential","greedy","candidate","selection",",","we","could","incorporate","pending","points","X_baseline","and","compute","the","incremental","qNEI","from","the","new","point",".","This","would","greatly","increase","efficiency","for","large","batches","."]
code: def __init__(
      self,
      model: Model,
      X_baseline: Tensor,
      sampler: Optional[MCSampler] = None,
      objective: Optional[MCAcquisitionObjective] = None,
      posterior_transform: Optional[PosteriorTransform] = None,
      X_pending: Optional[Tensor] = None,
      prune_baseline: bool = True,
      cache_root: bool = True,
      constraints: Optional[List[Callable[[Tensor], Tensor]]] = None,
      eta: Union[Tensor, float] = 1e-3,
      marginalize_dim: Optional[int] = None,
  ) -> None:
      r"""q-Noisy Expected Improvement.
  
      Args:
          model: A fitted model.
          X_baseline: A `batch_shape x r x d`-dim Tensor of `r` design points
              that have already been observed. These points are considered as
              the potential best design point.
          sampler: The sampler used to draw base samples. See `MCAcquisitionFunction`
              more details.
          objective: The MCAcquisitionObjective under which the samples are
              evaluated. Defaults to `IdentityMCObjective()`.
              NOTE: `ConstrainedMCObjective` for outcome constraints is deprecated in
              favor of passing the `constraints` directly to this constructor.
          posterior_transform: A PosteriorTransform (optional).
          X_pending: A `batch_shape x 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. Concatenated into `X` upon
              forward call. Copied and set to have no gradient.
          prune_baseline: If True, remove points in `X_baseline` that are
              highly unlikely to be the best point. This can significantly
              improve performance and is generally recommended. In order to
              customize pruning parameters, instead manually call
              `botorch.acquisition.utils.prune_inferior_points` on `X_baseline`
              before instantiating the acquisition function.
          cache_root: A boolean indicating whether to cache the root
              decomposition over `X_baseline` and use low-rank updates.
          constraints: A list of constraint callables which map a Tensor of posterior
              samples of dimension `sample_shape x batch-shape x q x m`-dim to a
              `sample_shape x batch-shape x q`-dim Tensor. The associated constraints
              are considered satisfied if the output is less than zero.
          eta: Temperature parameter(s) governing the smoothness of the sigmoid
              approximation to the constraint indicators. For more details, on this
              parameter, see the docs of `compute_smoothed_feasibility_indicator`.
          marginalize_dim: The dimension to marginalize over.
  
      TODO: similar to qNEHVI, when we are using sequential greedy candidate
      selection, we could incorporate pending points X_baseline and compute
      the incremental qNEI from the new point. This would greatly increase
      efficiency for large batches.
      """
      super().__init__(
          model=model,
          sampler=sampler,
          objective=objective,
          posterior_transform=posterior_transform,
          X_pending=X_pending,
          constraints=constraints,
          eta=eta,
      )
      CachedCholeskyMCSamplerMixin.__init__(
          self, model=model, cache_root=cache_root, sampler=sampler
      )
      if prune_baseline:
          X_baseline = prune_inferior_points(
              model=model,
              X=X_baseline,
              objective=objective,
              posterior_transform=posterior_transform,
              constraints=self._constraints,
              marginalize_dim=marginalize_dim,
          )
      self.register_buffer("X_baseline", X_baseline)
      # registering buffers for _get_samples_and_objectives in the next `if` block
      self.register_buffer("baseline_samples", None)
      self.register_buffer("baseline_obj", None)
      if self._cache_root:
          self.q_in = -1
          # set baseline samples
          with torch.no_grad():  # this is _get_samples_and_objectives(X_baseline)
              posterior = self.model.posterior(
                  X_baseline,
                  posterior_transform=self.posterior_transform,
              )
              # Note: The root decomposition is cached in two different places. It
              # may be confusing to have two different caches, but this is not
              # trivial to change since each is needed for a different reason:
              # - LinearOperator caching to `posterior.mvn` allows for reuse within
              #  this function, which may be helpful if the same root decomposition
              #  is produced by the calls to `self.base_sampler` and
              #  `self._cache_root_decomposition`.
              # - self._baseline_L allows a root decomposition to be persisted outside
              #   this method.
              baseline_samples = self.get_posterior_samples(posterior)
              baseline_obj = self.objective(
                  baseline_samples, X=X_baseline
              )
  
          # We make a copy here because we will write an attribute `base_samples`
          # to `self.base_sampler.base_samples`, and we don't want to mutate
          # `self.sampler`.
          self.base_sampler = deepcopy(self.sampler)
          self.baseline_samples = baseline_samples
          self.baseline_obj = baseline_obj
          self.register_buffer(
              "_baseline_best_f",
              self._compute_best_feasible_objective(
                  samples=baseline_samples, obj=baseline_obj
              ),
          )
          self._baseline_L = self._compute_root_decomposition(
              posterior=posterior
          )
  
imports: from __future__ import annotations
  import math
  from abc import ABC, abstractmethod
  from copy import deepcopy
  from functools import partial
  from typing import Callable, List, Optional, Protocol, Tuple, Union
  import torch
  from botorch.acquisition.acquisition import AcquisitionFunction, MCSamplerMixin
  from botorch.acquisition.cached_cholesky import CachedCholeskyMCSamplerMixin
  from botorch.acquisition.objective import ConstrainedMCObjective, IdentityMCObjective, MCAcquisitionObjective, PosteriorTransform
  from botorch.acquisition.utils import compute_best_feasible_objective, prune_inferior_points, repeat_to_match_aug_dim
  from botorch.exceptions.errors import UnsupportedError
  from botorch.models.model import Model
  from botorch.sampling.base import MCSampler
  from botorch.utils.objective import compute_smoothed_feasibility_indicator
  from botorch.utils.transforms import concatenate_pending_points, match_batch_shape, t_batch_mode_transform
  from torch import Tensor
label_desc: Use image node_id 1 to create a new qNoisyExpectedImprovement object from inherited base classes: SampleReducingMCAcquisitionFunction, CachedCholeskyMCSamplerMixin with example: obj = qNoisyExpectedImprovement(model, X_baseline, sampler, objective, posterior_transform, X_pending, prune_baseline, cache_root, constraints, eta, marginalize_dim)