ax.modelbridge¶

Generation Strategy¶

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

Bases: 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() → 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 current_step: GenerationStep¶

Current generation step.

property experiment: Experiment¶

Experiment, currently set on this generation strategy.

gen(experiment: Experiment, data: Optional[Data] = None, n: int = 1, pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None, **kwargs: Any) → 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: Optional[GeneratorRun]¶

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

property model: Optional[ModelBridge]¶

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[int]¶

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

property name: str¶

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

property num_can_complete_this_step: int¶

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: int¶

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: int¶

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

property trial_indices_by_step: Dict[int, Set[int]]¶

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: Optional[pandas.DataFrame]¶

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: bool¶

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[ModelSpec], best_model_selector: Optional[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.

property cv_results: Optional[List[CVResult]]¶

cv_results from self.model_spec_to_gen_from for convenience

property diagnostics: Optional[Dict[str, Dict[str, float]]]¶

diagnostics from self.model_spec_to_gen_from for convenience

fit(experiment: Experiment, data: Data, search_space: Optional[SearchSpace] = None, optimization_config: Optional[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.

property fitted_model: ModelBridge¶

fitted_model from self.model_spec_to_gen_from for convenience

property fixed_features: Optional[ObservationFeatures]¶

fixed_features from self.model_spec_to_gen_from for convenience

gen(n: Optional[int] = None, pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None) → 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).

property model_cv_kwargs: Optional[Dict[str, Any]]¶

model_cv_kwargs from self.model_spec_to_gen_from for convenience

property model_enum: ModelRegistryBase¶

model_enum from self.model_spec_to_gen_from for convenience

property model_gen_kwargs: Optional[Dict[str, Any]]¶

model_gen_kwargs from self.model_spec_to_gen_from for convenience

property model_kwargs: Optional[Dict[str, Any]]¶

model_kwargs from self.model_spec_to_gen_from for convenience

property model_spec_to_gen_from: ModelSpec¶

Returns the cached _model_spec_to_gen_from or gets it from _pick_fitted_model_to_gen_from and then caches and returns it

model_specs: List[ModelSpec]¶

update(experiment: Experiment, new_data: Data) → None[source]¶

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

class ax.modelbridge.generation_node.GenerationStep(model: Union[ModelRegistryBase, Callable[[...], 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: GenerationNode, 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[ModelRegistryBase, Callable[[...], ModelBridge]]¶

model_gen_kwargs: Optional[Dict[str, Any]] = None¶

model_kwargs: Optional[Dict[str, Any]] = None¶

property model_name: str¶

property model_spec: ModelSpec¶

num_trials: int¶

should_deduplicate: bool = False¶

use_update: bool = False¶

Registry¶

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.

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

Bases: 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[ModelBridge]¶

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

property model_class: Type[Model]¶

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.

bridge_class: Type[ModelBridge]¶

Alias for field number 0

model_class: Type[Model]¶

Alias for field number 1

not_saved_model_kwargs: Optional[List[str]]¶

Alias for field number 4

standard_bridge_kwargs: Optional[Dict[str, Any]]¶

Alias for field number 3

transforms: List[Type[Transform]]¶

Alias for field number 2

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

Bases: 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'¶

CONTEXT_SACBO = 'Contextual_SACBO'¶

EMPIRICAL_BAYES_THOMPSON = 'EB'¶

FACTORIAL = 'Factorial'¶

FULLYBAYESIAN = 'FullyBayesian'¶

FULLYBAYESIANMOO = 'FullyBayesianMOO'¶

FULLYBAYESIANMOO_MTGP = 'FullyBayesianMOO_MTGP'¶

FULLYBAYESIAN_MTGP = 'FullyBayesian_MTGP'¶

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: GeneratorRun, experiment: Experiment, data: Data, models_enum: Type[ModelRegistryBase], after_gen: bool = True) → 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 – 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.

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: Experiment, data: Data, search_space: Optional[SearchSpace] = None, dtype: dtype = torch.float64, device: device = device(type='cpu')) → 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: ~typing.Optional[~ax.core.search_space.SearchSpace] = None, cost_intercept: float = 0.01, dtype: ~torch.dtype = torch.float64, device: ~torch.device = device(type='cpu'), transforms: ~typing.List[~typing.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: ~typing.Optional[~typing.Dict[str, ~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]]]] = None, **kwargs: ~typing.Any) → 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: ~typing.Optional[~ax.core.search_space.SearchSpace] = None, cost_intercept: float = 0.01, dtype: ~torch.dtype = torch.float64, device: ~torch.device = device(type='cpu'), transforms: ~typing.List[~typing.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: ~typing.Optional[~typing.Dict[str, ~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]]]] = None, **kwargs: ~typing.Any) → TorchModelBridge[source]¶

