code: def test_sample_points_around_best(self):
      tkwargs = {"device": self.device}
      _bounds = torch.ones(2, 2)
      _bounds[1] = 2
      for dtype in (torch.float, torch.double):
          tkwargs["dtype"] = dtype
          bounds = _bounds.to(**tkwargs)
          X_train = 1 + torch.rand(20, 2, **tkwargs)
          model = MockModel(
              MockPosterior(
                  mean=(2 * X_train + 1).sum(dim=-1, keepdim=True)
              )
          )
          # test NEI with X_baseline
          acqf = qNoisyExpectedImprovement(
              model,
              X_baseline=X_train,
              prune_baseline=False,
              cache_root=False,
          )
          with mock.patch(
              "botorch.optim.initializers.sample_perturbed_subset_dims"
          ) as mock_subset_dims:
              X_rnd = sample_points_around_best(
                  acq_function=acqf,
                  n_discrete_points=4,
                  sigma=1e-3,
                  bounds=bounds,
              )
              mock_subset_dims.assert_not_called()
          self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
          self.assertTrue((X_rnd >= 1).all())
          self.assertTrue((X_rnd <= 2).all())
          # test model that returns a batched mean
          model = MockModel(
              MockPosterior(
                  mean=(2 * X_train + 1)
                  .sum(dim=-1, keepdim=True)
                  .unsqueeze(0)
              )
          )
          acqf = qNoisyExpectedImprovement(
              model,
              X_baseline=X_train,
              prune_baseline=False,
              cache_root=False,
          )
          X_rnd = sample_points_around_best(
              acq_function=acqf,
              n_discrete_points=4,
              sigma=1e-3,
              bounds=bounds,
          )
          self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
          self.assertTrue((X_rnd >= 1).all())
          self.assertTrue((X_rnd <= 2).all())
  
          # test EI without X_baseline
          acqf = qExpectedImprovement(model, best_f=0.0)
  
          with warnings.catch_warnings(
              record=True
          ) as w, settings.debug(True):
              X_rnd = sample_points_around_best(
                  acq_function=acqf,
                  n_discrete_points=4,
                  sigma=1e-3,
                  bounds=bounds,
              )
              self.assertEqual(len(w), 1)
              self.assertTrue(
                  issubclass(w[-1].category, BotorchWarning)
              )
              self.assertIsNone(X_rnd)
  
          # set train inputs
          model.train_inputs = (X_train,)
          X_rnd = sample_points_around_best(
              acq_function=acqf,
              n_discrete_points=4,
              sigma=1e-3,
              bounds=bounds,
          )
          self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
          self.assertTrue((X_rnd >= 1).all())
          self.assertTrue((X_rnd <= 2).all())
  
          # test an acquisition function that has no posterior_transform
          # and maximize=False
          pm = PosteriorMean(model, maximize=False)
          self.assertIsNone(pm.posterior_transform)
          self.assertFalse(pm.maximize)
          X_rnd = sample_points_around_best(
              acq_function=pm,
              n_discrete_points=4,
              sigma=0,
              bounds=bounds,
              best_pct=1e-8,  # ensures that we only use best value
          )
          idx = (-model.posterior(X_train).mean).argmax()
          self.assertTrue(
              (X_rnd == X_train[idx : idx + 1]).all(dim=-1).all()
          )
  
          # test acquisition function that has no model
          ff = FixedFeatureAcquisitionFunction(
              pm, d=2, columns=[0], values=[0]
          )
          # set X_baseline for testing purposes
          ff.X_baseline = X_train
          with warnings.catch_warnings(
              record=True
          ) as w, settings.debug(True):
              X_rnd = sample_points_around_best(
                  acq_function=ff,
                  n_discrete_points=4,
                  sigma=1e-3,
                  bounds=bounds,
              )
              self.assertEqual(len(w), 1)
              self.assertTrue(
                  issubclass(w[-1].category, BotorchWarning)
              )
              self.assertIsNone(X_rnd)
  
          # test constraints with NEHVI
          constraints = [lambda Y: Y[..., 0]]
          ref_point = torch.zeros(2, **tkwargs)
          # test cases when there are and are not any feasible points
          for any_feas in (True, False):
              Y_train = torch.stack(
                  [
                      torch.linspace(
                          -0.5, 0.5, X_train.shape[0], **tkwargs
                      )
                      if any_feas
                      else torch.ones(X_train.shape[0], **tkwargs),
                      X_train.sum(dim=-1),
                  ],
                  dim=-1,
              )
              moo_model = MockModel(
                  MockPosterior(mean=Y_train, samples=Y_train)
              )
              acqf = qNoisyExpectedHypervolumeImprovement(
                  moo_model,
                  ref_point=ref_point,
                  X_baseline=X_train,
                  constraints=constraints,
                  cache_root=False,
                  sampler=IIDNormalSampler(
                      sample_shape=torch.Size([2])
                  ),
              )
              X_rnd = sample_points_around_best(
                  acq_function=acqf,
                  n_discrete_points=4,
                  sigma=0.0,
                  bounds=bounds,
              )
              self.assertTrue(X_rnd.shape, torch.Size([4, 2]))
              # this should be true since sigma=0
              # and we should only be returning feasible points
              violation = constraints[0](Y_train)
              neg_violation = -violation.clamp_min(0.0)
              feas = neg_violation == 0
              eq_mask = (
                  X_train.unsqueeze(1) == X_rnd.unsqueeze(0)
              ).all(dim=-1)
              if feas.any():
                  # determine
                  # create n_train x n_rnd tensor of booleans
                  eq_mask = (
                      X_train.unsqueeze(1) == X_rnd.unsqueeze(0)
                  ).all(dim=-1)
                  # check that all X_rnd correspond to feasible points
                  self.assertEqual(eq_mask[feas].sum(), 4)
              else:
                  idcs = torch.topk(neg_violation, k=2).indices
                  self.assertEqual(eq_mask[idcs].sum(), 4)
              self.assertTrue((X_rnd >= 1).all())
              self.assertTrue((X_rnd <= 2).all())
          # test that subset_dims is called if d>=20
          X_train = 1 + torch.rand(10, 20, **tkwargs)
          model = MockModel(
              MockPosterior(
                  mean=(2 * X_train + 1).sum(dim=-1, keepdim=True)
              )
          )
          bounds = torch.ones(2, 20, **tkwargs)
          bounds[1] = 2
          # test NEI with X_baseline
          acqf = qNoisyExpectedImprovement(
              model,
              X_baseline=X_train,
              prune_baseline=False,
              cache_root=False,
          )
          with mock.patch(
              "botorch.optim.initializers.sample_perturbed_subset_dims",
              wraps=sample_perturbed_subset_dims,
          ) as mock_subset_dims:
              X_rnd = sample_points_around_best(
                  acq_function=acqf,
                  n_discrete_points=5,
                  sigma=1e-3,
                  bounds=bounds,
              )
          self.assertEqual(X_rnd.shape, torch.Size([5, 20]))
          self.assertTrue((X_rnd >= 1).all())
          self.assertTrue((X_rnd <= 2).all())
          mock_subset_dims.assert_called_once()
          # test tiny prob_perturb to make sure we perturb at least one dimension
          X_rnd = sample_points_around_best(
              acq_function=acqf,
              n_discrete_points=5,
              sigma=1e-3,
              bounds=bounds,
              prob_perturb=1e-8,
          )
          self.assertTrue(
              (
                  (X_rnd.unsqueeze(0) == X_train.unsqueeze(1)).all(
                      dim=-1
                  )
              ).sum()
              == 0
          )
  
