ax.modelbridge

Generation Strategy, Registry, and Factory

Generation Strategy

class ax.modelbridge.generation_strategy.GenerationStrategy(steps: List[ax.modelbridge.generation_node.GenerationStep], name: Optional[str] = None)[source]

Bases: ax.utils.common.base.Base

GenerationStrategy describes which model should be used to generate new points for which trials, enabling and automating use of different models throughout the optimization process. For instance, it allows to use one model for the initialization trials, and another one for all subsequent trials. In the general case, this allows to automate use of an arbitrary number of models to generate an arbitrary numbers of trials described in the trials_per_model argument.

Parameters

  • steps – A list of GenerationStep describing steps of this strategy.

  • name – An optional name for this generaiton strategy. If not specified, strategy’s name will be names of its steps’ models joined with ‘+’.

clone_reset() → ax.modelbridge.generation_strategy.GenerationStrategy[source]

Copy this generation strategy without it’s state.

current_generator_run_limit() → Tuple[int, bool][source]

How many generator runs can this generation strategy generate right now, assuming each one of them becomes its own trial, and whether optimization is completed.

NOTE: This method might move the generation strategy to the next step, which is safe, as the next call to gen will just pick up from there.

Returns: a two-item tuple of:

  • the number of generator runs that can currently be produced, with -1 meaning unlimited generator runs,

  • whether optimization is completed and the generation strategy cannot generate any more generator runs at all.

property experiment

Experiment, currently set on this generation strategy.

gen(experiment: ax.core.experiment.Experiment, data: Optional[ax.core.data.Data] = None, n: int = 1, pending_observations: Optional[Dict[str, List[ax.core.observation.ObservationFeatures]]] = None, **kwargs: Any) → ax.core.generator_run.GeneratorRun[source]

Produce the next points in the experiment. Additional kwargs passed to this method are propagated directly to the underlying model’s gen, along with the model_gen_kwargs set on the current generation step.

NOTE: Each generator run returned from this function must become a single trial on the experiment to comply with assumptions made in generation strategy. Do not split one generator run produced from generation strategy into multiple trials (never making a generator run into a trial is allowed).

Parameters

  • experiment – Experiment, for which the generation strategy is producing a new generator run in the course of gen, and to which that generator run will be added as trial(s). Information stored on the experiment (e.g., trial statuses) is used to determine which model will be used to produce the generator run returned from this method.

  • data – Optional data to be passed to the underlying model’s gen, which is called within this method and actually produces the resulting generator run. By default, data is all data on the experiment if use_update is False and only the new data since the last call to this method if use_update is True.

  • n – Integer representing how many arms should be in the generator run produced by this method. NOTE: Some underlying models may ignore the n and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that can differ from n.

  • pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.

property last_generator_run

Latest generator run produced by this generation strategy. Returns None if no generator runs have been produced yet.

property model

Current model in this strategy. Returns None if no model has been set yet (i.e., if no generator runs have been produced from this GS).

property model_transitions

List of trial indices where a transition happened from one model to another.

property name

Name of this generation strategy. Defaults to a combination of model names provided in generation steps.

property num_can_complete_this_step

Number of trials for the current step in generation strategy that can be completed (so are not in status FAILED or ABANDONED). Used to keep track of how many generator runs (that become trials) can be produced from the current generation step.

NOTE: This includes COMPLETED trials.

property num_completed_this_step

Number of trials in status COMPLETED or EARLY_STOPPED for the current generation step of this strategy. We include early stopped trials because their data will be used in the model, so they are completed from the model’s point of view and should count towards that total.

property num_running_trials_this_step

Number of trials in status RUNNING for the current generation step of this strategy.

property trial_indices_by_step

Find trials in experiment that are not mapped to a generation step yet and add them to the mapping of trials by generation step.

property trials_as_df

Puts information on individual trials into a data frame for easy viewing. For example: Gen. Step | Model | Trial Index | Trial Status | Arm Parameterizations 0 | Sobol | 0 | RUNNING | {“0_0”:{“x”:9.17…}}

property uses_non_registered_models

Whether this generation strategy involves models that are not registered and therefore cannot be stored.

Generation Node

class ax.modelbridge.generation_node.GenerationNode(model_specs: List[ax.modelbridge.model_spec.ModelSpec], best_model_selector: Optional[ax.modelbridge.cross_validation.BestModelSelector] = None)[source]

Bases: object

Base class for generation node, capable of fitting one or more model specs under the hood and generating candidates from them.

fit(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, **kwargs: Any) → None[source]

Fits the specified models to the given experiment + data using the model kwargs set on each corresponding model spec and the kwargs passed to this method.

NOTE: Local kwargs take precedence over the ones stored in ModelSpec.model_kwargs.

gen(n: Optional[int] = None, pending_observations: Optional[Dict[str, List[ax.core.observation.ObservationFeatures]]] = None) → ax.core.generator_run.GeneratorRun[source]

Picks a fitted model, from which to generate candidates (via self._pick_fitted_model_to_gen_from) and generates candidates from it. Uses the model_gen_kwargs set on the selected ModelSpec alongside any kwargs passed in to this function (with local kwargs) taking precedent.

Parameters

  • n – Optional nteger representing how many arms should be in the generator run produced by this method. When this is None, n will be determined by the ModelSpec that we are generating from.

  • pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.

NOTE: Models must have been fit prior to calling gen. NOTE: Some underlying models may ignore the n argument and produce a

model-determined number of arms. In that case this method will also output a generator run with number of arms (that can differ from n).

model_specs: List[ax.modelbridge.model_spec.ModelSpec]
update(experiment: ax.core.experiment.Experiment, new_data: ax.core.data.Data) → None[source]

Updates the specified models on the given experiment + new data.

class ax.modelbridge.generation_node.GenerationStep(model: Union[ax.modelbridge.registry.ModelRegistryBase, Callable[[…], ax.modelbridge.base.ModelBridge]], num_trials: int, model_kwargs: Optional[Dict[str, Any]] = None, model_gen_kwargs: Optional[Dict[str, Any]] = None, min_trials_observed: int = 0, max_parallelism: Optional[int] = None, use_update: bool = False, enforce_num_trials: bool = True, should_deduplicate: bool = False, index: int = - 1)[source]

Bases: ax.modelbridge.generation_node.GenerationNode, ax.utils.common.base.SortableBase

One step in the generation strategy, corresponds to a single model. Describes the model, how many trials will be generated with this model, what minimum number of observations is required to proceed to the next model, etc.

NOTE: Model can be specified either from the model registry (ax.modelbridge.registry.Models or using a callable model constructor. Only models from the registry can be saved, and thus optimization can only be resumed if interrupted when using models from the registry.

Parameters

  • model – A member of Models enum or a callable returning an instance of ModelBridge with an instantiated underlying Model. Refer to ax/modelbridge/factory.py for examples of such callables.

  • num_trials – How many trials to generate with the model from this step. If set to -1, trials will continue to be generated from this model as long as generation_strategy.gen is called (available only for the last of the generation steps).

  • min_trials_observed – How many trials must be completed before the generation strategy can proceed to the next step. Defaults to 0. If num_trials of a given step have been generated but min_trials_ observed have not been completed, a call to generation_strategy.gen will fail with a DataRequiredError.

  • max_parallelism – How many trials generated in the course of this step are allowed to be run (i.e. have trial.status of RUNNING) simultaneously. If max_parallelism trials from this step are already running, a call to generation_strategy.gen will fail with a MaxParallelismReached Exception, indicating that more trials need to be completed before generating and running next trials.

  • use_update – Whether to use model_bridge.update instead or reinstantiating model + bridge on every call to gen within a single generation step. NOTE: use of update on stateful models that do not implement _get_state may result in inability to correctly resume a generation strategy from a serialized state.

  • enforce_num_trials – Whether to enforce that only num_trials are generated from the given step. If False and num_trials have been generated, but min_trials_observed have not been completed, generation_strategy.gen will continue generating trials from the current step, exceeding num_ trials for it. Allows to avoid DataRequiredError, but delays proceeding to next generation step.

  • model_kwargs – Dictionary of kwargs to pass into the model constructor on instantiation. E.g. if model is Models.SOBOL, kwargs will be applied as Models.SOBOL(**model_kwargs); if model is get_sobol, get_sobol( **model_kwargs). NOTE: if generation strategy is interrupted and resumed from a stored snapshot and its last used model has state saved on its generator runs, model_kwargs is updated with the state dict of the model, retrieved from the last generator run of this generation strategy.

  • model_gen_kwargs – Each call to generation_strategy.gen performs a call to the step’s model’s gen under the hood; model_gen_kwargs will be passed to the model’s gen like so: model.gen(**model_gen_kwargs).

  • index – Index of this generation step, for use internally in Generation Strategy. Do not assign as it will be reassigned when instantiating GenerationStrategy with a list of its steps.

  • should_deduplicate – Whether to deduplicate the parameters of proposed arms against those of previous arms via rejection sampling. If this is True, the generation strategy will discard generator runs produced from the generation step that has should_deduplicate=True if they contain arms already present on the experiment and replace them with new generator runs. If no generator run with entirely unique arms could be produced in 5 attempts, a GenerationStrategyRepeatedPoints error will be raised, as we assume that the optimization converged when the model can no longer suggest unique arms.

enforce_num_trials: bool = True
index: int = -1
max_parallelism: Optional[int] = None
min_trials_observed: int = 0
model: Union[ax.modelbridge.registry.ModelRegistryBase, Callable[[…], ax.modelbridge.base.ModelBridge]]
model_gen_kwargs: Optional[Dict[str, Any]] = None
model_kwargs: Optional[Dict[str, Any]] = None
property model_name
property model_spec
num_trials: int
should_deduplicate: bool = False
use_update: bool = False

Registry

class ax.modelbridge.registry.ModelRegistryBase(value)[source]

Bases: enum.Enum

Base enum that provides instrumentation of __call__ on enum values, for enums that link their values to ModelSetup-s like Models.

property model_bridge_class

Type of ModelBridge used for the given model+bridge setup.

property model_class

Type of Model used for the given model+bridge setup.

view_defaults() → Tuple[Dict[str, Any], Dict[str, Any]][source]

Obtains the default keyword arguments for the model and the modelbridge specified through the Models enum, for ease of use in notebook environment, since models and bridges cannot be inspected directly through the enum.

Returns

A tuple of default keyword arguments for the model and the model bridge.

view_kwargs() → Tuple[Dict[str, Any], Dict[str, Any]][source]

Obtains annotated keyword arguments that the model and the modelbridge (corresponding to a given member of the Models enum) constructors expect.

Returns

A tuple of annotated keyword arguments for the model and the model bridge.

class ax.modelbridge.registry.ModelSetup(bridge_class: Type[ModelBridge], model_class: Type[Model], transforms: List[Type[Transform]], standard_bridge_kwargs: Optional[Dict[str, Any]] = None, not_saved_model_kwargs: Optional[List[str]] = None)[source]

Bases: tuple

A model setup defines a coupled combination of a model, a model bridge, standard set of transforms, and standard model bridge keyword arguments. This coupled combination yields a given standard modeling strategy in Ax, such as BoTorch GP+EI, a Thompson sampler, or a Sobol quasirandom generator.

property bridge_class

Alias for field number 0

property model_class

Alias for field number 1

property not_saved_model_kwargs

Alias for field number 4

property standard_bridge_kwargs

Alias for field number 3

property transforms

Alias for field number 2

class ax.modelbridge.registry.Models(value)[source]

Bases: ax.modelbridge.registry.ModelRegistryBase

Registry of available models.

Uses MODEL_KEY_TO_MODEL_SETUP to retrieve settings for model and model bridge, by the key stored in the enum value.

To instantiate a model in this enum, simply call an enum member like so: Models.SOBOL(search_space=search_space) or Models.GPEI(experiment=experiment, data=data). Keyword arguments specified to the call will be passed into the model or the model bridge constructors according to their keyword.

For instance, Models.SOBOL(search_space=search_space, scramble=False) will instantiate a RandomModelBridge(search_space=search_space) with a SobolGenerator(scramble=False) underlying model.

ALEBO = 'ALEBO'
ALEBO_INITIALIZER = 'ALEBO_Initializer'
BOTORCH = 'BO'
BOTORCH_MODULAR = 'BoTorch'
BO_MIXED = 'BO_MIXED'
EMPIRICAL_BAYES_THOMPSON = 'EB'
FACTORIAL = 'Factorial'
FULLYBAYESIAN = 'FullyBayesian'
FULLYBAYESIANMOO = 'FullyBayesianMOO'
GPEI = 'GPEI'
GPKG = 'GPKG'
GPMES = 'GPMES'
MOO = 'MOO'
MOO_MODULAR = 'MOO_Modular'
SOBOL = 'Sobol'
ST_MTGP = 'ST_MTGP'
ST_MTGP_NEHVI = 'ST_MTGP_NEHVI'
THOMPSON = 'Thompson'
UNIFORM = 'Uniform'
ax.modelbridge.registry.get_model_from_generator_run(generator_run: ax.core.generator_run.GeneratorRun, experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, models_enum: Optional[Type[ax.modelbridge.registry.ModelRegistryBase]] = None, after_gen: bool = True) → ax.modelbridge.base.ModelBridge[source]

Reinstantiate a model from model key and kwargs stored on a given generator run, with the given experiment and the data to initialize the model with.

Note: requires that the model that was used to get the generator run, is part of the Models registry enum.

Parameters

  • generator_run – A GeneratorRun created by the model we are looking to reinstantiate.

  • experiment – The experiment for which the model is reinstantiated.

  • data – Data, with which to reinstantiate the model.

  • models_enum – Optional subclass of Models registry, from which to obtain the settings of the model. Useful only if the generator run was created via a model that could not be included into the main registry, but can still be represented as a ModelSetup and was added to a registry that extends Models.

  • after_gen – Whether to reinstantiate the model in the state, in which it was after it created this generator run, as opposed to before. Defaults to True, useful when reinstantiating the model to resume optimization, rather than to recreate its state at the time of generation. TO recreate state at the time of generation, set to False.

ax.modelbridge.registry.logger = <Logger ax.modelbridge.registry (DEBUG)>

Module containing a registry of standard models (and generators, samplers etc.) such as Sobol generator, GP+EI, Thompson sampler, etc.

Use of Models enum allows for serialization and reinstantiation of models and generation strategies from generator runs they produced. To reinstantiate a model from generator run, use get_model_from_generator_run utility from this module.

Factory

ax.modelbridge.factory.DEFAULT_EHVI_BATCH_LIMIT = 5

Module containing functions that generate standard models, such as Sobol, GP+EI, etc.

Note: a special case here is a composite generator, which requires an additional GenerationStrategy and is able to delegate work to multiple models (for instance, to a random model to generate the first trial, and to an optimization model for subsequent trials).

ax.modelbridge.factory.get_GPEI(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu')) → ax.modelbridge.torch.TorchModelBridge[source]

Instantiates a GP model that generates points with EI.

ax.modelbridge.factory.get_GPKG(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, cost_intercept: float = 0.01, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), transforms: List[Type[ax.modelbridge.transforms.base.Transform]] = [<class 'ax.modelbridge.transforms.remove_fixed.RemoveFixed'>, <class 'ax.modelbridge.transforms.choice_encode.OrderedChoiceEncode'>, <class 'ax.modelbridge.transforms.one_hot.OneHot'>, <class 'ax.modelbridge.transforms.int_to_float.IntToFloat'>, <class 'ax.modelbridge.transforms.log.Log'>, <class 'ax.modelbridge.transforms.logit.Logit'>, <class 'ax.modelbridge.transforms.unit_x.UnitX'>, <class 'ax.modelbridge.transforms.ivw.IVW'>, <class 'ax.modelbridge.transforms.derelativize.Derelativize'>, <class 'ax.modelbridge.transforms.standardize_y.StandardizeY'>], transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, **kwargs: Any) → ax.modelbridge.torch.TorchModelBridge[source]

Instantiates a GP model that generates points with KG.

ax.modelbridge.factory.get_GPMES(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, cost_intercept: float = 0.01, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), transforms: List[Type[ax.modelbridge.transforms.base.Transform]] = [<class 'ax.modelbridge.transforms.remove_fixed.RemoveFixed'>, <class 'ax.modelbridge.transforms.choice_encode.OrderedChoiceEncode'>, <class 'ax.modelbridge.transforms.one_hot.OneHot'>, <class 'ax.modelbridge.transforms.int_to_float.IntToFloat'>, <class 'ax.modelbridge.transforms.log.Log'>, <class 'ax.modelbridge.transforms.logit.Logit'>, <class 'ax.modelbridge.transforms.unit_x.UnitX'>, <class 'ax.modelbridge.transforms.ivw.IVW'>, <class 'ax.modelbridge.transforms.derelativize.Derelativize'>, <class 'ax.modelbridge.transforms.standardize_y.StandardizeY'>], transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, **kwargs: Any) → ax.modelbridge.torch.TorchModelBridge[source]