Instantiates a GP model that generates points with MES.

ax.modelbridge.factory.get_MOO_EHVI(experiment: Experiment, data: Data, search_space: Optional[SearchSpace] = None, dtype: dtype = torch.float64, device: Optional[device] = None, optimization_config: Optional[OptimizationConfig] = None) → TorchModelBridge[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 should be included in the optimization_config or experiment.optimization_config.

ax.modelbridge.factory.get_MOO_NEHVI(experiment: Experiment, data: Data, search_space: Optional[SearchSpace] = None, dtype: dtype = torch.float64, device: Optional[device] = None, status_quo_features: Optional[ObservationFeatures] = None, use_input_warping: bool = False, optimization_config: Optional[OptimizationConfig] = None) → TorchModelBridge[source]¶

Instantiates a multi-objective model using qNEHVI.

ax.modelbridge.factory.get_MOO_PAREGO(experiment: Experiment, data: Data, search_space: Optional[SearchSpace] = None, dtype: dtype = torch.float64, device: device = device(type='cpu'), optimization_config: Optional[OptimizationConfig] = None) → TorchModelBridge[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: Experiment, data: Data, search_space: Optional[SearchSpace] = None, dtype: dtype = torch.float64, device: device = device(type='cpu'), optimization_config: Optional[OptimizationConfig] = None) → TorchModelBridge[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: Experiment, data: Data, search_space: Optional[SearchSpace] = None, trial_index: Optional[int] = None, device: device = device(type='cpu'), dtype: dtype = torch.float64) → 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: Experiment, data: Data, search_space: Optional[SearchSpace] = None, dtype: dtype = torch.float64, device: Optional[device] = None, trial_index: Optional[int] = None, optimization_config: Optional[OptimizationConfig] = None) → 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: Experiment, data: Data, trial_index: Optional[int] = None, search_space: Optional[SearchSpace] = None, dtype: dtype = torch.float64, device: device = device(type='cpu'), optimization_config: Optional[OptimizationConfig] = None) → TorchModelBridge[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: ~typing.Optional[~ax.core.search_space.SearchSpace] = None, dtype: ~torch.dtype = torch.float64, device: ~torch.device = device(type='cpu'), transforms: ~typing.List[~typing.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: ~typing.Optional[~typing.Dict[str, ~typing.Dict[str, ~typing.Optional[~typing.Union[int, float, str, ~botorch.acquisition.acquisition.AcquisitionFunction, ~typing.Dict[str, ~typing.Any], ~ax.core.optimization_config.OptimizationConfig]]]]] = None, model_constructor: ~typing.Callable[[~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[~torch.Tensor], ~typing.List[int], ~typing.List[int], ~typing.List[str], ~typing.Optional[~typing.Dict[str, ~torch.Tensor]], ~typing.Any], ~botorch.models.model.Model] = <function get_and_fit_model>, model_predictor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function predict_from_model>, acqf_constructor: ~typing.Callable[[~botorch.models.model.Model, ~torch.Tensor, ~typing.Optional[~typing.Tuple[~torch.Tensor, ~torch.Tensor]], ~typing.Optional[~torch.Tensor], ~typing.Optional[~torch.Tensor], ~typing.Any], ~botorch.acquisition.acquisition.AcquisitionFunction] = <function get_NEI>, acqf_optimizer: ~typing.Callable[[~botorch.acquisition.acquisition.AcquisitionFunction, ~torch.Tensor, int, ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.List[~typing.Tuple[~torch.Tensor, ~torch.Tensor, float]]], ~typing.Optional[~typing.Dict[int, float]], ~typing.Optional[~typing.Callable[[~torch.Tensor], ~torch.Tensor]], ~typing.Any], ~typing.Tuple[~torch.Tensor, ~torch.Tensor]] = <function scipy_optimizer>, refit_on_cv: bool = False, refit_on_update: bool = True, optimization_config: ~typing.Optional[~ax.core.optimization_config.OptimizationConfig] = None) → TorchModelBridge[source]¶