imports: import warnings
  from contextlib import ExitStack
  from itertools import product
  from random import random
  from typing import Optional
  from unittest import mock
  import torch
  from botorch import settings
  from botorch.acquisition.analytic import PosteriorMean
  from botorch.acquisition.fixed_feature import FixedFeatureAcquisitionFunction
  from botorch.acquisition.knowledge_gradient import qKnowledgeGradient
  from botorch.acquisition.monte_carlo import qExpectedImprovement, qNoisyExpectedImprovement
  from botorch.acquisition.multi_objective.hypervolume_knowledge_gradient import qHypervolumeKnowledgeGradient, qMultiFidelityHypervolumeKnowledgeGradient
  from botorch.acquisition.multi_objective.monte_carlo import qNoisyExpectedHypervolumeImprovement
  from botorch.exceptions import BadInitialCandidatesWarning, SamplingWarning
  from botorch.exceptions.errors import BotorchTensorDimensionError, UnsupportedError
  from botorch.exceptions.warnings import BotorchWarning
  from botorch.models import SingleTaskGP
  from botorch.models.model_list_gp_regression import ModelListGP
  from botorch.optim import initialize_q_batch, initialize_q_batch_nonneg
  from botorch.optim.initializers import gen_batch_initial_conditions, gen_one_shot_hvkg_initial_conditions, gen_one_shot_kg_initial_conditions, gen_value_function_initial_conditions, sample_perturbed_subset_dims, sample_points_around_best, sample_q_batches_from_polytope, sample_truncated_normal_perturbations, transform_constraints, transform_inter_point_constraint, transform_intra_point_constraint
  from botorch.sampling.normal import IIDNormalSampler
  from botorch.utils.sampling import draw_sobol_samples, manual_seed
  from botorch.utils.testing import _get_max_violation_of_bounds, _get_max_violation_of_constraints, BotorchTestCase, MockAcquisitionFunction, MockModel, MockPosterior
label_desc: Use image node_id 3 for calling the TestSampleAroundBest obj's underlying member method code with example usage: obj.test_sample_points_around_best() without return types