Instantiates a GP model that generates points with MES.

ax.modelbridge.factory.get_MOO_EHVI(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, objective_thresholds: Optional[List[ax.core.outcome_constraint.ObjectiveThreshold]] = None, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: Optional[torch.device] = None) → ax.modelbridge.multi_objective_torch.MultiObjectiveTorchModelBridge[source]

Instantiates a multi-objective model that generates points with EHVI.

Requires objective_thresholds, a list of ax.core.ObjectiveThresholds, for every objective being optimized. An arm only improves hypervolume if it is strictly better than all objective thresholds.

objective_thresholds can be passed in the optimization_config or passed directly here.

ax.modelbridge.factory.get_MOO_NEHVI(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, objective_thresholds: Optional[List[ax.core.outcome_constraint.ObjectiveThreshold]] = None, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, use_input_warping: bool = False, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None) → ax.modelbridge.multi_objective_torch.MultiObjectiveTorchModelBridge[source]

Instantiates a multi-objective model using qNEHVI.

ax.modelbridge.factory.get_MOO_PAREGO(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, objective_thresholds: Optional[List[ax.core.outcome_constraint.ObjectiveThreshold]] = None, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu')) → ax.modelbridge.multi_objective_torch.MultiObjectiveTorchModelBridge[source]

Instantiates a multi-objective model that generates points with ParEGO.

qParEGO optimizes random augmented chebyshev scalarizations of the multiple objectives. This allows it to explore non-convex pareto frontiers.

ax.modelbridge.factory.get_MOO_RS(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, objective_thresholds: Optional[List[ax.core.outcome_constraint.ObjectiveThreshold]] = None, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu')) → ax.modelbridge.multi_objective_torch.MultiObjectiveTorchModelBridge[source]

Instantiates a Random Scalarization multi-objective model.

Chooses a different random linear scalarization of the objectives for generating each new candidate arm. This will only explore the convex hull of the pareto frontier.

ax.modelbridge.factory.get_MTGP(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, trial_index: Optional[int] = None) → ax.modelbridge.torch.TorchModelBridge[source]

Instantiates a Multi-task Gaussian Process (MTGP) model that generates points with EI.

If the input experiment is a MultiTypeExperiment then a Multi-type Multi-task GP model will be instantiated. Otherwise, the model will be a Single-type Multi-task GP.

ax.modelbridge.factory.get_MTGP_NEHVI(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, objective_thresholds: Optional[List[ax.core.outcome_constraint.ObjectiveThreshold]] = None, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), trial_index: Optional[int] = None) → ax.modelbridge.torch.TorchModelBridge[source]

Instantiates a Multi-task Gaussian Process (MTGP) model that generates points with qNEHVI.

If the input experiment is a MultiTypeExperiment then a Multi-type Multi-task GP model will be instantiated. Otherwise, the model will be a Single-type Multi-task GP.

ax.modelbridge.factory.get_MTGP_PAREGO(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, trial_index: Optional[int] = None, objective_thresholds: Optional[List[ax.core.outcome_constraint.ObjectiveThreshold]] = None, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu')) → ax.modelbridge.multi_objective_torch.MultiObjectiveTorchModelBridge[source]

Instantiates a multi-objective, multi-task model that uses qParEGO.

qParEGO optimizes random augmented chebyshev scalarizations of the multiple objectives. This allows it to explore non-convex pareto frontiers.

ax.modelbridge.factory.get_botorch(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), transforms: List[Type[ax.modelbridge.transforms.base.Transform]] = [<class 'ax.modelbridge.transforms.remove_fixed.RemoveFixed'>, <class 'ax.modelbridge.transforms.choice_encode.OrderedChoiceEncode'>, <class 'ax.modelbridge.transforms.one_hot.OneHot'>, <class 'ax.modelbridge.transforms.int_to_float.IntToFloat'>, <class 'ax.modelbridge.transforms.log.Log'>, <class 'ax.modelbridge.transforms.logit.Logit'>, <class 'ax.modelbridge.transforms.unit_x.UnitX'>, <class 'ax.modelbridge.transforms.ivw.IVW'>, <class 'ax.modelbridge.transforms.derelativize.Derelativize'>, <class 'ax.modelbridge.transforms.standardize_y.StandardizeY'>], transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, model_constructor: Callable[[List[torch.Tensor], List[torch.Tensor], List[torch.Tensor], List[int], List[int], List[str], Optional[Dict[str, torch.Tensor]], Any], botorch.models.model.Model] = <function get_and_fit_model>, model_predictor: Callable[[botorch.models.model.Model, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]] = <function predict_from_model>, acqf_constructor: Callable[[botorch.models.model.Model, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[torch.Tensor], Optional[torch.Tensor], Any], botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEI>, acqf_optimizer: Callable[[botorch.acquisition.acquisition.AcquisitionFunction, torch.Tensor, int, Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[List[Tuple[torch.Tensor, torch.Tensor, float]]], Optional[Dict[int, float]], Optional[Callable[[torch.Tensor], torch.Tensor]], Any], Tuple[torch.Tensor, torch.Tensor]] = <function scipy_optimizer>, refit_on_cv: bool = False, refit_on_update: bool = True, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None) → ax.modelbridge.torch.TorchModelBridge[source]

Instantiates a BotorchModel.

ax.modelbridge.factory.get_empirical_bayes_thompson(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, num_samples: int = 10000, min_weight: Optional[float] = None, uniform_weights: bool = False) → ax.modelbridge.discrete.DiscreteModelBridge[source]

Instantiates an empirical Bayes / Thompson sampling model.

ax.modelbridge.factory.get_factorial(search_space: ax.core.search_space.SearchSpace) → ax.modelbridge.discrete.DiscreteModelBridge[source]

Instantiates a factorial generator.

ax.modelbridge.factory.get_sobol(search_space: ax.core.search_space.SearchSpace, seed: Optional[int] = None, deduplicate: bool = False, init_position: int = 0, scramble: bool = True, fallback_to_sample_polytope: bool = False) → ax.modelbridge.random.RandomModelBridge[source]

Instantiates a Sobol sequence quasi-random generator.

Parameters

  • search_space – Sobol generator search space.

  • kwargs – Custom args for sobol generator.

Returns

RandomModelBridge, with SobolGenerator as model.

ax.modelbridge.factory.get_thompson(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, num_samples: int = 10000, min_weight: Optional[float] = None, uniform_weights: bool = False) → ax.modelbridge.discrete.DiscreteModelBridge[source]

Instantiates a Thompson sampling model.

ax.modelbridge.factory.get_uniform(search_space: ax.core.search_space.SearchSpace, deduplicate: bool = False, seed: Optional[int] = None) → ax.modelbridge.random.RandomModelBridge[source]

Instantiate uniform generator.

Parameters

  • search_space – Uniform generator search space.

  • kwargs – Custom args for uniform generator.

Returns

RandomModelBridge, with UniformGenerator as model.

ModelSpec

class ax.modelbridge.model_spec.FactoryFunctionModelSpec(model_enum: Union[ax.modelbridge.registry.ModelRegistryBase, NoneType] = None, model_kwargs: Union[Dict[str, Any], NoneType] = None, model_gen_kwargs: Union[Dict[str, Any], NoneType] = None, model_cv_kwargs: Union[Dict[str, Any], NoneType] = None, fixed_features: Union[ax.core.observation.ObservationFeatures, NoneType] = None, _fitted_model: Union[ax.modelbridge.base.ModelBridge, NoneType] = None, _cv_results: Union[List[ax.modelbridge.cross_validation.CVResult], NoneType] = None, _diagnostics: Union[Dict[str, Dict[str, float]], NoneType] = None, factory_function: Union[Callable[…, ax.modelbridge.base.ModelBridge], NoneType] = None)[source]

Bases: ax.modelbridge.model_spec.ModelSpec

factory_function: Optional[Callable[[…], ax.modelbridge.base.ModelBridge]] = None
fit(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, search_space: Optional[ax.core.search_space.SearchSpace] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, **model_kwargs: Any) → None[source]

Fits the specified model on the given experiment + data using the model kwargs set on the model spec, alongside any passed down as kwargs to this function (local kwargs take precedent)

model_enum: Optional[ax.modelbridge.registry.ModelRegistryBase] = None
property model_key

Key string to identify the model used by this ModelSpec.

class ax.modelbridge.model_spec.ModelSpec(model_enum: ax.modelbridge.registry.ModelRegistryBase, model_kwargs: Union[Dict[str, Any], NoneType] = None, model_gen_kwargs: Union[Dict[str, Any], NoneType] = None, model_cv_kwargs: Union[Dict[str, Any], NoneType] = None, fixed_features: Union[ax.core.observation.ObservationFeatures, NoneType] = None, _fitted_model: Union[ax.modelbridge.base.ModelBridge, NoneType] = None, _cv_results: Union[List[ax.modelbridge.cross_validation.CVResult], NoneType] = None, _diagnostics: Union[Dict[str, Dict[str, float]], NoneType] = None)[source]

Bases: ax.utils.common.base.Base

cross_validate() → Tuple[Optional[List[ax.modelbridge.cross_validation.CVResult]], Optional[Dict[str, Dict[str, float]]]][source]

Call cross_validate, compute_diagnostics and cache the results If the model cannot be cross validated, warn and return None

property cv_results

Cached CV results from self.cross_validate() if it has been successfully called

property diagnostics

Cached CV diagnostics from self.cross_validate() if it has been successfully called

fit(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, **model_kwargs: Any) → None[source]

Fits the specified model on the given experiment + data using the model kwargs set on the model spec, alongside any passed down as kwargs to this function (local kwargs take precedent)

property fitted_model

Returns the fitted Ax model, asserting fit() was called

fixed_features: Optional[ax.core.observation.ObservationFeatures] = None
gen(**model_gen_kwargs: Any) → ax.core.generator_run.GeneratorRun[source]

Generates candidates from the fitted model, using the model gen kwargs set on the model spec, alongside any passed as kwargs to this function (local kwargs take precedent)

NOTE: Model must have been fit prior to calling gen()

Parameters

  • n – Integer representing how many arms should be in the generator run produced by this method. NOTE: Some underlying models may ignore the n and produce a model-determined number of arms. In that case this method will also output a generator run with number of arms that can differ from n.

  • pending_observations – A map from metric name to pending observations for that metric, used by some models to avoid resuggesting points that are currently being evaluated.

model_cv_kwargs: Optional[Dict[str, Any]] = None
model_enum: ax.modelbridge.registry.ModelRegistryBase
model_gen_kwargs: Optional[Dict[str, Any]] = None
property model_key

Key string to identify the model used by this ModelSpec.

model_kwargs: Optional[Dict[str, Any]] = None
update(experiment: ax.core.experiment.Experiment, new_data: ax.core.data.Data) → None[source]

Updates the current fitted model on the given experiment + new data

Model must have been fit prior to calling update()

Model Bridges

Base Model Bridge

class ax.modelbridge.base.BaseGenArgs(search_space: ax.core.search_space.SearchSpace, optimization_config: ax.core.optimization_config.OptimizationConfig, pending_observations: Dict[str, List[ax.core.observation.ObservationFeatures]], fixed_features: ax.core.observation.ObservationFeatures)[source]

Bases: object

fixed_features: ax.core.observation.ObservationFeatures
optimization_config: ax.core.optimization_config.OptimizationConfig
pending_observations: Dict[str, List[ax.core.observation.ObservationFeatures]]
search_space: ax.core.search_space.SearchSpace
class ax.modelbridge.base.ModelBridge(search_space: ax.core.search_space.SearchSpace, model: Any, transforms: Optional[List[Type[ax.modelbridge.transforms.base.Transform]]] = None, experiment: Optional[ax.core.experiment.Experiment] = None, data: Optional[ax.core.data.Data] = None, transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, fit_out_of_design: bool = False, fit_abandoned: bool = False)[source]

Bases: abc.ABC

The main object for using models in Ax.

ModelBridge specifies 3 methods for using models:

  • predict: Make model predictions. This method is not optimized for speed and so should be used primarily for plotting or similar tasks and not inside an optimization loop.

  • gen: Use the model to generate new candidates.

  • cross_validate: Do cross validation to assess model predictions.

ModelBridge converts Ax types like Data and Arm to types that are meant to be consumed by the models. The data sent to the model will depend on the implementation of the subclass, which will specify the actual API for external model.

This class also applies a sequence of transforms to the input data and problem specification which can be used to ensure that the external model receives appropriate inputs.

Subclasses will implement what is here referred to as the “terminal transform,” which is a transform that changes types of the data and problem specification.

cross_validate(cv_training_data: List[ax.core.observation.Observation], cv_test_points: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Make a set of cross-validation predictions.

Parameters

  • cv_training_data – The training data to use for cross validation.

  • cv_test_points – The test points at which predictions will be made.

Returns

A list of predictions at the test points.

evaluate_acquisition_function(observation_features: List[ax.core.observation.ObservationFeatures], search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: numpy.ndarray, objective_thresholds: Optional[numpy.ndarray] = None, outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]] = None, fixed_features: Optional[Dict[int, float]] = None, pending_observations: Optional[List[numpy.ndarray]] = None, acq_options: Optional[Dict[str, Any]] = None) → List[float][source]

Evaluate the acquisition function for given set of observation features.

Parameters

  • observation_features – A list of observation features, representing parameterizations, for which to evaluate the acquisition function.

  • search_space_digest – A dataclass used to compactly represent a search space.

  • objective_weights – The objective is to maximize a weighted sum of the columns of f(x). These are the weights.

  • objective_thresholds – The m-dim tensor of objective thresholds. There is one for each modeled metric.

  • outcome_constraints – A tuple of (A, b). For k outcome constraints and m outputs at f(x), A is (k x m) and b is (k x 1) such that A f(x) <= b. (Not used by single task models)

  • linear_constraints – A tuple of (A, b). For k linear constraints on d-dimensional x, A is (k x d) and b is (k x 1) such that A x <= b.

  • fixed_features – A map {feature_index: value} for features that should be held fixed during the evaluation.

  • pending_observations – A list of m (k_i x d) feature tensors X for m outcomes and k_i pending observations for outcome i.

  • acq_options – Keyword arguments used to contruct the acquisition function.

Returns

A list of acquisition function values, in the same order as the input observation features.

feature_importances(metric_name: str) → Dict[str, float][source]
gen(n: int, search_space: Optional[ax.core.search_space.SearchSpace] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, pending_observations: Optional[Dict[str, List[ax.core.observation.ObservationFeatures]]] = None, fixed_features: Optional[ax.core.observation.ObservationFeatures] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → ax.core.generator_run.GeneratorRun[source]
Parameters

  • n – Number of points to generate

  • search_space – Search space

  • optimization_config – Optimization config

  • pending_observations – A map from metric name to pending observations for that metric.

  • fixed_features – An ObservationFeatures object containing any features that should be fixed at specified values during generation.

  • model_gen_options – A config dictionary that is passed along to the model.

get_training_data() → List[ax.core.observation.Observation][source]

A copy of the (untransformed) data with which the model was fit.

property metric_names

Metric names present in training data.

property model_space

SearchSpace used to fit model.

predict(observation_features: List[ax.core.observation.ObservationFeatures]) → Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]][source]

