ax¶
- class ax.Arm(parameters: dict[str, None | str | bool | float | int], name: str | None = 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: bool = False) Arm [source]
Create a copy of this arm.
- Parameters:
clear_name – whether this cloned copy should set its name to None instead of the name of the arm being cloned. Defaults to False.
- property has_name: bool
Return true if arm’s name is not None.
- static md5hash(parameters: Mapping[str, None | str | bool | float | int]) str [source]
Return unique identifier for arm’s parameters.
- Parameters:
parameters – Parameterization; mapping of param name to value.
- Returns:
Hash of arm’s parameters.
- property name: str
Get arm name. Throws if name is None.
- property name_or_short_signature: str
Returns arm name if exists; else last 8 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).
- property parameters: dict[str, None | str | bool | float | int]
Get mapping from parameter names to values.
- property signature: str
Get unique representation of a arm.
- class ax.BatchTrial(experiment: core.experiment.Experiment, generator_run: GeneratorRun | None = None, generator_runs: list[GeneratorRun] | None = None, trial_type: str | None = None, optimize_for_power: bool | None = False, ttl_seconds: int | None = None, index: int | None = None, lifecycle_stage: LifecycleStage | None = None)[source]
Batched trial that has multiple attached arms, meant to be deployed and evaluated together, and possibly arm weights, which are a measure of how much of the total resources allocated to evaluating a batch should go towards evaluating the specific arm. For instance, for field experiments the weights could describe the fraction of the total experiment population assigned to the different treatment arms. Interpretation of the weights is defined in Runner.
NOTE: A BatchTrial is not just a trial with many arms; it is a trial, for which it is important that the arms are evaluated simultaneously, e.g. in an A/B test where the evaluation results are subject to nonstationarity. For cases where multiple arms are evaluated separately and independently of each other, use multiple Trial objects with a single arm each.
- Parameters:
experiment – Experiment, to which this trial is attached
generator_run – GeneratorRun, associated with this trial. This can a also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set.
generator_runs – GeneratorRuns, associated with this trial. This can a also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set. This cannot be combined with the generator_run argument.
trial_type – Type of this trial, if used in MultiTypeExperiment.
optimize_for_power – Whether to optimize the weights of arms in this trial such that the experiment’s power to detect effects of certain size is as high as possible. Refer to documentation of BatchTrial.set_status_quo_and_optimize_power for more detail.
ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.
index – If specified, the trial’s index will be set accordingly. This should generally not be specified, as in the index will be automatically determined based on the number of existing trials. This is only used for the purpose of loading from storage.
lifecycle_stage – The stage of the experiment lifecycle that this trial represents
- property abandoned_arm_names: set[str]
Set of names of arms that have been abandoned within this trial.
- property arm_weights: MutableMapping[Arm, float]
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.
- attach_batch_trial_data(raw_data: dict[str, dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]], sample_sizes: dict[str, int] | None = None, metadata: dict[str, str | int] | None = None) None [source]
Attaches data to the trial
- Parameters:
raw_data – Map from arm name to metric outcomes.
sample_sizes – Dict from arm name to sample size.
metadata – Additional metadata to track about this run. importantly the start_date and end_date
complete_trial – Whether to mark trial as complete after attaching data. Defaults to False.
- clone() BatchTrial [source]
Clone the trial and attach it to the current experiment.
- clone_to(experiment: core.experiment.Experiment | None = None, include_sq: bool = True) BatchTrial [source]
Clone the trial and attach it to a specified experiment. If None provided, attach it to the current experiment.
- Parameters:
experiment – The experiment to which the cloned trial will belong. If unspecified, uses the current experiment.
include_sq – Whether to include status quo in the cloned trial.
- Returns:
A new instance of the trial.
- property experiment: core.experiment.Experiment
The experiment this batch belongs to.
- property generator_run_structs: list[GeneratorRunStruct]
List of generator run structs attached to this trial.
Struct holds generator_run object and the weight with which it was added.
- property generator_runs: list[GeneratorRun]
All generator runs associated with this trial.
- property index: int
The index of this batch within the experiment’s batch list.
- property is_factorial: bool
Return true if the trial’s arms are a factorial design with no linked factors.
- mark_arm_abandoned(arm_name: str, reason: str | None = None) BatchTrial [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.
NOTE: Abandoned arms are considered to be ‘pending points’ in experiment after their abandonment to avoid Ax models suggesting the same arm again as a new candidate. Abandoned arms are also excluded from model training data unless
fit_abandoned
option is specified to model bridge.- Parameters:
arm_name – The name of the arm to abandon.
reason – The reason for abandoning the arm.
- Returns:
The batch instance.
- normalized_arm_weights(total: float = 1, trunc_digits: int | None = None) MutableMapping[Arm, float] [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 – The total weight to which to normalize. Default is 1, in which case arm weights can be interpreted as probabilities.
trunc_digits – 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.
- Returns:
Mapping from arms to the new set of weights.
- run() BatchTrial [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.
- Returns:
The trial instance.
- class ax.ChoiceParameter(name: str, parameter_type: ParameterType, values: list[None | str | bool | float | int], is_ordered: bool | None = None, is_task: bool = False, is_fidelity: bool = False, target_value: None | str | bool | float | int = None, sort_values: bool | None = None, dependents: dict[None | str | bool | float | int, list[str]] | None = None)[source]
Parameter object that specifies a discrete set of values.
- Parameters:
name – Name of the parameter.
parameter_type – Enum indicating the type of parameter value (e.g. string, int).
values – List of allowed values for the parameter.
is_ordered – If False, the parameter is a categorical variable. Defaults to False if parameter_type is STRING and
values
is longer than 2, else True.is_task – Treat the parameter as a task parameter for modeling.
is_fidelity – Whether this parameter is a fidelity parameter.
target_value – Target value of this parameter if it’s a fidelity or task parameter.
sort_values – Whether to sort
values
before encoding. Defaults to False ifparameter_type
is STRING, else True.dependents – Optional mapping for parameters in hierarchical search spaces; format is { value -> list of dependent parameter names }.
- add_values(values: list[None | str | bool | float | int]) ChoiceParameter [source]
Add input list to the set of allowed values for parameter.
Cast all input values to the parameter type.
- Parameters:
values – Values being added to the allowed list.
- property domain_repr: str
Returns a string representation of the domain.
- class ax.ComparisonOp(value)[source]
Class for enumerating comparison operations.
- class ax.Data(df: DataFrame | None = None, description: str | None = None)[source]
Class storing numerical data for an experiment.
The dataframe is retrieved via the df property. The data can be stored to an external store for future use by attaching it to an experiment using experiment.attach_data() (this requires a description to be set.)
- df
DataFrame with underlying data, and required columns. For BaseData, the required columns are “arm_name”, “metric_name”, “mean”, and “sem”, the latter two of which must be numeric.
- description
Human-readable description of data.
- clone() Data [source]
Returns a new Data object with the same underlying dataframe.
- filter(trial_indices: Iterable[int] | None = None, metric_names: Iterable[str] | None = None) Data [source]
Construct a new object with the subset of rows corresponding to the provided trial indices AND metric names. If either trial_indices or metric_names are not provided, that dimension will not be filtered.
- classmethod from_multiple_data(data: Iterable[Data], subset_metrics: Iterable[str] | None = None) Data [source]
Combines multiple objects into one (with the concatenated underlying dataframe).
- Parameters:
data – Iterable of Ax objects of this class to combine.
subset_metrics – If specified, combined object will only contain metrics, names of which appear in this iterable, in the underlying dataframe.
- class ax.Experiment(search_space: SearchSpace, name: str | None = None, optimization_config: OptimizationConfig | None = None, tracking_metrics: list[Metric] | None = None, runner: Runner | None = None, status_quo: Arm | None = None, description: str | None = None, is_test: bool = False, experiment_type: str | None = None, properties: dict[str, Any] | None = None, default_data_type: DataType | None = None, auxiliary_experiments_by_purpose: None | dict[AuxiliaryExperimentPurpose, list[AuxiliaryExperiment]] = None)[source]
Base class for defining an experiment.
- add_tracking_metric(metric: Metric) Experiment [source]
Add a new metric to the experiment.
- Parameters:
metric – Metric to be added.
- add_tracking_metrics(metrics: list[Metric]) Experiment [source]
Add a list of new metrics to the experiment.
If any of the metrics are already defined on the experiment, we raise an error and don’t add any of them to the experiment
- Parameters:
metrics – Metrics to be added.
- property arms_by_signature: dict[str, Arm]
The arms belonging to this experiment, by their signature.
- property arms_by_signature_for_deduplication: dict[str, Arm]
The arms belonging to this experiment that should be used for deduplication in
GenerationStrategy
, by their signature.In its current form, this includes all arms except for those that are associated with a
FAILED
trial. - TheCANDIDATE
,STAGED
,RUNNING
, andABANDONED
arms are included as pending points during generation, so they should be less likelyto get suggested by the model again.
The
EARLY_STOPPED
andCOMPLETED
trials were already evaluated, so
the model will have data for these and is unlikely to suggest them again.
- attach_data(data: Data, combine_with_last_data: bool = False, overwrite_existing_data: bool = False) int [source]
Attach data to experiment. Stores data in experiment._data_by_trial, to be looked up via experiment.lookup_data_for_trial.
- Parameters:
data – Data object to store.
combine_with_last_data –
By default, when attaching data, it’s identified by its timestamp, and experiment.lookup_data_for_trial returns data by most recent timestamp. Sometimes, however, we want to combine the data from multiple calls to attach_data into one dataframe. This might be because:
We attached data for some metrics at one point and data for
the rest of the metrics later on. - We attached data for some fidelity at one point and data for another fidelity later one.
To achieve that goal, set combine_with_last_data to True. In this case, we will take the most recent previously attached data, append the newly attached data to it, attach a new Data object with the merged result, and delete the old one. Afterwards, calls to lookup_data_for_trial will return this new combined data object. This operation will also validate that the newly added data does not contain observations for metrics that already have observations at the same fidelity in the most recent data.
overwrite_existing_data – By default, we keep around all data that has ever been attached to the experiment. However, if we know that the incoming data contains all the information we need for a given trial, we can replace the existing data for that trial, thereby reducing the amount we need to store in the database.
- Returns:
Timestamp of storage in millis.
- attach_fetch_results(results: Mapping[int, Mapping[str, Result[Data, MetricFetchE]]], combine_with_last_data: bool = False, overwrite_existing_data: bool = False) int | None [source]
UNSAFE: Prefer to use attach_data directly instead.
Attach fetched data results to the Experiment so they will not have to be fetched again. Returns the timestamp from attachment, which is used as a dict key for _data_by_trial.
NOTE: Any Errs in the results passed in will silently be dropped! This will cause the Experiment to fail to find them in the _data_by_trial cache and attempt to refetch at fetch time. If this is not your intended behavior you MUST resolve your results first and use attach_data directly instead.
- attach_trial(parameterizations: list[dict[str, None | str | bool | float | int]], arm_names: list[str] | None = None, ttl_seconds: int | None = None, run_metadata: dict[str, Any] | None = None, optimize_for_power: bool = False) tuple[dict[str, dict[str, None | str | bool | float | int]], int] [source]
Attach a new trial with the given parameterization to the experiment.
- Parameters:
parameterizations – List of parameterization for the new trial. If only one is provided a single-arm Trial is created. If multiple arms are provided a BatchTrial is created.
arm_names – Names of arm(s) in the new trial.
ttl_seconds – If specified, will consider the trial failed after this many seconds. Used to detect dead trials that were not marked failed properly.
run_metadata – Metadata to attach to the trial.
optimize_for_power – For BatchTrial only. Whether to optimize the weights of arms in this trial such that the experiment’s power to detect effects of certain size is as high as possible. Refer to documentation of BatchTrial.set_status_quo_and_optimize_power for more detail.
- Returns:
Tuple of arm name to parameterization dict, and trial index from newly created trial.
- clone_with(search_space: SearchSpace | None = None, name: str | None = None, optimization_config: OptimizationConfig | None = None, tracking_metrics: list[Metric] | None = None, runner: Runner | None = None, status_quo: Arm | None = None, description: str | None = None, is_test: bool | None = None, properties: dict[str, Any] | None = None, trial_indices: list[int] | None = None, data: Data | None = None) Experiment [source]
Return a copy of this experiment with some attributes replaced.
NOTE: This method only retains the latest data attached to the experiment. This is the same data that would be accessed using common APIs such as
Experiment.lookup_data()
.- Parameters:
search_space – New search space. If None, it uses the cloned search space of the original experiment.
name – New experiment name. If None, it adds cloned_experiment_ prefix to the original experiment name.
optimization_config – New optimization config. If None, it clones the same optimization_config from the orignal experiment.
tracking_metrics – New list of metrics to track. If None, it clones the tracking metrics already attached to the main experiment.
runner – New runner. If None, it clones the existing runner.
status_quo – New status quo arm. If None, it clones the existing status quo.
description – New description. If None, it uses the same description.
is_test – Whether the cloned experiment should be considered a test. If None, it uses the same value.
properties – New properties dictionary. If None, it uses a copy of the same properties.
trial_indices – If specified, only clones the specified trials. If None, clones all trials.
data – If specified, attach this data to the cloned experiment. If None, clones the latest data attached to the original experiment if the experiment has any data.
- property completed_trials: list[BaseTrial]
the list of all trials for which data has arrived or is expected to arrive.
- property data_by_trial: dict[int, OrderedDict[int, Data]]
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.
- property default_trial_type: str | None
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.
- fetch_data(metrics: list[Metric] | None = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) Data [source]
Fetches data for all trials on this experiment and for either the specified metrics or all metrics currently on the experiment, if metrics argument is not specified.
NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whether a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.
NOTE: This can be lossy (ex. a MapData could get implicitly cast to a Data and lose rows) if Experiment.default_data_type is misconfigured!
- Parameters:
metrics – If provided, fetch data for these metrics; otherwise, fetch data for all metrics defined on the experiment.
kwargs – keyword args to pass to underlying metrics’ fetch data functions.
- Returns:
Data for the experiment.
- fetch_data_results(metrics: list[Metric] | None = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) dict[int, dict[str, Result[Data, MetricFetchE]]] [source]
Fetches data for all trials on this experiment and for either the specified metrics or all metrics currently on the experiment, if metrics argument is not specified.
If a metric fetch fails, the Exception will be captured in the MetricFetchResult along with a message.
NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whether a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.
- Parameters:
metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.
kwargs – keyword args to pass to underlying metrics’ fetch data functions.
- Returns:
A nested Dictionary from trial_index => metric_name => result
- fetch_trials_data(trial_indices: Iterable[int], metrics: list[Metric] | None = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) Data [source]
Fetches data for specific trials on the experiment.
NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whetner a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.
NOTE: This can be lossy (ex. a MapData could get implicitly cast to a Data and lose rows) if Experiment.default_data_type is misconfigured!
- Parameters:
trial_indices – Indices of trials, for which to fetch data.
metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.
kwargs – Keyword args to pass to underlying metrics’ fetch data functions.
- Returns:
Data for the specific trials on the experiment.
- fetch_trials_data_results(trial_indices: Iterable[int], metrics: list[Metric] | None = None, combine_with_last_data: bool = False, overwrite_existing_data: bool = False, **kwargs: Any) dict[int, dict[str, Result[Data, MetricFetchE]]] [source]
Fetches data for specific trials on the experiment.
If a metric fetch fails, the Exception will be captured in the MetricFetchResult along with a message.
NOTE: For metrics that are not available while trial is running, the data may be retrieved from cache on the experiment. Data is cached on the experiment via calls to experiment.attach_data and whether a given metric class is available while trial is running is determined by the boolean returned from its is_available_while_running class method.
- Parameters:
trial_indices – Indices of trials, for which to fetch data.
metrics – If provided, fetch data for these metrics instead of the ones defined on the experiment.
kwargs – keyword args to pass to underlying metrics’ fetch data functions.
- Returns:
A nested Dictionary from trial_index => metric_name => result
- get_trials_by_indices(trial_indices: Iterable[int]) list[BaseTrial] [source]
Grabs trials on this experiment by their indices.
- property has_name: bool
Return true if experiment’s name is not None.
- property immutable_search_space_and_opt_config: bool
Boolean representing whether search space and metrics on this experiment are mutable (by default they are).
NOTE: For experiments with immutable search spaces and metrics, generator runs will not store copies of search space and metrics, which improves storage layer performance. Not keeping copies of those on generator runs also disables keeping track of changes to search space and metrics, thereby necessitating that those attributes be immutable on experiment.
- property is_moo_problem: bool
Whether the experiment’s optimization config contains multiple objectives.
- property is_test: bool
Get whether the experiment is a test.
- lookup_data(trial_indices: Iterable[int] | None = None) Data [source]
Lookup stored data for trials on this experiment.
For each trial, returns latest data object present for this trial. Returns empty data if no data is present. In particular, this method will not fetch data from metrics - to do that, use fetch_data() instead.
- Parameters:
trial_indices – Indices of trials for which to fetch data. If omitted, lookup data for all trials on the experiment.
- Returns:
Data for the trials on the experiment.
- lookup_data_for_trial(trial_index: int) tuple[Data, int] [source]
Lookup stored data for a specific trial.
Returns latest data object and its storage timestamp present for this trial. Returns empty data and -1 if no data is present. In particular, this method will not fetch data from metrics - to do that, use fetch_data() instead.
- Parameters:
trial_index – The index of the trial to lookup data for.
- Returns:
The requested data object, and its storage timestamp in milliseconds.
- lookup_data_for_ts(timestamp: int) Data [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 – Timestamp in millis at which data was stored.
- Returns:
Data object with all data stored at the timestamp.
- property metric_config_summary_df: DataFrame
Creates a dataframe with information about each metric in the experiment. The resulting dataframe has one row per metric, and the following columns:
Name: the name of the metric.
Type: the metric subclass (e.g., Metric, BraninMetric).
Goal: the goal for this for this metric, based on the optimization config (minimize, maximize, constraint or track).
Bound: the bound of this metric (e.g., “<=10.0”) if it is being used as part of an ObjectiveThreshold or OutcomeConstraint.
Lower is Better: whether the user prefers this metric to be lower, if provided.
- property name: str
Get experiment name. Throws if name is None.
- new_batch_trial(generator_run: GeneratorRun | None = None, generator_runs: list[GeneratorRun] | None = None, trial_type: str | None = None, optimize_for_power: bool | None = False, ttl_seconds: int | None = None, lifecycle_stage: LifecycleStage | None = None) BatchTrial [source]
Create a new batch trial associated with this experiment.
- Parameters:
generator_run – GeneratorRun, associated with this trial. This can a also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set.
generator_runs – GeneratorRuns, associated with this trial. This can a also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set. This cannot be combined with the generator_run argument.
trial_type – Type of this trial, if used in MultiTypeExperiment.
optimize_for_power – Whether to optimize the weights of arms in this trial such that the experiment’s power to detect effects of certain size is as high as possible. Refer to documentation of BatchTrial.set_status_quo_and_optimize_power for more detail.
ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.
lifecycle_stage – The stage of the experiment lifecycle that this trial represents
- new_trial(generator_run: GeneratorRun | None = None, trial_type: str | None = None, ttl_seconds: int | None = None) Trial [source]
Create a new trial associated with this experiment.
- Parameters:
generator_run – GeneratorRun, associated with this trial. Trial has only one arm attached to it and this generator_run must therefore contain one arm. This arm can also be set later through add_arm or add_generator_run, but a trial’s associated generator run is immutable once set.
trial_type – Type of this trial, if used in MultiTypeExperiment.
ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.
- property num_abandoned_arms: int
How many arms attached to this experiment are abandoned.
- property num_trials: int
How many trials are associated with this experiment.
- property optimization_config: OptimizationConfig | None
The experiment’s optimization config.
- remove_tracking_metric(metric_name: str) Experiment [source]
Remove a metric that already exists on the experiment.
- Parameters:
metric_name – Unique name of metric to remove.
- reset_runners(runner: Runner) None [source]
Replace all candidate trials runners.
- Parameters:
runner – New runner to replace with.
- runner_for_trial(trial: BaseTrial) Runner | None [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.
- property search_space: SearchSpace
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.
- stop_trial_runs(trials: list[BaseTrial], reasons: list[str | None] | None = None) None [source]
Stops the jobs that execute given trials.
Used if, for example, TTL for a trial was specified and expired, or poor early results suggest the trial is not worth running to completion.
Override default implementation on the
Runner
if its desirable to stop trials in bulk.- Parameters:
trials – Trials to be stopped.
reasons – A list of strings describing the reasons for why the trials are to be stopped (in the same order).
- property sum_trial_sizes: int
Sum of numbers of arms attached to each trial in this experiment.
- supports_trial_type(trial_type: str | None) bool [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.
- property time_created: datetime
Creation time of the experiment.
- property trial_indices_by_status: dict[TrialStatus, set[int]]
Indices of trials associated with the experiment, grouped by trial status.
- property trial_indices_expecting_data: set[int]
Set of indices of trials, statuses of which indicate that we expect these trials to have data, either already or in the future.
- property trials: dict[int, BaseTrial]
The trials associated with the experiment.
NOTE: If some trials on this experiment specify their TTL, RUNNING trials will be checked for whether their TTL elapsed during this call. Found past- TTL trials will be marked as FAILED.
- property trials_by_status: dict[TrialStatus, list[BaseTrial]]
Trials associated with the experiment, grouped by trial status.
- property trials_expecting_data: list[BaseTrial]
the list of all trials for which data has arrived or is expected to arrive.
- update_tracking_metric(metric: Metric) Experiment [source]
Redefine a metric that already exists on the experiment.
- Parameters:
metric – New metric definition.
- validate_trials(trials: Iterable[BaseTrial]) None [source]
Raise ValueError if any of the trials in the input are not from this experiment.
- warm_start_from_old_experiment(old_experiment: Experiment, copy_run_metadata_keys: list[str] | None = None, trial_statuses_to_copy: list[TrialStatus] | None = None, search_space_check_membership_raise_error: bool = True) list[Trial] [source]
Copy all completed trials with data from an old Ax expeirment to this one. This function checks that the parameters of each trial are members of the current experiment’s search_space.
NOTE: Currently only handles experiments with 1-arm
Trial
-s, notBatchTrial
-s as there has not yet been need for support of the latter.- Parameters:
old_experiment – The experiment from which to transfer trials and data
copy_run_metadata_keys – A list of keys denoting which items to copy over from each trial’s run_metadata. Defaults to
old_experiment.runner.run_metadata_report_keys
.trial_statuses_to_copy – All trials with a status in this list will be copied. By default, copies all
RUNNING
,COMPLETED
,ABANDONED
, andEARLY_STOPPED
trials.search_space_check_membership_raise_error – Whether to raise an exception if the warm started trials being imported fall outside of the defined search space.
- Returns:
List of trials successfully copied from old_experiment to this one
- class ax.FixedParameter(name: str, parameter_type: ParameterType, value: None | str | bool | float | int, is_fidelity: bool = False, target_value: None | str | bool | float | int = None, dependents: dict[None | str | bool | float | int, list[str]] | None = None)[source]
Parameter object that specifies a single fixed value.
- property domain_repr: str
Returns a string representation of the domain.
- class ax.GeneratorRun(arms: list[Arm], weights: list[float] | None = None, optimization_config: OptimizationConfig | None = None, search_space: SearchSpace | None = None, model_predictions: tuple[dict[str, list[float]], dict[str, dict[str, list[float]]]] | None = None, best_arm_predictions: tuple[Arm, tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None = None, type: str | None = None, fit_time: float | None = None, gen_time: float | None = None, model_key: str | None = None, model_kwargs: dict[str, Any] | None = None, bridge_kwargs: dict[str, Any] | None = None, gen_metadata: dict[str, Any] | None = None, model_state_after_gen: dict[str, Any] | None = None, generation_step_index: int | None = None, candidate_metadata_by_arm_signature: None | dict[str, dict[str, Any] | None] = None, generation_node_name: str | None = 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.- add_arm(arm: Arm, weight: float = 1.0) None [source]
Adds an arm to this generator run. This should not be used to mutate generator runs that are attached to trials.
- Parameters:
arm – The arm to add.
weight – The weight to associate with the arm.
- property arm_weights: MutableMapping[Arm, float]
Mapping from arms to weights (order matches order in arms property).
- property best_arm_predictions: tuple[Arm, tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] | None
Best arm in this run (according to the optimization config) and its optional respective model predictions.
- property candidate_metadata_by_arm_signature: dict[str, dict[str, Any] | None] | None
Retrieves model-produced candidate metadata as a mapping from arm name (for the arm the candidate became when added to experiment) to the metadata dict.
- clone() GeneratorRun [source]
Return a deep copy of a GeneratorRun.
- property index: int | None
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.
- property model_predictions: tuple[dict[str, list[float]], dict[str, dict[str, list[float]]]] | None
Means and covariances for the arms in this run recorded at the time the run was executed.
- property model_predictions_by_arm: dict[str, tuple[dict[str, float], dict[str, dict[str, float]] | None]] | None
Model predictions for each arm in this run, at the time the run was executed.
- property optimization_config: OptimizationConfig | None
The optimization config used during generation of this run.
- property param_df: DataFrame
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
- property search_space: SearchSpace | None
The search used during generation of this run.
- property time_created: datetime
Creation time of the batch.
- class ax.Metric(name: str, lower_is_better: bool | None = None, properties: dict[str, Any] | None = None)[source]
Base class for representing metrics.
The fetch_trial_data method is the essential method to override when subclassing, which specifies how to retrieve a Metric, for a given trial.
- A Metric must return a Data object, which requires (at minimum) the following:
https://ax.dev/api/_modules/ax/core/data.html#Data.required_columns
- lower_is_better
Flag for metrics which should be minimized.
- properties
Properties specific to a particular metric.
- bulk_fetch_experiment_data(experiment: core.experiment.Experiment, metrics: list[Metric], trials: list[core.base_trial.BaseTrial] | None = None, **kwargs: Any) dict[int, dict[str, MetricFetchResult]] [source]
Fetch multiple metrics data for multiple trials on an experiment, using instance attributes of the metrics.
Returns Dict of metric_name => Result Default behavior calls fetch_trial_data for each metric. Subclasses should override this to trial data computation for multiple metrics.
- bulk_fetch_trial_data(trial: core.base_trial.BaseTrial, metrics: list[Metric], **kwargs: Any) dict[str, MetricFetchResult] [source]
Fetch multiple metrics data for one trial, using instance attributes of the metrics.
Returns Dict of metric_name => Result Default behavior calls fetch_trial_data for each metric. Subclasses should override this to perform trial data computation for multiple metrics.
- clone() Metric [source]
Create a copy of this Metric.
- data_constructor
alias of
Data
- fetch_data_prefer_lookup(experiment: core.experiment.Experiment, metrics: list[Metric], trials: list[core.base_trial.BaseTrial] | None = None, **kwargs: Any) tuple[dict[int, dict[str, MetricFetchResult]], bool] [source]
Fetch or lookup (with fallback to fetching) data for given metrics, depending on whether they are available while running. Return a tuple containing the data, along with a boolean that will be True if new data was fetched, and False if all data was looked up from cache.
If metric is available while running, its data can change (and therefore we should always re-fetch it). If metric is available only upon trial completion, its data does not change, so we can look up that data on the experiment and only fetch the data that is not already attached to the experiment.
NOTE: If fetching data for a metrics class that is only available upon trial completion, data fetched in this function (data that was not yet available on experiment) will be attached to experiment.
- classmethod fetch_experiment_data_multi(experiment: core.experiment.Experiment, metrics: Iterable[Metric], trials: Iterable[core.base_trial.BaseTrial] | None = None, **kwargs: Any) dict[int, dict[str, MetricFetchResult]] [source]
Fetch multiple metrics data for an experiment.
Returns Dict of trial_index => (metric_name => Result) Default behavior calls fetch_trial_data_multi for each trial. Subclasses should override to batch data computation across trials + metrics.
- property fetch_multi_group_by_metric: type[Metric]
Metric class, with which to group this metric in Experiment._metrics_by_class, which is used to combine metrics on experiment into groups and then fetch their data via Metric.fetch_trial_data_multi for each group.
NOTE: By default, this property will just return the class on which it is defined; however, in some cases it is useful to group metrics by their superclass, in which case this property should return that superclass.
- fetch_trial_data(trial: core.base_trial.BaseTrial, **kwargs: Any) MetricFetchResult [source]
Fetch data for one trial.
- classmethod fetch_trial_data_multi(trial: core.base_trial.BaseTrial, metrics: Iterable[Metric], **kwargs: Any) dict[str, MetricFetchResult] [source]
Fetch multiple metrics data for one trial.
Returns Dict of metric_name => Result Default behavior calls fetch_trial_data for each metric. Subclasses should override this to trial data computation for multiple metrics.
- classmethod is_available_while_running() bool [source]
Whether metrics of this class are available while the trial is running. Metrics that are not available while the trial is running are assumed to be available only upon trial completion. For such metrics, data is assumed to never change once the trial is completed.
NOTE: If this method returns False, data-fetching via experiment.fetch_data will return the data cached on the experiment (for the metrics of the given class) whenever its available. Data is cached on experiment when attached via experiment.attach_data.
- property name: str
Get name of metric.
- classmethod period_of_new_data_after_trial_completion() timedelta [source]
Period of time metrics of this class are still expecting new data to arrive after trial completion. This is useful for metrics whose results are processed by some sort of data pipeline, where the pipeline will continue to land additional data even after the trial is completed.
If the metric is not available after trial completion, this method will return timedelta(0). Otherwise, it should return the maximum amount of time that the metric may have new data arrive after the trial is completed.
NOTE: This property will not prevent new data from attempting to be refetched for completed trials when calling experiment.fetch_data(). Its purpose is to prevent experiment.fetch_data() from being called in Scheduler and anywhere else it is checked.
- class ax.Models(value)[source]
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.BOTORCH(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.
NOTE: If you deprecate a model, please add its replacement to ax.storage.json_store.decoder._DEPRECATED_MODEL_TO_REPLACEMENT to ensure backwards compatibility of the storage layer.
- class ax.MultiObjective(objectives: list[Objective] | None = None, **extra_kwargs: Any)[source]
Class for an objective composed of a multiple component objectives.
The Acquisition function determines how the objectives are weighted.
- objectives
List of objectives.
- clone() MultiObjective [source]
Create a copy of the objective.
- property metric: Metric
Override base method to error.
- property metrics: list[Metric]
Get the objective metrics.
- property objectives: list[Objective]
Get the objectives.
- class ax.MultiObjectiveOptimizationConfig(objective: MultiObjective | ScalarizedObjective, outcome_constraints: list[OutcomeConstraint] | None = None, objective_thresholds: list[ObjectiveThreshold] | None = None, risk_measure: RiskMeasure | None = None)[source]
An optimization configuration for multi-objective optimization, which comprises multiple objective, outcome constraints, objective thresholds, and an optional risk measure.
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.
ObjectiveThresholds should be present for every objective. A good rule of thumb is to set them 10% below the minimum acceptable value for each metric.
- property all_constraints: list[OutcomeConstraint]
Get all constraints and thresholds.
- clone_with_args(objective: ~ax.core.objective.MultiObjective | ~ax.core.objective.ScalarizedObjective | None = None, outcome_constraints: None | list[~ax.core.outcome_constraint.OutcomeConstraint] = [OutcomeConstraint( >= 0%)], objective_thresholds: None | list[~ax.core.outcome_constraint.ObjectiveThreshold] = [ObjectiveThreshold( <= 0%)], risk_measure: ~ax.core.risk_measures.RiskMeasure | None = RiskMeasure(risk_measure=, options={})) MultiObjectiveOptimizationConfig [source]
Make a copy of this optimization config.
- property objective: MultiObjective | ScalarizedObjective
Get objective.
- property objective_thresholds: list[ObjectiveThreshold]
Get objective thresholds.
- class ax.Objective(metric: Metric, minimize: bool | None = None)[source]
Base class for representing an objective.
- minimize
If True, minimize metric.
- clone() Objective [source]
Create a copy of the objective.
- get_unconstrainable_metrics() list[Metric] [source]
Return a list of metrics that are incompatible with OutcomeConstraints.
- property metric: Metric
Get the objective metric.
- property metrics: list[Metric]
Get a list of objective metrics.
- class ax.ObjectiveThreshold(metric: Metric, bound: float, relative: bool = True, op: ComparisonOp | None = None)[source]
Class for representing Objective Thresholds.
An objective threshold represents the threshold for an objective metric to contribute to hypervolume calculations. A list containing the objective threshold for each metric collectively form a reference point.
Objective thresholds may bound the metric from above or from below. The bound can be expressed as an absolute measurement or relative to the status quo (if applicable).
The direction of the bound is inferred from the Metric’s lower_is_better attribute.
- metric
Metric to constrain.
- 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. That is, the bound’s value will be
(1 + sign * bound/100.0) * status_quo_metric_value
, where sign is the sign of status_quo_metric_value, ensuring that a positive relative bound gives rise to an absolute upper bound, even if the status-quo arm has a negative metric value. This requires specification of a status-quo arm inExperiment
.
- op
automatically inferred, but manually overwritable. specifies whether metric should be greater or equal to, or less than or equal to, some bound.
- clone() ObjectiveThreshold [source]
Create a copy of this ObjectiveThreshold.
- class ax.OptimizationConfig(objective: Objective, outcome_constraints: list[OutcomeConstraint] | None = None, risk_measure: RiskMeasure | None = None)[source]
An optimization configuration, which comprises an objective, outcome constraints and an optional risk measure.
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.
- property all_constraints: list[OutcomeConstraint]
Get outcome constraints.
- clone() OptimizationConfig [source]
Make a copy of this optimization config.
- clone_with_args(objective: ~ax.core.objective.Objective | None = None, outcome_constraints: None | list[~ax.core.outcome_constraint.OutcomeConstraint] = [OutcomeConstraint( >= 0%)], risk_measure: ~ax.core.risk_measures.RiskMeasure | None = RiskMeasure(risk_measure=, options={})) OptimizationConfig [source]
Make a copy of this optimization config.
- property objective: Objective
Get objective.
- property outcome_constraints: list[OutcomeConstraint]
Get outcome constraints.
- class ax.OptimizationLoop(experiment: Experiment, evaluation_function: Callable[[dict[str, None | str | bool | float | int]], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]] | Callable[[dict[str, None | str | bool | float | int], float | None], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]], total_trials: int = 20, arms_per_trial: int = 1, random_seed: int | None = None, wait_time: int = 0, run_async: bool = False, generation_strategy: GenerationStrategy | None = None)[source]
Managed optimization loop, in which Ax oversees deployment of trials and gathering data.
- full_run() OptimizationLoop [source]
Runs full optimization loop as defined in the provided optimization plan.
- get_best_point() tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None] [source]
Obtains the best point encountered in the course of this optimization.
- get_current_model() ModelBridge | None [source]
Obtain the most recently used model in optimization.
- static with_evaluation_function(parameters: list[dict[str, None | str | bool | float | int | Sequence[None | str | bool | float | int] | dict[str, list[str]]]], evaluation_function: Callable[[dict[str, None | str | bool | float | int]], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]] | Callable[[dict[str, None | str | bool | float | int], float | None], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]], experiment_name: str | None = None, objective_name: str | None = None, minimize: bool = False, parameter_constraints: list[str] | None = None, outcome_constraints: list[str] | None = None, total_trials: int = 20, arms_per_trial: int = 1, wait_time: int = 0, random_seed: int | None = None, generation_strategy: GenerationStrategy | None = None) OptimizationLoop [source]
Constructs a synchronous OptimizationLoop using an evaluation function.
- classmethod with_runners_and_metrics(parameters: list[dict[str, None | str | bool | float | int | Sequence[None | str | bool | float | int] | dict[str, list[str]]]], path_to_runner: str, paths_to_metrics: list[str], experiment_name: str | None = None, objective_name: str | None = None, minimize: bool = False, parameter_constraints: list[str] | None = None, outcome_constraints: list[str] | None = None, total_trials: int = 20, arms_per_trial: int = 1, wait_time: int = 0, random_seed: int | None = None) OptimizationLoop [source]
Constructs an asynchronous OptimizationLoop using Ax runners and metrics.
- class ax.OrderConstraint(lower_parameter: Parameter, upper_parameter: Parameter)[source]
Constraint object for specifying one parameter to be smaller than another.
- clone() OrderConstraint [source]
Clone.
- clone_with_transformed_parameters(transformed_parameters: dict[str, Parameter]) OrderConstraint [source]
Clone, but replace parameters with transformed versions.
- property constraint_dict: dict[str, float]
Weights on parameters for linear constraint representation.
- property lower_parameter: Parameter
Parameter with lower value.
- property upper_parameter: Parameter
Parameter with higher value.
- class ax.OutcomeConstraint(metric: Metric, op: ComparisonOp, bound: float, relative: bool = 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
[default
True
] Whether the provided bound value is relative to some status-quo arm’s metric value. If False,bound
is interpreted as an absolute number, elsebound
specifies percent-difference from the observed metric value on the status-quo arm. That is, the bound’s value will be(1 + sign * bound/100.0) * status_quo_metric_value
, where sign is the sign of status_quo_metric_value. This ensures that a positive relative bound gives rise to an absolute upper bound, even if the status-quo arm has a negative metric value. This requires specification of a status-quo arm inExperiment
.
- clone() OutcomeConstraint [source]
Create a copy of this OutcomeConstraint.
- class ax.Parameter[source]
-
- abstract property domain_repr: str
Returns a string representation of the domain.
- class ax.ParameterConstraint(constraint_dict: dict[str, float], bound: float)[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 sum_i(w_i * v_i) <= b where:
w is the vector of parameter weights. v is a vector of parameter values. b is the specified bound.
- property bound: float
Get bound of the inequality of the constraint.
- check(parameter_dict: dict[str, int | float]) bool [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 – Map from parameter name to parameter value.
- Returns:
Whether the constraint is satisfied.
- clone() ParameterConstraint [source]
Clone.
- clone_with_transformed_parameters(transformed_parameters: dict[str, Parameter]) ParameterConstraint [source]
Clone, but replaced parameters with transformed versions.
- class ax.ParameterType(value)[source]
An enumeration.
- class ax.RangeParameter(name: str, parameter_type: ParameterType, lower: float, upper: float, log_scale: bool = False, logit_scale: bool = False, digits: int | None = None, is_fidelity: bool = False, target_value: None | str | bool | float | int = None)[source]
Parameter object that specifies a range of values.
- property digits: int | None
Number of digits to round values to for float type.
Upper and lower bound are re-cast after this property is changed.
- property domain_repr: str
Returns a string representation of the domain.
- is_valid_type(value: None | str | bool | float | int) bool [source]
Same as default except allows floats whose value is an int for Int parameters.
- property log_scale: bool
Whether the parameter’s random values should be sampled from log space.
- property logit_scale: bool
Whether the parameter’s random values should be sampled from logit space.
- property lower: float | int
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.
- update_range(lower: float | None = None, upper: float | None = None) RangeParameter [source]
Set the range to the given values.
If lower or upper is not provided, it will be left at its current value.
- Parameters:
lower – New value for the lower bound.
upper – New value for the upper bound.
- property upper: float | int
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.
- validate(value: None | str | bool | float | int, tol: float = 1.5e-07) bool [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 – Value being checked.
tol – Absolute tolerance for floating point comparisons.
- Returns:
True if valid, False otherwise.
- class ax.Runner[source]
Abstract base class for custom runner classes
- poll_available_capacity() int [source]
Checks how much available capacity there is to schedule trial evaluations. Required for runners used with Ax
Scheduler
.NOTE: This method might be difficult to implement in some systems. Returns -1 if capacity of the system is “unlimited” or “unknown” (meaning that the
Scheduler
should be trying to schedule as many trials as is possible without violating scheduler settings). There is no need to artificially force this method to limit capacity;Scheduler
has other limitations in place to limit number of trials running at once, like theSchedulerOptions.max_pending_trials
setting, or more granular control in the form of the max_parallelism setting in each of the GenerationStep`s of a `GenerationStrategy).- Returns:
An integer, representing how many trials there is available capacity for; -1 if capacity is “unlimited” or not possible to know in advance.
- poll_exception(trial: core.base_trial.BaseTrial) str [source]
Returns the exception from a trial.
- Parameters:
trial – Trial to get exception for.
- Returns:
Exception string.
- poll_trial_status(trials: Iterable[core.base_trial.BaseTrial]) dict[core.base_trial.TrialStatus, set[int]] [source]
Checks the status of any non-terminal trials and returns their indices as a mapping from TrialStatus to a list of indices. Required for runners used with Ax
Scheduler
.NOTE: Does not need to handle waiting between polling calls while trials are running; this function should just perform a single poll.
- Parameters:
trials – Trials to poll.
- Returns:
A dictionary mapping TrialStatus to a list of trial indices that have the respective status at the time of the polling. This does not need to include trials that at the time of polling already have a terminal (ABANDONED, FAILED, COMPLETED) status (but it may).
- abstract run(trial: core.base_trial.BaseTrial) dict[str, Any] [source]
Deploys a trial based on custom runner subclass implementation.
- Parameters:
trial – The trial to deploy.
- Returns:
Dict of run metadata from the deployment process.
- property run_metadata_report_keys: list[str]
A list of keys of the metadata dict returned by run() that are relevant outside the runner-internal impolementation. These can e.g. be reported in Scheduler.report_results().
- run_multiple(trials: Iterable[core.base_trial.BaseTrial]) dict[int, dict[str, Any]] [source]
Runs a single evaluation for each of the given trials. Useful when deploying multiple trials at once is more efficient than deploying them one-by-one. Used in Ax
Scheduler
.NOTE: By default simply loops over run_trial. Should be overwritten if deploying multiple trials in batch is preferable.
- Parameters:
trials – Iterable of trials to be deployed, each containing arms with parameterizations to be evaluated. Can be a Trial if contains only one arm or a BatchTrial if contains multiple arms.
- Returns:
Dict of trial index to the run metadata of that trial from the deployment process.
- property staging_required: bool
Whether the trial goes to staged or running state once deployed.
- stop(trial: core.base_trial.BaseTrial, reason: str | None = None) dict[str, Any] [source]
Stop a trial based on custom runner subclass implementation.
Optional method.
- Parameters:
trial – The trial to stop.
reason – A message containing information why the trial is to be stopped.
- Returns:
A dictionary of run metadata from the stopping process.
- class ax.SearchSpace(parameters: Sequence[Parameter], parameter_constraints: list[ParameterConstraint] | None = 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: 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 to cast.
- Returns:
New casted arm.
- check_all_parameters_present(parameterization: dict[str, None | str | bool | float | int], raise_error: bool = False) bool [source]
Whether a given parameterization contains all the parameters in the search space.
- Parameters:
parameterization – Dict from parameter name to value to validate.
raise_error – If true parameterization does not belong, raises an error with detailed explanation of why.
- Returns:
Whether the parameterization is contained in the search space.
- check_membership(parameterization: dict[str, None | str | bool | float | int], raise_error: bool = False, check_all_parameters_present: bool = True) bool [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:
parameterization – Dict from parameter name to value to validate.
raise_error – If true parameterization does not belong, raises an error with detailed explanation of why.
check_all_parameters_present – Ensure that parameterization specifies values for all parameters as expected by the search space.
- Returns:
Whether the parameterization is contained in the search space.
- check_types(parameterization: dict[str, None | str | bool | float | int], allow_none: bool = True, allow_extra_params: bool = True, raise_error: bool = False) bool [source]
Checks that the given parameterization’s types match the search space.
- Parameters:
parameterization – Dict from parameter name to value to validate.
allow_none – Whether None is a valid parameter value.
allow_extra_params – If parameterization can have params not in search space.
raise_error – If true and parameterization does not belong, raises an error with detailed explanation of why.
- Returns:
Whether the parameterization has valid types.
- construct_arm(parameters: dict[str, None | str | bool | float | int] | None = None, name: str | None = None) Arm [source]
Construct new arm using given parameters and name. Any missing parameters fallback to the experiment defaults, represented as None.
- out_of_design_arm() 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.
- Returns:
New arm w/ null parameter values.
- property summary_df: DataFrame
Creates a dataframe with information about each parameter in the given search space. The resulting dataframe has one row per parameter, and the following columns:
Name: the name of the parameter.
Type: the parameter subclass (Fixed, Range, Choice).
Domain: the parameter’s domain (e.g., “range=[0, 1]” or “values=[‘a’, ‘b’]”).
Datatype: the datatype of the parameter (int, float, str, bool).
Flags: flags associated with the parameter, if any.
Target Value: the target value of the parameter, if applicable.
Dependent Parameters: for parameters in hierarchical search spaces,
mapping from parameter value -> list of dependent parameter names.
- class ax.SumConstraint(parameters: list[Parameter], is_upper_bound: bool, bound: float)[source]
Constraint on the sum of parameters being greater or less than a bound.
- clone() SumConstraint [source]
Clone.
To use the same constraint, we need to reconstruct the original bound. We do this by re-applying the original bound weighting.
- clone_with_transformed_parameters(transformed_parameters: dict[str, Parameter]) SumConstraint [source]
Clone, but replace parameters with transformed versions.
- property constraint_dict: dict[str, float]
Weights on parameters for linear constraint representation.
- property is_upper_bound: bool
Whether the bound is an upper or lower bound on the sum.
- property op: ComparisonOp
Whether the sum is constrained by a <= or >= inequality.
- class ax.Trial(experiment: core.experiment.Experiment, generator_run: GeneratorRun | None = None, trial_type: str | None = None, ttl_seconds: int | None = None, index: int | None = None)[source]
Trial that only has one attached arm and no arm weights.
- Parameters:
experiment – Experiment, to which this trial is attached.
generator_run – GeneratorRun, associated with this trial. Trial has only one generator run (of just one 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.
trial_type – Type of this trial, if used in MultiTypeExperiment.
ttl_seconds – If specified, trials will be considered failed after this many seconds since the time the trial was ran, unless the trial is completed before then. Meant to be used to detect ‘dead’ trials, for which the evaluation process might have crashed etc., and which should be considered failed after their ‘time to live’ has passed.
index – If specified, the trial’s index will be set accordingly. This should generally not be specified, as in the index will be automatically determined based on the number of existing trials. This is only used for the purpose of loading from storage.
- property arms: list[Arm]
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
- property arms_by_name: dict[str, Arm]
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
- clone_to(experiment: core.experiment.Experiment | None = None) Trial [source]
Clone the trial and attach it to the specified experiment. If no experiment is provided, the original experiment will be used.
- Parameters:
experiment – The experiment to which the cloned trial will belong. If unspecified, uses the current experiment.
- Returns:
A new instance of the trial.
- property generator_run: GeneratorRun | None
Generator run attached to this trial.
- property generator_runs: list[GeneratorRun]
All generator runs associated with this trial.
- get_metric_mean(metric_name: str) float [source]
Metric mean for the arm attached to this trial, retrieved from the latest data available for the metric for the trial.
- property objective_mean: float
Objective mean for the arm attached to this trial, retrieved from the latest data available for the objective for the trial.
Note: the retrieved objective is the experiment-level objective at the time of the call to objective_mean, which is not necessarily the objective that was set at the time the trial was created or ran.
- update_trial_data(raw_data: dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]], metadata: dict[str, str | int] | None = None, sample_size: int | None = None, combine_with_last_data: bool = False) str [source]
Utility method that attaches data to a trial and returns an update message.
- Parameters:
raw_data – Evaluation data for the trial. Can be a mapping from metric name to a tuple of mean and SEM, just a tuple of mean and SEM if only one metric in optimization, or just the mean if SEM is unknown (then Ax will infer observation noise level). Can also be a list of (fidelities, mapping from metric name to a tuple of mean and SEM).
metadata – Additional metadata to track about this run, optional.
sample_size – Number of samples collected for the underlying arm, optional.
combine_with_last_data – Whether to combine the given data with the data that was previously attached to the trial. See Experiment.attach_data for a detailed explanation.
- Returns:
A string message summarizing the update.
- ax.optimize(parameters: list[dict[str, None | str | bool | float | int | Sequence[None | str | bool | float | int] | dict[str, list[str]]]], evaluation_function: Callable[[dict[str, None | str | bool | float | int]], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]] | Callable[[dict[str, None | str | bool | float | int], float | None], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]] | int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None] | list[tuple[dict[str, None | str | bool | float | int], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]] | list[tuple[dict[str, Hashable], dict[str, int | float | floating | integer | tuple[int | float | floating | integer, int | float | floating | integer | None]]]]], experiment_name: str | None = None, objective_name: str | None = None, minimize: bool = False, parameter_constraints: list[str] | None = None, outcome_constraints: list[str] | None = None, total_trials: int = 20, arms_per_trial: int = 1, random_seed: int | None = None, generation_strategy: GenerationStrategy | None = None) tuple[dict[str, None | str | bool | float | int], tuple[dict[str, float], dict[str, dict[str, float]] | None] | None, Experiment, ModelBridge | None] [source]
Construct and run a full optimization loop.