Instantiates a BotorchModel.

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

Instantiates an empirical Bayes / Thompson sampling model.

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

Instantiates a factorial generator.

ax.modelbridge.factory.get_sobol(search_space: SearchSpace, seed: Optional[int] = None, deduplicate: bool = False, init_position: int = 0, scramble: bool = True, fallback_to_sample_polytope: bool = False) → 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: Experiment, data: Data, search_space: Optional[SearchSpace] = None, num_samples: int = 10000, min_weight: Optional[float] = None, uniform_weights: bool = False) → DiscreteModelBridge[source]¶

Instantiates a Thompson sampling model.

ax.modelbridge.factory.get_uniform(search_space: SearchSpace, deduplicate: bool = False, seed: Optional[int] = None) → 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: 'Optional[ModelRegistryBase]' = None, model_kwargs: 'Optional[Dict[str, Any]]' = None, model_gen_kwargs: 'Optional[Dict[str, Any]]' = None, model_cv_kwargs: 'Optional[Dict[str, Any]]' = None, _fitted_model: 'Optional[ModelBridge]' = None, _cv_results: 'Optional[List[CVResult]]' = None, _diagnostics: 'Optional[CVDiagnostics]' = None, factory_function: 'Optional[TModelFactory]' = None)[source]¶

Bases: ModelSpec

factory_function: Optional[Callable[[...], ModelBridge]] = None¶

