class_docstr: **All classes in 'botorch_modular' directory are under
  construction, incomplete, and should be treated as alpha
  versions only.**
  
  Ax wrapper for BoTorch `AcquisitionFunction`, subcomponent
  of `BoTorchModel` and is not meant to be used outside of it.
  
  Args:
      surrogates: Dict of name => Surrogate model pairs, with which this acquisition
          function will be used.
      search_space_digest: A SearchSpaceDigest object containing metadata
          about the search space (e.g. bounds, parameter types).
      torch_opt_config: A TorchOptConfig object containing optimization
          arguments (e.g., objective weights, constraints).
      botorch_acqf_class: Type of BoTorch `AcquistitionFunction` that
          should be used. Subclasses of `Acquisition` often specify
          these via `default_botorch_acqf_class` attribute, in which
          case specifying one here is not required.
      options: Optional mapping of kwargs to the underlying `Acquisition
          Function` in BoTorch.
class_docstr_tok: ["*","*","All","classes","in","'botorch_modular","'","directory","are","under","construction",",","incomplete",",","and","should","be","treated","as","alpha","versions","only",".","*","*","Ax","wrapper","for","BoTorch","`","AcquisitionFunction","`",",","subcomponent","of","`","BoTorchModel","`","and","is","not","meant","to","be","used","outside","of","it",".","Args",":","surrogates",":","Dict","of","name","=",">","Surrogate","model","pairs",",","with","which","this","acquisition","function","will","be","used",".","search_space_digest",":","A","SearchSpaceDigest","object","containing","metadata","about","the","search","space","(","e.g",".","bounds",",","parameter","types",")",".","torch_opt_config",":","A","TorchOptConfig","object","containing","optimization","arguments","(","e.g.",",","objective","weights",",","constraints",")",".","botorch_acqf_class",":","Type","of","BoTorch","`","AcquistitionFunction","`","that","should","be","used",".","Subclasses","of","`","Acquisition","`","often","specify","these","via","`","default_botorch_acqf_class","`","attribute",",","in","which","case","specifying","one","here","is","not","required",".","options",":","Optional","mapping","of","kwargs","to","the","underlying","`","Acquisition","Function","`","in","BoTorch","."]
docstr: Generate a set of candidates via multi-start optimization. Obtains
  candidates and their associated acquisition function values.
  
  Args:
      n: The number of candidates to generate.
      search_space_digest: A ``SearchSpaceDigest`` object containing search space
          properties, e.g. ``bounds`` for optimization.
      inequality_constraints: A list of tuples (indices, coefficients, rhs),
          with each tuple encoding an inequality constraint of the form
          ``sum_i (X[indices[i]] * coefficients[i]) >= rhs``.
      fixed_features: A map `{feature_index: value}` for features that
          should be fixed to a particular value during generation.
      rounding_func: A function that post-processes an optimization
          result appropriately. This is typically passed down from
          `ModelBridge` to ensure compatibility of the candidates with
          with Ax transforms. For additional post processing, use
          `post_processing_func` option in `optimizer_options`.
      optimizer_options: Options for the optimizer function, e.g. ``sequential``
          or ``raw_samples``. This can also include a `post_processing_func`
          which is applied to the candidates before the `rounding_func`.
          `post_processing_func` can be used to support more customized options
          that typically only exist in MBM, such as BoTorch transforms.
          See the docstring of `TorchOptConfig` for more information on passing
          down these options while constructing a generation strategy.
  
  Returns:
      A two-element tuple containing an `n x d`-dim tensor of generated candidates
      and a tensor with the associated acquisition value.