Make model predictions (mean and covariance) for the given observation features.

Predictions are made for all outcomes. If an out-of-design observation can successfully be transformed, the predicted value will be returned. Othwerise, we will attempt to find that observation in the training data and return the raw value.

Parameters

observation_features – observation features

Returns

2-element tuple containing

  • Dictionary from metric name to list of mean estimates, in same order as observation_features.

  • Nested dictionary with cov[‘metric1’][‘metric2’] a list of cov(metric1@x, metric2@x) for x in observation_features.

property status_quo

Observation corresponding to status quo, if any.

property training_in_design

For each observation in the training data, a bool indicating if it is in-design for the model.

transform_observation_data(observation_data: List[ax.core.observation.ObservationData]) → Any[source]

Applies transforms to given observation features and returns them in the model space.

Parameters

observation_features – ObservationFeatures to be transformed.

Returns

Transformed values. This could be e.g. a torch Tensor, depending on the ModelBridge subclass.

transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → Any[source]

Applies transforms to given observation features and returns them in the model space.

Parameters

observation_features – ObservationFeatures to be transformed.

Returns

Transformed values. This could be e.g. a torch Tensor, depending on the ModelBridge subclass.

transform_optimization_config(optimization_config: ax.core.optimization_config.OptimizationConfig, fixed_features: ax.core.observation.ObservationFeatures) → Any[source]

Applies transforms to given optimization config.

Parameters

  • optimization_config – OptimizationConfig to transform.

  • fixed_features – features which should not be transformed.

Returns

Transformed values. This could be e.g. a torch Tensor, depending on the ModelBridge subclass.

update(new_data: ax.core.data.Data, experiment: ax.core.experiment.Experiment) → None[source]

Update the model bridge and the underlying model with new data. This method should be used instead of fit, in cases where the underlying model does not need to be re-fit from scratch, but rather updated.

Note: update expects only new data (obtained since the model initialization or last update) to be passed in, not all data in the experiment.

Parameters

  • new_data – Data from the experiment obtained since the last call to update.

  • experiment – Experiment, in which this data was obtained.

ax.modelbridge.base.clamp_observation_features(observation_features: List[ax.core.observation.ObservationFeatures], search_space: ax.core.search_space.SearchSpace) → List[ax.core.observation.ObservationFeatures][source]
ax.modelbridge.base.gen_arms(observation_features: List[ax.core.observation.ObservationFeatures], arms_by_signature: Optional[Dict[str, ax.core.arm.Arm]] = None) → Tuple[List[ax.core.arm.Arm], Optional[Dict[str, Optional[Dict[str, Any]]]]][source]

Converts observation features to a tuple of arms list and candidate metadata dict, where arm signatures are mapped to their respective candidate metadata.

ax.modelbridge.base.unwrap_observation_data(observation_data: List[ax.core.observation.ObservationData]) → Tuple[Dict[str, List[float]], Dict[str, Dict[str, List[float]]]][source]

Converts observation data to the format for model prediction outputs. That format assumes each observation data has the same set of metrics.

Array Model Bridge

class ax.modelbridge.array.ArrayModelBridge(search_space: ax.core.search_space.SearchSpace, model: Any, transforms: Optional[List[Type[ax.modelbridge.transforms.base.Transform]]] = None, experiment: Optional[ax.core.experiment.Experiment] = None, data: Optional[ax.core.data.Data] = None, transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, fit_out_of_design: bool = False, fit_abandoned: bool = False)[source]

Bases: ax.modelbridge.base.ModelBridge

A model bridge for using array-based models.

Requires that all non-task parameters have been transformed to RangeParameters.

If there are any (non-task) discrete parameters (e.g. as obtained via a ChoiceEncode transform), those need to be of integer type with parameter space normalized to {0, 1, …, num_choices-1}. The num_choices information is passed to the model and optimization needs to take this into account and return only candidates that take values in this parameter space (specifically, there is no relaxation and no rounding is applied).

All other parameters need to be of float type on a regular (non-log) scale.

This will convert all parameter types to float and put data into arrays.

feature_importances(metric_name: str) → Dict[str, float][source]
model: Any
model_best_point(search_space: Optional[ax.core.search_space.SearchSpace] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, pending_observations: Optional[Dict[str, List[ax.core.observation.ObservationFeatures]]] = None, fixed_features: Optional[ax.core.observation.ObservationFeatures] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → Optional[Tuple[ax.core.arm.Arm, Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]][source]
outcomes: List[str]
parameters: List[str]
class ax.modelbridge.array.ArrayModelGenArgs(search_space_digest: ax.core.search_space.SearchSpaceDigest, objective_weights: numpy.ndarray, outcome_constraints: Union[Tuple[numpy.ndarray, numpy.ndarray], NoneType], linear_constraints: Union[Tuple[numpy.ndarray, numpy.ndarray], NoneType], fixed_features: Union[Dict[int, float], NoneType], pending_observations: Union[List[numpy.ndarray], NoneType], rounding_func: Callable[[numpy.ndarray], numpy.ndarray], extra_model_gen_kwargs: Dict[str, Any])[source]

Bases: object

extra_model_gen_kwargs: Dict[str, Any]
fixed_features: Optional[Dict[int, float]]
linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]]
objective_weights: numpy.ndarray
outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]]
pending_observations: Optional[List[numpy.ndarray]]
rounding_func: Callable[[numpy.ndarray], numpy.ndarray]
search_space_digest: ax.core.search_space.SearchSpaceDigest
ax.modelbridge.array.validate_optimization_config(optimization_config: ax.core.optimization_config.OptimizationConfig, outcomes: List[str]) → None[source]

Validate optimization config against model fitted outcomes.

Parameters

  • optimization_config – Config to validate.

  • outcomes – List of metric names w/ valid model fits.

Raises

ValueError if –

  1. Relative constraints are found 2. Optimization metrics are not present in model fitted outcomes.

Discrete Model Bridge

class ax.modelbridge.discrete.DiscreteModelBridge(search_space: ax.core.search_space.SearchSpace, model: Any, transforms: Optional[List[Type[ax.modelbridge.transforms.base.Transform]]] = None, experiment: Optional[ax.core.experiment.Experiment] = None, data: Optional[ax.core.data.Data] = None, transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, fit_out_of_design: bool = False, fit_abandoned: bool = False)[source]

Bases: ax.modelbridge.base.ModelBridge

A model bridge for using models based on discrete parameters.

Requires that all parameters have been transformed to ChoiceParameters.

model: ax.models.discrete_base.DiscreteModel
outcomes: List[str]
parameters: List[str]
search_space: Optional[ax.core.search_space.SearchSpace]

NumPy Model Bridge