fit(experiment: Experiment, data: Data, search_space: Optional[SearchSpace] = None, optimization_config: Optional[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[ModelRegistryBase] = None¶

property model_key: str¶

Key string to identify the model used by this ModelSpec.

class ax.modelbridge.model_spec.ModelSpec(model_enum: 'ModelRegistryBase', model_kwargs: 'Optional[Dict[str, Any]]' = None, model_gen_kwargs: 'Optional[Dict[str, Any]]' = None, model_cv_kwargs: 'Optional[Dict[str, Any]]' = None, _fitted_model: 'Optional[ModelBridge]' = None, _cv_results: 'Optional[List[CVResult]]' = None, _diagnostics: 'Optional[CVDiagnostics]' = None)[source]¶

Bases: Base

copy() → ModelSpec[source]¶

ModelSpec is both a spec and an object that performs actions. Copying is useful to avoid changes to a singleton model spec.

cross_validate() → Tuple[Optional[List[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: Optional[List[CVResult]]¶

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

property diagnostics: Optional[Dict[str, Dict[str, float]]]¶

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

fit(experiment: Experiment, 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: ModelBridge¶

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

property fixed_features: Optional[ObservationFeatures]¶

Fixed generation features to pass into the Model’s .gen function.

gen(**model_gen_kwargs: Any) → 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: ModelRegistryBase¶

model_gen_kwargs: Optional[Dict[str, Any]] = None¶

property model_key: str¶

Key string to identify the model used by this ModelSpec.

model_kwargs: Optional[Dict[str, Any]] = None¶

update(experiment: Experiment, new_data: Data) → None[source]¶

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

Model must have been fit prior to calling update()

class ax.modelbridge.model_spec.ModelSpecJSONEncoder(*, skipkeys=False, ensure_ascii=True, check_circular=True, allow_nan=True, sort_keys=False, indent=None, separators=None, default=None)[source]¶

Bases: JSONEncoder

Generic encoder to avoid JSON errors in ModelSpec.__repr__

default(o: Any) → str[source]¶

Implement this method in a subclass such that it returns a serializable object for o, or calls the base implementation (to raise a TypeError).

For example, to support arbitrary iterators, you could implement default like this:

def default(self, o):
    try:
        iterable = iter(o)
    except TypeError:
        pass
    else:
        return list(iterable)
    # Let the base class default method raise the TypeError
    return JSONEncoder.default(self, o)

Base Model Bridge¶

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

Bases: object

fixed_features: ObservationFeatures¶

optimization_config: Optional[OptimizationConfig]¶

pending_observations: Dict[str, List[ObservationFeatures]]¶

search_space: SearchSpace¶

class ax.modelbridge.base.GenResults(observation_features: List[ax.core.observation.ObservationFeatures], weights: List[float], best_observation_features: Union[ax.core.observation.ObservationFeatures, NoneType] = None, gen_metadata: Dict[str, Any] = <factory>)[source]¶

Bases: object

best_observation_features: Optional[ObservationFeatures] = None¶

gen_metadata: Dict[str, Any]¶

observation_features: List[ObservationFeatures]¶

weights: List[float]¶

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

Bases: 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[Observation], cv_test_points: List[ObservationFeatures]) → List[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.

feature_importances(metric_name: str) → Dict[str, float][source]¶

gen(n: int, search_space: Optional[SearchSpace] = None, optimization_config: Optional[OptimizationConfig] = None, pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None, fixed_features: Optional[ObservationFeatures] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None) → GeneratorRun[source]¶

Generate new points from the underlying model according to search_space, optimization_config and other parameters.

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.

Returns:

A GeneratorRun object that contains the generated points and other metadata.

get_training_data() → List[Observation][source]¶

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

property metric_names: Set[str]¶

Metric names present in training data.

property model_space: SearchSpace¶

SearchSpace used to fit model.

predict(observation_features: List[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: Optional[Observation]¶

Observation corresponding to status quo, if any.

property training_in_design: List[bool]¶

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

transform_observation_features(observation_features: List[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_observations(observations: List[Observation]) → 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: OptimizationConfig, fixed_features: 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: Data, 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[ObservationFeatures], search_space: SearchSpace) → List[ObservationFeatures][source]¶

ax.modelbridge.base.gen_arms(observation_features: List[ObservationFeatures], arms_by_signature: Optional[Dict[str, Arm]] = None) → Tuple[List[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[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.

Discrete Model Bridge¶

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

Bases: ModelBridge

A model bridge for using models based on discrete parameters.

Requires that all parameters have been transformed to ChoiceParameters.

model: DiscreteModel¶

outcomes: List[str]¶

parameters: List[str]¶

search_space: Optional[SearchSpace]¶

Random Model Bridge¶

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

Bases: 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: RandomModel

parameters: List[str]

update(new_data: Data, 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: Experiment, search_space: SearchSpace, data: Data, model: TorchModel, transforms: List[Type[Transform]], transform_configs: Optional[Dict[str, Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]]] = None, torch_dtype: Optional[dtype] = None, torch_device: Optional[device] = None, status_quo_name: Optional[str] = None, status_quo_features: Optional[ObservationFeatures] = None, optimization_config: Optional[OptimizationConfig] = None, fit_out_of_design: bool = False, fit_abandoned: bool = False, default_model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None)[source]¶

Bases: ModelBridge

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.

evaluate_acquisition_function(observation_features: Union[List[ObservationFeatures], List[List[ObservationFeatures]]], search_space: Optional[SearchSpace] = None, optimization_config: Optional[OptimizationConfig] = None, pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None, fixed_features: Optional[ObservationFeatures] = None, acq_options: Optional[Dict[str, Any]] = None) → List[float][source]¶

Evaluate the acquisition function for given set of observation features.

Parameters:

• observation_features – Either a list or a list of lists of observation features, representing parameterizations, for which to evaluate the acquisition function. If a single list is passed, the acquisition function is evaluated for each observation feature. If a list of lists is passed each element (itself a list of observation features) represents a batch of points for which to evaluate the joint acquisition value.

• search_space – Search space for fitting the model.

• optimization_config – Optimization config defining how to optimize the model.

• 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.

• 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]¶

infer_objective_thresholds(search_space: Optional[SearchSpace] = None, optimization_config: Optional[OptimizationConfig] = None, fixed_features: Optional[ObservationFeatures] = None) → List[ObjectiveThreshold][source]¶

Infer objective thresholds.

This method is only applicable for Multi-Objective optimization problems.

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.

model: Optional[TorchModel] = None¶

model_best_point(search_space: Optional[SearchSpace] = None, optimization_config: Optional[OptimizationConfig] = None, pending_observations: Optional[Dict[str, List[ObservationFeatures]]] = None, fixed_features: Optional[ObservationFeatures] = None, model_gen_options: Optional[Dict[str, Optional[Union[int, float, str, AcquisitionFunction, Dict[str, Any], OptimizationConfig]]]] = None) → Optional[Tuple[Arm, Optional[Tuple[Dict[str, float], Optional[Dict[str, Dict[str, float]]]]]]][source]¶

outcomes: List[str]¶

parameters: List[str]¶