code: def optimize(
      self,
      n: int,
      search_space_digest: SearchSpaceDigest,
      inequality_constraints: Optional[
          List[Tuple[Tensor, Tensor, float]]
      ] = None,
      fixed_features: Optional[Dict[int, float]] = None,
      rounding_func: Optional[Callable[[Tensor], Tensor]] = None,
      optimizer_options: Optional[Dict[str, Any]] = None,
  ) -> Tuple[Tensor, Tensor]:
      """Generate a set of candidates via multi-start optimization. Obtains
      candidates and their associated acquisition function values.
  
      Args:
          n: The number of candidates to generate.
          search_space_digest: A ``SearchSpaceDigest`` object containing search space
              properties, e.g. ``bounds`` for optimization.
          inequality_constraints: A list of tuples (indices, coefficients, rhs),
              with each tuple encoding an inequality constraint of the form
              ``sum_i (X[indices[i]] * coefficients[i]) >= rhs``.
          fixed_features: A map `{feature_index: value}` for features that
              should be fixed to a particular value during generation.
          rounding_func: A function that post-processes an optimization
              result appropriately. This is typically passed down from
              `ModelBridge` to ensure compatibility of the candidates with
              with Ax transforms. For additional post processing, use
              `post_processing_func` option in `optimizer_options`.
          optimizer_options: Options for the optimizer function, e.g. ``sequential``
              or ``raw_samples``. This can also include a `post_processing_func`
              which is applied to the candidates before the `rounding_func`.
              `post_processing_func` can be used to support more customized options
              that typically only exist in MBM, such as BoTorch transforms.
              See the docstring of `TorchOptConfig` for more information on passing
              down these options while constructing a generation strategy.
  
      Returns:
          A two-element tuple containing an `n x d`-dim tensor of generated candidates
          and a tensor with the associated acquisition value.
      """
      # NOTE: Could make use of `optimizer_class` when it's added to BoTorch
      # instead of calling `optimizer_acqf` or `optimize_acqf_discrete` etc.
      _tensorize = partial(
          torch.tensor, dtype=self.dtype, device=self.device
      )
      ssd = search_space_digest
      bounds = _tensorize(ssd.bounds).t()
  
      # Prepare arguments for optimizer
      optimizer_options_with_defaults = optimizer_argparse(
          self.acqf,
          bounds=bounds,
          q=n,
          optimizer_options=optimizer_options,
      )
      post_processing_func = get_post_processing_func(
          rounding_func=rounding_func,
          optimizer_options=optimizer_options_with_defaults,
      )
      discrete_features = sorted(
          ssd.ordinal_features + ssd.categorical_features
      )
      if fixed_features is not None:
          for i in fixed_features:
              if not 0 <= i < len(ssd.feature_names):
                  raise ValueError(f"Invalid fixed_feature index: {i}")
  
      # 1. Handle the fully continuous search space.
      if (
          optimizer_options_with_defaults.pop(
              "force_use_optimize_acqf", False
          )
          or not discrete_features
      ):
          return optimize_acqf(
              acq_function=self.acqf,
              bounds=bounds,
              q=n,
              inequality_constraints=inequality_constraints,
              fixed_features=fixed_features,
              post_processing_func=post_processing_func,
              **optimizer_options_with_defaults,
          )
  
      # 2. Handle search spaces with discrete features.
      discrete_choices = mk_discrete_choices(
          ssd=ssd, fixed_features=fixed_features
      )
  
      # 2a. Handle the fully discrete search space.
      if len(discrete_choices) == len(ssd.feature_names):
          X_observed = self.X_observed
          if self.X_pending is not None:
              X_observed = torch.cat(
                  [X_observed, self.X_pending], dim=0
              )
  
          # Special handling for search spaces with a large number of choices
          total_choices = reduce(
              operator.mul,
              [float(len(c)) for c in discrete_choices.values()],
          )
          if total_choices > MAX_CHOICES_ENUMERATE:
              discrete_choices = [
                  torch.tensor(c, device=self.device, dtype=self.dtype)
                  for c in discrete_choices.values()
              ]
              return optimize_acqf_discrete_local_search(
                  acq_function=self.acqf,
                  q=n,
                  discrete_choices=discrete_choices,
                  inequality_constraints=inequality_constraints,
                  X_avoid=X_observed,
                  **optimizer_options_with_defaults,
              )
  
          # Enumerate all possible choices
          all_choices = (
              discrete_choices[i] for i in range(len(discrete_choices))
          )
          all_choices = _tensorize(tuple(product(*all_choices)))
  
          # This can be vectorized, but using a for-loop to avoid memory issues
          for x in X_observed:
              all_choices = all_choices[
                  (all_choices - x).abs().max(dim=-1).values
                  > DUPLICATE_TOL
              ]
  
          # Filter out candidates that violate the constraints
          # TODO: It will be more memory-efficient to do this filtering before
          # converting the generator into a tensor. However, if we run into memory
          # issues we are likely better off being smarter in how we optimize the
          # acquisition function.
          inequality_constraints = inequality_constraints or []
          is_feasible = torch.ones(
              all_choices.shape[0], dtype=torch.bool
          )
          for inds, weights, bound in inequality_constraints:
              is_feasible &= (all_choices[..., inds] * weights).sum(
                  dim=-1
              ) >= bound
          all_choices = all_choices[is_feasible]
  
          num_choices = all_choices.size(dim=0)
          if num_choices == 0:
              raise SearchSpaceExhausted(
                  "No more feasible choices in a fully discrete search space."
              )
          if num_choices < n:
              warnings.warn(
                  (
                      f"Requested n={n} candidates from fully discrete search "
                      f"space, but only {num_choices} possible choices remain."
                  ),
                  AxWarning,
              )
              n = num_choices
  
          discrete_opt_options = optimizer_argparse(
              self.acqf,
              bounds=bounds,
              q=n,
              optimizer_options=optimizer_options,
              optimizer_is_discrete=True,
          )
          return optimize_acqf_discrete(
              acq_function=self.acqf,
              q=n,
              choices=all_choices,
              **discrete_opt_options,
          )
  
      # 2b. Handle mixed search spaces that have discrete and continuous features.
      return optimize_acqf_mixed(
          acq_function=self.acqf,
          bounds=bounds,
          q=n,
          # For now we just enumerate all possible discrete combinations. This is not
          # scalable and and only works for a reasonably small number of choices. A
          # slowdown warning is logged in `enumerate_discrete_combinations` if needed.
          fixed_features_list=enumerate_discrete_combinations(
              discrete_choices=discrete_choices
          ),
          inequality_constraints=inequality_constraints,
          post_processing_func=post_processing_func,
          **optimizer_options_with_defaults,
      )
  