class ax.modelbridge.numpy.NumpyModelBridge(search_space: ax.core.search_space.SearchSpace, model: Any, transforms: Optional[List[Type[ax.modelbridge.transforms.base.Transform]]] = None, experiment: Optional[ax.core.experiment.Experiment] = None, data: Optional[ax.core.data.Data] = None, transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, fit_out_of_design: bool = False, fit_abandoned: bool = False)[source]

Bases: ax.modelbridge.array.ArrayModelBridge

A model bridge for using numpy array-based models.

This model bridge interfaces with NumpyModel.

Requires that all parameters have been transformed to RangeParameters or FixedParameters with float type and no log scale.

model: ax.models.numpy_base.NumpyModel
outcomes: List[str]
parameters: List[str]

Random Model Bridge

class ax.modelbridge.random.RandomModelBridge(search_space: ax.core.search_space.SearchSpace, model: Any, transforms: Optional[List[Type[ax.modelbridge.transforms.base.Transform]]] = None, experiment: Optional[ax.core.experiment.Experiment] = None, data: Optional[ax.core.data.Data] = None, transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, fit_out_of_design: bool = False, fit_abandoned: bool = False)[source]

Bases: ax.modelbridge.base.ModelBridge

A model bridge for using purely random ‘models’. Data and optimization configs are not required.

This model bridge interfaces with RandomModel.

model

A RandomModel used to generate candidates (note: this an awkward use of the word ‘model’).

Type

ax.models.random.base.RandomModel

parameters

Params found in search space on modelbridge init.

Type

List[str]

model: ax.models.random.base.RandomModel
parameters: List[str]
update(new_data: ax.core.data.Data, experiment: ax.core.experiment.Experiment) → None[source]

Update the model bridge and the underlying model with new data. This method should be used instead of fit, in cases where the underlying model does not need to be re-fit from scratch, but rather updated.

Note: update expects only new data (obtained since the model initialization or last update) to be passed in, not all data in the experiment.

Parameters

  • new_data – Data from the experiment obtained since the last call to update.

  • experiment – Experiment, in which this data was obtained.

Torch Model Bridge

class ax.modelbridge.torch.TorchModelBridge(experiment: ax.core.experiment.Experiment, search_space: ax.core.search_space.SearchSpace, data: ax.core.data.Data, model: ax.models.torch_base.TorchModel, transforms: List[Type[ax.modelbridge.transforms.base.Transform]], transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, torch_dtype: Optional[torch.dtype] = None, torch_device: Optional[torch.device] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, fit_out_of_design: bool = False, default_model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.array.ArrayModelBridge

A model bridge for using torch-based models.

Specifies an interface that is implemented by TorchModel. In particular, model should have methods fit, predict, and gen. See TorchModel for the API for each of these methods.

Requires that all parameters have been transformed to RangeParameters or FixedParameters with float type and no log scale.

This class converts Ax parameter types to torch tensors before passing them to the model.

model: Optional[ax.models.torch_base.TorchModel]

Multi-Objective Torch Model Bridge

class ax.modelbridge.multi_objective_torch.MultiObjectiveTorchModelBridge(experiment: ax.core.experiment.Experiment, search_space: ax.core.search_space.SearchSpace, data: ax.core.data.Data, model: ax.models.torch_base.TorchModel, transforms: List[Type[ax.modelbridge.transforms.base.Transform]], transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]]] = None, torch_dtype: Optional[torch.dtype] = None, torch_device: Optional[torch.device] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ax.core.observation.ObservationFeatures] = None, optimization_config: Optional[ax.core.optimization_config.MultiObjectiveOptimizationConfig] = None, fit_out_of_design: bool = False, objective_thresholds: Optional[List[ax.core.outcome_constraint.ObjectiveThreshold]] = None, default_model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.torch.TorchModelBridge

A model bridge for using multi-objective torch-based models.

Specifies an interface that is implemented by MultiObjectiveTorchModel. In particular, model should have methods fit, predict, and gen. See MultiObjectiveTorchModel for the API for each of these methods.

Requires that all parameters have been transformed to RangeParameters or FixedParameters with float type and no log scale.

This class converts Ax parameter types to torch tensors before passing them to the model.

gen(n: int, search_space: Optional[ax.core.search_space.SearchSpace] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, pending_observations: Optional[Dict[str, List[ax.core.observation.ObservationFeatures]]] = None, fixed_features: Optional[ax.core.observation.ObservationFeatures] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None) → ax.core.generator_run.GeneratorRun[source]
Parameters

  • n – Number of points to generate

  • search_space – Search space

  • optimization_config – Optimization config

  • pending_observations – A map from metric name to pending observations for that metric.

  • fixed_features – An ObservationFeatures object containing any features that should be fixed at specified values during generation.

  • model_gen_options – A config dictionary that is passed along to the model.

infer_objective_thresholds(search_space: Optional[ax.core.search_space.SearchSpace] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, fixed_features: Optional[ax.core.observation.ObservationFeatures] = None) → List[ax.core.outcome_constraint.ObjectiveThreshold][source]

Infer objective thresholds.

This method uses the model-estimated Pareto frontier over the in-sample points to infer absolute (not relativized) objective thresholds.

This uses a heuristic that sets the objective threshold to be a scaled nadir point, where the nadir point is scaled back based on the range of each objective across the current in-sample Pareto frontier.

untransform_objective_thresholds(objective_thresholds: torch.Tensor, objective_weights: torch.Tensor, bounds: List[Tuple[Union[int, float], Union[int, float]]], fixed_features: Optional[Dict[int, float]]) → List[ax.core.outcome_constraint.ObjectiveThreshold][source]

Utilities

General Utilities

ax.modelbridge.modelbridge_utils.array_to_observation_data(f: numpy.ndarray, cov: numpy.ndarray, outcomes: List[str]) → List[ax.core.observation.ObservationData][source]

Convert arrays of model predictions to a list of ObservationData.

Parameters

  • f – An (n x m) array

  • cov – An (n x m x m) array

  • outcomes – A list of d outcome names

Returns: A list of n ObservationData

ax.modelbridge.modelbridge_utils.extract_objective_thresholds(objective_thresholds: List[ax.core.outcome_constraint.ObjectiveThreshold], objective: ax.core.objective.Objective, outcomes: List[str]) → Optional[numpy.ndarray][source]

Extracts objective thresholds’ values, in the order of outcomes.

Will return None if no objective thresholds, otherwise the extracted tensor will be the same length as outcomes.

If one objective threshold is specified, they must be specified for every metric in the objective.

Outcomes that are not part of an objective will be given a threshold of 0 in this tensor, under the assumption that its value will not be used. Note that setting it to 0 for an outcome that is part of the objective would be incorrect, hence we validate that all objective metrics are represented.

Parameters

  • objective_thresholds – Objective thresholds to extract values from.

  • objective – The corresponding Objective, for validation purposes.

  • outcomes – n-length list of names of metrics.

Returns

(n,) array of thresholds

ax.modelbridge.modelbridge_utils.extract_objective_weights(objective: ax.core.objective.Objective, outcomes: List[str]) → numpy.ndarray[source]

Extract a weights for objectives.

Weights are for a maximization problem.

Give an objective weight to each modeled outcome. Outcomes that are modeled but not part of the objective get weight 0.

In the single metric case, the objective is given either +/- 1, depending on the minimize flag.

In the multiple metric case, each objective is given the input weight, multiplied by the minimize flag.

Parameters

  • objective – Objective to extract weights from.

  • outcomes – n-length list of names of metrics.

Returns

n-length list of weights.

ax.modelbridge.modelbridge_utils.extract_outcome_constraints(outcome_constraints: List[ax.core.outcome_constraint.OutcomeConstraint], outcomes: List[str]) → Optional[Tuple[numpy.ndarray, numpy.ndarray]][source]
ax.modelbridge.modelbridge_utils.extract_parameter_constraints(parameter_constraints: List[ax.core.parameter_constraint.ParameterConstraint], param_names: List[str]) → Optional[Tuple[numpy.ndarray, numpy.ndarray]][source]

Extract parameter constraints.

ax.modelbridge.modelbridge_utils.extract_search_space_digest(search_space: ax.core.search_space.SearchSpace, param_names: List[str]) → ax.core.search_space.SearchSpaceDigest[source]

Extract basic parameter prpoerties from a search space.

ax.modelbridge.modelbridge_utils.get_fixed_features(fixed_features: ax.core.observation.ObservationFeatures, param_names: List[str]) → Optional[Dict[int, float]][source]

Reformat a set of fixed_features.

ax.modelbridge.modelbridge_utils.get_pareto_frontier_and_configs(modelbridge: modelbridge_module.array.ArrayModelBridge, observation_features: List[ObservationFeatures], observation_data: Optional[List[ObservationData]] = None, objective_thresholds: Optional[TRefPoint] = None, optimization_config: Optional[MultiObjectiveOptimizationConfig] = None, arm_names: Optional[List[Optional[str]]] = None, use_model_predictions: bool = True, transform_outcomes_and_configs: bool = True) → Tuple[List[Observation], Tensor, Tensor, Optional[Tensor]][source]

Helper that applies transforms and calls frontier_evaluator.

Returns the frontier_evaluator configs in addition to the Pareto observations.

Parameters

  • modelbridge – Modelbridge used to predict metrics outcomes.

  • observation_features – Observation features to consider for the Pareto frontier.

  • observation_data – Data for computing the Pareto front, unless observation_features are provided and model_predictions is True.

  • objective_thresholds – Metric values bounding the region of interest in the objective outcome space; used to override objective thresholds specified in optimization_config, if necessary.

  • optimization_config – Multi-objective optimization config.

  • arm_names – Arm names for each observation in observation_features.

  • use_model_predictions – If True, will use model predictions at observation_features to compute Pareto front. If False, will use observation_data directly to compute Pareto front, ignoring observation_features.

  • transform_outcomes_and_configs – If True, will transform the optimization config, observation features and observation data, before calling frontier_evaluator, then will untransform all of the above before returning the observations.

Returns: Four-item tuple of:

  • frontier_observations: Observations of points on the pareto frontier,

  • f: n x m tensor representation of the Pareto frontier values where n is the length of frontier_observations and m is the number of metrics,

  • obj_w: m tensor of objective weights,

  • obj_t: m tensor of objective thresholds corresponding to Y, or None if no objective thresholds used.

ax.modelbridge.modelbridge_utils.get_pending_observation_features(experiment: ax.core.experiment.Experiment, include_failed_as_pending: bool = False) → Optional[Dict[str, List[ax.core.observation.ObservationFeatures]]][source]

Computes a list of pending observation features (corresponding to arms that have been generated and deployed in the course of the experiment, but have not been completed with data or to arms that have been abandoned or belong to abandoned trials).

NOTE: Pending observation features are passed to the model to instruct it to not generate the same points again.

Parameters

  • experiment – Experiment, pending features on which we seek to compute.

  • include_failed_as_pending – Whether to include failed trials as pending (for example, to avoid the model suggesting them again).

Returns

An optional mapping from metric names to a list of observation features, pending for that metric (i.e. do not have evaluation data for that metric). If there are no pending features for any of the metrics, return is None.

ax.modelbridge.modelbridge_utils.get_pending_observation_features_based_on_trial_status(experiment: ax.core.experiment.Experiment) → Optional[Dict[str, List[ax.core.observation.ObservationFeatures]]][source]

A faster analogue of get_pending_observation_features that makes assumptions about trials in experiment in order to speed up extraction of pending points.

Assumptions:

  • All arms in all trials in STAGED, RUNNING and ABANDONED statuses are to be considered pending for all outcomes.

  • All arms in all trials in other statuses are to be considered not pending for all outcomes.

This entails:

  • No actual data-fetching for trials to determine whether arms in them are pending for specific outcomes.

  • Even if data is present for some outcomes in RUNNING trials, their arms will still be considered pending for those outcomes.

NOTE: This function should not be used to extract pending features in field experiments, where arms in running trials should not be considered pending if there is data for those arms.

Parameters

experiment – Experiment, pending features on which we seek to compute.

Returns

An optional mapping from metric names to a list of observation features, pending for that metric (i.e. do not have evaluation data for that metric). If there are no pending features for any of the metrics, return is None.

ax.modelbridge.modelbridge_utils.hypervolume(modelbridge: modelbridge_module.array.ArrayModelBridge, observation_features: List[ObservationFeatures], objective_thresholds: Optional[TRefPoint] = None, observation_data: Optional[List[ObservationData]] = None, optimization_config: Optional[MultiObjectiveOptimizationConfig] = None, selected_metrics: Optional[List[str]] = None, use_model_predictions: bool = True) → float[source]

Helper function that computes (feasible) hypervolume.

Parameters

  • modelbridge – The modelbridge.

  • observation_features – The observation features for the in-sample arms.

  • objective_thresholds – The objective thresholds to be used for computing the hypervolume. If None, these are extracted from the optimization config.

  • observation_data – The observed outcomes for the in-sample arms.

  • optimization_config – The optimization config specifying the objectives, objectives thresholds, and outcome constraints.

  • selected_metrics – A list of objective metric names specifying which objectives to use in hypervolume computation. By default, all objectives are used.

  • use_model_predictions – A boolean indicating whether to use model predictions for determining the in-sample Pareto frontier instead of the raw observed values.

