Ax API Reference¶
-
class
ax.
Arm
(parameters, name=None)[source]¶ Base class for defining arms.
Randomization in experiments assigns units to a given arm. Thus, the arm encapsulates the parametrization needed by the unit.
-
clone
(clear_name=False)[source]¶ Create a copy of this arm.
Parameters: clear_name ( bool
) – whether this cloned copy should set its name to None instead of the name of the arm being cloned. Defaults to False.Return type: Arm
-
static
md5hash
(parameters)[source]¶ Return unique identifier for arm’s parameters.
Parameters: parameters ( Dict
[str
,Union
[str
,bool
,float
,int
,None
]]) – Parameterization; mapping of param name to value.Return type: str
Returns: Hash of arm’s parameters.
-
name_or_short_signature
¶ Returns arm name if exists; else last 4 characters of the hash.
Used for presentation of candidates (e.g. plotting and tables), where the candidates do not yet have names (since names are automatically set upon addition to a trial).
Return type: str
-
-
class
ax.
BatchTrial
(experiment, generator_run=None, trial_type=None)[source]¶ -
-
arm_weights
¶ The set of arms and associated weights for the trial.
These are constructed by merging the arms and weights from each generator run that is attached to the trial.
Return type: Optional
[Mutablemapping
[Arm
,float
]]
-
clone
()[source]¶ Clone the trial.
Return type: BatchTrial
Returns: A new instance of the trial.
-
experiment
¶ The experiment this batch belongs to.
Return type: Experiment
-
generator_run_structs
¶ List of generator run structs attached to this trial.
Struct holds generator_run object and the weight with which it was added.
Return type: List
[GeneratorRunStruct
]
-
is_factorial
¶ Return true if the trial’s arms are a factorial design with no linked factors.
Return type: bool
-
mark_arm_abandoned
(arm_name, reason=None)[source]¶ Mark a arm abandoned.
Usually done after deployment when one arm causes issues but user wants to continue running other arms in the batch.
Parameters: Return type: Returns: The batch instance.
-
normalized_arm_weights
(total=1, trunc_digits=None)[source]¶ Returns arms with a new set of weights normalized to the given total.
This method is useful for many runners where we need to normalize weights to a certain total without mutating the weights attached to a trial.
Parameters: - total (
float
) – The total weight to which to normalize. Default is 1, in which case arm weights can be interpreted as probabilities. - trunc_digits (
Optional
[int
]) – The number of digits to keep. If the resulting total weight is not equal to total, re-allocate weight in such a way to maintain relative weights as best as possible.
Return type: Returns: Mapping from arms to the new set of weights.
- total (
-
run
()[source]¶ Deploys the trial according to the behavior on the runner.
The runner returns a run_metadata dict containining metadata of the deployment process. It also returns a deployed_name of the trial within the system to which it was deployed. Both these fields are set on the trial.
Return type: BatchTrial
Returns: The trial instance.
-
-
class
ax.
ChoiceParameter
(name, parameter_type, values, is_ordered=False, is_task=False)[source]¶ Parameter object that specifies a discrete set of values.
-
add_values
(values)[source]¶ Add input list to the set of allowed values for parameter.
Cast all input values to the parameter type.
Parameters: values ( List
[Union
[str
,bool
,float
,int
,None
]]) – Values being added to the allowed list.Return type: ChoiceParameter
-
-
class
ax.
Data
(df=None, description=None)[source]¶ Class storing data for an experiment.
The dataframe is retrieved via the df property. The data can be stored to gluster for future use by attaching it to an experiment using experiment.add_data() (this requires a description to be set.)
-
df
¶ DataFrame with underlying data, and required columns.
-
description
¶ Human-readable description of data.
-
static
column_data_types
()[source]¶ Type specification for all supported columns.
Return type: Dict
[str
,Type
[+CT_co]]
-
df_hash
¶ Compute hash of pandas DataFrame.
This first serializes the DataFrame and computes the md5 hash on the resulting string. Note that this may cause performance issue for very large DataFrames.
Parameters: df – The DataFrame for which to compute the hash. - Returns
- str: The hash of the DataFrame.
Return type: str
-
-
class
ax.
Experiment
(search_space, name=None, optimization_config=None, tracking_metrics=None, runner=None, status_quo=None, description=None, is_test=False)[source]¶ Base class for defining an experiment.
-
add_tracking_metric
(metric)[source]¶ Add a new metric to the experiment.
Parameters: metric ( Metric
) – Metric to be added.Return type: Experiment
-
arms_by_signature
¶ The arms belonging to this experiment, by their signature.
Return type: Dict
[str
,Arm
]
-
attach_data
(data)[source]¶ Attach data to experiment.
Parameters: data ( Data
) – Data object to store.Return type: int
Returns: Timestamp of storage in millis.
-
data_by_trial
¶ Data stored on the experiment, indexed by trial index and storage time.
First key is trial index and second key is storage time in milliseconds. For a given trial, data is ordered by storage time, so first added data will appear first in the list.
Return type: Dict
[int
,OrderedDict
]
-
default_trial_type
¶ Default trial type assigned to trials in this experiment.
In the base experiment class this is always None. For experiments with multiple trial types, use the MultiTypeExperiment class.
Return type: Optional
[str
]
-
fetch_data
(metrics=None, **kwargs)[source]¶ Fetches data for all metrics and trials on this experiment.
Parameters: Return type: Returns: Data for the experiment.
-
is_simple_experiment
¶ Whether this experiment is a regular Experiment or the subclassing SimpleExperiment.
-
lookup_data_for_trial
(trial_index)[source]¶ Lookup stored data for a specific trial.
Returns latest data object present for this trial. Returns empty data if no data present.
Parameters: trial_index ( int
) – The index of the trial to lookup data for.Return type: Data
Returns: Requested data object.
-
lookup_data_for_ts
(timestamp)[source]¶ Collect data for all trials stored at this timestamp.
Useful when many trials’ data was fetched and stored simultaneously and user wants to retrieve same collection of data later.
Can also be used to lookup specific data for a single trial when storage time is known.
Parameters: timestamp ( int
) – Timestamp in millis at which data was stored.Return type: Data
Returns: Data object with all data stored at the timestamp.
-
new_batch_trial
(generator_run=None, trial_type=None)[source]¶ Create a new batch trial associated with this experiment.
Return type: BatchTrial
-
new_trial
(generator_run=None, trial_type=None)[source]¶ Create a new trial associated with this experiment.
Return type: Trial
-
optimization_config
¶ The experiment’s optimization config.
Return type: Optional
[OptimizationConfig
]
-
remove_tracking_metric
(metric_name)[source]¶ Remove a metric that already exists on the experiment.
Parameters: metric_name ( str
) – Unique name of metric to remove.Return type: Experiment
-
runner_for_trial
(trial)[source]¶ The default runner to use for a given trial.
In the base experiment class, this is always the default experiment runner. For experiments with multiple trial types, use the MultiTypeExperiment class.
Return type: Optional
[Runner
]
-
search_space
¶ The search space for this experiment.
When setting a new search space, all parameter names and types must be preserved. However, if no trials have been created, all modifications are allowed.
Return type: SearchSpace
-
supports_trial_type
(trial_type)[source]¶ Whether this experiment allows trials of the given type.
The base experiment class only supports None. For experiments with multiple trial types, use the MultiTypeExperiment class.
Return type: bool
-
update_tracking_metric
(metric)[source]¶ Redefine a metric that already exists on the experiment.
Parameters: metric ( Metric
) – New metric definition.Return type: Experiment
-
-
class
ax.
FixedParameter
(name, parameter_type, value)[source]¶ Parameter object that specifies a single fixed value.
-
class
ax.
GeneratorRun
(arms, weights=None, optimization_config=None, search_space=None, model_predictions=None, best_arm_predictions=None, type=None, fit_time=None, gen_time=None)[source]¶ An object that represents a single run of a generator.
This object is created each time the
gen
method of a generator is called. It stores the arms and (optionally) weights that were generated by the run. When we add a generator run to a trial, its arms and weights will be merged with those from previous generator runs that were already attached to the trial.-
arm_weights
¶ Mapping from arms to weights (order matches order in arms property).
Return type: Mutablemapping
[Arm
,float
]
-
clone
()[source]¶ Return a deep copy of a GeneratorRun.
Return type: GeneratorRun
-
index
¶ The index of this generator run within a trial’s list of generator run structs. This field is set when the generator run is added to a trial.
Return type: Optional
[int
]
-
optimization_config
¶ The optimization config used during generation of this run.
Return type: Optional
[OptimizationConfig
]
-
param_df
¶ Constructs a Pandas dataframe with the parameter values for each arm.
Useful for inspecting the contents of a generator run.
Returns: a dataframe with the generator run’s arms. Return type: pd.DataFrame
-
search_space
¶ The search used during generation of this run.
Return type: Optional
[SearchSpace
]
-
-
class
ax.
Metric
(name, lower_is_better=None)[source]¶ Base class for representing metrics.
-
lower_is_better
¶ Flag for metrics which should be minimized.
-
fetch_experiment_data
(experiment, **kwargs)[source]¶ Fetch this metric’s data for an experiment.
Default behavior is to fetch data from all trials expecting data and concatenate the results.
Return type: Data
-
classmethod
fetch_experiment_data_multi
(experiment, metrics, **kwargs)[source]¶ Fetch multiple metrics data for an experiment.
Default behavior calls fetch_experiment_data for each metric. Subclasses should override this to batch data computation for multiple metrics.
Return type: Data
-
-
class
ax.
Models
[source]¶ Registry of available factory functions.
-
BOTORCH
(data, search_space=None, dtype=torch.float64, device=device(type='cpu'), transforms=[<class 'ax.modelbridge.transforms.out_of_design.OutOfDesign'>, <class 'ax.modelbridge.transforms.remove_fixed.RemoveFixed'>, <class 'ax.modelbridge.transforms.ordered_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.unit_x.UnitX'>, <class 'ax.modelbridge.transforms.ivw.IVW'>, <class 'ax.modelbridge.transforms.derelativize.Derelativize'>, <class 'ax.modelbridge.transforms.standardize_y.StandardizeY'>], model_constructor=<function get_and_fit_model>, model_predictor=<function predict_from_model>, acqf_constructor=<function get_NEI>, acqf_optimizer=<function scipy_optimizer>, refit_on_cv=False, refit_on_update=True)¶ Instantiates a BotorchModel.
Return type: TorchModelBridge
-
EMPIRICAL_BAYES_THOMPSON
(data, search_space=None, num_samples=10000, min_weight=None, uniform_weights=False)¶ Instantiates an empirical Bayes / Thompson sampling model.
Return type: DiscreteModelBridge
-
FACTORIAL
()¶ Instantiates a factorial generator.
Return type: DiscreteModelBridge
-
GPEI
(data, search_space=None, dtype=torch.float64, device=device(type='cpu'))¶ Instantiates a GP model that generates points with EI.
Return type: TorchModelBridge
-
SOBOL
(seed=None, deduplicate=False, init_position=0, scramble=True)¶ Instantiates a Sobol sequence quasi-random generator.
Parameters: - search_space (
SearchSpace
) – Sobol generator search space. - kwargs – Custom args for sobol generator.
Return type: RandomModelBridge
Returns: RandomModelBridge, with SobolGenerator as model.
- search_space (
-
THOMPSON
(data, search_space=None, num_samples=10000, min_weight=None, uniform_weights=False)¶ Instantiates a Thompson sampling model.
Return type: DiscreteModelBridge
-
UNIFORM
(deduplicate=False, seed=None)¶ Instantiate uniform generator.
Parameters: - search_space (
SearchSpace
) – Uniform generator search space. - kwargs – Custom args for uniform generator.
Return type: RandomModelBridge
Returns: RandomModelBridge, with UniformGenerator as model.
- search_space (
-
-
class
ax.
Objective
(metric, minimize=False)[source]¶ Base class for representing an objective.
-
minimize
¶ If True, minimize metric.
-
-
class
ax.
OptimizationConfig
(objective, outcome_constraints=None)[source]¶ An optimization configuration, which comprises an objective and outcome constraints.
There is no minimum or maximum number of outcome constraints, but an individual metric can have at most two constraints–which is how we represent metrics with both upper and lower bounds.
-
clone
()[source]¶ Make a copy of this optimization config.
Return type: OptimizationConfig
-
outcome_constraints
¶ Get outcome constraints.
Return type: List
[OutcomeConstraint
]
-
-
class
ax.
OptimizationLoop
(experiment, total_trials=20, arms_per_trial=1, wait_time=0, run_async=False)[source]¶ Managed optimization loop, in which Ax oversees deployment of trials and gathering data.
-
full_run
()[source]¶ Runs full optimization loop as defined in the provided optimization plan.
Return type: OptimizationLoop
-
get_best_point
()[source]¶ Obtains the best point encountered in the course of this optimization.
Return type: Tuple
[Dict
[str
,Union
[str
,bool
,float
,int
,None
]],Optional
[Tuple
[Dict
[str
,float
],Optional
[Dict
[str
,Dict
[str
,float
]]]]]]
-
get_current_model
()[source]¶ Obtain the most recently used model in optimization.
Return type: Optional
[ModelBridge
]
-
static
with_evaluation_function
(parameters, evaluation_function, experiment_name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, total_trials=20, arms_per_trial=1, wait_time=0)[source]¶ Constructs a synchronous OptimizationLoop using an evaluation function.
Return type: OptimizationLoop
-
classmethod
with_runners_and_metrics
(parameters, path_to_runner, paths_to_metrics, experiment_name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, total_trials=20, arms_per_trial=1, wait_time=0)[source]¶ Constructs an asynchronous OptimizationLoop using Ax runners and metrics.
Return type: OptimizationLoop
-
-
class
ax.
OrderConstraint
(lower_parameter, upper_parameter)[source]¶ Constraint object for specifying one parameter to be smaller than another.
-
clone
()[source]¶ Clone.
Return type: OrderConstraint
-
-
class
ax.
OutcomeConstraint
(metric, op, bound, relative=True)[source]¶ Base class for representing outcome constraints.
Outcome constraints may of the form metric >= bound or metric <= bound, where the bound can be expressed as an absolute measurement or relative to the status quo (if applicable).
-
metric
¶ Metric to constrain.
-
op
¶ Specifies whether metric should be greater or equal to, or less than or equal to, some bound.
-
bound
¶ The bound in the constraint.
-
relative
¶ Whether you want to bound on an absolute or relative scale. If relative, bound is the acceptable percent change.
-
clone
()[source]¶ Create a copy of this OutcomeConstraint.
Return type: OutcomeConstraint
-
-
class
ax.
ParameterConstraint
(constraint_dict, bound)[source]¶ Base class for linear parameter constraints.
Constraints are expressed using a map from parameter name to weight followed by a bound.
- The constraint is satisfied if w * v <= b where:
- w is the vector of parameter weights. v is a vector of parameter values. b is the specified bound. * is the dot product operator.
-
check
(parameter_dict)[source]¶ Whether or not the set of parameter values satisfies the constraint.
Does a weighted sum of the parameter values based on the constraint_dict and checks that the sum is less than the bound.
Parameters: parameter_dict ( Dict
[str
,Union
[int
,float
]]) – Map from parameter name to parameter value.Return type: bool
Returns: Whether the constraint is satisfied.
-
clone
()[source]¶ Clone.
Return type: ParameterConstraint
-
class
ax.
RangeParameter
(name, parameter_type, lower, upper, log_scale=False, digits=None)[source]¶ Parameter object that specifies a continuous numerical range of values.
-
digits
¶ Number of digits to round values to for float type.
Upper and lower bound are re-cast after this property is changed.
Return type: Optional
[int
]
-
is_valid_type
(value)[source]¶ Same as default except allows floats whose value is an int for Int parameters.
Return type: bool
-
log_scale
¶ Whether to sample in log space when drawing random values of the parameter.
Return type: bool
-
lower
¶ Lower bound of the parameter range.
Value is cast to parameter type upon set and also validated to ensure the bound is strictly less than upper bound.
Return type: float
-
update_range
(lower=None, upper=None)[source]¶ Set the range to the given values.
If lower or upper is not provided, it will be left at its current value.
Parameters: Return type:
-
upper
¶ Upper bound of the parameter range.
Value is cast to parameter type upon set and also validated to ensure the bound is strictly greater than lower bound.
Return type: float
-
validate
(value)[source]¶ Returns True if input is a valid value for the parameter.
Checks that value is of the right type and within the valid range for the parameter. Returns False if value is None.
Parameters: value ( Union
[str
,bool
,float
,int
,None
]) – Value being checked.Return type: bool
Returns: True if valid, False otherwise.
-
-
class
ax.
Runner
[source]¶ Abstract base class for custom runner classes
-
class
ax.
SearchSpace
(parameters, parameter_constraints=None)[source]¶ Base object for SearchSpace object.
Contains a set of Parameter objects, each of which have a name, type, and set of valid values. The search space also contains a set of ParameterConstraint objects, which can be used to define restrictions across parameters (e.g. p_a < p_b).
-
cast_arm
(arm)[source]¶ Cast parameterization of given arm to the types in this SearchSpace.
For each parameter in given arm, cast it to the proper type specified in this search space. Throws if there is a mismatch in parameter names. This is mostly useful for int/float, which user can be sloppy with when hand written.
Parameters: arm ( Arm
) – Arm to cast.Return type: Arm
Returns: New casted arm.
-
check_membership
(parameterization, raise_error=False)[source]¶ Whether the given parameterization belongs in the search space.
Checks that the given parameter values have the same name/type as search space parameters, are contained in the search space domain, and satisfy the parameter constraints.
Parameters: Return type: Returns: Whether the parameterization is contained in the search space.
-
check_types
(parameterization, allow_none=True, raise_error=False)[source]¶ Checks that the given parameterization’s types match the search space.
Checks that the names of the parameterization match those specified in the search space, and the given values are of the correct type.
Parameters: Return type: Returns: Whether the parameterization has valid types.
-
out_of_design_arm
()[source]¶ Create a default out-of-design arm.
An out of design arm contains values for some parameters which are outside of the search space. In the modeling conversion, these parameters are all stripped down to an empty dictionary, since the point is already outside of the modeled space.
Return type: Arm
Returns: New arm w/ null parameter values.
-
-
class
ax.
SimpleExperiment
(search_space, name=None, objective_name=None, evaluation_function=<function unimplemented_evaluation_function>, minimize=False, outcome_constraints=None, status_quo=None)[source]¶ Simplified experiment class with defaults.
Parameters: - search_space (
SearchSpace
) – parameter space - name (
Optional
[str
]) – name of this experiment - objective_name (
Optional
[str
]) – which of the metrics computed by the evaluation function is the objective - evaluation_function (
Callable
[[Dict
[str
,Union
[str
,bool
,float
,int
,None
]],Optional
[float
]],Union
[Dict
[str
,Tuple
[float
,float
]],Tuple
[float
,float
],float
]]) – function that evaluates mean and standard error for a parameter configuration. This function should accept a dictionary of parameter names to parameter values (TParametrization) and optionally a weight, and return a dictionary of metric names to a tuple of means and standard errors (TEvaluationOutcome). The function can also return a single tuple, in which case we assume the metric is the objective. - minimize (
bool
) – whether the objective should be minimized, defaults to False - outcome_constraints (
Optional
[List
[OutcomeConstraint
]]) – constraints on the outcome, if any - status_quo (
Optional
[Arm
]) – Arm representing existing “control” arm
-
add_tracking_metric
(metric)[source]¶ Add a new metric to the experiment.
Parameters: metric ( Metric
) – Metric to be added.Return type: SimpleExperiment
-
eval
()[source]¶ Evaluate all arms in the experiment with the evaluation function passed as argument to this SimpleExperiment.
Return type: Data
-
eval_trial
(trial)[source]¶ Evaluate trial arms with the evaluation function of this experiment.
Parameters: trial ( BaseTrial
) – trial, whose arms to evaluate.Return type: Data
-
evaluation_function
¶ Get the evaluation function.
Return type: Callable
[[Dict
[str
,Union
[str
,bool
,float
,int
,None
]],Optional
[float
]],Union
[Dict
[str
,Tuple
[float
,float
]],Tuple
[float
,float
],float
]]
-
fetch_data
(metrics=None, **kwargs)[source]¶ Fetches data for all metrics and trials on this experiment.
Parameters: Return type: Returns: Data for the experiment.
-
has_evaluation_function
¶ Whether this SimpleExperiment has a valid evaluation function attached.
Return type: bool
-
is_simple_experiment
¶ Whether this experiment is a regular Experiment or the subclassing SimpleExperiment.
-
update_tracking_metric
(metric)[source]¶ Redefine a metric that already exists on the experiment.
Parameters: metric ( Metric
) – New metric definition.Return type: SimpleExperiment
- search_space (
-
class
ax.
SumConstraint
(parameters, is_upper_bound, bound)[source]¶ Constraint on the sum of parameters being greater or less than a bound.
-
clone
()[source]¶ Clone.
Return type: SumConstraint
-
constraint_dict
¶ Weights on parameters for linear constraint representation.
Return type: Dict
[str
,float
]
-
op
¶ Whether the sum is constrained by a <= or >= inequality.
Return type: ComparisonOp
-
-
class
ax.
Trial
(experiment, generator_run=None, trial_type=None)[source]¶ Trial that only has one attached arm and no arm weights.
Parameters: - experiment (
Experiment
) – experiment, to which this trial is attached - generator_run (
Optional
[GeneratorRun
]) – generator_run associated with this trial. Trial has only one generator run (and thus arm) attached to it. This can also be set later through add_arm or add_generator_run, but a trial’s associated genetor run is immutable once set.
-
arms
¶ All arms attached to this trial.
Returns: - list of a single arm
- attached to this trial if there is one, else None.
Return type: arms
-
arms_by_name
¶ Dictionary of all arms attached to this trial with their names as keys.
Returns: - dictionary of a single
- arm name to arm if one is attached to this trial, else None.
Return type: arms
-
generator_run
¶ Generator run attached to this trial.
Return type: Optional
[GeneratorRun
]
- experiment (
-
ax.
optimize
(parameters, evaluation_function, experiment_name=None, objective_name=None, minimize=False, parameter_constraints=None, outcome_constraints=None, total_trials=20, arms_per_trial=1, wait_time=0)[source]¶ Construct and run a full optimization loop.
Return type: Tuple
[Dict
[str
,Union
[str
,bool
,float
,int
,None
]],Optional
[Tuple
[Dict
[str
,float
],Optional
[Dict
[str
,Dict
[str
,float
]]]]],Experiment
,Optional
[ModelBridge
]]
-
ax.
save
(experiment, filepath)¶ Save experiment to file.
- Convert Ax experiment to JSON-serializable dictionary.
- Write to file.
Return type: None
-
ax.
load
(filepath)¶ Load experiment from file.
- Read file.
- Convert dictionary to Ax experiment instance.
Return type: Experiment