imports: from __future__ import annotations
  import dataclasses
  import functools
  import operator
  import warnings
  from functools import partial, reduce
  from itertools import product
  from logging import Logger
  from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Type
  import torch
  from ax.core.search_space import SearchSpaceDigest
  from ax.exceptions.core import AxWarning, SearchSpaceExhausted
  from ax.models.model_utils import enumerate_discrete_combinations, mk_discrete_choices
  from ax.models.torch.botorch_modular.optimizer_argparse import optimizer_argparse
  from ax.models.torch.botorch_modular.surrogate import Surrogate
  from ax.models.torch.botorch_modular.utils import _tensor_difference, get_post_processing_func
  from ax.models.torch.botorch_moo_defaults import infer_objective_thresholds
  from ax.models.torch.utils import _get_X_pending_and_observed, get_botorch_objective_and_transform, subset_model
  from ax.models.torch_base import TorchOptConfig
  from ax.utils.common.base import Base
  from ax.utils.common.constants import Keys
  from ax.utils.common.logger import get_logger
  from ax.utils.common.typeutils import not_none
  from botorch.acquisition.acquisition import AcquisitionFunction
  from botorch.acquisition.input_constructors import get_acqf_input_constructor
  from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
  from botorch.acquisition.objective import MCAcquisitionObjective, PosteriorTransform
  from botorch.acquisition.risk_measures import RiskMeasureMCObjective
  from botorch.models.model import Model, ModelDict
  from botorch.optim.optimize import optimize_acqf, optimize_acqf_discrete, optimize_acqf_discrete_local_search, optimize_acqf_mixed
  from botorch.utils.constraints import get_outcome_constraint_transforms
  from torch import Tensor
label_desc: Use image node_id 7 for calling the Acquisition obj's underlying member method code with example usage: obj.optimize(n, search_space_digest, inequality_constraints, fixed_features, rounding_func, optimizer_options) and returns: optimize_acqf_mixed, optimize_acqf, optimize_acqf_discrete, optimize_acqf_discrete_local_search