Returns

The (feasible) hypervolume.

ax.modelbridge.modelbridge_utils.observation_data_to_array(outcomes: List[str], observation_data: List[ax.core.observation.ObservationData]) → Tuple[numpy.ndarray, numpy.ndarray][source]

Convert a list of Observation data to arrays.

Parameters

observation_data – A list of n ObservationData

Returns

An array of n ObservationData, each containing

  • f: An (n x m) array

  • cov: An (n x m x m) array

ax.modelbridge.modelbridge_utils.observation_features_to_array(parameters: List[str], obsf: List[ax.core.observation.ObservationFeatures]) → numpy.ndarray[source]

Convert a list of Observation features to arrays.

ax.modelbridge.modelbridge_utils.observed_hypervolume(modelbridge: modelbridge_module.array.ArrayModelBridge, objective_thresholds: Optional[TRefPoint] = None, optimization_config: Optional[MultiObjectiveOptimizationConfig] = None, selected_metrics: Optional[List[str]] = None) → float[source]

Calculate hypervolume of a pareto frontier based on observed data.

Given observed data, return the hypervolume of the pareto frontier formed from those outcomes.

Parameters

  • modelbridge – Modelbridge that holds previous training data.

  • objective_thresholds – point defining the origin of hyperrectangles that can contribute to hypervolume.

  • observation_features – observation features to predict. Model’s training data used by default if unspecified.

  • optimization_config – Optimization config

  • selected_metrics – If specified, hypervolume will only be evaluated on the specified subset of metrics. Otherwise, all metrics will be used.

Returns

(float) calculated hypervolume.

ax.modelbridge.modelbridge_utils.observed_pareto_frontier(modelbridge: modelbridge_module.array.ArrayModelBridge, objective_thresholds: Optional[TRefPoint] = None, optimization_config: Optional[MultiObjectiveOptimizationConfig] = None) → List[Observation][source]

Generate a pareto frontier based on observed data. Given observed data (sourced from model training data), return points on the Pareto frontier as Observation-s.

Parameters

  • modelbridge – Modelbridge that holds previous training data.

  • objective_thresholds – Metric values bounding the region of interest in the objective outcome space; used to override objective thresholds in the optimization config, if needed.

  • optimization_config – Multi-objective optimization config.

Returns

Data representing points on the pareto frontier.

ax.modelbridge.modelbridge_utils.pareto_frontier(modelbridge: modelbridge_module.array.ArrayModelBridge, observation_features: List[ObservationFeatures], observation_data: Optional[List[ObservationData]] = None, objective_thresholds: Optional[TRefPoint] = None, optimization_config: Optional[MultiObjectiveOptimizationConfig] = None, arm_names: Optional[List[Optional[str]]] = None, use_model_predictions: bool = True) → List[Observation][source]

Compute the list of points on the Pareto frontier as Observation-s in the untransformed search space.

Parameters

  • modelbridge – Modelbridge used to predict metrics outcomes.

  • observation_features – Observation features to consider for the Pareto frontier.

  • observation_data – Data for computing the Pareto front, unless observation_features are provided and model_predictions is True.

  • objective_thresholds – Metric values bounding the region of interest in the objective outcome space; used to override objective thresholds specified in optimization_config, if necessary.

  • optimization_config – Multi-objective optimization config.

  • arm_names – Arm names for each observation in observation_features.

  • use_model_predictions – If True, will use model predictions at observation_features to compute Pareto front. If False, will use observation_data directly to compute Pareto front, ignoring observation_features.

Returns: Points on the Pareto frontier as Observation-s.

ax.modelbridge.modelbridge_utils.parse_observation_features(X: numpy.ndarray, param_names: List[str], candidate_metadata: Optional[List[Optional[Dict[str, Any]]]] = None) → List[ax.core.observation.ObservationFeatures][source]

Re-format raw model-generated candidates into ObservationFeatures.

Parameters

  • param_names – List of param names.

  • X – Raw np.ndarray of candidate values.

  • candidate_metadata – Model’s metadata for candidates it produced.

Returns

List of candidates, represented as ObservationFeatures.

ax.modelbridge.modelbridge_utils.pending_observations_as_array(pending_observations: Dict[str, List[ax.core.observation.ObservationFeatures]], outcome_names: List[str], param_names: List[str]) → Optional[List[numpy.ndarray]][source]

Re-format pending observations.

Parameters

  • pending_observations – List of raw numpy pending observations.

  • outcome_names – List of outcome names.

  • param_names – List fitted param names.

Returns

Filtered pending observations data, by outcome and param names.

ax.modelbridge.modelbridge_utils.predicted_hypervolume(modelbridge: modelbridge_module.array.ArrayModelBridge, objective_thresholds: Optional[TRefPoint] = None, observation_features: Optional[List[ObservationFeatures]] = None, optimization_config: Optional[MultiObjectiveOptimizationConfig] = None, selected_metrics: Optional[List[str]] = None) → float[source]

Calculate hypervolume of a pareto frontier based on the posterior means of given observation features.

Given a model and features to evaluate calculate the hypervolume of the pareto frontier formed from their predicted outcomes.

Parameters

  • modelbridge – Modelbridge used to predict metrics outcomes.

  • objective_thresholds – point defining the origin of hyperrectangles that can contribute to hypervolume.

  • observation_features – observation features to predict. Model’s training data used by default if unspecified.

  • optimization_config – Optimization config

  • selected_metrics – If specified, hypervolume will only be evaluated on the specified subset of metrics. Otherwise, all metrics will be used.

Returns

calculated hypervolume.

ax.modelbridge.modelbridge_utils.predicted_pareto_frontier(modelbridge: modelbridge_module.array.ArrayModelBridge, objective_thresholds: Optional[TRefPoint] = None, observation_features: Optional[List[ObservationFeatures]] = None, optimization_config: Optional[MultiObjectiveOptimizationConfig] = None) → List[Observation][source]

Generate a Pareto frontier based on the posterior means of given observation features. Given a model and optionally features to evaluate (will use model training data if not specified), use the model to predict which points lie on the Pareto frontier.

Parameters

  • modelbridge – Modelbridge used to predict metrics outcomes.

  • observation_features – Observation features to predict, if provided and use_model_predictions is True.

  • objective_thresholds – Metric values bounding the region of interest in the objective outcome space; used to override objective thresholds specified in optimization_config, if necessary.

  • optimization_config – Multi-objective optimization config.

Returns

Observations representing points on the Pareto frontier.

ax.modelbridge.modelbridge_utils.transform_callback(param_names: List[str], transforms: MutableMapping[str, ax.modelbridge.transforms.base.Transform]) → Callable[[numpy.ndarray], numpy.ndarray][source]

A closure for performing the round trip transformations.

The function round points by de-transforming points back into the original space (done by applying transforms in reverse), and then re-transforming them. This function is specifically for points which are formatted as numpy arrays. This function is passed to _model_gen.

Parameters

  • param_names – Names of parameters to transform.

  • transforms – Ordered set of transforms which were applied to the points.

Returns

a function with for performing the roundtrip transform.

ax.modelbridge.modelbridge_utils.validate_and_apply_final_transform(objective_weights: numpy.ndarray, outcome_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]], linear_constraints: Optional[Tuple[numpy.ndarray, numpy.ndarray]], pending_observations: Optional[List[numpy.ndarray]], objective_thresholds: Optional[numpy.ndarray] = None, final_transform: Callable[[numpy.ndarray], torch.Tensor] = <built-in method tensor of type object>) → Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[Tuple[torch.Tensor, torch.Tensor]], Optional[List[torch.Tensor]], Optional[torch.Tensor]][source]

Cross Validation

class ax.modelbridge.cross_validation.AssessModelFitResult(good_fit_metrics_to_fisher_score: Dict[str, float], bad_fit_metrics_to_fisher_score: Dict[str, float])[source]

Bases: tuple

Container for model fit assessment results

property bad_fit_metrics_to_fisher_score

Alias for field number 1

property good_fit_metrics_to_fisher_score

Alias for field number 0

class ax.modelbridge.cross_validation.BestModelSelector[source]

Bases: object

best_diagnostic(diagnostics: List[Dict[str, Dict[str, float]]]) → int[source]

Return the index of the best diagnostic.

class ax.modelbridge.cross_validation.CVResult(observed: ax.core.observation.Observation, predicted: ax.core.observation.ObservationData)[source]

Bases: tuple

Container for cross validation results.

property observed

Alias for field number 0

property predicted

Alias for field number 1

class ax.modelbridge.cross_validation.CallableEnum(value)[source]

Bases: enum.Enum

An enumeration.

class ax.modelbridge.cross_validation.DiagnosticCriterion(value)[source]

Bases: ax.modelbridge.cross_validation.CallableEnum

An enumeration.

MIN: Callable[[Iterable[numbers.Number]], numbers.Number] = functools.partial(<function amin>)
class ax.modelbridge.cross_validation.MetricAggregation(value)[source]

Bases: ax.modelbridge.cross_validation.CallableEnum

An enumeration.

MEAN: Callable[[Iterable[numbers.Number]], numbers.Number] = functools.partial(<function mean>)
class ax.modelbridge.cross_validation.SingleDiagnosticBestModelSelector(diagnostic: str, metric_aggregation: ax.modelbridge.cross_validation.MetricAggregation, criterion: ax.modelbridge.cross_validation.DiagnosticCriterion)[source]

Bases: ax.modelbridge.cross_validation.BestModelSelector

Choose the best model using a single cross-validation diagnostic.

The input is a list of CVDiagnostics, each corresponding to one model. The specified diagnostic is extracted from each of the CVDiagnostics, its values (each of which corresponds to a separate metric) are aggregated with the aggregation function, the best one is determined with the criterion, and the index of the best diagnostic result is returned.

Example:

::
s = SingleDiagnosticBestModelSelector(

diagnostic = ‘Fisher exact test p’, criterion = DiagnosticCriterion.MIN, metric_aggregation = MetricAggregation.MEAN,

) best_diagnostic_index = s.best_diagnostic(diagnostics)

Parameters

  • diagnostic (str) – The name of the diagnostic to use, which should be a key in CVDiagnostic.

  • metric_aggregation (MetricAggregation) – Callable applied to the values of the diagnostic for a single model to produce a single number.

  • criterion – Callable used to determine which of the (aggregated) diagnostics is the best.

best_diagnostic(diagnostics: List[Dict[str, Dict[str, float]]]) → int[source]

Return the index of the best diagnostic.

ax.modelbridge.cross_validation.assess_model_fit(diagnostics: Dict[str, Dict[str, float]], significance_level: float = 0.1) → ax.modelbridge.cross_validation.AssessModelFitResult[source]

Assess model fit for given diagnostics results.

It determines if a model fit is good or bad based on Fisher exact test p

Parameters

diagnostics – Output of compute_diagnostics

Returns

Two dictionaries, one for good metrics, one for bad metrics, each mapping metric name to p-value

ax.modelbridge.cross_validation.compute_diagnostics(result: List[ax.modelbridge.cross_validation.CVResult]) → Dict[str, Dict[str, float]][source]

Computes diagnostics for given cross validation results.

It provides a dictionary with values for the following diagnostics, for each metric:

  • ‘Mean prediction CI’: the average width of the CIs at each of the CV predictions, relative to the observed mean.

  • ‘MAPE’: mean absolute percentage error of the estimated mean relative to the observed mean.

  • ‘Total raw effect’: the percent change from the smallest observed mean to the largest observed mean.

  • ‘Correlation coefficient’: the Pearson correlation of the estimated and observed means.

  • ‘Rank correlation’: the Spearman correlation of the estimated and observed means.

  • ‘Fisher exact test p’: we test if the model is able to distinguish the bottom half of the observations from the top half, using Fisher’s exact test and the observed/estimated means. A low p value indicates that the model has some ability to identify good arms. A high p value indicates that the model cannot identify arms better than chance, or that the observations are too noisy to be able to tell.

Each of these is returned as a dictionary from metric name to value for that metric.

Parameters

result – Output of cross_validate

Returns

A dictionary keyed by diagnostic name with results as described above.

ax.modelbridge.cross_validation.cross_validate(model: ax.modelbridge.base.ModelBridge, folds: int = - 1, test_selector: Optional[Callable] = None) → List[ax.modelbridge.cross_validation.CVResult][source]

Cross validation for model predictions.

Splits the model’s training data into train/test folds and makes out-of-sample predictions on the test folds.

Train/test splits are made based on arm names, so that repeated observations of a arm will always be in the train or test set together.

The test set can be limited to a specific set of observations by passing in a test_selector callable. This function should take in an Observation and return a boolean indiciating if it should be used in the test set or not. For example, we can limit the test set to arms with trial 0 with test_selector = lambda obs: obs.features.trial_index == 0 If not provided, all observations will be available for the test set.

