class_docstr: Tests that dispatching utilities correctly select generation strategies.
code: def test_choose_generation_strategy(self) -> None:
      expected_transforms = [Winsorize] + Cont_X_trans + Y_trans
      expected_transform_configs = {
          "Winsorize": {"derelativize_with_raw_status_quo": False},
          "Derelativize": {"use_raw_status_quo": False},
      }
      with self.subTest("GPEI"):
          sobol_gpei = choose_generation_strategy(
              search_space=get_branin_search_space()
          )
          self.assertEqual(sobol_gpei._steps[0].model, Models.SOBOL)
          self.assertEqual(sobol_gpei._steps[0].num_trials, 5)
          self.assertEqual(
              sobol_gpei._steps[1].model, Models.BOTORCH_MODULAR
          )
          expected_model_kwargs: Dict[str, Any] = {
              "torch_device": None,
              "transforms": expected_transforms,
              "transform_configs": expected_transform_configs,
          }
          self.assertEqual(
              sobol_gpei._steps[1].model_kwargs, expected_model_kwargs
          )
          device = torch.device("cpu")
          sobol_gpei = choose_generation_strategy(
              search_space=get_branin_search_space(),
              verbose=True,
              torch_device=device,
          )
          expected_model_kwargs["torch_device"] = device
          self.assertEqual(
              sobol_gpei._steps[1].model_kwargs, expected_model_kwargs
          )
          run_branin_experiment_with_generation_strategy(
              generation_strategy=sobol_gpei
          )
      with self.subTest("max initialization trials"):
          sobol_gpei = choose_generation_strategy(
              search_space=get_branin_search_space(),
              max_initialization_trials=2,
          )
          self.assertEqual(sobol_gpei._steps[0].model, Models.SOBOL)
          self.assertEqual(sobol_gpei._steps[0].num_trials, 2)
          self.assertEqual(
              sobol_gpei._steps[1].model, Models.BOTORCH_MODULAR
          )
      with self.subTest(
          "num_initialization_trials > max_initialization_trials"
      ):
          sobol_gpei = choose_generation_strategy(
              search_space=get_branin_search_space(),
              max_initialization_trials=2,
              num_initialization_trials=3,
          )
          self.assertEqual(sobol_gpei._steps[0].model, Models.SOBOL)
          self.assertEqual(sobol_gpei._steps[0].num_trials, 3)
          self.assertEqual(
              sobol_gpei._steps[1].model, Models.BOTORCH_MODULAR
          )
      with self.subTest(
          "num_initialization_trials > max_initialization_trials"
      ):
          sobol_gpei = choose_generation_strategy(
              search_space=get_branin_search_space(),
              max_initialization_trials=2,
              num_initialization_trials=3,
          )
          self.assertEqual(sobol_gpei._steps[0].model, Models.SOBOL)
          self.assertEqual(sobol_gpei._steps[0].num_trials, 3)
          self.assertEqual(
              sobol_gpei._steps[1].model, Models.BOTORCH_MODULAR
          )
      with self.subTest("MOO"):
          optimization_config = MultiObjectiveOptimizationConfig(
              objective=MultiObjective(objectives=[])
          )
          sobol_gpei = choose_generation_strategy(
              search_space=get_branin_search_space(),
              optimization_config=optimization_config,
          )
          self.assertEqual(sobol_gpei._steps[0].model, Models.SOBOL)
          self.assertEqual(sobol_gpei._steps[0].num_trials, 5)
          self.assertEqual(
              sobol_gpei._steps[1].model, Models.BOTORCH_MODULAR
          )
          model_kwargs = not_none(sobol_gpei._steps[1].model_kwargs)
          self.assertEqual(
              set(model_kwargs.keys()),
              {"torch_device", "transforms", "transform_configs"},
          )
          self.assertGreater(len(model_kwargs["transforms"]), 0)
      with self.subTest("Sobol (we can try every option)"):
          sobol = choose_generation_strategy(
              search_space=get_factorial_search_space(), num_trials=1000
          )
          self.assertEqual(sobol._steps[0].model, Models.SOBOL)
          self.assertEqual(len(sobol._steps), 1)
      with self.subTest("Sobol (because of too many categories)"):
          sobol_large = choose_generation_strategy(
              search_space=get_large_factorial_search_space(),
              verbose=True,
          )
          self.assertEqual(sobol_large._steps[0].model, Models.SOBOL)
          self.assertEqual(len(sobol_large._steps), 1)
      with self.subTest(
          "Sobol (because of too many categories) with saasbo"
      ):
          with self.assertLogs(
              choose_generation_strategy.__module__, logging.WARNING
          ) as logger:
              sobol_large = choose_generation_strategy(
                  search_space=get_large_factorial_search_space(),
                  verbose=True,
                  use_saasbo=True,
              )
              self.assertTrue(
                  any(
                      "SAASBO is incompatible with Sobol" in output
                      for output in logger.output
                  ),
                  logger.output,
              )
          self.assertEqual(sobol_large._steps[0].model, Models.SOBOL)
          self.assertEqual(len(sobol_large._steps), 1)
      with self.subTest("SOBOL due to too many unordered choices"):
          # Search space with more unordered choices than ordered parameters.
          sobol = choose_generation_strategy(
              search_space=get_search_space_with_choice_parameters(
                  num_ordered_parameters=5,
                  num_unordered_choices=100,
              )
          )
          self.assertEqual(sobol._steps[0].model, Models.SOBOL)
          self.assertEqual(len(sobol._steps), 1)
      with self.subTest(
          "GPEI with more unordered choices than ordered parameters"
      ):
          # Search space with more unordered choices than ordered parameters.
          sobol_gpei = choose_generation_strategy(
              search_space=get_search_space_with_choice_parameters(
                  num_ordered_parameters=5,
                  num_unordered_choices=10,
              )
          )
          self.assertEqual(
              sobol_gpei._steps[1].model, Models.BOTORCH_MODULAR
          )
      with self.subTest(
          "GPEI despite many unordered 2-value parameters"
      ):
          gs = choose_generation_strategy(
              search_space=get_large_factorial_search_space(
                  num_levels=2, num_parameters=10
              ),
          )
          self.assertEqual(gs._steps[0].model, Models.SOBOL)
          self.assertEqual(gs._steps[1].model, Models.BOTORCH_MODULAR)
      with self.subTest("GPEI-Batched"):
          sobol_gpei_batched = choose_generation_strategy(
              search_space=get_branin_search_space(),
              use_batch_trials=True,
          )
          self.assertEqual(sobol_gpei_batched._steps[0].num_trials, 1)
      with self.subTest("BO_MIXED (purely categorical)"):
          bo_mixed = choose_generation_strategy(
              search_space=get_factorial_search_space()
          )
          self.assertEqual(bo_mixed._steps[0].model, Models.SOBOL)
          self.assertEqual(bo_mixed._steps[0].num_trials, 6)
          self.assertEqual(bo_mixed._steps[1].model, Models.BO_MIXED)
          expected_model_kwargs = {
              "torch_device": None,
              "transforms": [Winsorize] + Mixed_transforms + Y_trans,
              "transform_configs": expected_transform_configs,
          }
          self.assertEqual(
              bo_mixed._steps[1].model_kwargs, expected_model_kwargs
          )
      with self.subTest("BO_MIXED (mixed search space)"):
          ss = get_branin_search_space(with_choice_parameter=True)
          # pyre-fixme[16]: `Parameter` has no attribute `_is_ordered`.
          ss.parameters["x2"]._is_ordered = False
          bo_mixed_2 = choose_generation_strategy(search_space=ss)
          self.assertEqual(bo_mixed_2._steps[0].model, Models.SOBOL)
          self.assertEqual(bo_mixed_2._steps[0].num_trials, 5)
          self.assertEqual(bo_mixed_2._steps[1].model, Models.BO_MIXED)
          expected_model_kwargs = {
              "torch_device": None,
              "transforms": [Winsorize] + Mixed_transforms + Y_trans,
              "transform_configs": expected_transform_configs,
          }
          self.assertEqual(
              bo_mixed._steps[1].model_kwargs, expected_model_kwargs
          )
      with self.subTest(
          "BO_MIXED (mixed multi-objective optimization)"
      ):
          search_space = get_branin_search_space(
              with_choice_parameter=True
          )
          search_space.parameters["x2"]._is_ordered = False
          optimization_config = MultiObjectiveOptimizationConfig(
              objective=MultiObjective(objectives=[])
          )
          moo_mixed = choose_generation_strategy(
              search_space=search_space,
              optimization_config=optimization_config,
          )
          self.assertEqual(moo_mixed._steps[0].model, Models.SOBOL)
          self.assertEqual(moo_mixed._steps[0].num_trials, 5)
          self.assertEqual(moo_mixed._steps[1].model, Models.BO_MIXED)
          model_kwargs = not_none(moo_mixed._steps[1].model_kwargs)
          self.assertEqual(
              set(model_kwargs.keys()),
              {"torch_device", "transforms", "transform_configs"},
          )
          self.assertGreater(len(model_kwargs["transforms"]), 0)
      with self.subTest("SAASBO"):
          sobol_fullybayesian = choose_generation_strategy(
              search_space=get_branin_search_space(),
              use_batch_trials=True,
              num_initialization_trials=3,
              use_saasbo=True,
          )
          self.assertEqual(
              sobol_fullybayesian._steps[0].model, Models.SOBOL
          )
          self.assertEqual(sobol_fullybayesian._steps[0].num_trials, 3)
          self.assertEqual(
              sobol_fullybayesian._steps[1].model, Models.SAASBO
          )
      with self.subTest("SAASBO MOO"):
          sobol_fullybayesianmoo = choose_generation_strategy(
              search_space=get_branin_search_space(),
              use_batch_trials=True,
              num_initialization_trials=3,
              use_saasbo=True,
              optimization_config=MultiObjectiveOptimizationConfig(
                  objective=MultiObjective(objectives=[])
              ),
          )
          self.assertEqual(
              sobol_fullybayesianmoo._steps[0].model, Models.SOBOL
          )
          self.assertEqual(
              sobol_fullybayesianmoo._steps[0].num_trials, 3
          )
          self.assertEqual(
              sobol_fullybayesianmoo._steps[1].model,
              Models.SAASBO,
          )
      with self.subTest("SAASBO"):
          sobol_fullybayesian_large = choose_generation_strategy(
              search_space=get_large_ordinal_search_space(
                  n_ordinal_choice_parameters=5,
                  n_continuous_range_parameters=10,
              ),
              use_saasbo=True,
          )
          self.assertEqual(
              sobol_fullybayesian_large._steps[0].model, Models.SOBOL
          )
          self.assertEqual(
              sobol_fullybayesian_large._steps[0].num_trials, 30
          )
          self.assertEqual(
              sobol_fullybayesian_large._steps[1].model,
              Models.SAASBO,
          )
      with self.subTest("num_initialization_trials"):
          ss = get_large_factorial_search_space()
          for _, param in ss.parameters.items():
              param._is_ordered = True
          # 2 * len(ss.parameters) init trials are performed if num_trials is large
          gs_12_init_trials = choose_generation_strategy(
              search_space=ss, num_trials=100
          )
          self.assertEqual(
              gs_12_init_trials._steps[0].model, Models.SOBOL
          )
          self.assertEqual(gs_12_init_trials._steps[0].num_trials, 12)
          self.assertEqual(
              gs_12_init_trials._steps[1].model, Models.BOTORCH_MODULAR
          )
          # at least 5 initialization trials are performed
          gs_5_init_trials = choose_generation_strategy(
              search_space=ss, num_trials=0
          )
          self.assertEqual(
              gs_5_init_trials._steps[0].model, Models.SOBOL
          )
          self.assertEqual(gs_5_init_trials._steps[0].num_trials, 5)
          self.assertEqual(
              gs_5_init_trials._steps[1].model, Models.BOTORCH_MODULAR
          )
          # avoid spending >20% of budget on initialization trials if there are
          # more than 5 initialization trials
          gs_6_init_trials = choose_generation_strategy(
              search_space=ss, num_trials=30
          )
          self.assertEqual(
              gs_6_init_trials._steps[0].model, Models.SOBOL
          )
          self.assertEqual(gs_6_init_trials._steps[0].num_trials, 6)
          self.assertEqual(
              gs_6_init_trials._steps[1].model, Models.BOTORCH_MODULAR
          )
      with self.subTest("suggested_model_override"):
          sobol_gpei = choose_generation_strategy(
              search_space=get_branin_search_space()
          )
          self.assertEqual(
              sobol_gpei._steps[1].model, Models.BOTORCH_MODULAR
          )
          sobol_saasbo = choose_generation_strategy(
              search_space=get_branin_search_space(),
              suggested_model_override=Models.SAASBO,
          )
          self.assertEqual(sobol_saasbo._steps[1].model, Models.SAASBO)
  
imports: import logging
  import warnings
  from typing import Any, Dict
  import torch
  from ax.core.objective import MultiObjective
  from ax.core.optimization_config import MultiObjectiveOptimizationConfig
  from ax.modelbridge.dispatch_utils import _make_botorch_step, calculate_num_initialization_trials, choose_generation_strategy, DEFAULT_BAYESIAN_PARALLELISM
  from ax.modelbridge.factory import Cont_X_trans, Y_trans
  from ax.modelbridge.registry import Mixed_transforms, Models
  from ax.modelbridge.transforms.log_y import LogY
  from ax.modelbridge.transforms.winsorize import Winsorize
  from ax.models.winsorization_config import WinsorizationConfig
  from ax.utils.common.testutils import TestCase
  from ax.utils.common.typeutils import not_none
  from ax.utils.testing.core_stubs import get_branin_search_space, get_discrete_search_space, get_experiment, get_factorial_search_space, get_large_factorial_search_space, get_large_ordinal_search_space, get_search_space_with_choice_parameters, run_branin_experiment_with_generation_strategy
  from ax.utils.testing.mock import fast_botorch_optimize
label_desc: Use image node_id 1 for calling the TestDispatchUtils obj's underlying member method code with example usage: obj.test_choose_generation_strategy() without return types