Parameters

  • model – Fitted model (ModelBridge) to cross validate.

  • folds – Number of folds. Use -1 for leave-one-out, otherwise will be k-fold.

  • test_selector – Function for selecting observations for the test set.

Returns

A CVResult for each observation in the training data.

ax.modelbridge.cross_validation.cross_validate_by_trial(model: ax.modelbridge.base.ModelBridge, trial: int = - 1) → List[ax.modelbridge.cross_validation.CVResult][source]

Cross validation for model predictions on a particular trial.

Uses all of the data up until the specified trial to predict each of the arms that was launched in that trial. Defaults to the last trial.

Parameters

  • model – Fitted model (ModelBridge) to cross validate.

  • trial – Trial for which predictions are evaluated.

Returns

A CVResult for each observation in the training data.

ax.modelbridge.cross_validation.has_good_opt_config_model_fit(optimization_config: ax.core.optimization_config.OptimizationConfig, assess_model_fit_result: ax.modelbridge.cross_validation.AssessModelFitResult) → bool[source]

Assess model fit for given diagnostics results across the optimization config metrics

Bad fit criteria: Any objective metrics are poorly fit based on the Fisher exact test p (see assess_model_fit())

TODO[]: Incl. outcome constraints in assessment

Parameters

  • optimization_config – Objective/Outcome constraint metrics to assess

  • diagnostics – Output of compute_diagnostics

Returns

Two dictionaries, one for good metrics, one for bad metrics, each mapping metric name to p-value

Dispatch Utilities

ax.modelbridge.dispatch_utils.choose_generation_strategy(search_space: ax.core.search_space.SearchSpace, use_batch_trials: bool = False, enforce_sequential_optimization: bool = True, random_seed: Optional[int] = None, winsorization_config: Optional[Union[ax.modelbridge.transforms.winsorize.WinsorizationConfig, Dict[str, ax.modelbridge.transforms.winsorize.WinsorizationConfig]]] = None, no_winsorization: bool = False, no_bayesian_optimization: bool = False, num_trials: Optional[int] = None, num_initialization_trials: Optional[int] = None, max_parallelism_cap: Optional[int] = None, max_parallelism_override: Optional[int] = None, optimization_config: Optional[ax.core.optimization_config.OptimizationConfig] = None, should_deduplicate: bool = False, use_saasbo: bool = False, verbose: Optional[bool] = None, experiment: Optional[ax.core.experiment.Experiment] = None) → ax.modelbridge.generation_strategy.GenerationStrategy[source]

Select an appropriate generation strategy based on the properties of the search space and expected settings of the experiment, such as number of arms per trial, optimization algorithm settings, expected number of trials in the experiment, etc.

Parameters

  • search_space – SearchSpace, based on the properties of which to select the generation strategy.

  • use_batch_trials – Whether this generation strategy will be used to generate batched trials instead of 1-arm trials.

  • enforce_sequential_optimization – Whether to enforce that 1) the generation strategy needs to be updated with min_trials_observed observations for a given generation step before proceeding to the next one and 2) maximum number of trials running at once (max_parallelism) if enforced for the BayesOpt step. NOTE: max_parallelism_override and max_parallelism_cap settings will still take their effect on max parallelism even if enforce_sequential_optimization=False, so if those settings are specified, max parallelism will be enforced.

  • random_seed – Fixed random seed for the Sobol generator.

  • winsorize_botorch_model – Whether to apply the winsorization transform prior to applying other transforms for fitting the BoTorch model.

  • winsorization_config – Explicit winsorization settings, if winsorizing. Usually only upper_quantile_margin is set when minimizing, and only lower_quantile_margin when maximizing.

  • no_bayesian_optimization – If True, Bayesian optimization generation strategy will not be suggested and quasi-random strategy will be used.

  • num_trials – Total number of trials in the optimization, if known in advance.

  • num_initialization_trials – Specific number of initialization trials, if wanted. Typically, initialization trials are generated quasi-randomly.

  • max_parallelism_override – Integer, with which to override the default max parallelism setting for all steps in the generation strategy returned from this function. Each generation step has a max_parallelism value, which restricts how many trials can run simultaneously during a given generation step. By default, the parallelism setting is chosen as appropriate for the model in a given generation step. If max_parallelism_override is -1, no max parallelism will be enforced for any step of the generation strategy. Be aware that parallelism is limited to improve performance of Bayesian optimization, so only disable its limiting if necessary.

  • max_parallelism_cap – Integer cap on parallelism in this generation strategy. If specified, max_parallelism setting in each generation step will be set to the minimum of the default setting for that step and the value of this cap. max_parallelism_cap is meant to just be a hard limit on parallelism (e.g. to avoid overloading machine(s) that evaluate the experiment trials). Specify only if not specifying max_parallelism_override.

  • optimization_config – used to infer whether to use MOO and will be passed in to Winsorize via its transform_config in order to determine default winsorization behavior when necessary.

  • should_deduplicate – Whether to deduplicate the parameters of proposed arms against those of previous arms via rejection sampling. If this is True, the generation strategy will discard generator runs produced from the generation step that has should_deduplicate=True if they contain arms already present on the experiment and replace them with new generator runs. If no generator run with entirely unique arms could be produced in 5 attempts, a GenerationStrategyRepeatedPoints error will be raised, as we assume that the optimization converged when the model can no longer suggest unique arms.

  • use_saasbo – Whether to use SAAS prior for any GPEI generation steps.

  • verbose – Whether GP model should produce verbose logs. If not None, its value gets added to model_kwargs during generation_strategy construction. Defaults to True for SAASBO, else None. Verbose outputs are currently only available for SAASBO, so if verbose is not None for a different model type, it will be overridden to None with a warning.

  • experiment – If specified, _experiment attribute of the generation strategy will be set to this experiment (useful for associating a generation strategy with a given experiment before it’s first used to gen with that experiment).

Transforms

ax.modelbridge.transforms.base

class ax.modelbridge.transforms.base.Transform(search_space: Optional[SearchSpace], observation_features: List[ObservationFeatures], observation_data: List[ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[TConfig] = None)[source]

Bases: object

Defines the API for a transform that is applied to search_space, observation_features, observation_data, and optimization_config.

Transforms are used to adapt the search space and data into the types and structures expected by the model. When Transforms are used (for instance, in ModelBridge), it is always assumed that they may potentially mutate the transformed object in-place.

Forward transforms are defined for all four of those quantities. Reverse transforms are defined for observation_data and observation.

The forward transform for observation features must accept a partial observation with not all features recorded.

Forward and reverse transforms for observation data accept a list of observation features as an input, but they will not be mutated.

The forward transform for optimization config accepts the modelbridge and fixed features as inputs, but they will not be mutated.

This class provides an identify transform.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

This takes in observation_features, so that data transforms can be conditional on features, but observation_features are notmutated.

Parameters

  • observation_data – Observation data

  • observation_features – Corresponding observation features

Returns: transformed observation data

transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_optimization_config(optimization_config: OptimizationConfig, modelbridge: Optional[modelbridge_module.base.ModelBridge], fixed_features: ObservationFeatures) → OptimizationConfig[source]

Transform optimization config.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

optimization_config – The optimization config

Returns: transformed optimization config.

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Untransform observation data.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

  • observation_data – Observation data, in transformed space

  • observation_features – Corresponding observation features, in same space.

Returns: observation data in original space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.cast

class ax.modelbridge.transforms.cast.Cast(search_space: ax.core.search_space.SearchSpace, observation_features: Optional[List[ax.core.observation.ObservationFeatures]] = None, observation_data: Optional[List[ax.core.observation.ObservationData]] = None, modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Cast each param value to the respective parameter’s type/format and to a flattened version of the hierarchical search space, if applicable.

This is a default transform that should run across all models.

NOTE: In case where searh space is hierarchical and this transform is configured to flatten it:

  • All calls to Cast.transform_… transform Ax objects defined in terms of hierarchical search space, to their definitions in terms of flattened search space.

  • All calls to Cast.untransform_… cast Ax objects back to a hierarchical search space.

  • The hierarchical search space is seen as the “original” search space, and the flattened search space –– as “transformed”.

Transform is done in-place for casting types, but objects are copied during flattening of- and casting to the hierarchical search space.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features by adding parameter values that were removed during casting of observation features to hierarchical search space.

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Flattens the hierarchical search space and returns the flat SearchSpace if this transform is configured to flatten hierarchical search spaces. Does nothing if the search space is not hierarchical.

NOTE: All calls to Cast.transform_… transform Ax objects defined in terms of hierarchical search space, to their definitions in terms of flattened search space. All calls to Cast.untransform_… cast Ax objects back to a hierarchical search space.

Parameters

search_space – The search space to flatten.

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features by casting parameter values to their expected types and removing parameter values that are not applicable given the values of other parameters and the hierarchical structure of the search space.

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.cap_parameter

class ax.modelbridge.transforms.cap_parameter.CapParameter(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Cap parameter range(s) to given values. Expects a configuration of form { parameter_name -> new_upper_range_value }.

This transform only transforms the search space.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

ax.modelbridge.transforms.centered_unit_x

class ax.modelbridge.transforms.centered_unit_x.CenteredUnitX(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Map X to [-1, 1]^d for RangeParameter of type float and not log scale.

Currently does not support linear constraints, but could in the future be adjusted to transform them too, since this is a linear operation.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.choice_encode

class ax.modelbridge.transforms.choice_encode.ChoiceEncode(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Convert general ChoiceParameters to Float ChoiceParameters.

ChoiceParameters will be transformed to ChoiceParameters of float or int type. The resulting choice parameter will be considered ordered iff the original parameter is.

If the parameter type is numeric (int, float) and the parameter is orderd, then the values are normalized to the unit interval while retaining relative spacing. If the parameter type is unordered (categorical) or ordered but non-numeric, this transform uses an integer encoding.

In the inverse transform, parameters will be mapped back onto the original domain.

This transform does not transform task parameters (use TaskEncode for this).

Note that this behavior is different from that of OrderedChoiceEncode, which transforms (ordered) ChoiceParameters to integer RangeParameters (rather than ChoiceParameters).

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

class ax.modelbridge.transforms.choice_encode.OrderedChoiceEncode(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.choice_encode.ChoiceEncode

Convert ordered ChoiceParameters to integer RangeParameters.

Parameters will be transformed to an integer RangeParameters, mapped from the original choice domain to a contiguous range 0, 1, …, n_choices - 1 of integers. Does not transform task parameters.

In the inverse transform, parameters will be mapped back onto the original domain.

In order to encode all ChoiceParameters (not just ordered ChoiceParameters), use ChoiceEncode instead.

Transform is done in-place.

config: TConfig
encoded_parameters: Dict[str, Dict[TParamValue, TParamValue]]
encoded_parameters_inverse: Dict[str, ClosestLookupDict]
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

ax.modelbridge.transforms.choice_encode.transform_choice_values(p: ax.core.parameter.ChoiceParameter) → Tuple[numpy.ndarray, ax.core.parameter.ParameterType][source]

ax.modelbridge.transforms.convert_metric_names

class ax.modelbridge.transforms.convert_metric_names.ConvertMetricNames(search_space: Optional[ax.core.search_space.SearchSpace], observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Convert all metric names to canonical name as specified on a multi_type_experiment.

For example, a multi-type experiment may have an offline simulator which attempts to approximate observations from some online system. We want to map the offline metric names to the corresponding online ones so the model can associate them.

This is done by replacing metric names in the data with the corresponding online metric names.

In the inverse transform, data will be mapped back onto the original metric names. By default, this transform is turned off. It can be enabled by passing the “perform_untransform” flag to the config.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

This takes in observation_features, so that data transforms can be conditional on features, but observation_features are notmutated.

Parameters

  • observation_data – Observation data

  • observation_features – Corresponding observation features

Returns: transformed observation data

untransform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Untransform observation data.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

  • observation_data – Observation data, in transformed space

  • observation_features – Corresponding observation features, in same space.

Returns: observation data in original space.

ax.modelbridge.transforms.convert_metric_names.convert_mt_observations(observations: List[ax.core.observation.Observation], experiment: ax.core.multi_type_experiment.MultiTypeExperiment) → List[ax.core.observation.Observation][source]

Apply ConvertMetricNames transform to observations for a MT experiment.

ax.modelbridge.transforms.convert_metric_names.tconfig_from_mt_experiment(experiment: ax.core.multi_type_experiment.MultiTypeExperiment) → Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]][source]

Generate the TConfig for this transform given a multi_type_experiment.

Parameters

experiment – The experiment from which to generate the config.

Returns

The transform config to pass into the ConvertMetricNames constructor.

ax.modelbridge.transforms.derelativize

class ax.modelbridge.transforms.derelativize.Derelativize(search_space: Optional[SearchSpace], observation_features: List[ObservationFeatures], observation_data: List[ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[TConfig] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Changes relative constraints to not-relative constraints using a plug-in estimate of the status quo value.

If status quo is in-design, uses model estimate at status quo. If not, uses raw observation at status quo.

Will raise an error if status quo is in-design and model fails to predict for it, unless the flag “use_raw_status_quo” is set to True in the transform config, in which case it will fall back to using the observed value in the training data.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_optimization_config(optimization_config: ax.core.optimization_config.OptimizationConfig, modelbridge: Optional[modelbridge_module.base.ModelBridge], fixed_features: ax.core.observation.ObservationFeatures) → ax.core.optimization_config.OptimizationConfig[source]

Transform optimization config.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

optimization_config – The optimization config

Returns: transformed optimization config.

ax.modelbridge.transforms.int_range_to_choice

class ax.modelbridge.transforms.int_range_to_choice.IntRangeToChoice(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Convert a RangeParameter of type int to a ordered ChoiceParameter.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

ax.modelbridge.transforms.int_to_float

class ax.modelbridge.transforms.int_to_float.IntToFloat(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Convert a RangeParameter of type int to type float.

Uses either randomized_rounding or default python rounding, depending on ‘rounding’ flag.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.ivw

class ax.modelbridge.transforms.ivw.IVW(search_space: Optional[SearchSpace], observation_features: List[ObservationFeatures], observation_data: List[ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[TConfig] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

If an observation data contains multiple observations of a metric, they are combined using inverse variance weighting.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

This takes in observation_features, so that data transforms can be conditional on features, but observation_features are notmutated.

Parameters

  • observation_data – Observation data

  • observation_features – Corresponding observation features

Returns: transformed observation data

ax.modelbridge.transforms.ivw.ivw_metric_merge(obsd: ax.core.observation.ObservationData, conflicting_noiseless: str = 'warn') → ax.core.observation.ObservationData[source]

Merge multiple observations of a metric with inverse variance weighting.

Correctly updates the covariance of the new merged estimates: ybar1 = Sum_i w_i * y_i ybar2 = Sum_j w_j * y_j cov[ybar1, ybar2] = Sum_i Sum_j w_i * w_j * cov[y_i, y_j]

w_i will be infinity if any variance is 0. If one variance is 0., then the IVW estimate is the corresponding mean. If there are multiple measurements with 0 variance but means are all the same, then IVW estimate is that mean. If there are multiple measurements and means differ, behavior depends on argument conflicting_noiseless. “ignore” and “warn” will use the first of the measurements as the IVW estimate. “warn” will additionally log a warning. “raise” will raise an exception.

Parameters

  • obsd – An ObservationData object

  • conflicting_noiseless – “warn”, “ignore”, or “raise”

ax.modelbridge.transforms.inverse_gaussian_cdf_y

class ax.modelbridge.transforms.inverse_gaussian_cdf_y.InverseGaussianCdfY(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[base_modelbridge.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Apply inverse CDF transform to Y.

This means that we model uniform distributions as gaussian-distributed.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Map to inverse Gaussian CDF in place.

ax.modelbridge.transforms.log

class ax.modelbridge.transforms.log.Log(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Apply log base 10 to a float RangeParameter domain.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.log_y

class ax.modelbridge.transforms.log_y.LogY(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[base_modelbridge.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Apply (natural) log-transform to Y.

This essentially means that we are model the observations as log-normally distributed. If config specifies match_ci_width=True, use a matching procedure based on the width of the CIs, otherwise (the default), use the delta method,

Transform is applied only for the metrics specified in the transform config. Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

This takes in observation_features, so that data transforms can be conditional on features, but observation_features are notmutated.

Parameters

  • observation_data – Observation data

  • observation_features – Corresponding observation features

Returns: transformed observation data

transform_optimization_config(optimization_config: OptimizationConfig, modelbridge: Optional[base_modelbridge.ModelBridge], fixed_features: ObservationFeatures) → OptimizationConfig[source]

Transform optimization config.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

optimization_config – The optimization config

Returns: transformed optimization config.

untransform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Untransform observation data.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

  • observation_data – Observation data, in transformed space

  • observation_features – Corresponding observation features, in same space.

Returns: observation data in original space.

ax.modelbridge.transforms.log_y.lognorm_to_norm(mu_ln: numpy.ndarray, Cov_ln: numpy.ndarray) → Tuple[numpy.ndarray, numpy.ndarray][source]

Compute mean and covariance of a MVN from those of the associated log-MVN

If Y is log-normal with mean mu_ln and covariance Cov_ln, then X ~ N(mu_n, Cov_n) with

Cov_n_{ij} = log(1 + Cov_ln_{ij} / (mu_ln_{i} * mu_n_{j})) mu_n_{i} = log(mu_ln_{i}) - 0.5 * log(1 + Cov_ln_{ii} / mu_ln_{i}**2)

ax.modelbridge.transforms.log_y.match_ci_width(mean: numpy.ndarray, variance: numpy.ndarray, transform: Callable[[numpy.ndarray], numpy.ndarray], level: float = 0.95) → numpy.ndarray[source]
ax.modelbridge.transforms.log_y.norm_to_lognorm(mu_n: numpy.ndarray, Cov_n: numpy.ndarray) → Tuple[numpy.ndarray, numpy.ndarray][source]

Compute mean and covariance of a log-MVN from its MVN sufficient statistics

If X ~ N(mu_n, Cov_n) and Y = exp(X), then Y is log-normal with

mu_ln_{i} = exp(mu_n_{i}) + 0.5 * Cov_n_{ii} Cov_ln_{ij} = exp(mu_n_{i} + mu_n_{j} + 0.5 * (Cov_n_{ii} + Cov_n_{jj})) * (exp(Cov_n_{ij}) - 1)

ax.modelbridge.transforms.logit

class ax.modelbridge.transforms.logit.Logit(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Apply logit transfor to a float RangeParameter domain.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.one_hot

class ax.modelbridge.transforms.one_hot.OneHot(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Convert categorical parameters (unordered ChoiceParameters) to one-hot-encoded parameters.

Does not convert task parameters.

Parameters will be one-hot-encoded, yielding a set of RangeParameters, of type float, on [0, 1]. If there are two values, one single RangeParameter will be yielded, otherwise there will be a new RangeParameter for each ChoiceParameter value.

In the reverse transform, floats can be converted to a one-hot encoded vector using one of two methods:

Strict rounding: Choose the maximum value. With levels [‘a’, ‘b’, ‘c’] and

float values [0.2, 0.4, 0.3], the restored parameter would be set to ‘b’. Ties are broken randomly, so values [0.2, 0.4, 0.4] is randomly set to ‘b’ or ‘c’.

Randomized rounding: Sample from the distribution. Float values

[0.2, 0.4, 0.3] are transformed to ‘a’ w.p. 0.2/0.9, ‘b’ w.p. 0.4/0.9, or ‘c’ w.p. 0.3/0.9.

Type of rounding can be set using transform_config[‘rounding’] to either ‘strict’ or ‘randomized’. Defaults to strict.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

class ax.modelbridge.transforms.one_hot.OneHotEncoder(values: List[T])[source]

Bases: object

Joins the two encoders needed for OneHot transform.

property classes

Return number of classes discovered while fitting transform.

int_encoder: sklearn.preprocessing._label.LabelEncoder
inverse_transform(encoded_labels: List[T]) → List[T][source]

Inverse transorm a list of one hot encoded labels.

label_binarizer: sklearn.preprocessing._label.LabelBinarizer
transform(labels: List[T]) → numpy.ndarray[source]

One hot encode a list of labels.

ax.modelbridge.transforms.percentile_y

class ax.modelbridge.transforms.percentile_y.PercentileY(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Map Y values to percentiles based on their empirical CDF.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Map observation data to empirical CDF quantiles in place.

ax.modelbridge.transforms.power_transform_y

class ax.modelbridge.transforms.power_transform_y.PowerTransformY(search_space: SearchSpace, observation_features: List[ObservationFeatures], observation_data: List[ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[TConfig] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Transform the values to look as normally distributed as possible.

This fits a power transform to the data with the goal of making the transformed values look as normally distributed as possible. We use Yeo-Johnson (https://www.stat.umn.edu/arc/yjpower.pdf), which can handle both positive and negative values.

While the transform seems to be quite robust, it probably makes sense to apply a bit of winsorization and also standardize the inputs before applying the power transform. The power transform will automatically standardize the data so the data will remain standardized.

The transform can’t be inverted for all values, so we apply clipping to move values to the image of the transform. This behavior can be controlled via the clip_mean setting.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Winsorize observation data in place.

transform_optimization_config(optimization_config: OptimizationConfig, modelbridge: Optional[modelbridge_module.base.ModelBridge], fixed_features: ObservationFeatures) → OptimizationConfig[source]

Transform optimization config.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

optimization_config – The optimization config

Returns: transformed optimization config.

untransform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Winsorize observation data in place.

ax.modelbridge.transforms.remove_fixed

class ax.modelbridge.transforms.remove_fixed.RemoveFixed(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Remove fixed parameters.

Fixed parameters should not be included in the SearchSpace. This transform removes these parameters, leaving only tunable parameters.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.rounding

ax.modelbridge.transforms.rounding.contains_constrained_integer(search_space: ax.core.search_space.SearchSpace, transform_parameters: Set[str]) → bool[source]

Check if any integer parameters are present in parameter_constraints.

Order constraints are ignored since strict rounding preserves ordering.

ax.modelbridge.transforms.rounding.randomized_onehot_round(x: numpy.ndarray) → numpy.ndarray[source]

Randomized rounding of x to a one-hot vector. x should be 0 <= x <= 1.

ax.modelbridge.transforms.rounding.randomized_round(x: float) → int[source]

Randomized round of x

ax.modelbridge.transforms.rounding.randomized_round_parameters(parameters: Dict[str, Optional[Union[str, bool, float, int]]], transform_parameters: Set[str]) → Dict[str, Optional[Union[str, bool, float, int]]][source]
ax.modelbridge.transforms.rounding.strict_onehot_round(x: numpy.ndarray) → numpy.ndarray[source]

Round x to a one-hot vector by selecting the max element. Ties broken randomly.

ax.modelbridge.transforms.search_space_to_choice

class ax.modelbridge.transforms.search_space_to_choice.SearchSpaceToChoice(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Replaces the search space with a single choice parameter, whose values are the signatures of the arms observed in the data.

This transform is meant to be used with ThompsonSampler.

Choice parameter will be unordered unless config[“use_ordered”] specifies otherwise.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.standardize_y

class ax.modelbridge.transforms.standardize_y.StandardizeY(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[base_modelbridge.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Standardize Y, separately for each metric.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

This takes in observation_features, so that data transforms can be conditional on features, but observation_features are notmutated.

Parameters

  • observation_data – Observation data

  • observation_features – Corresponding observation features

Returns: transformed observation data

transform_optimization_config(optimization_config: ax.core.optimization_config.OptimizationConfig, modelbridge: Optional[base_modelbridge.ModelBridge], fixed_features: ax.core.observation.ObservationFeatures) → ax.core.optimization_config.OptimizationConfig[source]

Transform optimization config.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

optimization_config – The optimization config

Returns: transformed optimization config.

untransform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Untransform observation data.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

  • observation_data – Observation data, in transformed space

  • observation_features – Corresponding observation features, in same space.

Returns: observation data in original space.

ax.modelbridge.transforms.standardize_y.compute_standardization_parameters(Ys: DefaultDict[Union[str, Tuple[str, Optional[Union[str, bool, float, int]]]], List[float]]) → Tuple[Dict[Union[str, Tuple[str, str]], float], Dict[Union[str, Tuple[str, str]], float]][source]

Compute mean and std. dev of Ys.

ax.modelbridge.transforms.stratified_standardize_y

class ax.modelbridge.transforms.stratified_standardize_y.StratifiedStandardizeY(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Standardize Y, separately for each metric and for each value of a ChoiceParameter.

The name of the parameter by which to stratify the standardization can be specified in config[“parameter_name”]. If not specified, will use a task parameter if search space contains exactly 1 task parameter, and will raise an exception otherwise.

The stratification parameter must be fixed during generation if there are outcome constraints, in order to apply the standardization to the constraints.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

This takes in observation_features, so that data transforms can be conditional on features, but observation_features are notmutated.

Parameters

  • observation_data – Observation data

  • observation_features – Corresponding observation features

Returns: transformed observation data

transform_optimization_config(optimization_config: ax.core.optimization_config.OptimizationConfig, modelbridge: Optional[modelbridge_module.base.ModelBridge], fixed_features: ax.core.observation.ObservationFeatures) → ax.core.optimization_config.OptimizationConfig[source]

Transform optimization config.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

optimization_config – The optimization config

Returns: transformed optimization config.

untransform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Untransform observation data.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

  • observation_data – Observation data, in transformed space

  • observation_features – Corresponding observation features, in same space.

Returns: observation data in original space.

ax.modelbridge.transforms.task_encode

class ax.modelbridge.transforms.task_encode.TaskEncode(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.choice_encode.OrderedChoiceEncode

Convert task ChoiceParameters to integer-valued ChoiceParameters.

Parameters will be transformed to an integer ChoiceParameter with property is_task=True, mapping values from the original choice domain to a contiguous range integers 0, 1, …, n_choices-1.

In the inverse transform, parameters will be mapped back onto the original domain.

Transform is done in-place.

config: TConfig
encoded_parameters: Dict[str, Dict[TParamValue, TParamValue]]
encoded_parameters_inverse: Dict[str, ClosestLookupDict]
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

ax.modelbridge.transforms.trial_as_task

class ax.modelbridge.transforms.trial_as_task.TrialAsTask(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Convert trial to one or more task parameters.

How trial is mapped to parameter is specified with a map like {parameter_name: {trial_index: level name}}. For example, {“trial_param1”: {0: “level1”, 1: “level1”, 2: “level2”},} will create choice parameters “trial_param1” with is_task=True. Observations with trial 0 or 1 will have “trial_param1” set to “level1”, and those with trial 2 will have “trial_param1” set to “level2”. Multiple parameter names and mappings can be specified in this dict.

The trial level mapping can be specified in config[“trial_level_map”]. If not specified, defaults to a parameter with a level for every trial index.

For the reverse transform, if there are multiple mappings in the transform the trial will not be set.

Will raise if trial not specified for every point in the training data.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.unit_x

class ax.modelbridge.transforms.unit_x.UnitX(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Map X to [0, 1]^d for RangeParameter of type float and not log scale.

Uses bounds l <= x <= u, sets x_tilde_i = (x_i - l_i) / (u_i - l_i). Constraints wTx <= b are converted to gTx_tilde <= h, where g_i = w_i (u_i - l_i) and h = b - wTl.

Transform is done in-place.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features

Returns: transformed observation features

transform_search_space(search_space: ax.core.search_space.SearchSpace) → ax.core.search_space.SearchSpace[source]

Transform search space.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

search_space – The search space

Returns: transformed search space.

untransform_observation_features(observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationFeatures][source]

Untransform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

observation_features – Observation features in the transformed space

Returns: observation features in the original space

ax.modelbridge.transforms.unit_x.normalize_value(value: float, bounds: Tuple[float, float]) → float[source]

Transform bounds to [0,1], and apply the same transform to the value.

Note: if the value is outside of the bounds, then the value will be mapped

outside of [0,1].

ax.modelbridge.transforms.utils

class ax.modelbridge.transforms.utils.ClosestLookupDict(*args: Any, **kwargs: Any)[source]

Bases: dict

A dictionary with numeric keys that looks up the closest key.

ax.modelbridge.transforms.utils.get_data(observation_data: List[ax.core.observation.ObservationData], metric_names: Optional[List[str]] = None) → Dict[str, List[float]][source]

Extract all metrics if metric_names is None.

ax.modelbridge.transforms.utils.match_ci_width_truncated(mean: float, variance: float, transform: Callable[[float], float], level: float = 0.95, margin: float = 0.001, lower_bound: float = 0.0, upper_bound: float = 1.0, clip_mean: bool = False) → Tuple[float, float][source]

Estimate a transformed variance using the match ci width method.

See log_y transform for the original. Here, bounds are forced to lie within a [lower_bound, upper_bound] interval after transformation.

ax.modelbridge.winsorize

class ax.modelbridge.transforms.winsorize.WinsorizationConfig(lower_quantile_margin: float = 0.0, upper_quantile_margin: float = 0.0, lower_boundary: Optional[float] = None, upper_boundary: Optional[float] = None)[source]

Bases: object

Dataclass for storing Winsorization configuration parameters

Attributes: lower_quantile_margin: Winsorization will increase any metric value below this

quantile to this quantile’s value.

upper_quantile_margin: Winsorization will decrease any metric value above this

quantile to this quantile’s value. NOTE: this quantile will be inverted before any operations, e.g., a value of 0.2 will decrease values above the 80th percentile to the value of the 80th percentile.

lower_boundary: If this value is lesser than the metric value corresponding to

lower_quantile_margin, set metric values below lower_boundary to lower_boundary and leave larger values unaffected.

upper_boundary: If this value is greater than the metric value corresponding to

upper_quantile_margin, set metric values above upper_boundary to upper_boundary and leave smaller values unaffected.

lower_boundary: Optional[float] = None
lower_quantile_margin: float = 0.0
upper_boundary: Optional[float] = None
upper_quantile_margin: float = 0.0
class ax.modelbridge.transforms.winsorize.Winsorize(search_space: ax.core.search_space.SearchSpace, observation_features: List[ax.core.observation.ObservationFeatures], observation_data: List[ax.core.observation.ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[Dict[str, Optional[Union[int, float, str, botorch.acquisition.acquisition.AcquisitionFunction, Dict[str, Any]]]]] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Clip the mean values for each metric to lay within the limits provided in the config. The config can contain either or both of two keys: - "winsorization_config", corresponding to either a single

WinsorizationConfig, which, if provided will be used for all metrics; or a mapping Dict[str, WinsorizationConfig] between each metric name and its WinsorizationConfig.

  • "optimization_config", which can be used to determine default winsorization

    settings if "winsorization_config" does not provide them for a given metric.

For example, {"winsorization_config": WinsorizationConfig(lower_quantile_margin=0.3)} will specify the same 30% winsorization from below for all metrics, whereas ``` {

“winsorization_config”: {

“metric_1”: WinsorizationConfig(lower_quantile_margin=0.2), “metric_2”: WinsorizationConfig(upper_quantile_margin=0.1),

}

will winsorize 20% from below for metric_1 and 10% from above from metric_2. Additional metrics won’t be winsorized.

You can also determine the winsorization cutoffs automatically without having an OptimizationConfig by passing in AUTO_WINS_QUANTILE for the quantile you want to winsorize. For example, to automatically winsorize large values:

"m1": WinsorizationConfig(upper_quantile_margin=AUTO_WINS_QUANTILE).

This may be useful when fitting models in a notebook where there is no corresponding OptimizationConfig.

Additionally, you can pass in winsorization boundaries lower_boundary and upper_boundary``that specify a maximum allowable amount of winsorization. This is discouraged and will eventually be deprecated as we strongly encourage that users allow ``Winsorize to automatically infer these boundaries from the optimization config.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Winsorize observation data in place.

ax.modelbridge.transforms.relativize

class ax.modelbridge.transforms.relativize.Relativize(search_space: Optional[SearchSpace], observation_features: List[ObservationFeatures], observation_data: List[ObservationData], modelbridge: Optional[modelbridge_module.base.ModelBridge] = None, config: Optional[TConfig] = None)[source]

Bases: ax.modelbridge.transforms.base.Transform

Change the relative flag of the given relative optimization configuration to False. This is needed in order for the new opt config to pass ModelBridge that requires non-relativized opt config.

Also transforms absolute data and opt configs to relative.

Requires a modelbridge with a status quo set to work.

config: TConfig
modelbridge: Optional[modelbridge_module.base.ModelBridge]
transform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Transform observation features.

This is typically done in-place. This class implements the identity transform (does nothing).

This takes in observation_features, so that data transforms can be conditional on features, but observation_features are notmutated.

Parameters

  • observation_data – Observation data

  • observation_features – Corresponding observation features

Returns: transformed observation data

transform_optimization_config(optimization_config: OptimizationConfig, modelbridge: Optional[modelbridge_module.base.ModelBridge], fixed_features: ObservationFeatures) → OptimizationConfig[source]

Change the relative flag of the given relative optimization configuration to False. This is needed in order for the new opt config to pass ModelBridge that requires non-relativized opt config.

Parameters

opt_config – Optimization configuaration relative to status quo.

Returns

Optimization configuration relative to status quo with relative flag equal to false.

untransform_observation_data(observation_data: List[ax.core.observation.ObservationData], observation_features: List[ax.core.observation.ObservationFeatures]) → List[ax.core.observation.ObservationData][source]

Untransform observation data.

This is typically done in-place. This class implements the identity transform (does nothing).

Parameters

  • observation_data – Observation data, in transformed space

  • observation_features – Corresponding observation features, in same space.

Returns: observation data in original space.

Strategies

ax.modelbridge.strategies.alebo

class ax.modelbridge.strategies.alebo.ALEBOStrategy(D: int, d: int, init_size: int, name: str = 'ALEBO', dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), random_kwargs: Optional[Dict[str, Any]] = None, gp_kwargs: Optional[Dict[str, Any]] = None, gp_gen_kwargs: Optional[Dict[str, Any]] = None)[source]

Bases: ax.modelbridge.generation_strategy.GenerationStrategy

Generation strategy for Adaptive Linear Embedding BO.

Both quasirandom initialization and BO are done with the same random projection. All function evaluations are done within that projection.

Parameters

  • D – Dimensionality of high-dimensional space

  • d – Dimensionality of low-dimensional space

  • init_size – Size of random initialization

  • name – Name of strategy

  • dtype – torch dtype

  • device – torch device

  • random_kwargs – kwargs passed along to random model

  • gp_kwargs – kwargs passed along to GP model

  • gp_gen_kwargs – kwargs passed along to gen call on GP

clone_reset() → ax.modelbridge.strategies.alebo.ALEBOStrategy[source]

Copy without state.

gen_projection(d: int, D: int, dtype: torch.dtype, device: torch.device) → torch.Tensor[source]

Generate the projection matrix B as a (d x D) tensor

ax.modelbridge.strategies.alebo.get_ALEBO(experiment: ax.core.experiment.Experiment, search_space: ax.core.search_space.SearchSpace, data: ax.core.data.Data, B: torch.Tensor, **model_kwargs: Any) → ax.modelbridge.torch.TorchModelBridge[source]
ax.modelbridge.strategies.alebo.get_ALEBOInitializer(search_space: ax.core.search_space.SearchSpace, B: numpy.ndarray, **model_kwargs: Any) → ax.modelbridge.random.RandomModelBridge[source]
class ax.modelbridge.strategies.rembo.HeSBOStrategy(D: int, d: int, init_per_proj: int, k: int = 1, name: str = 'HeSBO', dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), gp_kwargs: Optional[Dict[str, Any]] = None)[source]

Bases: ax.modelbridge.strategies.rembo.REMBOStrategy

Generation strategy for HeSBO.

Parameters

  • D – Dimensionality of high-dimensional space

  • d – Dimensionality of low-dimensional space

  • k – Number of random projections

  • init_per_proj – Number of arms to use for random initialization of each of the k projections.

  • name – Name of strategy

  • dtype – torch dtype

  • device – torch device

  • gp_kwargs – kwargs sent along to the GP model

get_projection(D: int, d: int, dtype: torch.dtype, device: torch.device) → Tuple[torch.Tensor, List[Tuple[float, float]]][source]

Generate the projection matrix A as a (D x d) tensor

Also return the box bounds for the low-d space.

class ax.modelbridge.strategies.rembo.REMBOStrategy(D: int, d: int, init_per_proj: int, k: int = 4, name: str = 'REMBO', dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), gp_kwargs: Optional[Dict[str, Any]] = None)[source]

Bases: ax.modelbridge.generation_strategy.GenerationStrategy

Generation strategy for REMBO.

Both quasirandom initialization and BO are done with the same random projection. As is done in the REMBO paper, k independent optimizations are done, each with an independently generated projection.

Parameters

  • D – Dimensionality of high-dimensional space

  • d – Dimensionality of low-dimensional space

  • k – Number of random projections

  • init_per_proj – Number of arms to use for random initialization of each of the k projections.

  • name – Name of strategy

  • dtype – torch dtype

  • device – torch device

  • gp_kwargs – kwargs sent along to the GP model

clone_reset() → ax.modelbridge.strategies.rembo.REMBOStrategy[source]

Copy without state.

gen(experiment: ax.core.experiment.Experiment, data: Optional[ax.core.data.Data] = None, n: int = 1, **kwargs: Any) → ax.core.generator_run.GeneratorRun[source]

Generate new points, rotating through projections each time.

get_projection(D: int, d: int, dtype: torch.dtype, device: torch.device) → Tuple[torch.Tensor, List[Tuple[float, float]]][source]

Generate the projection matrix A as a (D x d) tensor

Also return the box bounds for the low-d space.

property model_transitions

Generator changes every iteration with rotating strategy

ax.modelbridge.strategies.rembo.get_REMBO(experiment: ax.core.experiment.Experiment, data: ax.core.data.Data, A: torch.Tensor, initial_X_d: torch.Tensor, bounds_d: List[Tuple[float, float]], search_space: Optional[ax.core.search_space.SearchSpace] = None, dtype: torch.dtype = torch.float64, device: torch.device = device(type='cpu'), **model_kwargs: Any) → ax.modelbridge.torch.TorchModelBridge[source]

Instantiates a BotorchModel.

ax.modelbridge.strategies.rembo.get_rembo_initializer(search_space: ax.core.search_space.SearchSpace, A: numpy.ndarray, bounds_d: List[Tuple[float, float]], **kwargs: Any) → ax.modelbridge.random.RandomModelBridge[source]

Instantiates a uniform random generator.

Parameters

  • search_space – Search space.

  • A – Projection matrix.

  • bounds_d – Bounds in low-d space.

  • kwargs – kwargs

Returns

RandomModelBridge, with REMBOInitializer as model.